Repository: xataio/pgstream
Pull request #98 (2024-12-04 14:58:42)

Add snapshot generator
This PR adds the logic to generate snapshots. It includes the interface and the postgres implementation. 

The postgres implementation relies on transaction snapshot ids which offer a consistent view of the database to read from, allowing to parallelise the work to snapshot schema tables by splitting the table pages into ranges based on the [ctid](https://www.postgresql.org/docs/current/ddl-system-columns.html#DDL-SYSTEM-COLUMNS-CTID) system column.

The snapshot generator concurrency can be configured, and the produced table rows can be processed by a configurable processor.

Integration and unit tests added.

Relates to #96 


Patch:
diff --git a/pkg/snapshot/generator/postgres/config.go b/pkg/snapshot/generator/postgres/config.go
new file mode 100644
index 0000000..2881627
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/config.go
@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+type Config struct {
+	// Postgres connection URL. Required.
+	URL string
+	// BatchPageSize represents the size of the table page range that will be
+	// processed concurrently by the table workers. Defaults to 1000.
+	BatchPageSize uint
+	// SchemaWorkers represents the number of tables the snapshot generator will
+	// process concurrently per schema. Defaults to 4.
+	SchemaWorkers uint
+	// TableWorkers represents the number of concurrent workers per table. Each
+	// worker will process a different page range in parallel. Defaults to 4.
+	TableWorkers uint
+}
+
+const (
+	defaultBatchPageSize = 1000
+	defaultTableWorkers  = 4
+	defaultSchemaWorkers = 4
+)
+
+func (c *Config) batchPageSize() uint {
+	if c.BatchPageSize > 0 {
+		return c.BatchPageSize
+	}
+	return defaultBatchPageSize
+}
+
+func (c *Config) schemaWorkers() uint {
+	if c.SchemaWorkers > 0 {
+		return c.SchemaWorkers
+	}
+	return defaultSchemaWorkers
+}
+
+func (c *Config) tableWorkers() uint {
+	if c.TableWorkers > 0 {
+		return c.TableWorkers
+	}
+	return defaultTableWorkers
+}
diff --git a/pkg/snapshot/generator/postgres/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/pg_snapshot_generator.go
new file mode 100644
index 0000000..3377eab
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/pg_snapshot_generator.go
@@ -0,0 +1,285 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/lib/pq"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"golang.org/x/sync/errgroup"
+)
+
+type SnapshotGenerator struct {
+	logger loglib.Logger
+	conn   pglib.Querier
+	mapper mapper
+
+	schemaWorkers uint
+	tableWorkers  uint
+	batchPageSize uint
+
+	// Function called for processing produced rows.
+	processRow snapshot.RowProcessor
+}
+
+type mapper interface {
+	TypeForOID(uint32) string
+}
+
+type pageRange struct {
+	start uint
+	end   uint
+}
+
+type Option func(sg *SnapshotGenerator)
+
+func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.RowProcessor, opts ...Option) (*SnapshotGenerator, error) {
+	conn, err := pglib.NewConnPool(ctx, cfg.URL)
+	if err != nil {
+		return nil, err
+	}
+
+	sg := &SnapshotGenerator{
+		logger:        loglib.NewNoopLogger(),
+		mapper:        pglib.NewMapper(),
+		conn:          conn,
+		processRow:    processRow,
+		batchPageSize: cfg.batchPageSize(),
+		tableWorkers:  cfg.tableWorkers(),
+		schemaWorkers: cfg.schemaWorkers(),
+	}
+
+	for _, opt := range opts {
+		opt(sg)
+	}
+
+	return sg, nil
+}
+
+func WithLogger(logger loglib.Logger) Option {
+	return func(sg *SnapshotGenerator) {
+		sg.logger = loglib.NewLogger(logger).WithFields(loglib.Fields{
+			loglib.ModuleField: "postgres_snapshot_generator",
+		})
+	}
+}
+
+func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+		snapshotID, err := sg.exportSnapshot(ctx, tx)
+		if err != nil {
+			return &snapshot.Errors{Snapshot: err}
+		}
+
+		tableChan := make(chan string, len(ss.TableNames))
+		// a map of table errors per worker to avoid race conditions
+		workerTableErrs := make([]map[string]error, sg.schemaWorkers)
+		wg := &sync.WaitGroup{}
+		// start as many go routines as configured concurrent workers per schema
+		for i := uint(0); i < sg.schemaWorkers; i++ {
+			wg.Add(1)
+			workerTableErrs[i] = make(map[string]error, len(ss.TableNames))
+			go sg.createSnapshotWorker(ctx, wg, ss, snapshotID, tableChan, workerTableErrs[i])
+		}
+
+		for _, table := range ss.TableNames {
+			tableChan <- table
+		}
+
+		close(tableChan)
+		wg.Wait()
+
+		return sg.collectSnapshotTableErrors(workerTableErrs)
+	}, snapshotTxOptions())
+}
+
+func (sg *SnapshotGenerator) Close() error {
+	return sg.conn.Close(context.Background())
+}
+
+func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, ss *snapshot.Snapshot, snapshotID string, tableChan <-chan string, tableErrMap map[string]error) {
+	defer wg.Done()
+	for table := range tableChan {
+		logFields := loglib.Fields{"schema": ss.SchemaName, "table": table, "snapshotID": snapshotID}
+		sg.logger.Info("snapshotting table", logFields)
+
+		if err := sg.snapshotTable(ctx, snapshotID, ss.SchemaName, table); 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) {
+				tableErrMap[table] = err
+			}
+		}
+	}
+}
+
+func (sg *SnapshotGenerator) collectSnapshotTableErrors(workerTableErrs []map[string]error) error {
+	var tableErrs []snapshot.TableError
+	for _, worker := range workerTableErrs {
+		for table, err := range worker {
+			tableErrs = append(tableErrs, snapshot.NewTableError(table, err))
+		}
+	}
+
+	if len(tableErrs) > 0 {
+		return &snapshot.Errors{Tables: tableErrs}
+	}
+
+	return nil
+}
+
+func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, schema, table string) error {
+	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+		if err := sg.setTransactionSnapshot(ctx, tx, snapshotID); err != nil {
+			return err
+		}
+
+		tablePageCount, err := sg.getTablePageCount(ctx, tx, schema, table)
+		if err != nil {
+			return err
+		}
+
+		sg.logger.Debug(fmt.Sprintf("table page count: %d", tablePageCount), loglib.Fields{
+			"schema": schema, "table": table, "snapshotID": snapshotID,
+		})
+
+		// 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, tablePageCount)
+		errGroup, ctx := errgroup.WithContext(ctx)
+		for i := uint(0); i < sg.tableWorkers; i++ {
+			errGroup.Go(func() error {
+				return sg.snapshotTableRangeWorker(ctx, snapshotID, schema, 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 <= tablePageCount; start += sg.batchPageSize {
+			rangeChan <- pageRange{
+				start: start,
+				end:   start + sg.batchPageSize,
+			}
+		}
+
+		// wait for all table ranges to complete
+		close(rangeChan)
+		return errGroup.Wait()
+	}, snapshotTxOptions())
+}
+
+func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snapshotID, schema, table string, pageRangeChan <-chan (pageRange)) error {
+	for pageRange := range pageRangeChan {
+		if err := sg.snapshotTableRange(ctx, snapshotID, schema, table, pageRange); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID, schema, table string, pageRange pageRange) error {
+	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+		if err := sg.setTransactionSnapshot(ctx, tx, snapshotID); err != nil {
+			return err
+		}
+		sg.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
+			"schema": schema, "table": table, "snapshotID": snapshotID,
+		})
+
+		query := fmt.Sprintf("SELECT * FROM %s.%s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'",
+			pq.QuoteIdentifier(schema), pq.QuoteIdentifier(table), pageRange.start, pageRange.end)
+		rows, err := tx.Query(ctx, query)
+		if err != nil {
+			return fmt.Errorf("querying table rows: %w", err)
+		}
+		defer rows.Close()
+
+		fieldDescriptions := rows.FieldDescriptions()
+		for rows.Next() {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				values, err := rows.Values()
+				if err != nil {
+					return fmt.Errorf("retrieving rows values: %w", err)
+				}
+
+				columns := sg.toSnapshotColumns(fieldDescriptions, values)
+				if len(columns) == 0 {
+					continue
+				}
+
+				if err := sg.processRow(ctx, &snapshot.Row{
+					Schema:  schema,
+					Table:   table,
+					Columns: columns,
+				}); err != nil {
+					return fmt.Errorf("processing snapshot row: %w", err)
+				}
+			}
+		}
+
+		return rows.Err()
+	}, snapshotTxOptions())
+}
+
+func (sg *SnapshotGenerator) toSnapshotColumns(fieldDescriptions []pgconn.FieldDescription, values []any) []snapshot.Column {
+	columns := make([]snapshot.Column, 0, len(fieldDescriptions))
+	for i, value := range values {
+		dataType := sg.mapper.TypeForOID(fieldDescriptions[i].DataTypeOID)
+		if dataType == "" {
+			sg.logger.Warn(nil, "unknown data type OID", loglib.Fields{"data_type_oid": fieldDescriptions[i].DataTypeOID})
+			continue
+		}
+		columns = append(columns, snapshot.Column{
+			Name:  fieldDescriptions[i].Name,
+			Type:  dataType,
+			Value: value,
+		})
+	}
+
+	return columns
+}
+
+func (sg *SnapshotGenerator) getTablePageCount(ctx context.Context, tx pglib.Tx, schemaName, tableName string) (uint, error) {
+	var pageCount uint
+	query := "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2"
+	if err := tx.QueryRow(ctx, query, tableName, schemaName).Scan(&pageCount); err != nil {
+		return 0, fmt.Errorf("getting page count for table %s.%s: %w", schemaName, tableName, err)
+	}
+
+	return pageCount, nil
+}
+
+func (sg *SnapshotGenerator) exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) {
+	var snapshotID string
+	if err := tx.QueryRow(ctx, "SELECT pg_export_snapshot()").Scan(&snapshotID); 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 snapshotTxOptions() pglib.TxOptions {
+	return pglib.TxOptions{
+		IsolationLevel: pglib.RepeatableRead,
+		AccessMode:     pglib.ReadOnly,
+	}
+}
diff --git a/pkg/snapshot/generator/snapshot_generator.go b/pkg/snapshot/generator/snapshot_generator.go
new file mode 100644
index 0000000..91868df
--- /dev/null
+++ b/pkg/snapshot/generator/snapshot_generator.go
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package generator
+
+import (
+	"context"
+
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+type SnapshotGenerator interface {
+	CreateSnapshot(ctx context.Context, snapshot *snapshot.Snapshot) error
+	Close() error
+}
diff --git a/pkg/snapshot/snapshot.go b/pkg/snapshot/snapshot.go
index 1ce0bb4..1789830 100644
--- a/pkg/snapshot/snapshot.go
+++ b/pkg/snapshot/snapshot.go
@@ -3,6 +3,7 @@
 package snapshot
 
 import (
+	"context"
 	"errors"
 )
 
@@ -17,6 +18,20 @@ type Request struct {
 	Errors   *Errors
 }
 
+type Row struct {
+	Schema  string
+	Table   string
+	Columns []Column
+}
+
+type Column struct {
+	Name  string
+	Type  string
+	Value any
+}
+
+type RowProcessor func(context.Context, *Row) error
+
 type Status string
 
 const (


Test patch:
diff --git a/pkg/snapshot/generator/postgres/helper_test.go b/pkg/snapshot/generator/postgres/helper_test.go
new file mode 100644
index 0000000..04a0826
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/helper_test.go
@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
+	"github.com/testcontainers/testcontainers-go/wait"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+type mockRowProcessor struct {
+	rowChan chan *snapshot.Row
+	once    sync.Once
+}
+
+func (mp *mockRowProcessor) process(ctx context.Context, row *snapshot.Row) error {
+	mp.rowChan <- row
+	return nil
+}
+
+func (mp *mockRowProcessor) close() {
+	mp.once.Do(func() { close(mp.rowChan) })
+}
+
+type cleanup func() error
+
+func setupPostgresContainer(ctx context.Context) (cleanup, string, error) {
+	waitForLogs := wait.
+		ForLog("database system is ready to accept connections").
+		WithOccurrence(2).
+		WithStartupTimeout(5 * time.Second)
+
+	ctr, err := postgres.RunContainer(ctx,
+		testcontainers.WithImage("debezium/postgres:14-alpine"),
+		testcontainers.WithWaitStrategy(waitForLogs),
+	)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to start postgres container: %w", err)
+	}
+
+	pgurl, err := ctr.ConnectionString(ctx, "sslmode=disable")
+	if err != nil {
+		return nil, "", fmt.Errorf("retrieving connection string for postgres container: %w", err)
+	}
+
+	return func() error {
+		return ctr.Terminate(ctx)
+	}, pgurl, nil
+}
+
+func execQuery(t *testing.T, ctx context.Context, pgurl, query string) {
+	conn, err := pglib.NewConn(ctx, pgurl)
+	require.NoError(t, err)
+
+	_, err = conn.Exec(ctx, query)
+	require.NoError(t, err)
+}
diff --git a/pkg/snapshot/generator/postgres/pg_snapshot_generator_integration_test.go b/pkg/snapshot/generator/postgres/pg_snapshot_generator_integration_test.go
new file mode 100644
index 0000000..a0334c5
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/pg_snapshot_generator_integration_test.go
@@ -0,0 +1,75 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+func Test_PostgresSnapshotGenerator(t *testing.T) {
+	if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" {
+		t.Skip("skipping integration test...")
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	pgCleanup, pgurl, err := setupPostgresContainer(ctx)
+	require.NoError(t, err)
+	defer pgCleanup()
+
+	// use a mock row processor to validate the generator produces the right
+	// rows
+	mockProcessor := &mockRowProcessor{
+		rowChan: make(chan *snapshot.Row),
+	}
+
+	generator, err := NewSnapshotGenerator(ctx, &Config{URL: pgurl}, mockProcessor.process)
+	require.NoError(t, err)
+	defer generator.Close()
+
+	// create a table and populate it with data
+	testTable := "snapshot_generator_integration_test"
+	execQuery(t, ctx, pgurl, fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s(id SERIAL PRIMARY KEY, name TEXT)", testTable))
+	execQuery(t, ctx, pgurl, fmt.Sprintf("INSERT INTO %s(name) VALUES('alice')", testTable))
+	execQuery(t, ctx, pgurl, fmt.Sprintf("INSERT INTO %s(name) VALUES('bob')", testTable))
+	execQuery(t, ctx, pgurl, fmt.Sprintf("INSERT INTO %s(name) VALUES('charlie')", testTable))
+
+	go func() {
+		err = generator.CreateSnapshot(ctx, &snapshot.Snapshot{
+			SchemaName: "public",
+			TableNames: []string{testTable},
+		})
+		require.NoError(t, err)
+		mockProcessor.close()
+	}()
+
+	rows := make([]*snapshot.Row, 0, 3)
+	for row := range mockProcessor.rowChan {
+		rows = append(rows, row)
+	}
+
+	wantRows := []*snapshot.Row{
+		newTestRow(testTable, 1, "alice"),
+		newTestRow(testTable, 2, "bob"),
+		newTestRow(testTable, 3, "charlie"),
+	}
+	require.Equal(t, wantRows, rows)
+}
+
+func newTestRow(tableName string, id int32, name string) *snapshot.Row {
+	return &snapshot.Row{
+		Schema: "public",
+		Table:  tableName,
+		Columns: []snapshot.Column{
+			{Name: "id", Type: "int4", Value: id},
+			{Name: "name", Type: "text", Value: name},
+		},
+	}
+}
diff --git a/pkg/snapshot/generator/postgres/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/pg_snapshot_generator_test.go
new file mode 100644
index 0000000..cb036b0
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/pg_snapshot_generator_test.go
@@ -0,0 +1,913 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/internal/log/zerolog"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
+	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/wal"
+)
+
+func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
+	t.Parallel()
+
+	testTable1 := "test-table-1"
+	testTable2 := "test-table-2"
+	testSnapshot := &snapshot.Snapshot{
+		SchemaName: "test-schema",
+		TableNames: []string{testTable1},
+	}
+
+	txOptions := pglib.TxOptions{
+		IsolationLevel: pglib.RepeatableRead,
+		AccessMode:     pglib.ReadOnly,
+	}
+
+	testSnapshotID := "test-snapshot-id"
+	testPageCount := uint(1)
+	testUUID := uuid.New().String()
+	testColumns := []snapshot.Column{
+		{Name: "id", Type: "uuid", Value: testUUID},
+		{Name: "name", Type: "text", Value: "alice"},
+	}
+
+	testRow := func(tableName string, columns []snapshot.Column) *snapshot.Row {
+		return &snapshot.Row{
+			Schema:  testSnapshot.SchemaName,
+			Table:   tableName,
+			Columns: columns,
+		}
+	}
+
+	errTest := errors.New("oh noes")
+
+	tests := []struct {
+		name          string
+		querier       pglib.Querier
+		snapshot      *snapshot.Snapshot
+		schemaWorkers uint
+
+		wantRows []*snapshot.Row
+		wantErr  error
+	}{
+		{
+			name: "ok",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSnapshot.SchemaName, testTable1, 0, 10), query)
+								require.Len(t, args, 0)
+								return &pgmocks.Rows{
+									CloseFn: func() {},
+									NextFn:  func(i uint) bool { return i == 1 },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription {
+										return []pgconn.FieldDescription{
+											{Name: "id", DataTypeOID: pgtype.UUIDOID},
+											{Name: "name", DataTypeOID: pgtype.TextOID},
+										}
+									},
+									ValuesFn: func() ([]any, error) {
+										return []any{testUUID, "alice"}, nil
+									},
+									ErrFn: func() error { return nil },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr:  nil,
+			wantRows: []*snapshot.Row{testRow(testTable1, testColumns)},
+		},
+		{
+			name: "ok - multiple tables and multiple workers",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					mockTx := pgmocks.Tx{
+						QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+							if query == "SELECT pg_export_snapshot()" {
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							}
+							if query == "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2" {
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							}
+							return &pgmocks.Row{
+								ScanFn: func(args ...any) error { return fmt.Errorf("unexpected call to QueryRowFn: %s", query) },
+							}
+						},
+						QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+							return &pgmocks.Rows{
+								CloseFn: func() {},
+								NextFn:  func(i uint) bool { return i == 1 },
+								FieldDescriptionsFn: func() []pgconn.FieldDescription {
+									return []pgconn.FieldDescription{
+										{Name: "id", DataTypeOID: pgtype.UUIDOID},
+										{Name: "name", DataTypeOID: pgtype.TextOID},
+									}
+								},
+								ValuesFn: func() ([]any, error) {
+									return []any{testUUID, "alice"}, nil
+								},
+								ErrFn: func() error { return nil },
+							}, nil
+						},
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+							require.Len(t, args, 0)
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSnapshot.SchemaName,
+				TableNames: []string{testTable1, testTable2},
+			},
+			schemaWorkers: 2,
+
+			wantErr:  nil,
+			wantRows: []*snapshot.Row{testRow(testTable1, testColumns), testRow(testTable2, testColumns)},
+		},
+		{
+			name: "ok - unsupported column type",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSnapshot.SchemaName, testTable1, 0, 10), query)
+								require.Len(t, args, 0)
+								return &pgmocks.Rows{
+									CloseFn: func() {},
+									NextFn:  func(i uint) bool { return i == 1 },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription {
+										return []pgconn.FieldDescription{
+											{Name: "id", DataTypeOID: pgtype.UUIDOID},
+											{Name: "name", DataTypeOID: pgtype.TextOID},
+											{Name: "unsupported", DataTypeOID: 99999},
+										}
+									},
+									ValuesFn: func() ([]any, error) {
+										return []any{testUUID, "alice", 1}, nil
+									},
+									ErrFn: func() error { return nil },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr:  nil,
+			wantRows: []*snapshot.Row{testRow(testTable1, testColumns)},
+		},
+		{
+			name: "ok - no data",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSnapshot.SchemaName, testTable1, 0, 10), query)
+								require.Len(t, args, 0)
+								return &pgmocks.Rows{
+									CloseFn:             func() {},
+									NextFn:              func(i uint) bool { return i == 0 },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription { return []pgconn.FieldDescription{} },
+									ValuesFn:            func() ([]any, error) { return []any{}, nil },
+									ErrFn:               func() error { return nil },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr:  nil,
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - exporting snapshot",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										return errTest
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Snapshot: fmt.Errorf("exporting snapshot: %w", errTest),
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - setting transaction snapshot before table page count",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, errTest
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("setting transaction snapshot: %v", errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - getting table page count",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										return errTest
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("getting page count for table test-schema.test-table-1: %v", errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - setting transaction snapshot for table range",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, errTest
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("setting transaction snapshot: %v", errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - querying range data",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSnapshot.SchemaName, testTable1, 0, 10), query)
+								require.Len(t, args, 0)
+								return nil, errTest
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("querying table rows: %v", errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - getting row values",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								return &pgmocks.Rows{
+									CloseFn:             func() {},
+									NextFn:              func(i uint) bool { return i == 1 },
+									ValuesFn:            func() ([]any, error) { return nil, errTest },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription { return []pgconn.FieldDescription{} },
+									ErrFn:               func() error { return nil },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("retrieving rows values: %v", errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - rows err",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(ctx context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, []any{testTable1, testSnapshot.SchemaName}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										pageCount, ok := args[0].(*uint)
+										require.True(t, ok, fmt.Sprintf("pageCount, expected *uint, got %T", args[0]))
+										*pageCount = testPageCount
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								return &pgmocks.Rows{
+									CloseFn:             func() {},
+									NextFn:              func(i uint) bool { return i == 1 },
+									ValuesFn:            func() ([]any, error) { return []any{}, nil },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription { return []pgconn.FieldDescription{} },
+									ErrFn:               func() error { return errTest },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: errTest.Error(),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+		{
+			name: "error - multiple tables and multiple workers",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					mockTx := pgmocks.Tx{
+						QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+							if query == "SELECT pg_export_snapshot()" {
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							}
+							if query == "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2" {
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										return errTest
+									},
+								}
+							}
+							return &pgmocks.Row{
+								ScanFn: func(args ...any) error { return fmt.Errorf("unexpected call to QueryRowFn: %s", query) },
+							}
+						},
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+							require.Len(t, args, 0)
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSnapshot.SchemaName,
+				TableNames: []string{testTable1, testTable2},
+			},
+			schemaWorkers: 2,
+
+			wantErr: &snapshot.Errors{
+				Tables: []snapshot.TableError{
+					{
+						Table:    testTable1,
+						ErrorMsg: fmt.Sprintf("getting page count for table %s.%s: %v", testSnapshot.SchemaName, testTable1, errTest),
+					},
+					{
+						Table:    testTable2,
+						ErrorMsg: fmt.Sprintf("getting page count for table %s.%s: %v", testSnapshot.SchemaName, testTable2, errTest),
+					},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			rowChan := make(chan *snapshot.Row, 10)
+			sg := SnapshotGenerator{
+				logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{
+					LogLevel: "debug",
+				})),
+				conn:   tc.querier,
+				mapper: pglib.NewMapper(),
+				processRow: func(ctx context.Context, e *snapshot.Row) error {
+					rowChan <- e
+					return nil
+				},
+				schemaWorkers: 1,
+				tableWorkers:  1,
+				batchPageSize: 10,
+			}
+
+			if tc.schemaWorkers != 0 {
+				sg.schemaWorkers = tc.schemaWorkers
+			}
+
+			s := testSnapshot
+			if tc.snapshot != nil {
+				s = tc.snapshot
+			}
+
+			err := sg.CreateSnapshot(context.Background(), s)
+			require.Equal(t, tc.wantErr, sortSnapshotTableErrors(err))
+			close(rowChan)
+
+			rows := []*snapshot.Row{}
+			for row := range rowChan {
+				rows = append(rows, row)
+			}
+			diff := cmp.Diff(rows, tc.wantRows,
+				cmpopts.IgnoreFields(wal.Data{}, "Timestamp"),
+				cmpopts.SortSlices(func(a, b *snapshot.Row) bool { return a.Table < b.Table }))
+			require.Empty(t, diff, fmt.Sprintf("got: \n%v, \nwant \n%v, \ndiff: \n%s", rows, tc.wantRows, diff))
+		})
+	}
+}
+
+func sortSnapshotTableErrors(err error) error {
+	var snapshotErrs *snapshot.Errors
+	if errors.As(err, &snapshotErrs) {
+		sort.Slice(snapshotErrs.Tables, func(i, j int) bool {
+			return snapshotErrs.Tables[i].Table < snapshotErrs.Tables[j].Table
+		})
+		return snapshotErrs
+	}
+	return err
+}