Repository: xataio/pgstream
Pull request #230 (2025-04-01 12:19:44)

[Snapshot] `pg_dump` snapshot fixes
This PR adds a few small bug fixes:
- Add quote identifiers to snapshot schema/tables
- Do not translate wildcard `*` tables for pgdump schema snapshots, since it prevents all resources in a given schema to be dumped by having a list of tables.


Patch:
diff --git a/internal/postgres/pg_utils.go b/internal/postgres/pg_utils.go
new file mode 100644
index 0000000..e2e0708
--- /dev/null
+++ b/internal/postgres/pg_utils.go
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "github.com/lib/pq"
+
+func QuoteIdentifier(s string) string {
+	return pq.QuoteIdentifier(s)
+}
+
+func QuoteQualifiedIdentifier(schema, table string) string {
+	return pq.QuoteIdentifier(schema) + "." + pq.QuoteIdentifier(table)
+}
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 618ca3f..083932b 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -9,7 +9,6 @@ import (
 	"sync"
 
 	"github.com/jackc/pgx/v5/pgconn"
-	"github.com/lib/pq"
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/snapshot"
@@ -186,8 +185,8 @@ func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID,
 			"schema": schema, "table": table, "snapshotID": snapshotID,
 		})
 
-		query := fmt.Sprintf("SELECT * FROM %s.%s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'",
-			pq.QuoteIdentifier(schema), pq.QuoteIdentifier(table), pageRange.start, pageRange.end)
+		query := fmt.Sprintf("SELECT * FROM %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'",
+			pglib.QuoteQualifiedIdentifier(schema, table), pageRange.start, pageRange.end)
 		rows, err := tx.Query(ctx, query)
 		if err != nil {
 			return fmt.Errorf("querying table rows: %w", err)
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 817be67..35062de 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
@@ -145,11 +145,18 @@ func (s *SnapshotGenerator) pgdumpOptions(ss *snapshot.Snapshot) pglib.PGDumpOpt
 		ConnectionString: s.sourceURL,
 		Format:           "c",
 		SchemaOnly:       true,
-		Schemas:          []string{ss.SchemaName},
+		Schemas:          []string{pglib.QuoteIdentifier(ss.SchemaName)},
 	}
 
+	const wildcard = "*"
 	for _, table := range ss.TableNames {
-		opts.Tables = append(opts.Tables, ss.SchemaName+"."+table)
+		if table == wildcard {
+			// wildcard means all tables in the schema, so no table filter
+			// required
+			opts.Tables = nil
+			break
+		}
+		opts.Tables = append(opts.Tables, pglib.QuoteQualifiedIdentifier(ss.SchemaName, table))
 	}
 
 	return opts
diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store.go b/pkg/snapshot/store/postgres/pg_snapshot_store.go
index 63a0a63..d56a94a 100644
--- a/pkg/snapshot/store/postgres/pg_snapshot_store.go
+++ b/pkg/snapshot/store/postgres/pg_snapshot_store.go
@@ -135,5 +135,5 @@ func (s *Store) createTable(ctx context.Context) error {
 }
 
 func snapshotsTable() string {
-	return fmt.Sprintf("%s.%s", pq.QuoteIdentifier(store.SchemaName), pq.QuoteIdentifier(store.TableName))
+	return postgres.QuoteQualifiedIdentifier(store.SchemaName, store.TableName)
 }
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 e4502d3..52735a5 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
@@ -51,22 +51,22 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 		return nil, err
 	}
 
+	var dataSnapshotGenerator generator.SnapshotGenerator
 	// postgres data snapshot generator
-	dataSnapshotGenerator, err := pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, processEventAdapter.ProcessRow, pgsnapshotgenerator.WithLogger(logger))
+	dataSnapshotGenerator, err = pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, processEventAdapter.ProcessRow, pgsnapshotgenerator.WithLogger(logger))
 	if err != nil {
 		return nil, err
 	}
-
-	var g generator.SnapshotGenerator
-	// snapshot generator aggregator
-	g = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
-
 	// snapshot table finder layer
-	g, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, g)
+	dataSnapshotGenerator, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, dataSnapshotGenerator)
 	if err != nil {
 		return nil, err
 	}
 
+	var g generator.SnapshotGenerator
+	// snapshot generator aggregator
+	g = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
+
 	if cfg.SnapshotStoreURL != "" {
 		// snapshot activity recorder layer
 		snapshotStore, err := pgsnapshotstore.New(ctx, cfg.SnapshotStoreURL)
diff --git a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
index 92197d1..d27e570 100644
--- a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
+++ b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
@@ -8,7 +8,7 @@ import (
 	"slices"
 	"strings"
 
-	"github.com/lib/pq"
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/pkg/wal"
 )
 
@@ -134,7 +134,7 @@ func (a *dmlAdapter) buildOnConflictQuery(d *wal.Data) string {
 
 		cols := make([]string, 0, len(d.Columns))
 		for _, col := range d.Columns {
-			cols = append(cols, fmt.Sprintf("%[1]s = EXCLUDED.%[1]s", pq.QuoteIdentifier(col.Name)))
+			cols = append(cols, fmt.Sprintf("%[1]s = EXCLUDED.%[1]s", pglib.QuoteIdentifier(col.Name)))
 		}
 		return fmt.Sprintf(" ON CONFLICT (%s) DO UPDATE SET %s", strings.Join(primaryKeyCols, ","), strings.Join(cols, ", "))
 	case onConflictDoNothing:
@@ -169,7 +169,7 @@ func (a *dmlAdapter) extractPrimaryKeyColumnNames(colIDs []string, cols []wal.Co
 }
 
 func quotedTableName(schemaName, tableName string) string {
-	return fmt.Sprintf("%s.%s", pq.QuoteIdentifier(schemaName), pq.QuoteIdentifier(tableName))
+	return pglib.QuoteQualifiedIdentifier(schemaName, tableName)
 }
 
 func parseOnConflictAction(action string) (onConflictAction, error) {


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 77ec633..a28232a 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
@@ -53,8 +53,40 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					ConnectionString: "source-url",
 					Format:           "c",
 					SchemaOnly:       true,
-					Schemas:          []string{testSchema},
-					Tables:           []string{testSchema + "." + testTable},
+					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
+					Tables:           []string{pglib.QuoteQualifiedIdentifier(testSchema, testTable)},
+				}, po)
+				return testDump, nil
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				require.Equal(t, pglib.PGRestoreOptions{
+					ConnectionString: "target-url",
+					SchemaOnly:       true,
+				}, po)
+				require.Equal(t, testDump, dump)
+				return "", nil
+			},
+
+			wantErr: nil,
+		},
+		{
+			name: "ok - wildcard",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"*"},
+			},
+			conn: &mocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+					require.Equal(t, "CREATE SCHEMA IF NOT EXISTS "+testSchema, query)
+					return pglib.CommandTag{}, nil
+				},
+			},
+			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+				require.Equal(t, pglib.PGDumpOptions{
+					ConnectionString: "source-url",
+					Format:           "c",
+					SchemaOnly:       true,
+					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
 				}, po)
 				return testDump, nil
 			},