Repository: xataio/pgstream
Pull request #395 (2025-06-06 16:11:43)

[Snapshot] Wait for all tables to finish processing data
This PR updates the postgres data snapshot generator to make sure the rows processor has finished before calling the next snapshot step. Otherwise, we can encounter situations where the index/constraint creation starts before all the data has been fully snapshoted, potentially leading to constraint issues since the data is not valid yet.

By closing the processor after the snapshot is done we guarantee that we wait for all the in-flight batches to complete before continuing.


Patch:
diff --git a/pkg/snapshot/generator/mocks/mocks_row_processor.go b/pkg/snapshot/generator/mocks/mocks_row_processor.go
new file mode 100644
index 0000000..ce20ff2
--- /dev/null
+++ b/pkg/snapshot/generator/mocks/mocks_row_processor.go
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	"context"
+
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+type RowsProcessor struct {
+	ProcessRowFn func(context.Context, *snapshot.Row) error
+	CloseFn      func() error
+}
+
+func (m *RowsProcessor) ProcessRow(ctx context.Context, row *snapshot.Row) error {
+	return m.ProcessRowFn(ctx, row)
+}
+
+func (m *RowsProcessor) Close() error {
+	if m.CloseFn != nil {
+		return m.CloseFn()
+	}
+	return nil
+}
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 30c8181..eafe73e 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -27,7 +27,7 @@ type SnapshotGenerator struct {
 	batchPageSize uint
 
 	// Function called for processing produced rows.
-	processRow             snapshot.RowProcessor
+	rowsProcessor          snapshot.RowsProcessor
 	tableSnapshotGenerator snapshotTableFn
 }
 
