From 2ad33911528a9d78ccb9c4f64b26a5b0e9ac5fb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?No=C3=A9mi=20V=C3=A1nyi?= <sitbackandwait@gmail.com>
Date: Wed, 8 Jul 2026 15:00:36 +0200
Subject: [PATCH 1/3] init

---
 .../postgres/data/ctid_table_reader.go        | 219 +++++++++++++++
 .../data/instrumented_table_reader.go         |  36 +++
 .../instrumented_table_snapshot_generator.go  |  32 ---
 .../postgres/data/pg_snapshot_generator.go    | 264 +++---------------
 .../data/pg_snapshot_generator_test.go        |  53 ++--
 .../generator/postgres/data/snapshot_tx.go    |  55 ++++
 .../generator/postgres/data/table_reader.go   |  29 ++
 7 files changed, 402 insertions(+), 286 deletions(-)
 create mode 100644 pkg/snapshot/generator/postgres/data/ctid_table_reader.go
 create mode 100644 pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
 delete mode 100644 pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
 create mode 100644 pkg/snapshot/generator/postgres/data/snapshot_tx.go
 create mode 100644 pkg/snapshot/generator/postgres/data/table_reader.go

diff --git a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
new file mode 100644
index 00000000..1c0b85f4
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
@@ -0,0 +1,219 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgtype"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/internal/progress"
+	synclib "github.com/xataio/pgstream/internal/sync"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/wal/processor"
+	"golang.org/x/sync/errgroup"
+)
+
+// ctidReader reads a schema's tables by ranging over a stable transaction
+// snapshot using the ctid. The transaction snapshot is exported once per schema
+// in beginSchema and imported by every page range transaction, which lets the
+// reader parallelise the work across page ranges while keeping a consistent
+// view of each table.
+type ctidReader struct {
+	conn         pglib.Querier
+	logger       loglib.Logger
+	adapter      *adapter
+	processor    processor.Processor
+	tableWorkers uint
+	batchBytes   uint64
+
+	// progress tracking, shared with the snapshot generator.
+	progressTracking bool
+	progressBars     *synclib.Map[string, progress.Bar]
+}
+
+// beginSchema opens the transaction that exports the shared transaction
+// snapshot and keeps it open for the duration of fn, so that every readTable
+// call can import it. The snapshot is only exported when the schema has at least
+// one ctid table to read.
+func (r *ctidReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error {
+	// 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.
+	// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION
+	return r.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+		session := &readSession{}
+		if len(st.tables) > 0 {
+			snapshotID, err := exportSnapshot(ctx, tx)
+			if err != nil {
+				return snapshot.NewSchemaErrors(st.schema, err)
+			}
+			session.snapshotID = snapshotID
+		}
+
+		return fn(ctx, session)
+	}, snapshotTxOptions())
+}
+
+func (r *ctidReader) readTable(ctx context.Context, session *readSession, table *table) error {
+	snapshotID := session.snapshotID
+	tableInfo, err := r.getTableInfo(ctx, table.schema, table.name, snapshotID)
+	if err != nil {
+		return err
+	}
+	if tableInfo.isEmpty() {
+		return nil
+	}
+	table.rowSize = tableInfo.avgRowBytes
+
+	// If one page range fails, we abort the entire table snapshot. The
+	// snapshot relies on the transaction snapshot id to ensure all workers
+	// have the same table view, which allows us to use the ctid to
+	// parallelise the work.
+	rangeChan := make(chan pageRange, tableInfo.pageCount)
+	errGroup, ctx := errgroup.WithContext(ctx)
+	for i := uint(0); i < r.tableWorkers; i++ {
+		errGroup.Go(func() error {
+			return r.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan)
+		})
+	}
+
+	// page count returned by postgres starts at 0, so we need to include it
+	// when creating the page ranges.
+	for start := uint(0); start <= uint(tableInfo.pageCount); start += tableInfo.batchPageSize {
+		rangeChan <- pageRange{
+			start: start,
+			end:   start + tableInfo.batchPageSize,
+		}
+	}
+
+	// wait for all table ranges to complete
+	close(rangeChan)
+	return errGroup.Wait()
+}
+
+func (r *ctidReader) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan pageRange) error {
+	for pageRange := range pageRangeChan {
+		if err := r.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+var pageRangeQuery = "SELECT * FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'"
+
+func (r *ctidReader) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error {
+	return execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error {
+		r.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
+			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
+		})
+
+		query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(table.schema, table.name), pageRange.start, pageRange.end)
+		rows, err := tx.Query(ctx, query)
+		if err != nil {
+			return fmt.Errorf("querying table rows: %w", err)
+		}
+		defer rows.Close()
+
+		// resolve the column metadata (names/types) and timestamp once per page
+		// range, since the field descriptions are identical for every row in the
+		// result set.
+		rowAdapter := r.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions())
+		rowCount := uint(0)
+		for rows.Next() {
+			rowCount++
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				values, err := rows.Values()
+				if err != nil {
+					return fmt.Errorf("retrieving rows values: %w", err)
+				}
+
+				event := rowAdapter.rowToWalEvent(values)
+				if event == nil {
+					continue
+				}
+
+				if err := r.processor.ProcessWALEvent(ctx, event); err != nil {
+					return fmt.Errorf("processing snapshot row: %w", err)
+				}
+			}
+		}
+
+		if r.progressTracking {
+			bar, found := r.progressBars.Get(table.schema)
+			if found {
+				bar.Add64(int64(rowCount) * table.rowSize)
+			}
+		}
+
+		r.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
+			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
+		})
+
+		return rows.Err()
+	})
+}
+
+const (
+	// use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices.
+	// pg_relation_size will return only the size of the table itself, without toast tables.
+	tableInfoQuery = `SELECT
+  (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes,
+  CASE
+	WHEN c.reltuples > 0 THEN
+		ROUND(pg_table_size(c.oid) / c.reltuples)
+	ELSE
+		0
+  END AS avg_row_size
+FROM
+  pg_class c
+  JOIN pg_namespace n ON n.oid = c.relnamespace
+WHERE
+  c.relname = $1
+  AND n.nspname = $2
+  AND c.relkind = 'r';`
+
+	// select the max page for the relation instead of using pg_class.relpages, it may not contain an accurate value if
+	// the table is small, the table has active inserts, or the database has not been vacuumed/analyzed recently.
+	maxPageQuery = `SELECT MAX(ctid) FROM ONLY %s;`
+)
+
+func (r *ctidReader) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) {
+	tableInfo := &tableInfo{}
+	err := execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error {
+		// make sure the schema and table names are unquoted since the system
+		// catalogs store unquoted names
+		err := tx.QueryRow(ctx,
+			[]any{&tableInfo.avgPageBytes, &tableInfo.avgRowBytes},
+			tableInfoQuery,
+			pglib.UnquoteIdentifier(tableName),
+			pglib.UnquoteIdentifier(schemaName))
+		if err != nil {
+			return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err)
+		}
+
+		var ctid pgtype.TID
+		if err := tx.QueryRow(ctx, []any{&ctid}, fmt.Sprintf(maxPageQuery, pglib.QuoteQualifiedIdentifier(schemaName, tableName))); err != nil {
+			return fmt.Errorf("getting max page for table %s.%s: %w", schemaName, tableName, err)
+		}
+		tableInfo.pageCount = int(ctid.BlockNumber)
+
+		tableInfo.calculateBatchPageSize(r.batchBytes)
+
+		r.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{
+			"schema": schemaName, "table": tableName, "snapshotID": snapshotID,
+		})
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return tableInfo, nil
+}
diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
new file mode 100644
index 00000000..dc12335b
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+
+	"github.com/xataio/pgstream/pkg/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type instrumentedTableReader struct {
+	tracer trace.Tracer
+	reader tableReader
+}
+
+func newInstrumentedTableReader(reader tableReader, i *otel.Instrumentation) *instrumentedTableReader {
+	return &instrumentedTableReader{
+		tracer: i.Tracer,
+		reader: reader,
+	}
+}
+
+func (i *instrumentedTableReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error {
+	return i.reader.beginSchema(ctx, st, fn)
+}
+
+func (i *instrumentedTableReader) readTable(ctx context.Context, session *readSession, table *table) (err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.ReadTable", trace.WithAttributes([]attribute.KeyValue{
+		{Key: "schema", Value: attribute.StringValue(table.schema)},
+		{Key: "table", Value: attribute.StringValue(table.name)},
+	}...))
+	defer otel.CloseSpan(span, err)
+	return i.reader.readTable(ctx, session, table)
+}
diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
deleted file mode 100644
index 6ed5e39f..00000000
--- a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
+++ /dev/null
@@ -1,32 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-
-package postgres
-
-import (
-	"context"
-
-	"github.com/xataio/pgstream/pkg/otel"
-	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/trace"
-)
-
-type instrumentedTableSnapshotGenerator struct {
-	tracer          trace.Tracer
-	snapshotTableFn snapshotTableFn
-}
-
-func newInstrumentedTableSnapshotGenerator(fn snapshotTableFn, i *otel.Instrumentation) *instrumentedTableSnapshotGenerator {
-	return &instrumentedTableSnapshotGenerator{
-		tracer:          i.Tracer,
-		snapshotTableFn: fn,
-	}
-}
-
-func (i *instrumentedTableSnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) (err error) {
-	ctx, span := otel.StartSpan(ctx, i.tracer, "tableSnapshotGenerator.SnapshotTable", trace.WithAttributes([]attribute.KeyValue{
-		{Key: "schema", Value: attribute.StringValue(table.schema)},
-		{Key: "table", Value: attribute.StringValue(table.name)},
-	}...))
-	defer otel.CloseSpan(span, err)
-	return i.snapshotTableFn(ctx, snapshotID, table)
-}
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 47e95abe..d415c940 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/jackc/pgx/v5/pgtype"
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pglibinstrumentation "github.com/xataio/pgstream/internal/postgres/instrumentation"
 	"github.com/xataio/pgstream/internal/progress"
