Merge pull request #152 from xataio/use-pgdumprestore-snapshot-generator

Use pg_dump/pg_restore snapshot generator in pg listener

Author: eminano

build/docker/postgres/Dockerfile

@@ -1,4 +1,4 @@
-FROM postgres:15.3
+FROM postgres:17.2
 
 RUN apt-get update \
-  && apt-get install -y postgresql-15-wal2json \
+  && apt-get install -y postgresql-17-wal2json \

cmd/config.go

@@ -10,6 +10,7 @@ import (
 	"github.com/xataio/pgstream/pkg/kafka"
 	pgschemalog "github.com/xataio/pgstream/pkg/schemalog/postgres"
 	pgsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/data"
+	"github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore"
 	"github.com/xataio/pgstream/pkg/stream"
 	"github.com/xataio/pgstream/pkg/tls"
 	kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka"
@@ -86,9 +87,6 @@ func parsePostgresListenerConfig() *stream.PostgresListenerConfig {
 	snapshotTables := viper.GetStringSlice("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES")
 	if len(snapshotTables) > 0 {
 		cfg.Snapshot = &pglistener.SnapshotConfig{
-			SchemaLogStore: pgschemalog.Config{
-				URL: pgURL,
-			},
 			SnapshotStoreURL: pgURL,
 			Generator: pgsnapshotgenerator.Config{
 				URL:           pgURL,
@@ -98,12 +96,30 @@ func parsePostgresListenerConfig() *stream.PostgresListenerConfig {
 			},
 			Tables:          snapshotTables,
 			SnapshotWorkers: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS"),
+			Schema:          parseSchemaSnapshotConfig(pgURL),
 		}
 	}
 
 	return cfg
 }
 
+func parseSchemaSnapshotConfig(pgurl string) pglistener.SchemaSnapshotConfig {
+	pgTargetURL := viper.GetString("PGSTREAM_POSTGRES_WRITER_TARGET_URL")
+	if pgTargetURL != "" {
+		return pglistener.SchemaSnapshotConfig{
+			DumpRestore: &pgdumprestore.Config{
+				SourcePGURL: pgurl,
+				TargetPGURL: pgTargetURL,
+			},
+		}
+	}
+	return pglistener.SchemaSnapshotConfig{
+		SchemaLogStore: &pgschemalog.Config{
+			URL: pgurl,
+		},
+	}
+}
+
 func parseKafkaListenerConfig() *stream.KafkaListenerConfig {
 	kafkaServers := viper.GetStringSlice("PGSTREAM_KAFKA_SERVERS")
 	kafkaTopic := viper.GetString("PGSTREAM_KAFKA_TOPIC_NAME")

pkg/snapshot/generator/postgres/data/pg_schema_table_parser.go

@@ -17,10 +17,9 @@ import (
 )
 
 type SnapshotGenerator struct {
-	logger      loglib.Logger
-	conn        pglib.Querier
-	mapper      mapper
-	tableParser tableParser
+	logger loglib.Logger
+	conn   pglib.Querier
+	mapper mapper
 
 	schemaWorkers uint
 	tableWorkers  uint
@@ -39,8 +38,6 @@ type pageRange struct {
 	end   uint
 }
 
-type tableParser func(ctx context.Context, snapshot *snapshot.Snapshot) error
-
 type Option func(sg *SnapshotGenerator)
 
 func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.RowProcessor, opts ...Option) (*SnapshotGenerator, error) {
@@ -57,7 +54,6 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.
 		batchPageSize: cfg.batchPageSize(),
 		tableWorkers:  cfg.tableWorkers(),
 		schemaWorkers: cfg.schemaWorkers(),
-		tableParser:   newSchemaTableParser(conn).parseSnapshotTables,
 	}
 
 	for _, opt := range opts {
@@ -76,9 +72,6 @@ func WithLogger(logger loglib.Logger) Option {
 }
 
 func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
-	if err := sg.tableParser(ctx, ss); err != nil {
-		return err
-	}
 	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
 		snapshotID, err := sg.exportSnapshot(ctx, tx)
 		if err != nil {

pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go

@@ -872,7 +872,6 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				schemaWorkers: 1,
 				tableWorkers:  1,
 				batchPageSize: 10,
-				tableParser:   func(ctx context.Context, snapshot *snapshot.Snapshot) error { return nil },
 			}
 
 			if tc.schemaWorkers != 0 {

pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package tablefinder
+
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/generator"
+)
+
+// SnapshotTableFinder is a decorator around a snapshot generator that will
+// explode the wildcard references in the snapshot requests and replace them by
+// all the schema tables in postgres.
+type SnapshotTableFinder struct {
+	wrapped generator.SnapshotGenerator
+	conn    pglib.Querier
+}
+
+const wildcard = "*"
+
+// NewSnapshotTableFinder will return the generator on input wrapped with a
+// table finder that will explode the wildcard references in the snapshot
+// request and translate them into all the postgres tables for the given schema.
+func NewSnapshotTableFinder(ctx context.Context, pgurl string, generator generator.SnapshotGenerator) (*SnapshotTableFinder, error) {
+	conn, err := pglib.NewConnPool(ctx, pgurl)
+	if err != nil {
+		return nil, err
+	}
+	return &SnapshotTableFinder{
+		wrapped: generator,
+		conn:    conn,
+	}, nil
+}
+
+func (s *SnapshotTableFinder) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+	if slices.Contains(ss.TableNames, wildcard) {
+		var err error
+		ss.TableNames, err = s.discoverAllSchemaTables(ctx, ss.SchemaName)
+		if err != nil {
+			return err
+		}
+	}
+	return s.wrapped.CreateSnapshot(ctx, ss)
+}
+
+func (s *SnapshotTableFinder) Close() error {
+	return s.conn.Close(context.Background())
+}
+
+func (s *SnapshotTableFinder) discoverAllSchemaTables(ctx context.Context, schema string) ([]string, error) {
+	const query = "SELECT tablename FROM pg_tables WHERE schemaname=$1"
+	rows, err := s.conn.Query(ctx, query, schema)
+	if err != nil {
+		return nil, fmt.Errorf("discovering all tables for schema %s: %w", schema, err)
+	}
+	defer rows.Close()
+
+	tableNames := []string{}
+	for rows.Next() {
+		var tableName string
+		if err := rows.Scan(&tableName); err != nil {
+			return nil, fmt.Errorf("scanning table name: %w", err)
+		}
+		tableNames = append(tableNames, tableName)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return tableNames, nil
+}