Repository: xataio/pgstream
Pull request #156 (2025-02-24 15:29:27)

Decouple snapshot from wal listener
This PR decouples the snapshot generator adapter from the postgres listener, so that it can be used for both the initial snapshot and the snapshot only mode that will be added in following PRs.

The existing adapter has been split into two parts, the `CreateSnapshot` adapter, and the `ProcessEvent` one, but the logic hasn't been changed, only moved (along with tests).

A `PGSTREAM_POSTGRES_LISTENER_INITIAL_SNAPSHOT_ENABLED` env variable is added to flag if the initial snapshot should be performed for the postgres listener to make it more straightforward to configure.


Patch:
diff --git a/README.md b/README.md
index 91aba06..2035e53 100644
--- a/README.md
+++ b/README.md
@@ -129,15 +129,16 @@ Here's a list of all the environment variables that can be used to configure the
 <details>
   <summary>Postgres Listener</summary>
 
-| Environment Variable                               | Default                    | Required | Description                                                                                                                                                                                                                                                                                                  |
-| -------------------------------------------------- | -------------------------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| PGSTREAM_POSTGRES_LISTENER_URL                     | N/A                        | Yes      | URL of the Postgres database to connect to for replication purposes.                                                                                                                                                                                                                                         |
-| PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME            | "pgstream\_<dbname>\_slot" | No       | Name of the Postgres replication slot name.                                                                                                                                                                                                                                                                  |
-| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES          | ""                         | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
-| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS  | 4                          | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
-| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS   | 4                          | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
-| 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.                                                                                                                                                                                                                            |
+| Environment Variable                                | Default                | Required | Description                                                                                                                                                                                                                                                                                                  |
+| --------------------------------------------------- | ---------------------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| PGSTREAM_POSTGRES_LISTENER_URL                      | N/A                    | Yes      | URL of the Postgres database to connect to for replication purposes.                                                                                                                                                                                                                                         |
+| PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME             | "pgstream_dbname_slot" | No       | Name of the Postgres replication slot name.                                                                                                                                                                                                                                                                  |
+| PGSTREAM_POSTGRES_LISTENER_INITIAL_SNAPSHOT_ENABLED | false                  | No       | Enables performing an initial snapshot of the Postgres database before starting to listen to the replication slot.                                                                                                                                                                                           |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES           | ""                     | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS   | 4                      | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS    | 4                      | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
+| 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.                                                                                                                                                                                                                            |
 
 </details>
 
diff --git a/cmd/config.go b/cmd/config.go
index 712eadd..3422f56 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -14,7 +14,8 @@ import (
 	"github.com/xataio/pgstream/pkg/stream"
 	"github.com/xataio/pgstream/pkg/tls"
 	kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka"
-	pglistener "github.com/xataio/pgstream/pkg/wal/listener/postgres"
+	"github.com/xataio/pgstream/pkg/wal/listener/snapshot/adapter"
+	snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder"
 	"github.com/xataio/pgstream/pkg/wal/processor/batch"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
@@ -84,36 +85,42 @@ func parsePostgresListenerConfig() *stream.PostgresListenerConfig {
 		},
 	}
 
-	snapshotTables := viper.GetStringSlice("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES")
-	if len(snapshotTables) > 0 {
-		cfg.Snapshot = &pglistener.SnapshotConfig{
-			SnapshotStoreURL: pgURL,
-			Generator: pgsnapshotgenerator.Config{
-				URL:           pgURL,
-				BatchPageSize: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_BATCH_PAGE_SIZE"),
-				SchemaWorkers: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS"),
-				TableWorkers:  viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS"),
-			},
-			Tables:          snapshotTables,
-			SnapshotWorkers: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS"),
-			Schema:          parseSchemaSnapshotConfig(pgURL),
-		}
+	initialSnapshotEnabled := viper.GetBool("PGSTREAM_POSTGRES_LISTENER_INITIAL_SNAPSHOT_ENABLED")
+	if initialSnapshotEnabled {
+		cfg.Snapshot = parseSnapshotListenerConfig(pgURL)
 	}
 
 	return cfg
 }
 