@@ -21,21 +20,20 @@ import (
 )
 
 type SnapshotGenerator struct {
-	logger  loglib.Logger
-	conn    pglib.Querier
-	adapter *adapter
+	logger    loglib.Logger
+	conn      pglib.Querier
+	processor processor.Processor
+	// reader encapsulates the strategy used to read a schema's tables (ctid
+	// range scan by default).
+	reader tableReader
+	// instrumentation is captured while applying options and used to wrap the
+	// reader once it has been built.
+	instrumentation *otel.Instrumentation
 
 	// workers per snapshot, parallelise the snapshot creation for each schema
 	snapshotWorkers uint
 	// workers per schema, parallelise the snapshot creation for each table
 	schemaWorkers uint
-	// workers per table, parallelise the snapshot creation for each page range
-	tableWorkers uint
-	batchBytes   uint64
-
-	// Function called for processing produced rows.
-	processor              processor.Processor
-	tableSnapshotGenerator snapshotTableFn
 
 	progressTracking   bool
 	progressBars       *synclib.Map[string, progress.Bar]
@@ -69,8 +67,6 @@ type table struct {
 	rowSize int64
 }
 
-type snapshotTableFn func(ctx context.Context, snapshotID string, table *table) error
-
 type Option func(sg *SnapshotGenerator)
 
 func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor.Processor, opts ...Option) (*SnapshotGenerator, error) {
@@ -87,19 +83,28 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor.
 		logger:          loglib.NewNoopLogger(),
 		conn:            conn,
 		processor:       processor,
-		batchBytes:      cfg.batchBytes(),
-		tableWorkers:    cfg.tableWorkers(),
 		schemaWorkers:   cfg.schemaWorkers(),
 		snapshotWorkers: cfg.snapshotWorkers(),
 	}
 
-	sg.tableSnapshotGenerator = sg.snapshotTable
-
 	for _, opt := range opts {
 		opt(sg)
 	}
 
-	sg.adapter = newAdapter(pglib.NewMapper(conn), sg.logger)
+	sg.reader = &ctidReader{
+		conn:             sg.conn,
+		logger:           sg.logger,
+		adapter:          newAdapter(pglib.NewMapper(conn), sg.logger),
+		processor:        sg.processor,
+		tableWorkers:     cfg.tableWorkers(),
+		batchBytes:       cfg.batchBytes(),
+		progressTracking: sg.progressTracking,
+		progressBars:     sg.progressBars,
+	}
+
+	if sg.instrumentation != nil {
+		sg.reader = newInstrumentedTableReader(sg.reader, sg.instrumentation)
+	}
 
 	return sg, nil
 }
@@ -121,8 +126,7 @@ func WithInstrumentation(i *otel.Instrumentation) Option {
 			panic(err)
 		}
 
-		ig := newInstrumentedTableSnapshotGenerator(sg.tableSnapshotGenerator, i)
-		sg.tableSnapshotGenerator = ig.snapshotTable
+		sg.instrumentation = i
 	}
 }
 