@@ -44,7 +44,7 @@ type snapshotTableFn func(ctx context.Context, snapshotID string, schema, table
 
 type Option func(sg *SnapshotGenerator)
 
-func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.RowProcessor, opts ...Option) (*SnapshotGenerator, error) {
+func NewSnapshotGenerator(ctx context.Context, cfg *Config, rowsProcessor snapshot.RowsProcessor, opts ...Option) (*SnapshotGenerator, error) {
 	conn, err := pglib.NewConnPool(ctx, cfg.URL)
 	if err != nil {
 		return nil, err
@@ -54,7 +54,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.
 		logger:        loglib.NewNoopLogger(),
 		mapper:        pglib.NewMapper(conn),
 		conn:          conn,
-		processRow:    processRow,
+		rowsProcessor: rowsProcessor,
 		batchPageSize: cfg.batchPageSize(),
 		tableWorkers:  cfg.tableWorkers(),
 		schemaWorkers: cfg.schemaWorkers(),
@@ -91,7 +91,12 @@ func WithInstrumentation(i *otel.Instrumentation) Option {
 	}
 }
 
-func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) (err error) {
+	defer func() {
+		// make sure we close the rows processor once the snapshot is completed.
+		// It will wait until all rows are processed before returning.
+		sg.rowsProcessor.Close()
+	}()
 	// use a transaction snapshot to ensure the table rows can be parallelised.
 	// The transaction snapshot is available for use only until the end of the
 	// transaction that exported it.
@@ -237,7 +242,7 @@ func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID,
 					continue
 				}
 
-				if err := sg.processRow(ctx, &snapshot.Row{
+				if err := sg.rowsProcessor.ProcessRow(ctx, &snapshot.Row{
 					Schema:  schema,
 					Table:   table,
 					Columns: columns,
diff --git a/pkg/snapshot/snapshot.go b/pkg/snapshot/snapshot.go
index 7766647..724798e 100644
--- a/pkg/snapshot/snapshot.go
+++ b/pkg/snapshot/snapshot.go
@@ -31,6 +31,11 @@ type Column struct {
 
 type RowProcessor func(context.Context, *Row) error
 
+type RowsProcessor interface {
+	ProcessRow(context.Context, *Row) error
+	Close() error
+}
+
 type Status string
 
 const (
diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go
index 1d053f6..2d4f852 100644
--- a/pkg/stream/stream_run.go
+++ b/pkg/stream/stream_run.go
@@ -121,7 +121,7 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, instrumentat
 			snapshotGenerator, err := snapshotbuilder.NewSnapshotGenerator(
 				ctx,
 				config.Listener.Postgres.Snapshot,
-				processor.ProcessWALEvent,
+				processor,
 				logger,
 				instrumentation)
 			if err != nil {
diff --git a/pkg/stream/stream_snapshot.go b/pkg/stream/stream_snapshot.go
index 11d3fc2..4ba59b5 100644
--- a/pkg/stream/stream_snapshot.go
+++ b/pkg/stream/stream_snapshot.go
@@ -45,7 +45,7 @@ func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrum
 	snapshotGenerator, err := snapshotbuilder.NewSnapshotGenerator(
 		ctx,
 		config.Listener.Snapshot,
-		processor.ProcessWALEvent,
+		processor,
 		logger,
 		instrumentation)
 	if err != nil {
diff --git a/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go
index 6f952f5..a915195 100644
--- a/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go
+++ b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go
@@ -13,19 +13,24 @@ import (
 )
 
 type ProcessEventAdapter struct {
-	processEvent listener.ProcessWalEvent
-	clock        clockwork.Clock
+	processor listener.Processor
+
+	clock clockwork.Clock
 }
 
-func NewProcessEventAdapter(processEvent listener.ProcessWalEvent) *ProcessEventAdapter {
+func NewProcessEventAdapter(p listener.Processor) *ProcessEventAdapter {
 	return &ProcessEventAdapter{
-		processEvent: processEvent,
-		clock:        clockwork.NewRealClock(),
+		processor: p,
+		clock:     clockwork.NewRealClock(),
 	}
 }
 
 func (a *ProcessEventAdapter) ProcessRow(ctx context.Context, row *snapshot.Row) error {
-	return a.processEvent(ctx, a.snapshotRowToWalEvent(row))
+	return a.processor.ProcessWALEvent(ctx, a.snapshotRowToWalEvent(row))
+}
+
+func (a *ProcessEventAdapter) Close() error {
+	return a.processor.Close()
 }
 
 func (a *ProcessEventAdapter) snapshotRowToWalEvent(row *snapshot.Row) *wal.Event {
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 835cc4f..963a4c9 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
@@ -48,8 +48,8 @@ var errSchemaSnapshotNotConfigured = errors.New("no schema snapshot has been con
 // │ └─────────┘  └─────────┘  └────────┘  └─────────┘ │
 // └───────────────────────────────────────────────────┘
 
-func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, processEvent listener.ProcessWalEvent, logger loglib.Logger, instrumentation *otel.Instrumentation) (listenersnapshot.Generator, error) {
-	processEventAdapter := adapter.NewProcessEventAdapter(processEvent)
+func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, p listener.Processor, logger loglib.Logger, instrumentation *otel.Instrumentation) (listenersnapshot.Generator, error) {
+	rowsProcessor := adapter.NewProcessEventAdapter(p)
 
 	var g generator.SnapshotGenerator
 	var err error
@@ -61,7 +61,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 	if instrumentation.IsEnabled() {
 		opts = append(opts, pgsnapshotgenerator.WithInstrumentation(instrumentation))
 	}
-	g, err = pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, processEventAdapter.ProcessRow, opts...)
+	g, err = pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, rowsProcessor, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -76,7 +76,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 	}
 
 	// postgres schema snapshot generator layer
-	g, err = newSchemaSnapshotGenerator(ctx, &cfg.Schema, g, processEventAdapter.ProcessRow, logger, instrumentation)
+	g, err = newSchemaSnapshotGenerator(ctx, &cfg.Schema, g, rowsProcessor.ProcessRow, logger, instrumentation)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/wal/listener/wal_listener.go b/pkg/wal/listener/wal_listener.go
index c62e242..1c9a1f5 100644
--- a/pkg/wal/listener/wal_listener.go
+++ b/pkg/wal/listener/wal_listener.go
@@ -15,3 +15,8 @@ type Listener interface {
 }
 
 type ProcessWalEvent func(context.Context, *wal.Event) error
+
+type Processor interface {
+	ProcessWALEvent(context.Context, *wal.Event) error
+	Close() error
+}
diff --git a/pkg/wal/processor/batch/wal_batch_sender.go b/pkg/wal/processor/batch/wal_batch_sender.go
index 92d1180..c14077c 100644
--- a/pkg/wal/processor/batch/wal_batch_sender.go
+++ b/pkg/wal/processor/batch/wal_batch_sender.go
@@ -200,17 +200,15 @@ func (s *Sender[T]) send(ctx context.Context) error {
 	return err
 }
 
+// Close will stop the sending, and wait for all ongoing batches to finish. It
+// is safe to call multiple times.
 func (s *Sender[T]) Close() {
-	s.cancelFn()
-	// wait for the send loop to finish
-	s.wg.Wait()
-	s.closeMsgChan()
-}
-
-// closeMsgChan closes the internal msg channel. It can be called multiple
-// times.
-func (s *Sender[T]) closeMsgChan() {
 	s.once.Do(func() {
+		s.logger.Trace("closing batch sender")
+		s.cancelFn()
+		// wait for the send loop to finish
+		s.wg.Wait()
 		close(s.msgChan)
+		s.logger.Trace("batch sender closed")
 	})
 }


Test patch:
diff --git a/pkg/snapshot/generator/postgres/data/helper_test.go b/pkg/snapshot/generator/postgres/data/helper_test.go
index 2fdab47..6ed6c22 100644
--- a/pkg/snapshot/generator/postgres/data/helper_test.go
+++ b/pkg/snapshot/generator/postgres/data/helper_test.go
@@ -18,13 +18,14 @@ type mockRowProcessor struct {
 	once    sync.Once
 }
 
-func (mp *mockRowProcessor) process(ctx context.Context, row *snapshot.Row) error {
+func (mp *mockRowProcessor) ProcessRow(ctx context.Context, row *snapshot.Row) error {
 	mp.rowChan <- row
 	return nil
 }
 
-func (mp *mockRowProcessor) close() {
+func (mp *mockRowProcessor) Close() error {
 	mp.once.Do(func() { close(mp.rowChan) })
+	return nil
 }
 
 func execQuery(t *testing.T, ctx context.Context, pgurl, query string) {
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_integration_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_integration_test.go
index ebf7e6c..204a85d 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_integration_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_integration_test.go
@@ -32,7 +32,7 @@ func Test_PostgresSnapshotGenerator(t *testing.T) {
 		rowChan: make(chan *snapshot.Row),
 	}
 
-	generator, err := NewSnapshotGenerator(ctx, &Config{URL: pgurl}, mockProcessor.process)
+	generator, err := NewSnapshotGenerator(ctx, &Config{URL: pgurl}, mockProcessor)
 	require.NoError(t, err)
 	defer generator.Close()
 
@@ -49,7 +49,7 @@ func Test_PostgresSnapshotGenerator(t *testing.T) {
 			TableNames: []string{testTable},
 		})
 		require.NoError(t, err)
-		mockProcessor.close()
+		mockProcessor.Close()
 	}()
 
 	rows := make([]*snapshot.Row, 0, 3)
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 093b72f..7e97ba5 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -19,6 +19,7 @@ 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/mocks"
 	"github.com/xataio/pgstream/pkg/wal"
 )
 
@@ -870,9 +871,11 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				})),
 				conn:   tc.querier,
 				mapper: pglib.NewMapper(tc.querier),
-				processRow: func(ctx context.Context, e *snapshot.Row) error {
-					rowChan <- e
-					return nil
+				rowsProcessor: &mocks.RowsProcessor{
+					ProcessRowFn: func(ctx context.Context, e *snapshot.Row) error {
+						rowChan <- e
+						return nil
+					},
 				},
 				schemaWorkers: 1,
 				tableWorkers:  1,
diff --git a/pkg/wal/processor/batch/wal_batch_sender_test.go b/pkg/wal/processor/batch/wal_batch_sender_test.go
index 35879c6..fa0ac12 100644
--- a/pkg/wal/processor/batch/wal_batch_sender_test.go
+++ b/pkg/wal/processor/batch/wal_batch_sender_test.go
@@ -145,7 +145,7 @@ func TestSender_SendMessage(t *testing.T) {
 			wg.Add(1)
 			go func() {
 				defer func() {
-					batchSender.closeMsgChan()
+					batchSender.Close()
 					wg.Done()
 				}()
 				err := batchSender.SendMessage(ctx, tc.msg)