Improve schema replacement for statistics restore

end_to_end/end_to_end_suite_test.go

@@ -969,18 +969,28 @@ var _ = Describe("backup and restore end to end tests", func() {
 				"DROP INDEX schema2.foo3_idx1")
 			testhelper.AssertQueryRuns(backupConn,
 				"ANALYZE schema2.foo3")
+			testhelper.AssertQueryRuns(backupConn,
+				`CREATE TABLE schema2.foo4 ("schema2." TEXT)`)
+			defer testhelper.AssertQueryRuns(backupConn,
+				"DROP TABLE schema2.foo4")
+			testhelper.AssertQueryRuns(backupConn,
+				`INSERT INTO schema2.foo4 VALUES ('schema2.'),('schema2.')`)
+			testhelper.AssertQueryRuns(backupConn,
+				"ANALYZE schema2.foo4")
 			output := gpbackup(gpbackupPath, backupHelperPath,
 				"--with-stats")
 			timestamp := getBackupTimestamp(string(output))
 
 			gprestore(gprestorePath, restoreHelperPath, timestamp,
 				"--redirect-db", "restoredb",
 				"--include-table", "schema2.foo3",
+				"--include-table", "schema2.foo4",
 				"--redirect-schema", "schema3",
 				"--with-stats")
 
 			schema3TupleCounts := map[string]int{
 				"schema3.foo3": 100,
+				"schema3.foo4": 2,
 			}
 			assertDataRestored(restoreConn, schema3TupleCounts)
 			assertPGClassStatsRestored(restoreConn, restoreConn, schema3TupleCounts)
@@ -992,6 +1002,14 @@ var _ = Describe("backup and restore end to end tests", func() {
 			actualStatisticCount := dbconn.MustSelectString(restoreConn,
 				`SELECT count(*) FROM pg_statistic WHERE starelid='schema3.foo3'::regclass::oid;`)
 			Expect(actualStatisticCount).To(Equal("1"))
+
+			actualStatisticCount = dbconn.MustSelectString(restoreConn,
+				`SELECT count(*) FROM pg_statistic WHERE starelid='schema3.foo4'::regclass::oid;`)
+			Expect(actualStatisticCount).To(Equal("1"))
+
+			stavaluesTableFoo4 := dbconn.MustSelectString(restoreConn,
+				`SELECT stavalues1 FROM pg_statistic WHERE starelid='schema3.foo4'::regclass::oid;`)
+			Expect(stavaluesTableFoo4).To(Equal("{schema2.}"))
 		})
 		It("runs gprestore with --redirect-schema to redirect data back to the original database which still contain the original tables", func() {
 			skipIfOldBackupVersionBefore("1.17.0")

restore/restore.go

@@ -410,11 +410,13 @@ func editStatementsRedirectSchema(statements []toc.StatementWithType, redirectSc
 		return
 	}
 
-	schemaMatch := `(?:".+?"|[^.]+?)` // matches either an unquoted schema with no dots or a quoted schema containing dots
+	schemaMatch := `(?:".+?"|[^."']+?)` // matches either an unquoted schema with no dots and no quotes or a quoted schema
 	// This expression matches a GRANT or REVOKE statement on any object and captures the old schema name
 	permissionsRE := regexp.MustCompile(fmt.Sprintf(`(?m)(^(?:REVOKE|GRANT) .+ ON .+?) (%s)((\..+)? (?:FROM|TO) .+)`, schemaMatch))
 	// This expression matches an ATTACH PARTITION statement and captures both the parent and child schema names
 	attachRE := regexp.MustCompile(fmt.Sprintf(`(ALTER TABLE(?: ONLY)?) (%[1]s)(\..+ ATTACH PARTITION) (%[1]s)(\..+)`, schemaMatch))
+	// This expression matches a '<schema>.<table>'::regclass::oid expression
+	regclassOidRE := regexp.MustCompile(fmt.Sprintf(`(')(%s)((\.[^']+)'\:\:regclass\:\:oid)`, schemaMatch))
 	for i := range statements {
 		oldSchema := fmt.Sprintf("%s.", statements[i].Schema)
 		newSchema := fmt.Sprintf("%s.", redirectSchema)
@@ -442,7 +444,7 @@ func editStatementsRedirectSchema(statements []toc.StatementWithType, redirectSc
 
 		// Statistic statements need two schema replacements
 		if strings.Contains(statement, "pg_statistic") {
-			statement = strings.ReplaceAll(statement, oldSchema, newSchema)
+			statement = regclassOidRE.ReplaceAllString(statement, fmt.Sprintf("'%s$3", redirectSchema))
 			replaced = true
 		}