Repository: xataio/pgstream
Pull request #418 (2025-06-24 13:48:41)

[Snapshot] Fix pg_dump multi schema support
This PR updates the pgdump/restore snapshot generator to not receive the exploded schema/tables values obtained by the schema/table discovery step, but rather the ones originally requested by the user, since the wildcard values both for tables and schemas, affect the options passed to the `pg_dump` command. 

If we receive all the explicit schemas/tables on input, we'll apply filters that are not necessary and will in some cases produce different outputs than the expected ones. By having the wildcard component, we can optimise and simplify the pg_dump options.

Fixes https://github.com/xataio/xata-actions-demo/actions/runs/15829106209/job/44616940684.


Patch:
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
index 7b4bc4c..8a97017 100644
--- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
@@ -221,7 +221,7 @@ func (s *SnapshotGenerator) restoreDump(ctx context.Context, schemaTables map[st
 	// will not be dumped, so it needs to be created explicitly (except for
 	// public schema)
 	for schema, tables := range schemaTables {
-		if len(tables) > 0 && schema != publicSchema {
+		if len(tables) > 0 && schema != publicSchema && schema != wildcard {
 			if err := s.createSchemaIfNotExists(ctx, schema); err != nil {
 				return err
 			}
@@ -296,7 +296,11 @@ func (s *SnapshotGenerator) pgdumpOptions(ctx context.Context, schemaTables map[
 		opts.NoOwner = true
 	}
 
-	if s.includeGlobalDBObjects {
+	switch {
+	case hasWildcardSchema(schemaTables):
+		// no need to filter schemas, since we are including all of them
+		opts.Schemas = nil
+	case s.includeGlobalDBObjects:
 		// instead of using the schema filter, we use the exclude schemas filter
 		// to make sure extensions and other database global objects are
 		// created. pg_dump will not include them when using the schema filter,
@@ -329,7 +333,10 @@ func (s *SnapshotGenerator) pgdumpOptions(ctx context.Context, schemaTables map[
 	return opts, nil
 }
 
-var selectTablesQuery = "SELECT tablename FROM pg_tables WHERE schemaname = '%s' AND tablename NOT IN (%s)"
+const (
+	selectTablesQuery       = "SELECT tablename FROM pg_tables WHERE tablename NOT IN (%s)"
+	selectSchemaTablesQuery = "SELECT tablename FROM pg_tables WHERE schemaname = '%s' AND tablename NOT IN (%s)"
+)
 
 func (s *SnapshotGenerator) pgdumpExcludedTables(ctx context.Context, schemaName string, includeTables []string) ([]string, error) {
 	conn, err := s.connBuilder(ctx, s.sourceURL)
@@ -345,8 +352,16 @@ func (s *SnapshotGenerator) pgdumpExcludedTables(ctx context.Context, schemaName
 		paramRefs = append(paramRefs, fmt.Sprintf("$%d", i+1))
 	}
 
+	var query string
+	switch schemaName {
+	case wildcard:
+		query = fmt.Sprintf(selectTablesQuery, strings.Join(paramRefs, ","))
+	default:
+		// if the schema is not wildcard, we need to filter by schema name
+		query = fmt.Sprintf(selectSchemaTablesQuery, schemaName, strings.Join(paramRefs, ","))
+	}
+
 	// get all tables in the schema that are not in the include list
-	query := fmt.Sprintf(selectTablesQuery, schemaName, strings.Join(paramRefs, ","))
 	rows, err := conn.Query(ctx, query, tableParams...)
 	if err != nil {
 		return nil, fmt.Errorf("retrieving tables from schema: %w", err)
@@ -369,7 +384,7 @@ func (s *SnapshotGenerator) pgdumpExcludedTables(ctx context.Context, schemaName
 	return excludeTables, nil
 }
 
-var selectSchemasQuery = "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN (%s)"
+const selectSchemasQuery = "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN (%s)"
 
 func (s *SnapshotGenerator) pgdumpExcludedSchemas(ctx context.Context, includeSchemas []string) ([]string, error) {
 	conn, err := s.connBuilder(ctx, s.sourceURL)
@@ -507,5 +522,9 @@ func (s *SnapshotGenerator) parseDump(d []byte) *dump {
 }
 
 func hasWildcardTable(tables []string) bool {
-	return slices.Contains(tables, "*")
+	return slices.Contains(tables, wildcard)
+}
+
+func hasWildcardSchema(schemaTables map[string][]string) bool {
+	return schemaTables[wildcard] != nil
 }
diff --git a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
index 2a7b43f..e1a0bd2 100644
--- a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
+++ b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
@@ -67,6 +67,17 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, p li
 		return nil, err
 	}
 
+	// snapshot table finder layer
+	finderOpts := []pgtablefinder.Option{}
+	if instrumentation.IsEnabled() {
+		finderOpts = append(finderOpts, pgtablefinder.WithInstrumentation(instrumentation))
+	}
+
+	g, err = pgtablefinder.NewSnapshotSchemaTableFinder(ctx, cfg.Generator.URL, g, finderOpts...)
+	if err != nil {
+		return nil, err
+	}
+
 	// postgres schema snapshot generator layer
 	g, err = newSchemaSnapshotGenerator(ctx, &cfg.Schema, g, rowsProcessor.ProcessRow, logger, instrumentation)
 	if err != nil {
@@ -86,17 +97,6 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, p li
 		g = generator.NewSnapshotRecorder(snapshotStore, g, cfg.Recorder.RepeatableSnapshots)
 	}
 
-	// snapshot table finder layer
-	finderOpts := []pgtablefinder.Option{}
-	if instrumentation.IsEnabled() {
-		finderOpts = append(finderOpts, pgtablefinder.WithInstrumentation(instrumentation))
-	}
-
-	g, err = pgtablefinder.NewSnapshotSchemaTableFinder(ctx, cfg.Generator.URL, g, finderOpts...)
-	if err != nil {
-		return nil, err
-	}
-
 	if instrumentation.IsEnabled() {
 		g, err = generatorinstrumentation.NewSnapshotGenerator(g, instrumentation)
 		if err != nil {


Test patch:
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
index debb464..e0a43ea 100644
--- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
@@ -58,7 +58,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						},
 						ErrFn: func() error { return nil },
 					}, nil
-				case fmt.Sprintf(selectTablesQuery, testSchema, "$1"):
+				case fmt.Sprintf(selectSchemaTablesQuery, testSchema, "$1"):
+					require.Equal(t, []any{testTable}, args)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				case fmt.Sprintf(selectTablesQuery, "$1"):
 					require.Equal(t, []any{testTable}, args)
 					return &mocks.Rows{
 						CloseFn: func() {},
@@ -232,10 +246,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			wantErr: nil,
 		},
 		{
-			name: "ok - wildcard",
+			name: "ok - wildcard table",
 			snapshot: &snapshot.Snapshot{
 				SchemaTables: map[string][]string{
-					testSchema: {"*"},
+					testSchema: {wildcard},
 				},
 			},
 			conn: validQuerier(),
@@ -272,6 +286,87 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			wantErr: nil,
 		},
+		{
+			name: "ok - wildcard schema",
+			snapshot: &snapshot.Snapshot{
+				SchemaTables: map[string][]string{
+					wildcard: {testTable},
+				},
+			},
+			conn: validQuerier(),
+			pgdumpFn: newMockPgdump(func(_ context.Context, i uint, po pglib.PGDumpOptions) ([]byte, error) {
+				switch i {
+				case 1:
+					require.Equal(t, pglib.PGDumpOptions{
+						ConnectionString: "source-url",
+						Format:           "p",
+						SchemaOnly:       true,
+						ExcludeTables:    []string{pglib.QuoteIdentifier(excludedTable)},
+					}, po)
+					return schemaDump, nil
+				case 2:
+					require.Equal(t, pglib.PGDumpOptions{
+						ConnectionString: "source-url",
+						Format:           "p",
+						DataOnly:         true,
+						Tables:           []string{testSequence},
+					}, po)
+					return sequenceDump, nil
+				default:
+					return nil, fmt.Errorf("unexpected call to pgdumpFn: %d", i)
+				}
+			}),
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				require.Equal(t, pglib.PGRestoreOptions{
+					ConnectionString: "target-url",
+					Format:           "p",
+				}, po)
+				require.Equal(t, append(schemaDump, sequenceDump...), dump)
+				return "", nil
+			},
+
+			wantErr: nil,
+		},
+		{
+			name: "ok - wildcard schema and table",
+			snapshot: &snapshot.Snapshot{
+				SchemaTables: map[string][]string{
+					wildcard: {wildcard},
+				},
+			},
+			conn: validQuerier(),
+			pgdumpFn: newMockPgdump(func(_ context.Context, i uint, po pglib.PGDumpOptions) ([]byte, error) {
+				switch i {
+				case 1:
+					require.Equal(t, pglib.PGDumpOptions{
+						ConnectionString: "source-url",
+						Format:           "p",
+						SchemaOnly:       true,
+					}, po)
+					return schemaDump, nil
+				case 2:
+					require.Equal(t, pglib.PGDumpOptions{
+						ConnectionString: "source-url",
+						Format:           "p",
+						DataOnly:         true,
+						Tables:           []string{testSequence},
+					}, po)
+					return sequenceDump, nil
+				default:
+					return nil, fmt.Errorf("unexpected call to pgdumpFn: %d", i)
+				}
+			}),
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				require.Equal(t, pglib.PGRestoreOptions{
+					ConnectionString: "target-url",
+					Format:           "p",
+				}, po)
+				require.Equal(t, append(schemaDump, sequenceDump...), dump)
+				return "", nil
+			},
+
+			wantErr: nil,
+		},
 		{
 			name: "ok - no tables in public schema",
 			snapshot: &snapshot.Snapshot{
@@ -341,7 +436,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 							},
 							ErrFn: func() error { return nil },
 						}, nil
-					case fmt.Sprintf(selectTablesQuery, testSchema, "$1"):
+					case fmt.Sprintf(selectSchemaTablesQuery, testSchema, "$1"):
 						require.Equal(t, []any{testTable}, args)
 						return &mocks.Rows{
 							CloseFn: func() {},
@@ -393,7 +488,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 							},
 							ErrFn: func() error { return nil },
 						}, nil
-					case fmt.Sprintf(selectTablesQuery, testSchema, "$1"):
+					case fmt.Sprintf(selectSchemaTablesQuery, testSchema, "$1"):
 						require.Equal(t, []any{testTable}, args)
 						return &mocks.Rows{
 							CloseFn: func() {},