@@ -188,18 +192,9 @@ func (sg *SnapshotGenerator) Close() error {
 }
 
 func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTables *schemaTables) error {
-	// 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.
-	// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION
-	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) (err error) {
-		snapshotID, err := sg.exportSnapshot(ctx, tx)
-		if err != nil {
-			return snapshot.NewSchemaErrors(schemaTables.schema, err)
-		}
-
+	return sg.reader.beginSchema(ctx, schemaTables, func(ctx context.Context, session *readSession) (err error) {
 		if sg.progressTracking {
-			if err := sg.addProgressBar(ctx, snapshotID, schemaTables); err != nil {
+			if err := sg.addProgressBar(ctx, session.snapshotID, schemaTables); err != nil {
 				return err
 			}
 			defer func() {
@@ -217,7 +212,7 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab
 		for i := uint(0); i < sg.schemaWorkers; i++ {
 			wg.Add(1)
 			workerTableErrs[i] = make(map[string]error, len(schemaTables.tables))
-			go sg.createSnapshotWorker(ctx, wg, snapshotID, tableChan, workerTableErrs[i])
+			go sg.createSnapshotWorker(ctx, wg, session, tableChan, workerTableErrs[i])
 		}
 
 		for _, tableName := range schemaTables.tables {
@@ -231,16 +226,16 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab
 		wg.Wait()
 
 		return sg.collectTableErrors(schemaTables.schema, workerTableErrs)
-	}, snapshotTxOptions())
+	})
 }
 
-func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, snapshotID string, tableChan <-chan *table, tableErrMap map[string]error) {
+func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, session *readSession, tableChan <-chan *table, tableErrMap map[string]error) {
 	defer wg.Done()
 	for t := range tableChan {
-		logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": snapshotID}
+		logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": session.snapshotID}
 		sg.logger.Debug("snapshotting table", logFields)
 
-		if err := sg.tableSnapshotGenerator(ctx, snapshotID, t); err != nil {
+		if err := sg.reader.readTable(ctx, session, t); err != nil {
 			sg.logger.Error(err, "snapshotting table", logFields)
 			// errors will get notified unless the table doesn't exist
 			if !errors.Is(err, pglib.ErrNoRows) {
@@ -288,108 +283,6 @@ func (sg *SnapshotGenerator) collectSchemaErrors(workerSchemaErrs map[string]err
 	return nil
 }
 
-func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) error {
-	tableInfo, err := sg.getTableInfo(ctx, table.schema, table.name, snapshotID)
-	if err != nil {
-		return err
-	}
-	if tableInfo.isEmpty() {
-		return nil
-	}
-	table.rowSize = tableInfo.avgRowBytes
-
-	// If one page range fails, we abort the entire table snapshot. The
-	// snapshot relies on the transaction snapshot id to ensure all workers
-	// have the same table view, which allows us to use the ctid to
-	// parallelise the work.
-	rangeChan := make(chan pageRange, tableInfo.pageCount)
-	errGroup, ctx := errgroup.WithContext(ctx)
-	for i := uint(0); i < sg.tableWorkers; i++ {
-		errGroup.Go(func() error {
-			return sg.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan)
-		})
-	}
-
-	// page count returned by postgres starts at 0, so we need to include it
-	// when creating the page ranges.
-	for start := uint(0); start <= uint(tableInfo.pageCount); start += tableInfo.batchPageSize {
-		rangeChan <- pageRange{
-			start: start,
-			end:   start + tableInfo.batchPageSize,
-		}
-	}
-
-	// wait for all table ranges to complete
-	close(rangeChan)
-	return errGroup.Wait()
-}
-
-func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan pageRange) error {
-	for pageRange := range pageRangeChan {
-		if err := sg.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-var pageRangeQuery = "SELECT * FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'"
-
-func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error {
-	return sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
-		sg.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
-		})
-
-		query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(table.schema, table.name), pageRange.start, pageRange.end)
-		rows, err := tx.Query(ctx, query)
-		if err != nil {
-			return fmt.Errorf("querying table rows: %w", err)
-		}
-		defer rows.Close()
-
-		// resolve the column metadata (names/types) and timestamp once per page
-		// range, since the field descriptions are identical for every row in the
-		// result set.
-		rowAdapter := sg.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions())
-		rowCount := uint(0)
-		for rows.Next() {
-			rowCount++
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-				values, err := rows.Values()
-				if err != nil {
-					return fmt.Errorf("retrieving rows values: %w", err)
-				}
-
-				event := rowAdapter.rowToWalEvent(values)
-				if event == nil {
-					continue
-				}
-
-				if err := sg.processor.ProcessWALEvent(ctx, event); err != nil {
-					return fmt.Errorf("processing snapshot row: %w", err)
-				}
-			}
-		}
-
-		if sg.progressTracking {
-			bar, found := sg.progressBars.Get(table.schema)
-			if found {
-				bar.Add64(int64(rowCount) * table.rowSize)
-			}
-		}
-
-		sg.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
-		})
-
-		return rows.Err()
-	})
-}
-
 func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID string, schemaTables *schemaTables) error {
 	totalBytes, err := sg.getSnapshotSchemaTotalBytes(ctx, snapshotID, schemaTables.schema, schemaTables.tables)
 	if err != nil {
@@ -409,64 +302,6 @@ func (sg *SnapshotGenerator) markProgressBarCompleted(schema string) {
 	sg.progressBars.Delete(schema)
 }
 
-const (
-	// use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices.
-	// pg_relation_size will return only the size of the table itself, without toast tables.
-	tableInfoQuery = `SELECT
-  (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes,
-  CASE
-	WHEN c.reltuples > 0 THEN
-		ROUND(pg_table_size(c.oid) / c.reltuples)
-	ELSE
-		0
-  END AS avg_row_size
-FROM
-  pg_class c
-  JOIN pg_namespace n ON n.oid = c.relnamespace
-WHERE
-  c.relname = $1
-  AND n.nspname = $2
-  AND c.relkind = 'r';`
-
-	// select the max page for the relation instead of using pg_class.relpages, it may not contain an accurate value if
-	// the table is small, the table has active inserts, or the database has not been vacuumed/analyzed recently.
-	maxPageQuery = `SELECT MAX(ctid) FROM ONLY %s;`
-)
-
-func (sg *SnapshotGenerator) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) {
-	tableInfo := &tableInfo{}
-	err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
-		// make sure the schema and table names are unquoted since the system
-		// catalogs store unquoted names
-		err := tx.QueryRow(ctx,
-			[]any{&tableInfo.avgPageBytes, &tableInfo.avgRowBytes},
-			tableInfoQuery,
-			pglib.UnquoteIdentifier(tableName),
-			pglib.UnquoteIdentifier(schemaName))
-		if err != nil {
-			return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err)
-		}
-
-		var ctid pgtype.TID
-		if err := tx.QueryRow(ctx, []any{&ctid}, fmt.Sprintf(maxPageQuery, pglib.QuoteQualifiedIdentifier(schemaName, tableName))); err != nil {
-			return fmt.Errorf("getting max page for table %s.%s: %w", schemaName, tableName, err)
-		}
-		tableInfo.pageCount = int(ctid.BlockNumber)
-
-		tableInfo.calculateBatchPageSize(sg.batchBytes)
-
-		sg.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{
-			"schema": schemaName, "table": tableName, "snapshotID": snapshotID,
-		})
-		return nil
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	return tableInfo, nil
-}
-
 const tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = ANY($2) AND c.relkind = 'r';`
 
 func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, snapshotID, schema string, tables []string) (int64, error) {
@@ -482,7 +317,7 @@ func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, sn
 		unquotedTables[i] = pglib.UnquoteIdentifier(table)
 	}
 
-	err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
+	err := execInSnapshotTx(ctx, sg.conn, snapshotID, func(tx pglib.Tx) error {
 		err := tx.QueryRow(ctx, []any{&totalBytes}, tablesBytesQuery, pglib.UnquoteIdentifier(schema), unquotedTables)
 		if err != nil {
 			return fmt.Errorf("retrieving total bytes for schema: %w", err)
@@ -493,41 +328,6 @@ func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, sn
 	return totalBytes, err
 }
 
-const exportSnapshotQuery = `SELECT pg_export_snapshot()`
-
-func (sg *SnapshotGenerator) exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) {
-	var snapshotID string
-	if err := tx.QueryRow(ctx, []any{&snapshotID}, exportSnapshotQuery); err != nil {
-		return "", fmt.Errorf("exporting snapshot: %w", err)
-	}
-	return snapshotID, nil
-}
-
-func (sg *SnapshotGenerator) setTransactionSnapshot(ctx context.Context, tx pglib.Tx, snapshotID string) error {
-	_, err := tx.Exec(ctx, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", snapshotID))
-	if err != nil {
-		return fmt.Errorf("setting transaction snapshot: %w", err)
-	}
-	return nil
-}
-
-func (sg *SnapshotGenerator) execInSnapshotTx(ctx context.Context, snapshotID string, fn func(tx pglib.Tx) error) error {
-	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
-		if err := sg.setTransactionSnapshot(ctx, tx, snapshotID); err != nil {
-			return err
-		}
-
-		return fn(tx)
-	}, snapshotTxOptions())
-}
-
-func snapshotTxOptions() pglib.TxOptions {
-	return pglib.TxOptions{
-		IsolationLevel: pglib.RepeatableRead,
-		AccessMode:     pglib.ReadOnly,
-	}
-}
-
 // calculateBatchPageSize will automatically determine the batch page size based
 // on the average page size and the configured batch bytes limit.
 func (t *tableInfo) calculateBatchPageSize(bytes uint64) {
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 bb3aeda4..99c13aee 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -1132,32 +1132,41 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			t.Parallel()
 
 			eventChan := make(chan *wal.Event, 10)
-			sg := SnapshotGenerator{
-				logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{
-					LogLevel: "debug",
-				})),
-				conn:    tc.querier,
-				adapter: newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
-				processor: &processormocks.Processor{
-					ProcessWALEventFn: func(ctx context.Context, e *wal.Event) error {
-						eventChan <- e
-						return nil
-					},
-					CloseFn: func() error {
-						return tc.processorCloseErr
-					},
+			logger := zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{
+				LogLevel: "debug",
+			}))
+			processor := &processormocks.Processor{
+				ProcessWALEventFn: func(ctx context.Context, e *wal.Event) error {
+					eventChan <- e
+					return nil
 				},
+				CloseFn: func() error {
+					return tc.processorCloseErr
+				},
+			}
+			progressBars := synclib.NewMap[string, progress.Bar]()
+			sg := SnapshotGenerator{
+				logger:           logger,
+				conn:             tc.querier,
+				processor:        processor,
 				schemaWorkers:    1,
-				tableWorkers:     1,
-				batchBytes:       1024 * 1024, // 1MB
 				snapshotWorkers:  1,
 				progressTracking: tc.progressBar != nil,
-				progressBars:     synclib.NewMap[string, progress.Bar](),
+				progressBars:     progressBars,
 				progressBarBuilder: func(totalBytes int64, description string) progress.Bar {
 					return tc.progressBar
 				},
+				reader: &ctidReader{
+					conn:             tc.querier,
+					logger:           logger,
+					adapter:          newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
+					processor:        processor,
+					tableWorkers:     1,
+					batchBytes:       1024 * 1024, // 1MB
+					progressTracking: tc.progressBar != nil,
+					progressBars:     progressBars,
+				},
 			}
-			sg.tableSnapshotGenerator = sg.snapshotTable
 
 			if tc.schemaWorkers != 0 {
 				sg.schemaWorkers = tc.schemaWorkers
@@ -1667,7 +1676,7 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) {
 				},
 			}
 
-			sg := SnapshotGenerator{
+			reader := &ctidReader{
 				logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{
 					LogLevel: "debug",
 				})),
@@ -1688,11 +1697,11 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) {
 				progressBars:     synclib.NewMap[string, progress.Bar](),
 			}
 
-			if sg.progressTracking {
-				sg.progressBars.Set(tc.table.schema, progressBar)
+			if reader.progressTracking {
+				reader.progressBars.Set(tc.table.schema, progressBar)
 			}
 
-			err := sg.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange)
+			err := reader.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange)
 			require.Equal(t, tc.wantErr, err)
 			close(eventChan)
 
diff --git a/pkg/snapshot/generator/postgres/data/snapshot_tx.go b/pkg/snapshot/generator/postgres/data/snapshot_tx.go
new file mode 100644
index 00000000..cec143b0
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/snapshot_tx.go
@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+)
+
+// snapshot_tx.go groups the helpers used to read from a shared transaction
+// snapshot. A transaction snapshot is exported once per schema and imported by
+// every reader transaction so that all workers observe the same stable view of
+// the database, which is what allows the ctid based reader to parallelise the
+// work.
+// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION
+
+const exportSnapshotQuery = `SELECT pg_export_snapshot()`
+
+func exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) {
+	var snapshotID string
+	if err := tx.QueryRow(ctx, []any{&snapshotID}, exportSnapshotQuery); err != nil {
+		return "", fmt.Errorf("exporting snapshot: %w", err)
+	}
+	return snapshotID, nil
+}
+
+func setTransactionSnapshot(ctx context.Context, tx pglib.Tx, snapshotID string) error {
+	_, err := tx.Exec(ctx, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", snapshotID))
+	if err != nil {
+		return fmt.Errorf("setting transaction snapshot: %w", err)
+	}
+	return nil
+}
+
+// execInSnapshotTx runs fn in a read only repeatable read transaction that
+// imports the given transaction snapshot, so it observes the same view of the
+// database as the transaction that exported it.
+func execInSnapshotTx(ctx context.Context, conn pglib.Querier, snapshotID string, fn func(tx pglib.Tx) error) error {
+	return conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+		if err := setTransactionSnapshot(ctx, tx, snapshotID); err != nil {
+			return err
+		}
+
+		return fn(tx)
+	}, snapshotTxOptions())
+}
+
+func snapshotTxOptions() pglib.TxOptions {
+	return pglib.TxOptions{
+		IsolationLevel: pglib.RepeatableRead,
+		AccessMode:     pglib.ReadOnly,
+	}
+}
diff --git a/pkg/snapshot/generator/postgres/data/table_reader.go b/pkg/snapshot/generator/postgres/data/table_reader.go
new file mode 100644
index 00000000..d06635e9
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/table_reader.go
@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "context"
+
+// readSession carries the per-schema state that a tableReader needs to snapshot
+// individual tables. For the ctid reader it holds the exported transaction
+// snapshot id shared by all the table read transactions. It is only valid for
+// the duration of the beginSchema callback that produced it.
+type readSession struct {
+	snapshotID string
+}
+
+// tableReader abstracts the strategy used to read a schema's tables during a
+// data snapshot. The ctid reader ranges over a stable transaction snapshot
+// using the ctid to parallelise the work; future implementations (e.g. a
+// primary-key keyset reader) can provide alternative behaviour behind the same
+// seam.
+type tableReader interface {
+	// beginSchema prepares the reader to snapshot the given schema tables and
+	// invokes fn with the readSession that must be used for every readTable call
+	// belonging to that schema. The session is only valid for the duration of
+	// fn.
+	beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error
+	// readTable snapshots a single table using the session provided by
+	// beginSchema.
+	readTable(ctx context.Context, session *readSession, table *table) error
+}

From 0cb8d3427193b0049b23099c856f9ce0676ad568 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?No=C3=A9mi=20V=C3=A1nyi?= <sitbackandwait@gmail.com>
Date: Mon, 20 Jul 2026 15:49:07 +0200
Subject: [PATCH 2/3] Sharpen tableReader seam: document readSession, group
 progress state

- Collapse the progressTracking bool + progressBars map clump, which was
  duplicated across the generator and the reader, into a single
  progressTracker value with set/advance/complete methods. A zero-value
  tracker is disabled and no-ops, so callers no longer guard each call.
---
 .../postgres/data/ctid_table_reader.go        | 15 ++---
 .../postgres/data/pg_snapshot_generator.go    | 36 ++++-------
 .../data/pg_snapshot_generator_test.go        | 43 +++++++-------
 .../postgres/data/progress_tracker.go         | 59 +++++++++++++++++++
 .../generator/postgres/data/table_reader.go   | 21 ++++---
 5 files changed, 108 insertions(+), 66 deletions(-)
 create mode 100644 pkg/snapshot/generator/postgres/data/progress_tracker.go

diff --git a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
index 1c0b85f4..3f2367ce 100644
--- a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
+++ b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
@@ -8,8 +8,6 @@ import (
 
 	"github.com/jackc/pgx/v5/pgtype"
 	pglib "github.com/xataio/pgstream/internal/postgres"
-	"github.com/xataio/pgstream/internal/progress"
-	synclib "github.com/xataio/pgstream/internal/sync"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/snapshot"
 	"github.com/xataio/pgstream/pkg/wal/processor"
@@ -29,9 +27,9 @@ type ctidReader struct {
 	tableWorkers uint
 	batchBytes   uint64
 
-	// progress tracking, shared with the snapshot generator.
-	progressTracking bool
-	progressBars     *synclib.Map[string, progress.Bar]
+	// progress is shared by value with the snapshot generator; the underlying
+	// bars map is shared by reference.
+	progress progressTracker
 }
 
 // beginSchema opens the transaction that exports the shared transaction
@@ -145,12 +143,7 @@ func (r *ctidReader) snapshotTableRange(ctx context.Context, snapshotID string,
 			}
 		}
 
-		if r.progressTracking {
-			bar, found := r.progressBars.Get(table.schema)
-			if found {
-				bar.Add64(int64(rowCount) * table.rowSize)
-			}
-		}
+		r.progress.advance(table.schema, int64(rowCount)*table.rowSize)
 
 		r.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
 			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index d415c940..20b8e00c 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -11,7 +11,6 @@ import (
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pglibinstrumentation "github.com/xataio/pgstream/internal/postgres/instrumentation"
 	"github.com/xataio/pgstream/internal/progress"
-	synclib "github.com/xataio/pgstream/internal/sync"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/otel"
 	"github.com/xataio/pgstream/pkg/snapshot"
@@ -35,8 +34,7 @@ type SnapshotGenerator struct {
 	// workers per schema, parallelise the snapshot creation for each table
 	schemaWorkers uint
 
-	progressTracking   bool
-	progressBars       *synclib.Map[string, progress.Bar]
+	progress           progressTracker
 	progressBarBuilder func(totalBytes int64, description string) progress.Bar
 }
 
@@ -92,14 +90,13 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor.
 	}
 
 	sg.reader = &ctidReader{
-		conn:             sg.conn,
-		logger:           sg.logger,
-		adapter:          newAdapter(pglib.NewMapper(conn), sg.logger),
-		processor:        sg.processor,
-		tableWorkers:     cfg.tableWorkers(),
-		batchBytes:       cfg.batchBytes(),
-		progressTracking: sg.progressTracking,
-		progressBars:     sg.progressBars,
+		conn:         sg.conn,
+		logger:       sg.logger,
+		adapter:      newAdapter(pglib.NewMapper(conn), sg.logger),
+		processor:    sg.processor,
+		tableWorkers: cfg.tableWorkers(),
+		batchBytes:   cfg.batchBytes(),
+		progress:     sg.progress,
 	}
 
 	if sg.instrumentation != nil {
@@ -132,8 +129,7 @@ func WithInstrumentation(i *otel.Instrumentation) Option {
 
 func WithProgressTracking() Option {
 	return func(sg *SnapshotGenerator) {
-		sg.progressTracking = true
-		sg.progressBars = synclib.NewMap[string, progress.Bar]()
+		sg.progress = newProgressTracker()
 		sg.progressBarBuilder = progress.NewBytesBar
 	}
 }
@@ -193,13 +189,13 @@ func (sg *SnapshotGenerator) Close() error {
 
 func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTables *schemaTables) error {
 	return sg.reader.beginSchema(ctx, schemaTables, func(ctx context.Context, session *readSession) (err error) {
-		if sg.progressTracking {
+		if sg.progress.enabled {
 			if err := sg.addProgressBar(ctx, session.snapshotID, schemaTables); err != nil {
 				return err
 			}
 			defer func() {
 				if err == nil {
-					sg.markProgressBarCompleted(schemaTables.schema)
+					sg.progress.complete(schemaTables.schema)
 				}
 			}()
 		}
@@ -290,18 +286,10 @@ func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID stri
 	}
 
 	bar := sg.progressBarBuilder(totalBytes, fmt.Sprintf("[cyan][%s][reset] Snapshotting data...", schemaTables.schema))
-	sg.progressBars.Set(schemaTables.schema, bar)
+	sg.progress.set(schemaTables.schema, bar)
 	return nil
 }
 
-func (sg *SnapshotGenerator) markProgressBarCompleted(schema string) {
-	bar, found := sg.progressBars.Get(schema)
-	if found {
-		bar.Close()
-	}
-	sg.progressBars.Delete(schema)
-}
-
 const tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = ANY($2) AND c.relkind = 'r';`
 
 func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, snapshotID, schema string, tables []string) (int64, error) {
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 99c13aee..8c67c22e 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -1144,27 +1144,28 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					return tc.processorCloseErr
 				},
 			}