-func parseSchemaSnapshotConfig(pgurl string) pglistener.SchemaSnapshotConfig {
+func parseSnapshotListenerConfig(pgURL string) *snapshotbuilder.SnapshotListenerConfig {
+	return &snapshotbuilder.SnapshotListenerConfig{
+		SnapshotStoreURL: pgURL,
+		Generator: pgsnapshotgenerator.Config{
+			URL:           pgURL,
+			BatchPageSize: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_BATCH_PAGE_SIZE"),
+			SchemaWorkers: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS"),
+			TableWorkers:  viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS"),
+		},
+		Adapter: adapter.SnapshotConfig{
+			Tables:          viper.GetStringSlice("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES"),
+			SnapshotWorkers: viper.GetUint("PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS"),
+		},
+		Schema: parseSchemaSnapshotConfig(pgURL),
+	}
+}
+
+func parseSchemaSnapshotConfig(pgurl string) snapshotbuilder.SchemaSnapshotConfig {
 	pgTargetURL := viper.GetString("PGSTREAM_POSTGRES_WRITER_TARGET_URL")
 	if pgTargetURL != "" {
-		return pglistener.SchemaSnapshotConfig{
+		return snapshotbuilder.SchemaSnapshotConfig{
 			DumpRestore: &pgdumprestore.Config{
 				SourcePGURL: pgurl,
 				TargetPGURL: pgTargetURL,
 			},
 		}
 	}
-	return pglistener.SchemaSnapshotConfig{
+	return snapshotbuilder.SchemaSnapshotConfig{
 		SchemaLogStore: &pgschemalog.Config{
 			URL: pgurl,
 		},
diff --git a/pg2pg.env b/pg2pg.env
index f98cf3d..55dd1d3 100644
--- a/pg2pg.env
+++ b/pg2pg.env
@@ -1,5 +1,6 @@
 # Listener config
 PGSTREAM_POSTGRES_LISTENER_URL="postgres://postgres:postgres@localhost?sslmode=disable"
+PGSTREAM_POSTGRES_LISTENER_INITIAL_SNAPSHOT_ENABLED=true
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS=4
diff --git a/pkg/stream/config.go b/pkg/stream/config.go
index 8f99914..bbb2d4c 100644
--- a/pkg/stream/config.go
+++ b/pkg/stream/config.go
@@ -8,7 +8,7 @@ import (
 
 	"github.com/xataio/pgstream/pkg/kafka"
 	kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka"
-	pglistener "github.com/xataio/pgstream/pkg/wal/listener/postgres"
+	snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
 	"github.com/xataio/pgstream/pkg/wal/processor/postgres"
@@ -32,7 +32,7 @@ type ListenerConfig struct {
 
 type PostgresListenerConfig struct {
 	Replication pgreplication.Config
-	Snapshot    *pglistener.SnapshotConfig
+	Snapshot    *snapshotbuilder.SnapshotListenerConfig
 }
 
 type KafkaListenerConfig struct {
diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go
index d060c9e..a887576 100644
--- a/pkg/stream/stream_run.go
+++ b/pkg/stream/stream_run.go
@@ -17,6 +17,7 @@ import (
 	"github.com/xataio/pgstream/pkg/wal/listener"
 	kafkalistener "github.com/xataio/pgstream/pkg/wal/listener/kafka"
 	pglistener "github.com/xataio/pgstream/pkg/wal/listener/postgres"
+	snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder"
 	"github.com/xataio/pgstream/pkg/wal/processor"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	processinstrumentation "github.com/xataio/pgstream/pkg/wal/processor/instrumentation"
@@ -258,7 +259,7 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, instrumentat
 		}
 		if config.Listener.Postgres.Snapshot != nil {
 			logger.Info("initial snapshot enabled")
-			snapshotGenerator, err := pglistener.NewSnapshotGeneratorAdapter(
+			snapshotGenerator, err := snapshotbuilder.NewSnapshotGenerator(
 				ctx,
 				config.Listener.Postgres.Snapshot,
 				processor.ProcessWALEvent,
diff --git a/pkg/wal/listener/postgres/config.go b/pkg/wal/listener/snapshot/adapter/config.go
similarity index 65%
rename from pkg/wal/listener/postgres/config.go
rename to pkg/wal/listener/snapshot/adapter/config.go
index 601a6c9..1365df8 100644
--- a/pkg/wal/listener/postgres/config.go
+++ b/pkg/wal/listener/snapshot/adapter/config.go
@@ -1,33 +1,23 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package postgres
+package adapter
 
 import (
 	"strings"
-
-	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
-	Schema           SchemaSnapshotConfig
+	Tables []string
 	// 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
 
+const publicSchema = "public"
+
 func (c *SnapshotConfig) schemaTableMap() map[string][]string {
 	schemaTableMap := make(map[string][]string, len(c.Tables))
 	for _, table := range c.Tables {
diff --git a/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go
new file mode 100644
index 0000000..6f952f5
--- /dev/null
+++ b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter.go
@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package adapter
+
+import (
+	"context"
+	"time"
+
+	"github.com/jonboulle/clockwork"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/wal"
+	"github.com/xataio/pgstream/pkg/wal/listener"
+)
+
+type ProcessEventAdapter struct {
+	processEvent listener.ProcessWalEvent
+	clock        clockwork.Clock
+}
+
+func NewProcessEventAdapter(processEvent listener.ProcessWalEvent) *ProcessEventAdapter {
+	return &ProcessEventAdapter{
+		processEvent: processEvent,
+		clock:        clockwork.NewRealClock(),
+	}
+}
+
+func (a *ProcessEventAdapter) ProcessRow(ctx context.Context, row *snapshot.Row) error {
+	return a.processEvent(ctx, a.snapshotRowToWalEvent(row))
+}
+
+func (a *ProcessEventAdapter) snapshotRowToWalEvent(row *snapshot.Row) *wal.Event {
+	if row == nil {
+		return nil
+	}
+
+	columns := make([]wal.Column, 0, len(row.Columns))
+	for _, col := range row.Columns {
+		columns = append(columns, a.snapshotColumnToWalColumn(col))
+	}
+	// use 0 since there's no LSN associated, but it can be used as the
+	// initial version downstream
+	const zeroLSN = "0/0"
+	return &wal.Event{
+		CommitPosition: wal.CommitPosition(zeroLSN),
+		Data: &wal.Data{
+			Action:    "I",
+			Timestamp: a.clock.Now().UTC().Format(time.RFC3339),
+			LSN:       zeroLSN,
+			Schema:    row.Schema,
+			Table:     row.Table,
+			Columns:   columns,
+		},
+	}
+}
+
+func (a *ProcessEventAdapter) snapshotColumnToWalColumn(col snapshot.Column) wal.Column {
+	return wal.Column{
+		Name:  col.Name,
+		Type:  col.Type,
+		Value: col.Value,
+	}
+}
diff --git a/pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter.go b/pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter.go
new file mode 100644
index 0000000..7cd5055
--- /dev/null
+++ b/pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter.go
@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package adapter
+
+import (
+	"context"
+
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/generator"
+	"golang.org/x/sync/errgroup"
+)
+
+// SnapshotGeneratorAdapter adapts a snapshot generator to work with WAL events
+type SnapshotGeneratorAdapter struct {
+	logger          loglib.Logger
+	generator       generator.SnapshotGenerator
+	schemaTables    map[string][]string
+	snapshotWorkers uint
+}
+
+type Option func(a *SnapshotGeneratorAdapter)
+
+func NewSnapshotGeneratorAdapter(cfg *SnapshotConfig, generator generator.SnapshotGenerator, opts ...Option) *SnapshotGeneratorAdapter {
+	a := &SnapshotGeneratorAdapter{
+		generator:       generator,
+		schemaTables:    cfg.schemaTableMap(),
+		logger:          loglib.NewNoopLogger(),
+		snapshotWorkers: cfg.snapshotWorkers(),
+	}
+
+	for _, opt := range opts {
+		opt(a)
+	}
+
+	return a
+}
+
+func WithLogger(logger loglib.Logger) Option {
+	return func(a *SnapshotGeneratorAdapter) {
+		a.logger = loglib.NewLogger(logger).WithFields(loglib.Fields{
+			loglib.ModuleField: "snapshot_generator_adapter",
+		})
+	}
+}
+
+func (s *SnapshotGeneratorAdapter) CreateSnapshot(ctx context.Context) error {
+	errGroup, ctx := errgroup.WithContext(ctx)
+	snapshotChan := make(chan *snapshot.Snapshot)
+	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
+				}
+			}
+			return nil
+		})
+	}
+	for schema, tables := range s.schemaTables {
+		snapshotChan <- &snapshot.Snapshot{
+			SchemaName: schema,
+			TableNames: tables,
+		}
+	}
+	close(snapshotChan)
+
+	return errGroup.Wait()
+}
+
+func (s *SnapshotGeneratorAdapter) Close() error {
+	s.logger.Info("closing snapshot generator")
+	return s.generator.Close()
+}
diff --git a/pkg/wal/listener/snapshot/builder/config.go b/pkg/wal/listener/snapshot/builder/config.go
new file mode 100644
index 0000000..75ba0e6
--- /dev/null
+++ b/pkg/wal/listener/snapshot/builder/config.go
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package builder
+
+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"
+	"github.com/xataio/pgstream/pkg/wal/listener/snapshot/adapter"
+)
+
+type SnapshotListenerConfig struct {
+	Generator        pgsnapshotgenerator.Config
+	Adapter          adapter.SnapshotConfig
+	SnapshotStoreURL string
+	Schema           SchemaSnapshotConfig
+}
+
+type SchemaSnapshotConfig struct {
+	SchemaLogStore *schemalogpg.Config
+	DumpRestore    *pgdumprestore.Config
+}
diff --git a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
similarity index 50%
rename from pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go
rename to pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
index 5284d6c..e4502d3 100644
--- a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter.go
+++ b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
@@ -1,14 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package postgres
+package builder
 
 import (
 	"context"
 	"errors"
 	"fmt"
-	"time"
 
-	"github.com/jonboulle/clockwork"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	schemalogpg "github.com/xataio/pgstream/pkg/schemalog/postgres"
@@ -19,33 +17,20 @@ import (
 	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"
-)
-
-// SnapshotGeneratorAdapter adapts a snapshot generator to work with WAL events
-type SnapshotGeneratorAdapter struct {
-	logger          loglib.Logger
-	generator       generator.SnapshotGenerator
-	processEvent    listenerProcessWalEvent
-	schemaTables    map[string][]string
-	snapshotWorkers uint
-	clock           clockwork.Clock
-}
-
-const (
-	publicSchema = "public"
+	"github.com/xataio/pgstream/pkg/wal/listener"
+	listenersnapshot "github.com/xataio/pgstream/pkg/wal/listener/snapshot"
+	"github.com/xataio/pgstream/pkg/wal/listener/snapshot/adapter"
 )
 
 var errSchemaSnapshotNotConfigured = errors.New("no schema snapshot has been configured")
 
-// NewSnapshotGeneratorAdapter builds a snapshot generator based on the
+// NewSnapshotGenerator 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          │
+// │               Snapshot Generator                  │
 // │ ┌─────────┐ ┌────────┐ ┌────────────────────────┐ │
 // │ │         │ │        │ │       Aggregator       │ │
 // │ │         │ │        │ │  ┌──────────────────┐  │ │
@@ -57,32 +42,27 @@ var errSchemaSnapshotNotConfigured = errors.New("no schema snapshot has been con
 // │ │         │ │        │ │  └──────────────────┘  │ │
 // │ └─────────┘ └────────┘ └────────────────────────┘ │
 // └───────────────────────────────────────────────────┘
-func NewSnapshotGeneratorAdapter(ctx context.Context, cfg *SnapshotConfig, processEvent listenerProcessWalEvent, logger loglib.Logger) (*SnapshotGeneratorAdapter, error) {
-	s := &SnapshotGeneratorAdapter{
-		processEvent:    processEvent,
-		schemaTables:    cfg.schemaTableMap(),
-		logger:          loglib.NewNoopLogger(),
-		snapshotWorkers: cfg.snapshotWorkers(),
-		clock:           clockwork.NewRealClock(),
-	}
 
+func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, processEvent listener.ProcessWalEvent, logger loglib.Logger) (listenersnapshot.Generator, error) {
+	processEventAdapter := adapter.NewProcessEventAdapter(processEvent)
 	// postgres schema snapshot generator
-	schemaSnapshotGenerator, err := s.newSchemaSnapshotGenerator(ctx, cfg.Schema, logger)
+	schemaSnapshotGenerator, err := newSchemaSnapshotGenerator(ctx, &cfg.Schema, processEventAdapter.ProcessRow, logger)
 	if err != nil {
 		return nil, err
 	}
 
 	// postgres data snapshot generator
-	dataSnapshotGenerator, err := pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, s.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
-	s.generator = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
+	g = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
 
 	// snapshot table finder layer
-	s.generator, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, s.generator)
+	g, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, g)
 	if err != nil {
 		return nil, err
 	}
@@ -93,81 +73,13 @@ func NewSnapshotGeneratorAdapter(ctx context.Context, cfg *SnapshotConfig, proce
 		if err != nil {
 			return nil, fmt.Errorf("create postgres snapshot store: %w", err)
 		}
-		s.generator = generator.NewSnapshotRecorder(snapshotStore, s.generator)
-	}
-
-	return s, nil
-}
-
-func (s *SnapshotGeneratorAdapter) CreateSnapshot(ctx context.Context) error {
-	errGroup, ctx := errgroup.WithContext(ctx)
-	snapshotChan := make(chan *snapshot.Snapshot)
-	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
-				}
-			}
-			return nil
-		})
-	}
-	for schema, tables := range s.schemaTables {
-		snapshotChan <- &snapshot.Snapshot{
-			SchemaName: schema,
-			TableNames: tables,
-		}
-	}
-	close(snapshotChan)
-
-	return errGroup.Wait()
-}
-
-func (s *SnapshotGeneratorAdapter) Close() error {
-	s.logger.Info("closing snapshot generator")
-	return s.generator.Close()
-}
-
-func (s *SnapshotGeneratorAdapter) processRow(ctx context.Context, row *snapshot.Row) error {
-	return s.processEvent(ctx, s.snapshotRowToWalEvent(row))
-}
-
-func (s *SnapshotGeneratorAdapter) snapshotRowToWalEvent(row *snapshot.Row) *wal.Event {
-	if row == nil {
-		return nil
-	}
-
-	columns := make([]wal.Column, 0, len(row.Columns))
-	for _, col := range row.Columns {
-		columns = append(columns, s.snapshotColumnToWalColumn(col))
+		g = generator.NewSnapshotRecorder(snapshotStore, g)
 	}
-	// use 0 since there's no LSN associated, but it can be used as the
-	// initial version downstream
-	const zeroLSN = "0/0"
-	return &wal.Event{
-		CommitPosition: wal.CommitPosition(zeroLSN),
-		Data: &wal.Data{
-			Action:    "I",
-			Timestamp: s.clock.Now().UTC().Format(time.RFC3339),
-			LSN:       zeroLSN,
-			Schema:    row.Schema,
-			Table:     row.Table,
-			Columns:   columns,
-		},
-	}
-}
 
-func (s *SnapshotGeneratorAdapter) snapshotColumnToWalColumn(col snapshot.Column) wal.Column {
-	return wal.Column{
-		Name:  col.Name,
-		Type:  col.Type,
-		Value: col.Value,
-	}
+	return adapter.NewSnapshotGeneratorAdapter(&cfg.Adapter, g, adapter.WithLogger(logger)), nil
 }
 
-func (s *SnapshotGeneratorAdapter) newSchemaSnapshotGenerator(ctx context.Context, cfg SchemaSnapshotConfig, logger loglib.Logger) (generator.SnapshotGenerator, error) {
+func newSchemaSnapshotGenerator(ctx context.Context, cfg *SchemaSnapshotConfig, processRow snapshot.RowProcessor, logger loglib.Logger) (generator.SnapshotGenerator, error) {
 	switch {
 	case cfg.SchemaLogStore != nil:
 		// postgres schemalog schema snapshot generator
@@ -178,7 +90,7 @@ func (s *SnapshotGeneratorAdapter) newSchemaSnapshotGenerator(ctx context.Contex
 			return nil, fmt.Errorf("create schema log postgres store: %w", err)
 		}
 		schemaLogStore = schemalog.NewStoreCache(schemaLogStore)
-		return schemalogsnapshotgenerator.NewSnapshotGenerator(schemaLogStore, s.processRow), nil
+		return schemalogsnapshotgenerator.NewSnapshotGenerator(schemaLogStore, processRow), nil
 	case cfg.DumpRestore != nil:
 		// postgres pgdump/pgrestore schema snapshot generator
 		return pgdumprestoregenerator.NewSnapshotGenerator(ctx, cfg.DumpRestore, pgdumprestoregenerator.WithLogger(logger))
diff --git a/pkg/wal/listener/snapshot/wal_snapshot_listener.go b/pkg/wal/listener/snapshot/wal_snapshot_listener.go
new file mode 100644
index 0000000..b03ce24
--- /dev/null
+++ b/pkg/wal/listener/snapshot/wal_snapshot_listener.go
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package snapshot
+
+import "context"
+
+type Generator interface {
+	CreateSnapshot(context.Context) error
+	Close() error
+}
diff --git a/pkg/wal/listener/wal_listener.go b/pkg/wal/listener/wal_listener.go
index 1719e4a..c62e242 100644
--- a/pkg/wal/listener/wal_listener.go
+++ b/pkg/wal/listener/wal_listener.go
@@ -2,10 +2,16 @@
 
 package listener
 
-import "context"
+import (
+	"context"
+
+	"github.com/xataio/pgstream/pkg/wal"
+)
 
 // Listener represents a process that listens to WAL events.
 type Listener interface {
 	Listen(ctx context.Context) error
 	Close() error
 }
+
+type ProcessWalEvent func(context.Context, *wal.Event) error


Test patch:
diff --git a/pkg/stream/integration/helper_test.go b/pkg/stream/integration/helper_test.go
index ee2bc87..4c6f742 100644
--- a/pkg/stream/integration/helper_test.go
+++ b/pkg/stream/integration/helper_test.go
@@ -22,7 +22,8 @@ import (
 	"github.com/xataio/pgstream/pkg/tls"
 	"github.com/xataio/pgstream/pkg/wal"
 	kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka"
-	pglistener "github.com/xataio/pgstream/pkg/wal/listener/postgres"
+	"github.com/xataio/pgstream/pkg/wal/listener/snapshot/adapter"
+	snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder"
 	"github.com/xataio/pgstream/pkg/wal/processor/batch"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
@@ -126,12 +127,14 @@ func testPostgresListenerCfgWithSnapshot(sourceURL, targetURL string, tables []s
 			Replication: pgreplication.Config{
 				PostgresURL: sourceURL,
 			},
-			Snapshot: &pglistener.SnapshotConfig{
+			Snapshot: &snapshotbuilder.SnapshotListenerConfig{
 				Generator: pgsnapshotgenerator.Config{
 					URL: sourceURL,
 				},
-				Tables: tables,
-				Schema: pglistener.SchemaSnapshotConfig{
+				Adapter: adapter.SnapshotConfig{
+					Tables: tables,
+				},
+				Schema: snapshotbuilder.SchemaSnapshotConfig{
 					DumpRestore: &pgdumprestore.Config{
 						SourcePGURL: sourceURL,
 						TargetPGURL: targetURL,
diff --git a/pkg/wal/listener/postgres/config_test.go b/pkg/wal/listener/snapshot/adapter/config_test.go
similarity index 97%
rename from pkg/wal/listener/postgres/config_test.go
rename to pkg/wal/listener/snapshot/adapter/config_test.go
index b8d52fa..00d4aa6 100644
--- a/pkg/wal/listener/postgres/config_test.go
+++ b/pkg/wal/listener/snapshot/adapter/config_test.go
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package postgres
+package adapter
 
 import (
 	"testing"
diff --git a/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter_test.go b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter_test.go
new file mode 100644
index 0000000..b69369b
--- /dev/null
+++ b/pkg/wal/listener/snapshot/adapter/wal_process_event_adapter_test.go
@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package adapter
+
+import (
+	"testing"
+	"time"
+
+	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/wal"
+)
+
+func TestProcessEventAdapter_snapshotRowToWalEvent(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	fakeClock := clockwork.NewFakeClockAt(now)
+	testTable := "table1"
+	zeroLSN := "0/0"
+
+	tests := []struct {
+		name string
+		row  *snapshot.Row
+
+		wantEvent *wal.Event
+	}{
+		{
+			name: "ok - nil row",
+			row:  nil,
+
+			wantEvent: nil,
+		},
+		{
+			name: "ok",
+			row: &snapshot.Row{
+				Schema: publicSchema,
+				Table:  testTable,
+				Columns: []snapshot.Column{
+					{Name: "id", Type: "int4", Value: 1},
+					{Name: "name", Type: "text", Value: "alice"},
+				},
+			},
+
+			wantEvent: &wal.Event{
+				CommitPosition: wal.CommitPosition(zeroLSN),
+				Data: &wal.Data{
+					Action:    "I",
+					Timestamp: fakeClock.Now().UTC().Format(time.RFC3339),
+					LSN:       zeroLSN,
+					Schema:    publicSchema,
+					Table:     testTable,
+					Columns: []wal.Column{
+						{Name: "id", Type: "int4", Value: 1},
+						{Name: "name", Type: "text", Value: "alice"},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			a := ProcessEventAdapter{
+				clock: fakeClock,
+			}
+			event := a.snapshotRowToWalEvent(tc.row)
+
+			require.Equal(t, tc.wantEvent, event)
+		})
+	}
+}
diff --git a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter_test.go b/pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter_test.go
similarity index 52%
rename from pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter_test.go
rename to pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter_test.go
index c0b5409..7d86396 100644
--- a/pkg/wal/listener/postgres/wal_pg_snapshot_generator_adapter_test.go
+++ b/pkg/wal/listener/snapshot/adapter/wal_snapshot_generator_adapter_test.go
@@ -1,20 +1,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package postgres
+package adapter
 
 import (
 	"context"
 	"errors"
 	"testing"
-	"time"
 
-	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
 	"github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/snapshot"
 	"github.com/xataio/pgstream/pkg/snapshot/generator"
 	generatormocks "github.com/xataio/pgstream/pkg/snapshot/generator/mocks"
-	"github.com/xataio/pgstream/pkg/wal"
 )
 
 func TestSnapshotGeneratorAdapter_CreateSnapshot(t *testing.T) {
@@ -68,9 +65,7 @@ func TestSnapshotGeneratorAdapter_CreateSnapshot(t *testing.T) {
 			ga := SnapshotGeneratorAdapter{
 				logger:          log.NewNoopLogger(),
 				generator:       tc.generator,
-				processEvent:    func(ctx context.Context, e *wal.Event) error { return nil },
 				schemaTables:    tc.schemaTables,
-				clock:           clockwork.NewFakeClock(),
 				snapshotWorkers: 1,
 			}
 			defer ga.Close()
@@ -80,66 +75,3 @@ func TestSnapshotGeneratorAdapter_CreateSnapshot(t *testing.T) {
 		})
 	}
 }
-
-func TestSnapshotGeneratorAdapter_snapshotRowToWalEvent(t *testing.T) {
-	t.Parallel()
-
-	now := time.Now()
-	fakeClock := clockwork.NewFakeClockAt(now)
-	testTable := "table1"
-	zeroLSN := "0/0"
-
-	tests := []struct {
-		name string
-		row  *snapshot.Row
-
-		wantEvent *wal.Event
-	}{
-		{
-			name: "ok - nil row",
-			row:  nil,
-
-			wantEvent: nil,
-		},
-		{
-			name: "ok",
-			row: &snapshot.Row{
-				Schema: publicSchema,
-				Table:  testTable,
-				Columns: []snapshot.Column{
-					{Name: "id", Type: "int4", Value: 1},
-					{Name: "name", Type: "text", Value: "alice"},
-				},
-			},
-
-			wantEvent: &wal.Event{
-				CommitPosition: wal.CommitPosition(zeroLSN),
-				Data: &wal.Data{
-					Action:    "I",
-					Timestamp: fakeClock.Now().UTC().Format(time.RFC3339),
-					LSN:       zeroLSN,
-					Schema:    publicSchema,
-					Table:     testTable,
-					Columns: []wal.Column{
-						{Name: "id", Type: "int4", Value: 1},
-						{Name: "name", Type: "text", Value: "alice"},
-					},
-				},
-			},
-		},
-	}
-
-	for _, tc := range tests {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			t.Parallel()
-
-			ga := SnapshotGeneratorAdapter{
-				clock: fakeClock,
-			}
-			event := ga.snapshotRowToWalEvent(tc.row)
-
-			require.Equal(t, tc.wantEvent, event)
-		})
-	}
-}