Repository: xataio/pgstream
Pull request #265 (2025-04-10 10:35:01)

Add option to repeat snapshots when using snapshot recorder
Add option to support repeatable snapshots when using the snapshot request recorder. Currently the behaviour is to filter out existing successfully completed snapshots to avoid performing them more than once, to fit the replication run with initial snapshot (usually it should only need to be performed once). However, for the snapshot mode only, we want to support having repeatable snapshots while keeping track of the requests and the status of all the runs. 


Patch:
diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go
index 536787a..447773f 100644
--- a/cmd/config/config_env.go
+++ b/cmd/config/config_env.go
@@ -79,8 +79,7 @@ func parseSnapshotListenerConfig() *snapshotbuilder.SnapshotListenerConfig {
 }
 
 func parseSnapshotConfig(pgURL, prefix string) *snapshotbuilder.SnapshotListenerConfig {
-	return &snapshotbuilder.SnapshotListenerConfig{
-		SnapshotStoreURL: viper.GetString(fmt.Sprintf("%s_SNAPSHOT_STORE_URL", prefix)),
+	cfg := &snapshotbuilder.SnapshotListenerConfig{
 		Generator: pgsnapshotgenerator.Config{
 			URL:           pgURL,
 			BatchPageSize: viper.GetUint(fmt.Sprintf("%s_SNAPSHOT_BATCH_PAGE_SIZE", prefix)),
@@ -93,6 +92,15 @@ func parseSnapshotConfig(pgURL, prefix string) *snapshotbuilder.SnapshotListener
 		},
 		Schema: parseSchemaSnapshotConfig(prefix, pgURL),
 	}
+
+	if storeURL := viper.GetString(fmt.Sprintf("%s_SNAPSHOT_STORE_URL", prefix)); storeURL != "" {
+		cfg.Recorder = &snapshotbuilder.SnapshotRecorderConfig{
+			RepeatableSnapshots: viper.GetBool(fmt.Sprintf("%s_SNAPSHOT_STORE_REPEATABLE", prefix)),
+			SnapshotStoreURL:    storeURL,
+		}
+	}
+
+	return cfg
 }
 
 func parseSchemaSnapshotConfig(prefix, pgurl string) snapshotbuilder.SchemaSnapshotConfig {
diff --git a/cmd/config/config_yaml.go b/cmd/config/config_yaml.go
index 6bbfee7..0a74deb 100644
--- a/cmd/config/config_yaml.go
+++ b/cmd/config/config_yaml.go
@@ -64,7 +64,8 @@ type SnapshotConfig struct {
 }
 
 type SnapshotRecorderConfig struct {
-	PostgresURL string `mapstructure:"postgres_url" yaml:"postgres_url"`
+	RepeatableSnapshots bool   `mapstructure:"repeatable_snapshots" yaml:"repeatable_snapshots"`
+	PostgresURL         string `mapstructure:"postgres_url" yaml:"postgres_url"`
 }
 
 type SnapshotDataConfig struct {
@@ -357,7 +358,10 @@ func (c *YAMLConfig) parseSnapshotConfig() (*snapshotbuilder.SnapshotListenerCon
 		if snapshotConfig.Recorder.PostgresURL == "" {
 			return nil, errInvalidSnapshotRecorderConfig
 		}
-		streamCfg.SnapshotStoreURL = snapshotConfig.Recorder.PostgresURL
+		streamCfg.Recorder = &snapshotbuilder.SnapshotRecorderConfig{
+			RepeatableSnapshots: snapshotConfig.Recorder.RepeatableSnapshots,
+			SnapshotStoreURL:    snapshotConfig.Recorder.PostgresURL,
+		}
 	}
 
 	switch snapshotConfig.Mode {
diff --git a/config_template.yaml b/config_template.yaml
index 235f0cc..fde5a95 100644
--- a/config_template.yaml
+++ b/config_template.yaml
@@ -6,6 +6,7 @@ source:
       mode: full # options are data_and, schema or data
       tables: ["test", "test_schema.test", "another_schema.*"] # tables to snapshot, can be a list of table names or a pattern
       recorder:
+        repeatable_snapshots: true # whether to repeat snapshots that have already been taken
         postgres_url: "postgresql://user:password@localhost:5432/mydatabase" # URL of the database where the snapshot status is recorded
       snapshot_workers: 4 # number of schemas to be snapshotted in parallel
       data: # when mode is full or data
diff --git a/docs/README.md b/docs/README.md
index 93d156c..c72ebb4 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -97,6 +97,8 @@ Here's a list of all the environment variables that can be used to configure the
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_BATCH_PAGE_SIZE  | 1000                   | No       | Size of the table page range which will be processed concurrently by the table workers from `PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS`.                                                                                                                                                              |
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS          | 1                      | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_USE_SCHEMALOG    | False                  | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres outputs.                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_STORE_URL        | ""                     | No       | Postgres URL for the database where the snapshot requests and their status will be tracked. A table `snapshot_requests` will be created under a `pgstream` schema.                                                                                                                                           |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_STORE_REPEATABLE | False                  | No       | Allow to repeat snapshots requests that have been already completed succesfully. By default snapshots that have been completed won't be repeated.                                                                                                                                                            |
 
 </details>
 
diff --git a/pkg/snapshot/generator/snapshot_generator_recorder.go b/pkg/snapshot/generator/snapshot_generator_recorder.go
index 5655363..3cbc3d1 100644
--- a/pkg/snapshot/generator/snapshot_generator_recorder.go
+++ b/pkg/snapshot/generator/snapshot_generator_recorder.go
@@ -14,8 +14,9 @@ import (
 // SnapshotRecorder is a decorator around a snapshot generator that will record
 // the snapshot request status.
 type SnapshotRecorder struct {
-	wrapped SnapshotGenerator
-	store   snapshotstore.Store
+	wrapped             SnapshotGenerator
+	store               snapshotstore.Store
+	repeatableSnapshots bool
 }
 
 const updateTimeout = time.Minute
@@ -23,10 +24,11 @@ const updateTimeout = time.Minute
 // NewSnapshotRecorder will return the generator on input wrapped with an
 // activity recorder that will keep track of the status of the snapshot
 // requests.
-func NewSnapshotRecorder(store snapshotstore.Store, generator SnapshotGenerator) *SnapshotRecorder {
+func NewSnapshotRecorder(store snapshotstore.Store, generator SnapshotGenerator, repeatableSnapshots bool) *SnapshotRecorder {
 	return &SnapshotRecorder{
-		wrapped: generator,
-		store:   store,
+		wrapped:             generator,
+		store:               store,
+		repeatableSnapshots: repeatableSnapshots,
 	}
 }
 
@@ -88,6 +90,11 @@ func (s *SnapshotRecorder) markSnapshotCompleted(req *snapshot.Request, err erro
 }
 
 func (s *SnapshotRecorder) filterOutExistingSnapshots(ctx context.Context, schema string, tables []string) ([]string, error) {
+	// if we want to be able to repeat snapshots, we don't filter out existing ones
+	if s.repeatableSnapshots {
+		return tables, nil
+	}
+
 	snapshotRequests, err := s.store.GetSnapshotRequestsBySchema(ctx, schema)
 	if err != nil {
 		return nil, fmt.Errorf("retrieving existing snapshots for schema: %w", err)
diff --git a/pkg/wal/listener/snapshot/builder/config.go b/pkg/wal/listener/snapshot/builder/config.go
index 75ba0e6..a2974cb 100644
--- a/pkg/wal/listener/snapshot/builder/config.go
+++ b/pkg/wal/listener/snapshot/builder/config.go
@@ -10,13 +10,18 @@ import (
 )
 
 type SnapshotListenerConfig struct {
-	Generator        pgsnapshotgenerator.Config
-	Adapter          adapter.SnapshotConfig
-	SnapshotStoreURL string
-	Schema           SchemaSnapshotConfig
+	Generator pgsnapshotgenerator.Config
+	Adapter   adapter.SnapshotConfig
+	Recorder  *SnapshotRecorderConfig
+	Schema    SchemaSnapshotConfig
 }
 
 type SchemaSnapshotConfig struct {
 	SchemaLogStore *schemalogpg.Config
 	DumpRestore    *pgdumprestore.Config
 }
+
+type SnapshotRecorderConfig struct {
+	RepeatableSnapshots bool
+	SnapshotStoreURL    string
+}
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 52735a5..f50a269 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,13 +67,13 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 	// snapshot generator aggregator
 	g = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
 
-	if cfg.SnapshotStoreURL != "" {
+	if cfg.Recorder != nil && cfg.Recorder.SnapshotStoreURL != "" {
 		// snapshot activity recorder layer
-		snapshotStore, err := pgsnapshotstore.New(ctx, cfg.SnapshotStoreURL)
+		snapshotStore, err := pgsnapshotstore.New(ctx, cfg.Recorder.SnapshotStoreURL)
 		if err != nil {
 			return nil, fmt.Errorf("create postgres snapshot store: %w", err)
 		}
-		g = generator.NewSnapshotRecorder(snapshotStore, g)
+		g = generator.NewSnapshotRecorder(snapshotStore, g, cfg.Recorder.RepeatableSnapshots)
 	}
 
 	return adapter.NewSnapshotGeneratorAdapter(&cfg.Adapter, g, adapter.WithLogger(logger)), nil


Test patch:
diff --git a/cmd/config/helper_test.go b/cmd/config/helper_test.go
index fbcc860..168b678 100644
--- a/cmd/config/helper_test.go
+++ b/cmd/config/helper_test.go
@@ -25,7 +25,9 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) {
 	assert.Equal(t, uint(1000), streamConfig.Listener.Postgres.Snapshot.Generator.BatchPageSize)
 	assert.NotNil(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore)
 	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore.CleanTargetDB)
-	assert.Equal(t, "postgresql://user:password@localhost:5432/mytargetdatabase", streamConfig.Listener.Postgres.Snapshot.SnapshotStoreURL)
+	assert.NotNil(t, streamConfig.Listener.Postgres.Snapshot.Recorder)
+	assert.Equal(t, "postgresql://user:password@localhost:5432/mytargetdatabase", streamConfig.Listener.Postgres.Snapshot.Recorder.SnapshotStoreURL)
+	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Recorder.RepeatableSnapshots)
 
 	assert.NotNil(t, streamConfig.Listener.Kafka)
 	assert.Equal(t, []string{"localhost:9092"}, streamConfig.Listener.Kafka.Reader.Conn.Servers)
diff --git a/cmd/config/test/test_config.env b/cmd/config/test/test_config.env
index 44c67f5..12b0114 100644
--- a/cmd/config/test/test_config.env
+++ b/cmd/config/test/test_config.env
@@ -5,6 +5,7 @@ PGSTREAM_POSTGRES_LISTENER_URL="postgresql://user:password@localhost:5432/mydata
 PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME="pgstream_mydatabase_slot"
 PGSTREAM_POSTGRES_LISTENER_INITIAL_SNAPSHOT_ENABLED=true
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_STORE_URL="postgresql://user:password@localhost:5432/mytargetdatabase"
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_STORE_REPEATABLE=true
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES="test test_schema.test another_schema.*"
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS=4
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS=4
diff --git a/cmd/config/test/test_config.yaml b/cmd/config/test/test_config.yaml
index f4dd321..97439e4 100644
--- a/cmd/config/test/test_config.yaml
+++ b/cmd/config/test/test_config.yaml
@@ -6,6 +6,7 @@ source:
       mode: full # options are data_and, schema or data
       tables: ["test", "test_schema.test", "another_schema.*"] # tables to snapshot, can be a list of table names or a pattern
       recorder:
+        repeatable_snapshots: true # whether to repeat snapshots that have already been taken
         postgres_url: "postgresql://user:password@localhost:5432/mytargetdatabase" # URL of the database where the snapshot status is recorded
       snapshot_workers: 4 # number of schemas to be snapshotted in parallel
       data: # when mode is full or data
diff --git a/pkg/snapshot/generator/snapshot_generator_recorder_test.go b/pkg/snapshot/generator/snapshot_generator_recorder_test.go
index 0c394fd..ed8eba7 100644
--- a/pkg/snapshot/generator/snapshot_generator_recorder_test.go
+++ b/pkg/snapshot/generator/snapshot_generator_recorder_test.go
@@ -347,7 +347,7 @@ func TestSnapshotRecorder_CreateSnapshot(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			sr := NewSnapshotRecorder(tc.store, tc.generator)
+			sr := NewSnapshotRecorder(tc.store, tc.generator, false)
 			defer sr.Close()
 
 			err := sr.CreateSnapshot(context.Background(), newTestSnapshot())
@@ -364,9 +364,10 @@ func TestSnapshotRecorder_filterOutExistingSnapshots(t *testing.T) {
 	errTest := errors.New("oh noes")
 
 	tests := []struct {
-		name   string
-		store  snapshotstore.Store
-		tables []string
+		name                string
+		store               snapshotstore.Store
+		tables              []string
+		repeatableSnapshots bool
 
 		wantTables []string
 		wantErr    error
@@ -415,6 +416,16 @@ func TestSnapshotRecorder_filterOutExistingSnapshots(t *testing.T) {
 			},
 			wantTables: []string{"table1"},
 		},
+		{
+			name: "ok - existing snapshots with repeatable snapshots",
+			store: &snapshotstoremocks.Store{
+				GetSnapshotRequestsBySchemaFn: func(ctx context.Context, s string) ([]*snapshot.Request, error) {
+					return nil, errors.New("unexpected call to GetSnapshotRequestsBySchemaFn")
+				},
+			},
+			repeatableSnapshots: true,
+			wantTables:          []string{"table1", "table2"},
+		},
 		{
 			name: "ok - existing wildcard snapshot with wildcard table",
 			store: &snapshotstoremocks.Store{
@@ -522,7 +533,8 @@ func TestSnapshotRecorder_filterOutExistingSnapshots(t *testing.T) {
 			t.Parallel()
 
 			sr := SnapshotRecorder{
-				store: tc.store,
+				store:               tc.store,
+				repeatableSnapshots: tc.repeatableSnapshots,
 			}
 
 			tables := testTables