-			progressBars := synclib.NewMap[string, progress.Bar]()
+			pt := progressTracker{
+				enabled: tc.progressBar != nil,
+				bars:    synclib.NewMap[string, progress.Bar](),
+			}
 			sg := SnapshotGenerator{
-				logger:           logger,
-				conn:             tc.querier,
-				processor:        processor,
-				schemaWorkers:    1,
-				snapshotWorkers:  1,
-				progressTracking: tc.progressBar != nil,
-				progressBars:     progressBars,
+				logger:          logger,
+				conn:            tc.querier,
+				processor:       processor,
+				schemaWorkers:   1,
+				snapshotWorkers: 1,
+				progress:        pt,
 				progressBarBuilder: func(totalBytes int64, description string) progress.Bar {
 					return tc.progressBar
 				},
 				reader: &ctidReader{
-					conn:             tc.querier,
-					logger:           logger,
-					adapter:          newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
-					processor:        processor,
-					tableWorkers:     1,
-					batchBytes:       1024 * 1024, // 1MB
-					progressTracking: tc.progressBar != nil,
-					progressBars:     progressBars,
+					conn:         tc.querier,
+					logger:       logger,
+					adapter:      newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
+					processor:    processor,
+					tableWorkers: 1,
+					batchBytes:   1024 * 1024, // 1MB
+					progress:     pt,
 				},
 			}
 
