Repository: xataio/pgstream
Pull request #152 (2025-02-19 13:44:13)

Use pg_dump/pg_restore snapshot generator in pg listener
This PR is a follow up of https://github.com/xataio/pgstream/pull/149. It uses the schema snapshot generator added in that PR in the pg listener initial snapshot. 
The table mapper logic from the pg data generator implementation in charge of discovering schema tables when wildcards are used, is extracted into a decorator so that the discovery can be used for both data and schema snapshot generators.


Patch:
diff --git a/build/docker/postgres/Dockerfile b/build/docker/postgres/Dockerfile
index 67ebbb2..0ad26b2 100644
--- a/build/docker/postgres/Dockerfile
+++ b/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 \
diff --git a/cmd/config.go b/cmd/config.go
index f26b5ad..712eadd 100644
--- a/cmd/config.go
+++ b/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")
diff --git a/pkg/snapshot/generator/postgres/data/pg_schema_table_parser.go b/pkg/snapshot/generator/postgres/data/pg_schema_table_parser.go
deleted file mode 100644
index a636a22..0000000
--- a/pkg/snapshot/generator/postgres/data/pg_schema_table_parser.go
+++ /dev/null
@@ -1,58 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-
-package postgres
-
-import (
-	"context"
-	"fmt"
-	"slices"
-
-	pglib "github.com/xataio/pgstream/internal/postgres"
-	"github.com/xataio/pgstream/pkg/snapshot"
-)
-
-type schemaTableParser struct {
-	conn pglib.Querier
-}
-
-func newSchemaTableParser(conn pglib.Querier) *schemaTableParser {
-	return &schemaTableParser{
-		conn: conn,
-	}
-}
-
-func (p *schemaTableParser) parseSnapshotTables(ctx context.Context, snapshot *snapshot.Snapshot) error {
-	if slices.Contains(snapshot.TableNames, "*") {
-		var err error
-		snapshot.TableNames, err = p.discoverAllSchemaTables(ctx, snapshot.SchemaName)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (p *schemaTableParser) discoverAllSchemaTables(ctx context.Context, schema string) ([]string, error) {
-	const query = "SELECT tablename FROM pg_tables WHERE schemaname=$1"
-	rows, err := p.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
-}
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 0c17a24..05d9c28 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.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 {
diff --git a/pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder.go b/pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder.go
new file mode 100644
index 0000000..3718149
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder.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
+}
diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go
index 470a534..d060c9e 100644
--- a/pkg/stream/stream_run.go
+++ b/pkg/stream/stream_run.go
@@ -11,7 +11,6 @@ import (
 	kafkainstrumentation "github.com/xataio/pgstream/pkg/kafka/instrumentation"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/otel"
-	pgsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/data"
 	"github.com/xataio/pgstream/pkg/wal/checkpointer"
 	kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka"
 	pgcheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/postgres"
@@ -263,7 +262,7 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, instrumentat
 				ctx,
 				config.Listener.Postgres.Snapshot,
 				processor.ProcessWALEvent,
-				pgsnapshotgenerator.WithLogger(logger))
+				logger)
 			if err != nil {
 				return err
 			}
diff --git a/pkg/wal/listener/postgres/config.go b/pkg/wal/listener/postgres/config.go
index 0a64ab6..601a6c9 100644
--- a/pkg/wal/listener/postgres/config.go
+++ b/pkg/wal/listener/postgres/config.go
@@ -7,19 +7,25 @@ import (
 
 	schemalogpg "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"
 )
 
 type SnapshotConfig struct {
 	Generator        pgsnapshotgenerator.Config
 	SnapshotStoreURL string
 	Tables           []string
-	SchemaLogStore   schemalogpg.Config
+	Schema           SchemaSnapshotConfig
 	// SnapshotWorkers represents the number of snapshots the generator will
 	// process concurrently. This doesn't affect the parallelism of the tables
 	// within each individual snapshot request. It defaults to 1.
 	SnapshotWorkers uint
 }
 
+type SchemaSnapshotConfig struct {
+	SchemaLogStore *schemalogpg.Config
+	DumpRestore    *pgdumprestore.Config
+}
+
 const defaultSnapshotWorkers = 1
 
 func (c *SnapshotConfig) schemaTableMap() map[string][]string {
diff --git a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go b/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go
index 0f9ab57..5284d6c 100644
--- a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go
+++ b/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go
@@ -4,6 +4,7 @@ package postgres
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -14,7 +15,9 @@ import (
 	"github.com/xataio/pgstream/pkg/snapshot"
 	"github.com/xataio/pgstream/pkg/snapshot/generator"
 	pgsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/data"
-	pgschemasnapshot "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/schemalog"
+	pgdumprestoregenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore"
+	schemalogsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/schemalog"
+	pgtablefinder "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/tablefinder"
 	pgsnapshotstore "github.com/xataio/pgstream/pkg/snapshot/store/postgres"
 	"github.com/xataio/pgstream/pkg/wal"
 	"golang.org/x/sync/errgroup"
@@ -34,7 +37,27 @@ const (
 	publicSchema = "public"
 )
 
-func NewSnapshotGeneratorAdapter(ctx context.Context, cfg *SnapshotConfig, processEvent listenerProcessWalEvent, opts ...pgsnapshotgenerator.Option) (*SnapshotGeneratorAdapter, error) {
+var errSchemaSnapshotNotConfigured = errors.New("no schema snapshot has been configured")
+
+// NewSnapshotGeneratorAdapter builds a snapshot generator based on the
+// configuration on input, adapting the process wal event to fit with the
+// snapshot package implementation. It will add an activity recorder layer if
+// configured which will keep track of snapshot requests and their status.
+//
+// ┌───────────────────────────────────────────────────┐
+// │               Snapshot Generator Adapter          │
+// │ ┌─────────┐ ┌────────┐ ┌────────────────────────┐ │
+// │ │         │ │        │ │       Aggregator       │ │
+// │ │         │ │        │ │  ┌──────────────────┐  │ │
+// │ │         │ │        │ │  │ Schema Snapshot  │  │ │
+// │ │Snapshot │ │ Table  │ │  └─────────┬────────┘  │ │
+// │ │Activity │─▶ Finder ├▶│            │           │ │
+// │ │Recorder │ │        │ │  ┌─────────▼────────┐  │ │
+// │ │         │ │        │ │  │  Data Snapshot   │  │ │
+// │ │         │ │        │ │  └──────────────────┘  │ │
+// │ └─────────┘ └────────┘ └────────────────────────┘ │
+// └───────────────────────────────────────────────────┘
+func NewSnapshotGeneratorAdapter(ctx context.Context, cfg *SnapshotConfig, processEvent listenerProcessWalEvent, logger loglib.Logger) (*SnapshotGeneratorAdapter, error) {
 	s := &SnapshotGeneratorAdapter{
 		processEvent:    processEvent,
 		schemaTables:    cfg.schemaTableMap(),
@@ -44,28 +67,34 @@ func NewSnapshotGeneratorAdapter(ctx context.Context, cfg *SnapshotConfig, proce
 	}
 
 	// postgres schema snapshot generator
-	var schemaLogStore schemalog.Store
-	var err error
-	schemaLogStore, err = schemalogpg.NewStore(ctx, cfg.SchemaLogStore)
+	schemaSnapshotGenerator, err := s.newSchemaSnapshotGenerator(ctx, cfg.Schema, logger)
 	if err != nil {
-		return nil, fmt.Errorf("create schema log postgres store: %w", err)
+		return nil, err
 	}
-	schemaLogStore = schemalog.NewStoreCache(schemaLogStore)
-	schemaSnapshotGenerator := pgschemasnapshot.NewSnapshotGenerator(schemaLogStore, s.processRow)
 
 	// postgres data snapshot generator
-	dataSnapshotGenerator, err := pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, s.processRow, opts...)
+	dataSnapshotGenerator, err := pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, s.processRow, pgsnapshotgenerator.WithLogger(logger))
 	if err != nil {
 		return nil, err
 	}
 
+	// snapshot generator aggregator
 	s.generator = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
-	// snapshot activity recorder layer
-	snapshotStore, err := pgsnapshotstore.New(ctx, cfg.SnapshotStoreURL)
+
+	// snapshot table finder layer
+	s.generator, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, s.generator)
 	if err != nil {
-		return nil, fmt.Errorf("create postgres snapshot store: %w", err)
+		return nil, err
+	}
+
+	if cfg.SnapshotStoreURL != "" {
+		// snapshot activity recorder layer
+		snapshotStore, err := pgsnapshotstore.New(ctx, cfg.SnapshotStoreURL)
+		if err != nil {
+			return nil, fmt.Errorf("create postgres snapshot store: %w", err)
+		}
+		s.generator = generator.NewSnapshotRecorder(snapshotStore, s.generator)
 	}
-	s.generator = generator.NewSnapshotRecorder(snapshotStore, s.generator)
 
 	return s, nil
 }
@@ -76,6 +105,7 @@ func (s *SnapshotGeneratorAdapter) CreateSnapshot(ctx context.Context) error {
 	for i := uint(0); i < s.snapshotWorkers; i++ {
 		errGroup.Go(func() error {
 			for snapshot := range snapshotChan {
+				s.logger.Info("creating snapshot for schema", loglib.Fields{"schema": snapshot.SchemaName, "tables": snapshot.TableNames})
 				if err := s.generator.CreateSnapshot(ctx, snapshot); err != nil {
 					s.logger.Error(err, "creating snapshot for schema", loglib.Fields{"schema": snapshot.SchemaName, "tables": snapshot.TableNames})
 					return err
@@ -136,3 +166,23 @@ func (s *SnapshotGeneratorAdapter) snapshotColumnToWalColumn(col snapshot.Column
 		Value: col.Value,
 	}
 }
+
+func (s *SnapshotGeneratorAdapter) newSchemaSnapshotGenerator(ctx context.Context, cfg SchemaSnapshotConfig, logger loglib.Logger) (generator.SnapshotGenerator, error) {
+	switch {
+	case cfg.SchemaLogStore != nil:
+		// postgres schemalog schema snapshot generator
+		var schemaLogStore schemalog.Store
+		var err error
+		schemaLogStore, err = schemalogpg.NewStore(ctx, *cfg.SchemaLogStore)
+		if err != nil {
+			return nil, fmt.Errorf("create schema log postgres store: %w", err)
+		}
+		schemaLogStore = schemalog.NewStoreCache(schemaLogStore)
+		return schemalogsnapshotgenerator.NewSnapshotGenerator(schemaLogStore, s.processRow), nil
+	case cfg.DumpRestore != nil:
+		// postgres pgdump/pgrestore schema snapshot generator
+		return pgdumprestoregenerator.NewSnapshotGenerator(ctx, cfg.DumpRestore, pgdumprestoregenerator.WithLogger(logger))
+	default:
+		return nil, errSchemaSnapshotNotConfigured
+	}
+}
diff --git a/pkg/wal/processor/batch/wal_batch.go b/pkg/wal/processor/batch/wal_batch.go
index 5c08672..f04ab6d 100644
--- a/pkg/wal/processor/batch/wal_batch.go
+++ b/pkg/wal/processor/batch/wal_batch.go
@@ -12,6 +12,8 @@ type Batch[T Message] struct {
 	totalBytes int
 }
 
+const zeroLSN = "0/0"
+
 func NewBatch[T Message](messages []T, positions []wal.CommitPosition) *Batch[T] {
 	return &Batch[T]{
 		messages:  messages,
@@ -33,7 +35,7 @@ func (b *Batch[T]) add(m *WALMessage[T]) {
 		b.totalBytes += m.message.Size()
 	}
 
-	if m.position != "" {
+	if m.position != "" && m.position != zeroLSN {
 		b.positions = append(b.positions, m.position)
 	}
 }
diff --git a/pkg/wal/processor/batch/wal_batch_sender.go b/pkg/wal/processor/batch/wal_batch_sender.go
index 10bb695..523f603 100644
--- a/pkg/wal/processor/batch/wal_batch_sender.go
+++ b/pkg/wal/processor/batch/wal_batch_sender.go
@@ -149,7 +149,7 @@ func (s *Sender[T]) Send(ctx context.Context) error {
 					}
 				}
 			case msg := <-s.msgChan:
-				if msgBatch.maxBatchBytesReached(s.maxBatchBytes, msg.message) {
+				if !msg.message.IsEmpty() && msgBatch.maxBatchBytesReached(s.maxBatchBytes, msg.message) {
 					if err := drainBatch(msgBatch); err != nil {
 						return err
 					}


Test patch:
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
index c8588a8..05b6530 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.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 {
diff --git a/pkg/snapshot/generator/postgres/data/pg_schema_table_parser_test.go b/pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder_test.go
similarity index 59%
rename from pkg/snapshot/generator/postgres/data/pg_schema_table_parser_test.go
rename to pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder_test.go
index 1d30e10..c158f98 100644
--- a/pkg/snapshot/generator/postgres/data/pg_schema_table_parser_test.go
+++ b/pkg/snapshot/generator/postgres/tablefinder/pg_snapshot_table_finder_test.go
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package postgres
+package tablefinder
 
 import (
 	"context"
@@ -11,32 +11,43 @@ import (
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
 	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/generator"
+	"github.com/xataio/pgstream/pkg/snapshot/generator/mocks"
 )
 
-func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
+func TestSnapshotTableFinder_CreateSnapshot(t *testing.T) {
 	t.Parallel()
 
 	testSchema := "test-schema"
 	errTest := errors.New("oh noes")
 
+	newTestSnapshot := func(tables []string) *snapshot.Snapshot {
+		return &snapshot.Snapshot{
+			SchemaName: testSchema,
+			TableNames: tables,
+		}
+	}
+
 	tests := []struct {
-		name     string
-		conn     *pgmocks.Querier
-		snapshot *snapshot.Snapshot
+		name      string
+		conn      *pgmocks.Querier
+		snapshot  *snapshot.Snapshot
+		generator generator.SnapshotGenerator
 
-		wantTables []string
-		wantErr    error
+		wantErr error
 	}{
 		{
-			name: "ok - no wildcard",
-			conn: &pgmocks.Querier{},
-			snapshot: &snapshot.Snapshot{
-				SchemaName: testSchema,
-				TableNames: []string{"table-a", "table-b", "table-c"},
+			name:     "ok - no wildcard",
+			conn:     &pgmocks.Querier{},
+			snapshot: newTestSnapshot([]string{"table-a", "table-b", "table-c"}),
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					require.Equal(t, newTestSnapshot([]string{"table-a", "table-b", "table-c"}), snapshot)
+					return nil
+				},
 			},
 
-			wantTables: []string{"table-a", "table-b", "table-c"},
-			wantErr:    nil,
+			wantErr: nil,
 		},
 		{
 			name: "ok - only wildcard",
@@ -60,11 +71,16 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 			},
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
-				TableNames: []string{"*"},
+				TableNames: []string{wildcard},
+			},
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					require.Equal(t, newTestSnapshot([]string{"table-1"}), snapshot)
+					return nil
+				},
 			},
 
-			wantTables: []string{"table-1"},
-			wantErr:    nil,
+			wantErr: nil,
 		},
 		{
 			name: "ok - with wildcard",
@@ -88,11 +104,16 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 			},
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
-				TableNames: []string{"table-a", "table-b", "*"},
+				TableNames: []string{"table-a", "table-b", wildcard},
+			},
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					require.Equal(t, newTestSnapshot([]string{"table-1"}), snapshot)
+					return nil
+				},
 			},
 
-			wantTables: []string{"table-1"},
-			wantErr:    nil,
+			wantErr: nil,
 		},
 		{
 			name: "error - querying schema tables",
@@ -105,11 +126,15 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 			},
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
-				TableNames: []string{"table-a", "table-b", "*"},
+				TableNames: []string{"table-a", "table-b", wildcard},
+			},
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					return errors.New("CreateSnapshotFn: should not be called")
+				},
 			},
 
-			wantTables: nil,
-			wantErr:    errTest,
+			wantErr: errTest,
 		},
 		{
 			name: "error - scanning row",
@@ -129,11 +154,15 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 			},
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
-				TableNames: []string{"table-a", "table-b", "*"},
+				TableNames: []string{"table-a", "table-b", wildcard},
+			},
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					return errors.New("CreateSnapshotFn: should not be called")
+				},
 			},
 
-			wantTables: nil,
-			wantErr:    errTest,
+			wantErr: errTest,
 		},
 		{
 			name: "error - rows error",
@@ -150,11 +179,15 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 			},
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
-				TableNames: []string{"table-a", "table-b", "*"},
+				TableNames: []string{"table-a", "table-b", wildcard},
+			},
+			generator: &mocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, snapshot *snapshot.Snapshot) error {
+					return errors.New("CreateSnapshotFn: should not be called")
+				},
 			},
 
-			wantTables: nil,
-			wantErr:    errTest,
+			wantErr: errTest,
 		},
 	}
 
@@ -163,10 +196,12 @@ func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			p := newSchemaTableParser(tc.conn)
-			err := p.parseSnapshotTables(context.Background(), tc.snapshot)
+			tableFinder := SnapshotTableFinder{
+				conn:    tc.conn,
+				wrapped: tc.generator,
+			}
+			err := tableFinder.CreateSnapshot(context.Background(), tc.snapshot)
 			require.ErrorIs(t, err, tc.wantErr)
-			require.Equal(t, tc.wantTables, tc.snapshot.TableNames)
 		})
 	}
 }