@@ -1693,12 +1694,14 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) {
 						return nil
 					},
 				},
-				progressTracking: tc.name == "ok - with progress tracking",
-				progressBars:     synclib.NewMap[string, progress.Bar](),
+				progress: progressTracker{
+					enabled: tc.name == "ok - with progress tracking",
+					bars:    synclib.NewMap[string, progress.Bar](),
+				},
 			}
 
-			if reader.progressTracking {
-				reader.progressBars.Set(tc.table.schema, progressBar)
+			if reader.progress.enabled {
+				reader.progress.set(tc.table.schema, progressBar)
 			}
 
 			err := reader.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange)
diff --git a/pkg/snapshot/generator/postgres/data/progress_tracker.go b/pkg/snapshot/generator/postgres/data/progress_tracker.go
new file mode 100644
index 00000000..3d4e52a1
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/progress_tracker.go
@@ -0,0 +1,59 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"github.com/xataio/pgstream/internal/progress"
+	synclib "github.com/xataio/pgstream/internal/sync"
+)
+
+// progressTracker owns the per-schema progress bars shared between the snapshot
+// generator, which creates and completes the bars, and the table readers, which
+// advance them as rows are processed. The enabled flag and the bars map always
+// travel together, so they are grouped here instead of being threaded as a pair
+// through every constructor.
+//
+// The bars map is a pointer, so copies of a progressTracker value share the same
+// underlying bars. A zero-value tracker is disabled and every method is a no-op,
+// which lets callers advance progress without guarding each call.
+type progressTracker struct {
+	enabled bool
+	bars    *synclib.Map[string, progress.Bar]
+}
+
+func newProgressTracker() progressTracker {
+	return progressTracker{
+		enabled: true,
+		bars:    synclib.NewMap[string, progress.Bar](),
+	}
+}
+
+// set registers the bar tracking the given schema's progress.
+func (p progressTracker) set(schema string, bar progress.Bar) {
+	if !p.enabled {
+		return
+	}
+	p.bars.Set(schema, bar)
+}
+
+// advance adds the given number of bytes to the given schema's bar, if one is
+// being tracked.
+func (p progressTracker) advance(schema string, bytes int64) {
+	if !p.enabled {
+		return
+	}
+	if bar, found := p.bars.Get(schema); found {
+		bar.Add64(bytes)
+	}
+}
+
+// complete closes the given schema's bar and stops tracking it.
+func (p progressTracker) complete(schema string) {
+	if !p.enabled {
+		return
+	}
+	if bar, found := p.bars.Get(schema); found {
+		bar.Close()
+	}
+	p.bars.Delete(schema)
+}
diff --git a/pkg/snapshot/generator/postgres/data/table_reader.go b/pkg/snapshot/generator/postgres/data/table_reader.go
index d06635e9..196289b9 100644
--- a/pkg/snapshot/generator/postgres/data/table_reader.go
+++ b/pkg/snapshot/generator/postgres/data/table_reader.go
@@ -5,25 +5,24 @@ package postgres
 import "context"
 
 // readSession carries the per-schema state that a tableReader needs to snapshot
-// individual tables. For the ctid reader it holds the exported transaction
-// snapshot id shared by all the table read transactions. It is only valid for
-// the duration of the beginSchema callback that produced it.
+// individual tables. It is only valid for the duration of the beginSchema
+// callback that produced it.
 type readSession struct {
+	// snapshotID is the exported transaction snapshot shared by all of a
+	// schema's table read transactions. It is empty when the schema has no
+	// tables to read, in which case readTable must not be called (an empty id
+	// would make SET TRANSACTION SNAPSHOT fail).
 	snapshotID string
 }
 
 // tableReader abstracts the strategy used to read a schema's tables during a
-// data snapshot. The ctid reader ranges over a stable transaction snapshot
-// using the ctid to parallelise the work; future implementations (e.g. a
-// primary-key keyset reader) can provide alternative behaviour behind the same
-// seam.
+// data snapshot.
 type tableReader interface {
 	// beginSchema prepares the reader to snapshot the given schema tables and
 	// invokes fn with the readSession that must be used for every readTable call
-	// belonging to that schema. The session is only valid for the duration of
-	// fn.
+	// belonging to that schema.
+	// The session is only valid for the duration of fn.
 	beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error
-	// readTable snapshots a single table using the session provided by
-	// beginSchema.
+	// readTable snapshots a single table using the session provided by beginSchema.
 	readTable(ctx context.Context, session *readSession, table *table) error
 }

From 85a2d502752fe1459223fe7392ae7ff788e54f79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?No=C3=A9mi=20V=C3=A1nyi?= <sitbackandwait@gmail.com>
Date: Wed, 12 Aug 2026 15:52:37 +0200
Subject: [PATCH 3/3] Turn the tableReader strategy into a factory of read
 sessions

readSession carried a snapshotID, which is a transaction snapshot detail
of the ctid strategy, and the generator reached into it: for its log
fields and to run its own getSnapshotSchemaTotalBytes. So the generator
still knew how the ctid reader works, and a strategy that doesn't export
a transaction snapshot would have had nothing to put in that field.

A tableReader now only opens a session per schema, and the session does
the reading, already bound to its schema and to the state its strategy
needs. ctidReader produces a ctidSession that owns the snapshotID, which
also drops the id from the signatures of readTable, getTableInfo and
snapshotTableRange, and makes it impossible to pair a session with the
wrong reader - the previous contract could only document that.

Alongside:

- Move getSnapshotSchemaTotalBytes onto ctidSession.totalBytes. The
  generator owns the progress bar, the strategy owns how the bytes
  behind it are measured.
- Wrap the session as well as the reader in instrumentedTableReader,
  since the reader hands out the sessions that do the work. beginSchema
  now opens a span too, so table spans have a schema parent.
- Extract rowSink: the row to wal event to processor to progress path is
  the same whichever rows a strategy decides to read, so it no longer
  sits inlined in snapshotTableRange next to the query and tx handling.
- Build the reader in newTableReader, so choosing a strategy and
  applying the decorators happens in one place instead of in
  NewSnapshotGenerator.
---
 .../postgres/data/ctid_table_reader.go        | 141 +++++++++++-------
 .../data/instrumented_table_reader.go         |  35 ++++-
 .../postgres/data/pg_snapshot_generator.go    |  65 +++-----
 .../data/pg_snapshot_generator_test.go        |  35 +++--
 .../generator/postgres/data/row_sink.go       |  66 ++++++++
 .../generator/postgres/data/table_reader.go   |  70 ++++++---
 6 files changed, 266 insertions(+), 146 deletions(-)
 create mode 100644 pkg/snapshot/generator/postgres/data/row_sink.go

diff --git a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
index d7c9e3ed..bb01a627 100644
--- a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
+++ b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go
@@ -12,7 +12,6 @@ import (
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/snapshot"
-	"github.com/xataio/pgstream/pkg/wal/processor"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -24,27 +23,25 @@ import (
 type ctidReader struct {
 	conn         pglib.Querier
 	logger       loglib.Logger
-	adapter      *adapter
-	processor    processor.Processor
+	sink         rowSink
 	tableWorkers uint
 	batchBytes   uint64
-
-	// progress is shared by value with the snapshot generator; the underlying
-	// bars map is shared by reference.
-	progress progressTracker
 }
 
 // beginSchema opens the transaction that exports the shared transaction
-// snapshot and keeps it open for the duration of fn, so that every readTable
-// call can import it. The snapshot is only exported when the schema has at least
-// one ctid table to read.
-func (r *ctidReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error {
+// snapshot and keeps it open for the duration of fn, so that every read the
+// session performs can import it. The snapshot is only exported when the schema
+// has at least one ctid table to read.
+func (r *ctidReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, readSession) error) error {
 	// 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.
 	// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION
 	return r.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
-		session := &readSession{}
+		session := &ctidSession{
+			reader:       r,
+			schemaTables: st,
+		}
 		if len(st.tables) > 0 {
 			snapshotID, err := exportSnapshot(ctx, tx)
 			if err != nil {
@@ -57,9 +54,31 @@ func (r *ctidReader) beginSchema(ctx context.Context, st *schemaTables, fn func(
 	}, snapshotTxOptions())
 }
 
-func (r *ctidReader) readTable(ctx context.Context, session *readSession, table *table) error {
-	snapshotID := session.snapshotID
-	tableInfo, err := r.getTableInfo(ctx, table.schema, table.name, snapshotID)
+// ctidSession reads the tables of a single schema over the transaction snapshot
+// exported by beginSchema. Every transaction it opens imports that snapshot, so
+// all the workers observe the same view of the schema, which is what makes it
+// safe to split a table into page ranges read in parallel.
+type ctidSession struct {
+	reader       *ctidReader
+	schemaTables *schemaTables
+	// snapshotID is empty when the schema has no tables to read, in which case
+	// no read is performed: an empty id would make SET TRANSACTION SNAPSHOT
+	// fail.
+	snapshotID string
+}
+
+func (s *ctidSession) logFields() loglib.Fields {
+	return loglib.Fields{"snapshotID": s.snapshotID}
+}
+
+// execInTx runs fn in a transaction that observes the same view of the database
+// as the transaction that exported this session's snapshot.
+func (s *ctidSession) execInTx(ctx context.Context, fn func(tx pglib.Tx) error) error {
+	return execInSnapshotTx(ctx, s.reader.conn, s.snapshotID, fn)
+}
+
+func (s *ctidSession) readTable(ctx context.Context, table *table) error {
+	tableInfo, err := s.getTableInfo(ctx, table.schema, table.name)
 	if err != nil {
 		return err
 	}
@@ -83,9 +102,9 @@ func (r *ctidReader) readTable(ctx context.Context, session *readSession, table
 	// parallelise the work.
 	rangeChan := make(chan pageRange, tableInfo.pageCount)
 	errGroup, ctx := errgroup.WithContext(ctx)
-	for i := uint(0); i < r.tableWorkers; i++ {
+	for i := uint(0); i < s.reader.tableWorkers; i++ {
 		errGroup.Go(func() error {
-			return r.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan)
+			return s.snapshotTableRangeWorker(ctx, table, rangeChan)
 		})
 	}
 
@@ -103,9 +122,9 @@ func (r *ctidReader) readTable(ctx context.Context, session *readSession, table
 	return errGroup.Wait()
 }
 
-func (r *ctidReader) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan pageRange) error {
+func (s *ctidSession) snapshotTableRangeWorker(ctx context.Context, table *table, pageRangeChan <-chan pageRange) error {
 	for pageRange := range pageRangeChan {
-		if err := r.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil {
+		if err := s.snapshotTableRange(ctx, table, pageRange); err != nil {
 			return err
 		}
 	}
@@ -128,10 +147,10 @@ func buildPageRangeQuery(t *table, r pageRange) string {
 	return fmt.Sprintf(pageRangeQuery, strings.Join(quotedColumns, ", "), quotedTable, r.start, r.end)
 }
 
-func (r *ctidReader) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error {
-	return execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error {
-		r.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
+func (s *ctidSession) snapshotTableRange(ctx context.Context, table *table, pageRange pageRange) error {
+	return s.execInTx(ctx, func(tx pglib.Tx) error {
+		s.reader.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
+			"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
 		})
 
 		query := buildPageRangeQuery(table, pageRange)
@@ -146,40 +165,16 @@ func (r *ctidReader) snapshotTableRange(ctx context.Context, snapshotID string,
 		}
 		defer rows.Close()
 
-		// resolve the column metadata (names/types) and timestamp once per page
-		// range, since the field descriptions are identical for every row in the
-		// result set.
-		rowAdapter := r.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions())
-		rowCount := uint(0)
-		for rows.Next() {
-			rowCount++
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-				values, err := rows.Values()
-				if err != nil {
-					return fmt.Errorf("retrieving rows values: %w", err)
-				}
-
-				event := rowAdapter.rowToWalEvent(values)
-				if event == nil {
-					continue
-				}
-
-				if err := r.processor.ProcessWALEvent(ctx, event); err != nil {
-					return fmt.Errorf("processing snapshot row: %w", err)
-				}
-			}
+		rowCount, err := s.reader.sink.emit(ctx, table, rows)
+		if err != nil {
+			return err
 		}
 
-		r.progress.advance(table.schema, int64(rowCount)*table.rowSize)
-
-		r.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
+		s.reader.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
+			"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
 		})
 
-		return rows.Err()
+		return nil
 	})
 }
 
@@ -213,11 +208,13 @@ WHERE
 	// select the max page for the relation instead of using pg_class.relpages, it may not contain an accurate value if
 	// the table is small, the table has active inserts, or the database has not been vacuumed/analyzed recently.
 	maxPageQuery = `SELECT MAX(ctid) FROM ONLY %s;`
+
+	tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = ANY($2) AND c.relkind = 'r';`
 )
 
-func (r *ctidReader) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) {
+func (s *ctidSession) getTableInfo(ctx context.Context, schemaName, tableName string) (*tableInfo, error) {
 	tableInfo := &tableInfo{}
-	err := execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error {
+	err := s.execInTx(ctx, func(tx pglib.Tx) error {
 		// make sure the schema and table names are unquoted since the system
 		// catalogs store unquoted names
 		err := tx.QueryRow(ctx,
@@ -235,10 +232,10 @@ func (r *ctidReader) getTableInfo(ctx context.Context, schemaName, tableName, sn
 		}
 		tableInfo.pageCount = int(ctid.BlockNumber)
 
-		tableInfo.calculateBatchPageSize(r.batchBytes)
+		tableInfo.calculateBatchPageSize(s.reader.batchBytes)
 
-		r.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{
-			"schema": schemaName, "table": tableName, "snapshotID": snapshotID,
+		s.reader.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{
+			"schema": schemaName, "table": tableName, "snapshotID": s.snapshotID,
 		})
 		return nil
 	})
@@ -248,3 +245,31 @@ func (r *ctidReader) getTableInfo(ctx context.Context, schemaName, tableName, sn
 
 	return tableInfo, nil
 }
+
+// totalBytes returns the on disk size of the session's schema tables, as seen
+// by the session's transaction snapshot.
+func (s *ctidSession) totalBytes(ctx context.Context) (int64, error) {
+	schema, tables := s.schemaTables.schema, s.schemaTables.tables
+
+	totalBytes := int64(0)
+	s.reader.logger.Debug("querying total bytes for schema", loglib.Fields{
+		"schema": schema, "tables": tables, "snapshotID": s.snapshotID,
+	})
+
+	// make sure the schema and table names are unquoted since the system
+	// catalogs store unquoted names
+	unquotedTables := make([]string, len(tables))
+	for i, table := range tables {
+		unquotedTables[i] = pglib.UnquoteIdentifier(table)
+	}
+
+	err := s.execInTx(ctx, func(tx pglib.Tx) error {
+		err := tx.QueryRow(ctx, []any{&totalBytes}, tablesBytesQuery, pglib.UnquoteIdentifier(schema), unquotedTables)
+		if err != nil {
+			return fmt.Errorf("retrieving total bytes for schema: %w", err)
+		}
+		return nil
+	})
+
+	return totalBytes, err
+}
diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
index dc12335b..597bb168 100644
--- a/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
+++ b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go
@@ -5,11 +5,16 @@ package postgres
 import (
 	"context"
 
+	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 )
 
+// instrumentedTableReader traces the reads of the reader it wraps. Since the
+// reader hands out the sessions that do the reading, it also wraps every
+// session it produces, so that the table reads are traced as children of the
+// schema they belong to.
 type instrumentedTableReader struct {
 	tracer trace.Tracer
 	reader tableReader
@@ -22,15 +27,37 @@ func newInstrumentedTableReader(reader tableReader, i *otel.Instrumentation) *in
 	}
 }
 
-func (i *instrumentedTableReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error {
-	return i.reader.beginSchema(ctx, st, fn)
+func (i *instrumentedTableReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, readSession) error) (err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.BeginSchema", trace.WithAttributes([]attribute.KeyValue{
+		{Key: "schema", Value: attribute.StringValue(st.schema)},
+	}...))
+	defer otel.CloseSpan(span, err)
+
+	return i.reader.beginSchema(ctx, st, func(ctx context.Context, session readSession) error {
+		return fn(ctx, &instrumentedReadSession{tracer: i.tracer, session: session})
+	})
+}
+
+type instrumentedReadSession struct {
+	tracer  trace.Tracer
+	session readSession
 }
 
-func (i *instrumentedTableReader) readTable(ctx context.Context, session *readSession, table *table) (err error) {
+func (i *instrumentedReadSession) readTable(ctx context.Context, table *table) (err error) {
 	ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.ReadTable", trace.WithAttributes([]attribute.KeyValue{
 		{Key: "schema", Value: attribute.StringValue(table.schema)},
 		{Key: "table", Value: attribute.StringValue(table.name)},
 	}...))
 	defer otel.CloseSpan(span, err)
-	return i.reader.readTable(ctx, session, table)
+	return i.session.readTable(ctx, table)
+}
+
+func (i *instrumentedReadSession) totalBytes(ctx context.Context) (bytes int64, err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.TotalBytes")
+	defer otel.CloseSpan(span, err)
+	return i.session.totalBytes(ctx)
+}
+
+func (i *instrumentedReadSession) logFields() loglib.Fields {
+	return i.session.logFields()
 }
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 01a7ce29..93d9e54a 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"sync"
 
 	pglib "github.com/xataio/pgstream/internal/postgres"
@@ -30,8 +31,8 @@ type SnapshotGenerator struct {
 	// reader encapsulates the strategy used to read a schema's tables (ctid
 	// range scan by default).
 	reader tableReader
-	// instrumentation is captured while applying options and used to wrap the
-	// reader once it has been built.
+	// instrumentation is captured while applying options and used to decorate
+	// the reader once it has been built.
 	instrumentation *otel.Instrumentation
 
 	// workers per snapshot, parallelise the snapshot creation for each schema
@@ -99,19 +100,8 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor.
 		opt(sg)
 	}
 
-	sg.reader = &ctidReader{
-		conn:         sg.conn,
-		logger:       sg.logger,
-		adapter:      newAdapter(pglib.NewMapper(conn), sg.logger),
-		processor:    sg.processor,
-		tableWorkers: cfg.tableWorkers(),
-		batchBytes:   cfg.batchBytes(),
-		progress:     sg.progress,
-	}
-
-	if sg.instrumentation != nil {
-		sg.reader = newInstrumentedTableReader(sg.reader, sg.instrumentation)
-	}
+	sink := newRowSink(pglib.NewMapper(conn), sg.processor, sg.logger, sg.progress)
+	sg.reader = newTableReader(sg.conn, sg.logger, sink, cfg, sg.instrumentation)
 
 	return sg, nil
 }
@@ -199,9 +189,9 @@ func (sg *SnapshotGenerator) Close() error {
 }
 
 func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTables *schemaTables) error {
-	return sg.reader.beginSchema(ctx, schemaTables, func(ctx context.Context, session *readSession) (err error) {
+	return sg.reader.beginSchema(ctx, schemaTables, func(ctx context.Context, session readSession) (err error) {
 		if sg.progress.enabled {
-			if err := sg.addProgressBar(ctx, session.snapshotID, schemaTables); err != nil {
+			if err := sg.addProgressBar(ctx, session, schemaTables); err != nil {
 				return err
 			}
 			defer func() {
@@ -272,13 +262,15 @@ func readableColumns(pinned, live []string) ([]string, bool) {
 	return readable, len(readable) == 0
 }
 
-func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, session *readSession, tableChan <-chan *table, tableErrMap map[string]error) {
+func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, session readSession, tableChan <-chan *table, tableErrMap map[string]error) {
 	defer wg.Done()
+	sessionFields := session.logFields()
 	for t := range tableChan {
-		logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": session.snapshotID}
+		logFields := loglib.Fields{"schema": t.schema, "table": t.name}
+		maps.Copy(logFields, sessionFields)
 		sg.logger.Debug("snapshotting table", logFields)
 
-		if err := sg.reader.readTable(ctx, session, t); err != nil {
+		if err := session.readTable(ctx, t); err != nil {
 			sg.logger.Error(err, "snapshotting table", logFields)
 			// errors will get notified unless the table doesn't exist
 			if !errors.Is(err, pglib.ErrNoRows) {
@@ -326,8 +318,11 @@ func (sg *SnapshotGenerator) collectSchemaErrors(workerSchemaErrs map[string]err
 	return nil
 }
 
-func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID string, schemaTables *schemaTables) error {
-	totalBytes, err := sg.getSnapshotSchemaTotalBytes(ctx, snapshotID, schemaTables.schema, schemaTables.tables)
+// addProgressBar sizes the schema's progress bar with the total bytes the read
+// session reports for it. How those bytes are measured is the reading
+// strategy's business; the generator only owns the bar.
+func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, session readSession, schemaTables *schemaTables) error {
+	totalBytes, err := session.totalBytes(ctx)
 	if err != nil {
 		return err
 	}
@@ -337,32 +332,6 @@ func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID stri
 	return nil
 }
 
-const tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = ANY($2) AND c.relkind = 'r';`
-
-func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, snapshotID, schema string, tables []string) (int64, error) {
-	totalBytes := int64(0)
-	sg.logger.Debug("querying total bytes for schema", loglib.Fields{
-		"schema": schema, "tables": tables, "snapshotID": snapshotID,
-	})
-
-	// make sure the schema and table names are unquoted since the system
-	// catalogs store unquoted names
-	unquotedTables := make([]string, len(tables))
-	for i, table := range tables {
-		unquotedTables[i] = pglib.UnquoteIdentifier(table)
-	}
-
-	err := execInSnapshotTx(ctx, sg.conn, snapshotID, func(tx pglib.Tx) error {
-		err := tx.QueryRow(ctx, []any{&totalBytes}, tablesBytesQuery, pglib.UnquoteIdentifier(schema), unquotedTables)
-		if err != nil {
-			return fmt.Errorf("retrieving total bytes for schema: %w", err)
-		}
-		return nil
-	})
-
-	return totalBytes, err
-}
-
 // calculateBatchPageSize will automatically determine the batch page size based
 // on the average page size and the configured batch bytes limit.
 func (t *tableInfo) calculateBatchPageSize(bytes uint64) {
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 a36268cd..843b37e3 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -302,7 +302,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
 								require.Equal(t, fmt.Sprintf(
 									`SELECT "id" FROM ONLY %s WHERE ctid BETWEEN '(0,0)' AND '(1,0)'`,
-									quotedSchemaTable1), query)
+									quotedSchemaTable1,
+								), query)
 								return &pgmocks.Rows{
 									CloseFn: func() {},
 									NextFn:  func(i uint) bool { return i == 1 },
@@ -1411,11 +1412,9 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				reader: &ctidReader{
 					conn:         tc.querier,
 					logger:       logger,
-					adapter:      newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
-					processor:    processor,
+					sink:         newRowSink(pglib.NewMapper(tc.querier), processor, loglib.NewNoopLogger(), pt),
 					tableWorkers: 1,
 					batchBytes:   1024 * 1024, // 1MB
-					progress:     pt,
 				},
 			}
 
@@ -2028,13 +2027,20 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) {
 				},
 			}
 
+			pt := progressTracker{
+				enabled: tc.name == "ok - with progress tracking",
+				bars:    synclib.NewMap[string, progress.Bar](),
+			}
+			if pt.enabled {
+				pt.set(tc.table.schema, progressBar)
+			}
+
 			reader := &ctidReader{
 				logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{
 					LogLevel: "debug",
 				})),
-				conn:    tc.querier,
-				adapter: newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()),
-				processor: &processormocks.Processor{
+				conn: tc.querier,
+				sink: newRowSink(pglib.NewMapper(tc.querier), &processormocks.Processor{
 					ProcessWALEventFn: func(ctx context.Context, walEvent *wal.Event) error {
 						if tc.processor != nil {
 							if err := tc.processor.ProcessWALEvent(ctx, walEvent); err != nil {
@@ -2044,18 +2050,15 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) {
 						eventChan <- walEvent
 						return nil
 					},
-				},
-				progress: progressTracker{
-					enabled: tc.name == "ok - with progress tracking",
-					bars:    synclib.NewMap[string, progress.Bar](),
-				},
+				}, loglib.NewNoopLogger(), pt),
 			}
-
-			if reader.progress.enabled {
-				reader.progress.set(tc.table.schema, progressBar)
+			session := &ctidSession{
+				reader:       reader,
+				schemaTables: &schemaTables{schema: tc.table.schema, tables: []string{tc.table.name}},
+				snapshotID:   testSnapshotID,
 			}
 
-			err := reader.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange)
+			err := session.snapshotTableRange(context.Background(), tc.table, tc.pageRange)
 			require.Equal(t, tc.wantErr, err)
 			close(eventChan)
 
diff --git a/pkg/snapshot/generator/postgres/data/row_sink.go b/pkg/snapshot/generator/postgres/data/row_sink.go
new file mode 100644
index 00000000..0bd952de
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/data/row_sink.go
@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/wal/processor"
+)
+
+// rowSink turns the rows a reading strategy queries into the wal events the
+// snapshot emits. It owns everything that happens to a row once it has been
+// read: adapting it to a wal event, handing it to the processor and reporting
+// the bytes read. A strategy is therefore only responsible for deciding which
+// rows to read, and every strategy consumes them the same way.
+type rowSink struct {
+	adapter   *adapter
+	processor processor.Processor
+	progress  progressTracker
+}
+
+func newRowSink(mapper mapper, processor processor.Processor, logger loglib.Logger, progress progressTracker) rowSink {
+	return rowSink{
+		adapter:   newAdapter(mapper, logger),
+		processor: processor,
+		progress:  progress,
+	}
+}
+
+// emit converts every row of the result set into a wal event and hands it to
+// the processor, reporting the bytes read against the table's schema once the
+// result set has been fully consumed. It returns the number of rows read.
+func (s rowSink) emit(ctx context.Context, table *table, rows pglib.Rows) (uint, error) {
+	// resolve the column metadata (names/types) and timestamp once per result
+	// set, since the field descriptions are identical for every row in it.
+	rowAdapter := s.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions())
+	rowCount := uint(0)
+	for rows.Next() {
+		rowCount++
+		select {
+		case <-ctx.Done():
+			return rowCount, ctx.Err()
+		default:
+			values, err := rows.Values()
+			if err != nil {
+				return rowCount, fmt.Errorf("retrieving rows values: %w", err)
+			}
+
+			event := rowAdapter.rowToWalEvent(values)
+			if event == nil {
+				continue
+			}
+
+			if err := s.processor.ProcessWALEvent(ctx, event); err != nil {
+				return rowCount, fmt.Errorf("processing snapshot row: %w", err)
+			}
+		}
+	}
+
+	s.progress.advance(table.schema, int64(rowCount)*table.rowSize)
+
+	return rowCount, rows.Err()
+}
diff --git a/pkg/snapshot/generator/postgres/data/table_reader.go b/pkg/snapshot/generator/postgres/data/table_reader.go
index 196289b9..7486c521 100644
--- a/pkg/snapshot/generator/postgres/data/table_reader.go
+++ b/pkg/snapshot/generator/postgres/data/table_reader.go
@@ -2,27 +2,57 @@
 
 package postgres
 
-import "context"
-
-// readSession carries the per-schema state that a tableReader needs to snapshot
-// individual tables. It is only valid for the duration of the beginSchema
-// callback that produced it.
-type readSession struct {
-	// snapshotID is the exported transaction snapshot shared by all of a
-	// schema's table read transactions. It is empty when the schema has no
-	// tables to read, in which case readTable must not be called (an empty id
-	// would make SET TRANSACTION SNAPSHOT fail).
-	snapshotID string
-}
+import (
+	"context"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/otel"
+)
 
 // tableReader abstracts the strategy used to read a schema's tables during a
-// data snapshot.
+// data snapshot. It does not read tables itself: it opens a read session per
+// schema, and the session does the reading. This keeps the per-schema state a
+// strategy needs (a shared transaction snapshot, in the ctid reader's case)
+// inside the strategy that understands it, instead of travelling through the
+// snapshot generator.
 type tableReader interface {
-	// beginSchema prepares the reader to snapshot the given schema tables and
-	// invokes fn with the readSession that must be used for every readTable call
-	// belonging to that schema.
-	// The session is only valid for the duration of fn.
-	beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error
-	// readTable snapshots a single table using the session provided by beginSchema.
-	readTable(ctx context.Context, session *readSession, table *table) error
+	// beginSchema opens a read session for the given schema tables and invokes
+	// fn with it. The session, and any resource backing it, is only valid for
+	// the duration of fn.
+	beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, readSession) error) error
+}
+
+// readSession reads the tables of a single schema. It is created by a
+// tableReader in beginSchema, already bound to the schema it reads and to
+// whatever per-schema state its strategy needs, and it is only valid for the
+// duration of the beginSchema callback that produced it.
+type readSession interface {
+	// readTable snapshots a single table of the session's schema.
+	readTable(ctx context.Context, table *table) error
+	// totalBytes returns the on disk size of the schema tables this session
+	// reads, so the caller can size its progress reporting.
+	totalBytes(ctx context.Context) (int64, error)
+	// logFields returns the session specific fields to attach to the log
+	// entries of the schema being read.
+	logFields() loglib.Fields
+}
+
+// newTableReader builds the strategy used to read the snapshot tables and wraps
+// it with any active decorator. Picking between strategies belongs here, so
+// that the snapshot generator only ever sees a tableReader.
+func newTableReader(conn pglib.Querier, logger loglib.Logger, sink rowSink, cfg *Config, instrumentation *otel.Instrumentation) tableReader {
+	var reader tableReader = &ctidReader{
+		conn:         conn,
+		logger:       logger,
+		sink:         sink,
+		tableWorkers: cfg.tableWorkers(),
+		batchBytes:   cfg.batchBytes(),
+	}
+
+	if instrumentation != nil {
+		reader = newInstrumentedTableReader(reader, instrumentation)
+	}
+
+	return reader
 }