diff --git a/internal/postgres/retrier/pg_querier_retrier.go b/internal/postgres/retrier/pg_querier_retrier.go index 3df6f4563..819927d5d 100644 --- a/internal/postgres/retrier/pg_querier_retrier.go +++ b/internal/postgres/retrier/pg_querier_retrier.go @@ -152,6 +152,13 @@ func (q *Querier) resetConn(ctx context.Context) error { } func (q *Querier) isRetriableError(err error) bool { + return IsRetriableError(err) +} + +// IsRetriableError reports whether retrying an operation that failed with err +// could succeed. Callers retrying at a coarser granularity than a single +// query use it so their rule cannot drift from this one. +func IsRetriableError(err error) bool { mappedErr := postgres.MapError(err) permissionDenied := &postgres.ErrPermissionDenied{} diff --git a/internal/sync/semaphore.go b/internal/sync/semaphore.go index 7aac104f6..3d343347c 100644 --- a/internal/sync/semaphore.go +++ b/internal/sync/semaphore.go @@ -17,3 +17,10 @@ type WeightedSemaphore interface { func NewWeightedSemaphore(size int64) *semaphore.Weighted { return semaphore.NewWeighted(size) } + +// CopyBudgetReserve leaves room for non-copy connections +const CopyBudgetReserve = 5 + +func CopyBudgetSize(maxConnections int32) int64 { + return max(1, int64(maxConnections)-CopyBudgetReserve) +} diff --git a/internal/sync/semaphore_test.go b/internal/sync/semaphore_test.go new file mode 100644 index 000000000..de7758341 --- /dev/null +++ b/internal/sync/semaphore_test.go @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 + +package sync + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCopyBudgetSize(t *testing.T) { + t.Parallel() + + // zero would block every copy forever + require.Equal(t, int64(1), CopyBudgetSize(1)) + require.Equal(t, int64(1), CopyBudgetSize(CopyBudgetReserve)) + require.Equal(t, int64(45), CopyBudgetSize(50)) +} diff --git a/pkg/snapshot/generator/postgres/data/config.go b/pkg/snapshot/generator/postgres/data/config.go index e4c88fd1e..119f3c9c2 100644 --- a/pkg/snapshot/generator/postgres/data/config.go +++ b/pkg/snapshot/generator/postgres/data/config.go @@ -2,6 +2,11 @@ package postgres +import ( + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/pkg/backoff" +) + type Config struct { // Postgres connection URL. Required. URL string @@ -30,6 +35,17 @@ type Config struct { // expect unmarshalled values. This setting is derived from the stream // configuration for postgres targets, not set by users. RawJSONValues bool + // derived from stream config + CopyPassthrough *CopyPassthroughConfig +} + +// the generator writes the target +type CopyPassthroughConfig struct { + TargetURL string + DisableTriggers bool + // 0 defers to the url + MaxConnections uint + RetryPolicy backoff.Config } const ( @@ -40,6 +56,17 @@ const ( defaultMaxConnections = 50 ) +// array_recv ignores user-defined OID mismatch +// needs target schema from source +const copyFormat = " WITH (FORMAT binary)" + +func (c *CopyPassthroughConfig) poolOptions() []pglib.PoolOption { + if c.MaxConnections == 0 { + return nil + } + return []pglib.PoolOption{pglib.WithMaxConnections(int32(c.MaxConnections))} +} + func (c *Config) batchBytes() uint64 { if c.BatchBytes > 0 { return c.BatchBytes diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough.go new file mode 100644 index 000000000..698e8e20a --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough.go @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "sync/atomic" + "time" + + pglib "github.com/xataio/pgstream/internal/postgres" + pglibinstrumentation "github.com/xataio/pgstream/internal/postgres/instrumentation" + pglibretrier "github.com/xataio/pgstream/internal/postgres/retrier" + synclib "github.com/xataio/pgstream/internal/sync" + "github.com/xataio/pgstream/pkg/backoff" + loglib "github.com/xataio/pgstream/pkg/log" + "github.com/xataio/pgstream/pkg/otel" + "golang.org/x/sync/errgroup" +) + +var ( + errMissingCopyPassthroughTarget = errors.New("copy passthrough requires a target postgres url") + errUnexpectedCopiedRows = errors.New("number of rows copied doesn't match the source rows") +) + +// copyPassthroughSnapshotter streams a page range from the source's COPY TO +// STDOUT into the target's COPY FROM STDIN, decoding nothing on the way. It +// writes to the target itself, so it takes on what the bypassed writer did: +// the target connection and its retry policy, trigger suppression, and a +// budget capping concurrent COPYs. +// +// It delegates the rows COPY cannot carry. +type copyPassthroughSnapshotter struct { + cfg *CopyPassthroughConfig + logger loglib.Logger + targetConn pglib.Querier + budget synclib.WeightedSemaphore + fallback rangeSnapshotter + backoffProvider backoff.Provider +} + +func newCopyPassthroughSnapshotter(ctx context.Context, cfg *CopyPassthroughConfig, logger loglib.Logger, + instrumentation *otel.Instrumentation, fallback rangeSnapshotter, +) (*copyPassthroughSnapshotter, error) { + if cfg.TargetURL == "" { + return nil, errMissingCopyPassthroughTarget + } + + poolOpts := cfg.poolOptions() + maxConnections, err := pglib.ConnPoolMaxConnections(cfg.TargetURL, poolOpts...) + if err != nil { + return nil, fmt.Errorf("resolving copy passthrough target connections: %w", err) + } + + // not the retrying querier: it replays the transaction closure, and the + // closure reads a pipe the failed attempt already drained. Retrying is + // done a page range at a time instead, where the pipe is rebuilt. + pool, err := pglib.NewConnPool(ctx, cfg.TargetURL, poolOpts...) + if err != nil { + return nil, fmt.Errorf("connecting to copy passthrough target: %w", err) + } + targetConn := pglib.Querier(pool) + + if instrumentation.IsEnabled() { + targetConn, err = pglibinstrumentation.NewQuerier(targetConn, instrumentation) + if err != nil { + return nil, errors.Join(fmt.Errorf("instrumenting copy passthrough target: %w", err), targetConn.Close(ctx)) + } + } + + if err := targetConn.Ping(ctx); err != nil { + return nil, errors.Join(fmt.Errorf("pinging copy passthrough target: %w", err), targetConn.Close(ctx)) + } + + return ©PassthroughSnapshotter{ + cfg: cfg, + logger: logger, + targetConn: targetConn, + budget: synclib.NewWeightedSemaphore(synclib.CopyBudgetSize(maxConnections)), + fallback: fallback, + backoffProvider: backoff.NewProvider(&cfg.RetryPolicy), + }, nil +} + +const targetGeneratedColumnsQuery = `SELECT a.attname::text +FROM pg_catalog.pg_attribute a + JOIN pg_catalog.pg_class c ON c.oid = a.attrelid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace +WHERE c.relname = $1 AND n.nspname = $2 AND a.attnum > 0 AND NOT a.attisdropped AND a.attgenerated <> '' +ORDER BY a.attnum` + +// the target rejects these, so the target decides +func (s *copyPassthroughSnapshotter) prepareTable(ctx context.Context, table *table) error { + if err := s.fallback.prepareTable(ctx, table); err != nil { + return err + } + + rows, err := s.targetConn.Query(ctx, targetGeneratedColumnsQuery, + pglib.UnquoteIdentifier(table.name), pglib.UnquoteIdentifier(table.schema)) + if err != nil { + return fmt.Errorf("getting target generated columns for %s.%s: %w", table.schema, table.name, err) + } + defer rows.Close() + + var generated []string + for rows.Next() { + var column string + if err := rows.Scan(&column); err != nil { + return fmt.Errorf("scanning target generated column: %w", err) + } + generated = append(generated, column) + } + if err := rows.Err(); err != nil { + return err + } + + table.generatedColumns = generated + return nil +} + +func (s *copyPassthroughSnapshotter) close(ctx context.Context) error { + return errors.Join(s.targetConn.Close(ctx), s.fallback.close(ctx)) +} + +func (s *copyPassthroughSnapshotter) snapshotRange(ctx context.Context, run runInSnapshotTx, table *table, r pageRange) (int64, error) { + if !table.hasCopyableColumns() { + return s.fallback.snapshotRange(ctx, run, table, r) + } + + // held outside the source tx, not inside it + if err := s.budget.Acquire(ctx, 1); err != nil { + return 0, fmt.Errorf("acquiring copy budget: %w", err) + } + defer s.budget.Release(1) + + rowCount, err := s.copyRange(ctx, run, table, r) + if err == nil || s.cfg.RetryPolicy.DisableRetries || !retriableCopyError(err) { + return rowCount, err + } + + // a page range is re-runnable: the target transaction rolled back, and the + // source is read from the same exported snapshot + err = s.backoffProvider(ctx).RetryNotify(func() error { + var retryErr error + rowCount, retryErr = s.copyRange(ctx, run, table, r) + if retryErr != nil && !retriableCopyError(retryErr) { + return fmt.Errorf("%w: %w", retryErr, backoff.ErrPermanent) + } + return retryErr + }, func(err error, d time.Duration) { + s.logger.Warn(err, "retrying copy passthrough page range", loglib.Fields{ + "schema": table.schema, "table": table.name, + "ctid_start": r.start, "ctid_end": r.end, "retry_delay": d.String(), + }) + }) + return rowCount, err +} + +// an integrity assertion is not a transient failure +func retriableCopyError(err error) bool { + return !errors.Is(err, backoff.ErrPermanent) && pglibretrier.IsRetriableError(err) +} + +func (s *copyPassthroughSnapshotter) copyRange(ctx context.Context, run runInSnapshotTx, table *table, r pageRange) (int64, error) { + var rowCount int64 + err := run(ctx, func(tx pglib.Tx) error { + var err error + rowCount, err = s.copyRangeInTx(ctx, tx, table, r) + return err + }) + return rowCount, err +} + +// the pipe bounds memory +func (s *copyPassthroughSnapshotter) copyRangeInTx(ctx context.Context, tx pglib.Tx, table *table, r pageRange) (int64, error) { + copyTable := table.withoutGeneratedColumns() + sourceSQL := buildCopyToSQL(copyTable, r) + targetSQL := buildCopyFromSQL(copyTable) + + s.logger.Trace("copy passthrough", loglib.Fields{ + "schema": table.schema, "table": table.name, + "source_sql": sourceSQL, "target_sql": targetSQL, + }) + + pr, pw := io.Pipe() + eg, egCtx := errgroup.WithContext(ctx) + + var rowsOut atomic.Int64 + eg.Go(func() error { + n, err := tx.CopyToWriter(egCtx, pw, sourceSQL) + if err != nil { + err = wrapPageRangeQueryError(err) + } else { + rowsOut.Store(n) + } + // unblocks a target awaiting rows + pw.CloseWithError(err) + return err + }) + + eg.Go(func() error { + err := s.targetConn.ExecInTx(egCtx, func(targetTx pglib.Tx) error { + if err := s.prepareTargetTx(egCtx, targetTx); err != nil { + return err + } + rowsIn, err := targetTx.CopyFromReader(egCtx, pr, targetSQL) + if err != nil { + return fmt.Errorf("copying rows into %s.%s (ctid %d-%d): %w", + table.schema, table.name, r.start, r.end, err) + } + if out := rowsOut.Load(); rowsIn != out { + return fmt.Errorf("%w: copied (%d), expected (%d): %w", + errUnexpectedCopiedRows, rowsIn, out, backoff.ErrPermanent) + } + return nil + }) + // unblocks a source mid-write + pr.CloseWithError(err) + return err + }) + + if err := eg.Wait(); err != nil { + return 0, err + } + + return rowsOut.Load(), nil +} + +// no COPY carries all-generated rows +func (t *table) hasCopyableColumns() bool { + return len(t.withoutGeneratedColumns().columns) > 0 +} + +// COPY rejects generated columns +func (t *table) withoutGeneratedColumns() *table { + if len(t.generatedColumns) == 0 { + return t + } + + generated := make(map[string]struct{}, len(t.generatedColumns)) + for _, column := range t.generatedColumns { + generated[column] = struct{}{} + } + + columns := make([]string, 0, len(t.columns)) + for _, column := range t.columns { + if _, isGenerated := generated[column]; !isGenerated { + columns = append(columns, column) + } + } + + copied := *t + copied.columns = columns + return &copied +} + +// reuses the decoding path's query +func buildCopyToSQL(t *table, r pageRange) string { + return fmt.Sprintf("COPY (%s) TO STDOUT%s", buildPageRangeQuery(t, r), copyFormat) +} + +// survives target column reordering +func buildCopyFromSQL(t *table) string { + target := pglib.QuoteQualifiedIdentifier(t.schema, t.name) + if len(t.columns) == 0 { + return fmt.Sprintf("COPY %s FROM STDIN%s", target, copyFormat) + } + + quotedColumns := make([]string, len(t.columns)) + for i, column := range t.columns { + quotedColumns[i] = pglib.QuoteRawIdentifier(column) + } + return fmt.Sprintf("COPY %s (%s) FROM STDIN%s", target, strings.Join(quotedColumns, ", "), copyFormat) +} + +// bounds the wait, does not cap the copy +const targetLockTimeout = "SET LOCAL lock_timeout = '30s'" + +// the bypassed writer did this +func (s *copyPassthroughSnapshotter) prepareTargetTx(ctx context.Context, tx pglib.Tx) error { + if _, err := tx.Exec(ctx, targetLockTimeout); err != nil { + return fmt.Errorf("setting lock timeout on postgres target: %w", err) + } + + if !s.cfg.DisableTriggers { + return nil + } + + if _, err := tx.Exec(ctx, "SET LOCAL session_replication_role = replica"); err != nil { + return fmt.Errorf("disabling triggers on postgres target: %w", err) + } + return nil +} diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_integration_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_integration_test.go new file mode 100644 index 000000000..57ff73547 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_integration_test.go @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + "os" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/internal/testcontainers" + "github.com/xataio/pgstream/pkg/snapshot" + "github.com/xataio/pgstream/pkg/wal" +) + +func Test_PostgresSnapshotGenerator_copyPassthrough(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var sourceURL, targetURL string + sourceCleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourceURL, testcontainers.Postgres17) + require.NoError(t, err) + defer sourceCleanup() + targetCleanup, err := testcontainers.SetupPostgresContainer(ctx, &targetURL, testcontainers.Postgres17) + require.NoError(t, err) + defer targetCleanup() + + const testTable = "copy_passthrough_test" + schema := fmt.Sprintf(`CREATE TABLE %s ( + id int primary key, + label text, + amount numeric, + ts timestamptz, + flags bool[], + payload bytea + )`, testTable) + execQuery(t, ctx, sourceURL, schema) + // stands in for the schema snapshot + execQuery(t, ctx, targetURL, schema) + + execQuery(t, ctx, sourceURL, fmt.Sprintf(`INSERT INTO %s VALUES + (1, 'plain', '10.25', '2024-01-02 03:04:05+00', '{t,f}', '\x00ff'), + (2, E'tab\there', '-0.001', '1999-12-31 23:59:59+00', '{}', '\x'), + (3, E'newline\nand\rreturn', '0', 'infinity', NULL, NULL), + (4, E'back\\slash', NULL, '-infinity', '{t}', '\xdeadbeef'), + (5, '', '1e10', '2024-06-01 12:00:00+00', '{f,f}', '\x0a0d09')`, + testTable)) + + // any event means it decoded instead + generator, err := NewSnapshotGenerator(ctx, &Config{ + URL: sourceURL, + CopyPassthrough: &CopyPassthroughConfig{TargetURL: targetURL}, + }, &failingProcessor{t: t}) + require.NoError(t, err) + defer generator.Close() + + require.NoError(t, generator.CreateSnapshot(ctx, &snapshot.Snapshot{ + SchemaTables: map[string][]string{"public": {testTable}}, + })) + + require.Equal(t, + readRowsAsText(t, ctx, sourceURL, testTable), + readRowsAsText(t, ctx, targetURL, testTable)) +} + +// pins what binary format relies on +func Test_PostgresSnapshotGenerator_copyPassthrough_userDefinedTypes(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var sourceURL, targetURL string + sourceCleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourceURL, testcontainers.Postgres17) + require.NoError(t, err) + defer sourceCleanup() + targetCleanup, err := testcontainers.SetupPostgresContainer(ctx, &targetURL, testcontainers.Postgres17) + require.NoError(t, err) + defer targetCleanup() + + const testTable = "copy_passthrough_udt_test" + schema := fmt.Sprintf(` + CREATE TYPE mood AS ENUM ('sad','ok','happy'); + CREATE TYPE addr AS (street text, num int); + CREATE DOMAIN pos AS int; + CREATE DOMAIN moodd AS mood; + CREATE TYPE numrng AS RANGE (subtype = numeric); + CREATE TABLE %s ( + id int primary key, + c_enum mood, + c_enumarr mood[], + c_comp addr, + c_comparr addr[], + c_domint pos, + c_domenum moodd, + c_rng numrng, + c_bltrng numrange, + c_bltmrng nummultirange + )`, testTable) + + execQuery(t, ctx, sourceURL, schema) + // the decoys shift every user-defined OID on the target + execQuery(t, ctx, targetURL, `CREATE TYPE decoy_enum AS ENUM ('a'); + CREATE TYPE decoy_comp AS (x int); + CREATE DOMAIN decoy_dom AS text; + CREATE TYPE decoy_rng AS RANGE (subtype = int4);`) + execQuery(t, ctx, targetURL, schema) + + require.NotEqual(t, + readTypeOID(t, ctx, sourceURL, "mood"), + readTypeOID(t, ctx, targetURL, "mood"), + "decoy types did not shift the target OIDs, so this test proves nothing") + + execQuery(t, ctx, sourceURL, fmt.Sprintf(`INSERT INTO %s VALUES + (1, 'happy', '{sad,ok}', ROW('main st', 2), ARRAY[ROW('main st', 2)::addr], 5, 'ok', '[1,2]', '[1,2]', '{[1,2]}'), + (2, 'sad', '{}', ROW(NULL, NULL), '{}', 0, NULL, 'empty', 'empty', '{}'), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL)`, + testTable)) + + generator, err := NewSnapshotGenerator(ctx, &Config{ + URL: sourceURL, + // no explicit format: this must exercise the binary default + CopyPassthrough: &CopyPassthroughConfig{TargetURL: targetURL}, + }, &failingProcessor{t: t}) + require.NoError(t, err) + defer generator.Close() + + require.NoError(t, generator.CreateSnapshot(ctx, &snapshot.Snapshot{ + SchemaTables: map[string][]string{"public": {testTable}}, + })) + + require.Equal(t, + readRowsAsText(t, ctx, sourceURL, testTable), + readRowsAsText(t, ctx, targetURL, testTable)) +} + +// no COPY carries all-generated rows +func Test_PostgresSnapshotGenerator_copyPassthrough_allGeneratedColumns(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var sourceURL, targetURL string + sourceCleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourceURL, testcontainers.Postgres17) + require.NoError(t, err) + defer sourceCleanup() + targetCleanup, err := testcontainers.SetupPostgresContainer(ctx, &targetURL, testcontainers.Postgres17) + require.NoError(t, err) + defer targetCleanup() + + const testTable = "copy_passthrough_all_generated_test" + schema := fmt.Sprintf("CREATE TABLE %s (only_generated int GENERATED ALWAYS AS (1) STORED)", testTable) + execQuery(t, ctx, sourceURL, schema) + execQuery(t, ctx, targetURL, schema) + execQuery(t, ctx, sourceURL, fmt.Sprintf("INSERT INTO %s DEFAULT VALUES", testTable)) + execQuery(t, ctx, sourceURL, fmt.Sprintf("INSERT INTO %s DEFAULT VALUES", testTable)) + + rows := &countingProcessor{} + generator, err := NewSnapshotGenerator(ctx, &Config{ + URL: sourceURL, + CopyPassthrough: &CopyPassthroughConfig{TargetURL: targetURL}, + }, rows) + require.NoError(t, err) + defer generator.Close() + + require.NoError(t, generator.CreateSnapshot(ctx, &snapshot.Snapshot{ + SchemaTables: map[string][]string{"public": {testTable}}, + })) + + // events prove it decoded instead + require.Equal(t, 2, rows.count()) +} + +type countingProcessor struct { + events atomic.Int64 +} + +func (p *countingProcessor) ProcessWALEvent(context.Context, *wal.Event) error { + p.events.Add(1) + return nil +} + +func (p *countingProcessor) Close() error { return nil } +func (p *countingProcessor) Name() string { return "countingProcessor" } +func (p *countingProcessor) count() int { return int(p.events.Load()) } + +func readTypeOID(t *testing.T, ctx context.Context, pgurl, typeName string) int { + t.Helper() + + conn, err := pglib.NewConn(ctx, pgurl) + require.NoError(t, err) + defer conn.Close(ctx) + + var oid int + require.NoError(t, conn.QueryRow(ctx, []any{&oid}, "SELECT $1::regtype::oid::int", typeName)) + return oid +} + +type failingProcessor struct { + t *testing.T +} + +func (p *failingProcessor) ProcessWALEvent(_ context.Context, event *wal.Event) error { + p.t.Errorf("copy passthrough emitted a wal event, so rows were decoded: %v", event) + return nil +} + +func (p *failingProcessor) Close() error { return nil } +func (p *failingProcessor) Name() string { return "failingProcessor" } + +// compares rows across two instances +func readRowsAsText(t *testing.T, ctx context.Context, pgurl, tableName string) []string { + t.Helper() + + conn, err := pglib.NewConn(ctx, pgurl) + require.NoError(t, err) + defer conn.Close(ctx) + + rows, err := conn.Query(ctx, fmt.Sprintf("SELECT (t.*)::text FROM %s t ORDER BY id", tableName)) + require.NoError(t, err) + defer rows.Close() + + var out []string + for rows.Next() { + var row string + require.NoError(t, rows.Scan(&row)) + out = append(out, row) + } + require.NoError(t, rows.Err()) + require.NotEmpty(t, out, "no rows read from %s", tableName) + return out +} diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_test.go new file mode 100644 index 000000000..25d9330c2 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_copy_passthrough_test.go @@ -0,0 +1,479 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "errors" + "io" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/require" + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/internal/postgres/mocks" + synclib "github.com/xataio/pgstream/internal/sync" + "github.com/xataio/pgstream/pkg/backoff" + loglib "github.com/xataio/pgstream/pkg/log" +) + +func TestBuildCopyToSQL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + table *table + + want string + }{ + { + name: "explicit columns", + table: &table{schema: "public", name: "users", columns: []string{"id", "name"}}, + want: `COPY (SELECT "id", "name" FROM ONLY "public"."users" WHERE ctid BETWEEN '(0,0)' AND '(10,0)') TO STDOUT WITH (FORMAT binary)`, + }, + { + name: "no columns falls back to star, as the decoding path does", + table: &table{schema: "public", name: "users"}, + want: `COPY (SELECT * FROM ONLY "public"."users" WHERE ctid BETWEEN '(0,0)' AND '(10,0)') TO STDOUT WITH (FORMAT binary)`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, buildCopyToSQL(tc.table, pageRange{start: 0, end: 10})) + }) + } +} + +func TestBuildCopyFromSQL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + table *table + + want string + }{ + { + name: "columns are named so target column order is irrelevant", + table: &table{schema: "public", name: "users", columns: []string{"id", "name"}}, + want: `COPY "public"."users" ("id", "name") FROM STDIN WITH (FORMAT binary)`, + }, + { + name: "no columns omits the column list", + table: &table{schema: "public", name: "users"}, + want: `COPY "public"."users" FROM STDIN WITH (FORMAT binary)`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, buildCopyFromSQL(tc.table)) + }) + } +} + +func TestCopyPassthroughSnapshotter_copyRange(t *testing.T) { + t.Parallel() + + errSource := errors.New("source copy failed") + errTarget := errors.New("target copy failed") + testTable := &table{schema: "public", name: "users", columns: []string{"id"}} + + newSourceTx := func(payload string, rows int64, sourceErr error) *mocks.Tx { + return &mocks.Tx{ + CopyToWriterFn: func(_ context.Context, w io.Writer, _ string) (int64, error) { + if sourceErr != nil { + return -1, sourceErr + } + if _, err := io.WriteString(w, payload); err != nil { + return -1, err + } + return rows, nil + }, + } + } + + newTargetConn := func(rows int64, targetErr error, got *string) *mocks.Querier { + return &mocks.Querier{ + ExecInTxFn: func(ctx context.Context, fn func(tx pglib.Tx) error) error { + return fn(&mocks.Tx{ + ExecFn: func(context.Context, uint, string, ...any) (pglib.CommandTag, error) { + return pglib.CommandTag{}, nil + }, + CopyFromReaderFn: func(_ context.Context, r io.Reader, _ string) (int64, error) { + if targetErr != nil { + return -1, targetErr + } + b, err := io.ReadAll(r) + if err != nil { + return -1, err + } + *got = string(b) + return rows, nil + }, + }) + }, + } + } + + tests := []struct { + name string + sourceTx *mocks.Tx + targetRow int64 + targetErr error + + wantRows int64 + wantPayload string + wantErr error + }{ + { + name: "rows stream through unchanged", + sourceTx: newSourceTx("1\n2\n3\n", 3, nil), + targetRow: 3, + wantRows: 3, + wantPayload: "1\n2\n3\n", + }, + { + name: "row count mismatch is reported", + sourceTx: newSourceTx("1\n2\n3\n", 3, nil), + targetRow: 2, + wantErr: errUnexpectedCopiedRows, + }, + { + name: "source failure surfaces", + sourceTx: newSourceTx("", 0, errSource), + targetRow: 0, + wantErr: errSource, + }, + { + name: "target failure surfaces and does not block the source", + sourceTx: newSourceTx("1\n2\n3\n", 3, nil), + targetErr: errTarget, + wantErr: errTarget, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var got string + s := newTestCopyPassthrough(newTargetConn(tc.targetRow, tc.targetErr, &got), nil) + + rows, err := s.copyRangeInTx(t.Context(), tc.sourceTx, testTable, pageRange{start: 0, end: 10}) + if tc.wantErr != nil { + require.ErrorIs(t, err, tc.wantErr) + return + } + + require.NoError(t, err) + require.Equal(t, tc.wantRows, rows) + require.Equal(t, tc.wantPayload, got) + }) + } +} + +func TestCopyPassthroughSnapshotter_prepareTargetTx(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + disableTriggers bool + + wantQueries []string + }{ + { + name: "triggers left alone by default", + disableTriggers: false, + wantQueries: []string{targetLockTimeout}, + }, + { + name: "triggers suppressed for the copy", + disableTriggers: true, + wantQueries: []string{ + targetLockTimeout, + "SET LOCAL session_replication_role = replica", + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var queries []string + tx := &mocks.Tx{ + ExecFn: func(_ context.Context, _ uint, query string, _ ...any) (pglib.CommandTag, error) { + queries = append(queries, query) + return pglib.CommandTag{}, nil + }, + } + + s := ©PassthroughSnapshotter{cfg: &CopyPassthroughConfig{DisableTriggers: tc.disableTriggers}} + require.NoError(t, s.prepareTargetTx(t.Context(), tx)) + require.Equal(t, tc.wantQueries, queries) + }) + } +} + +func TestTable_withoutGeneratedColumns(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + table *table + + wantColumns []string + }{ + { + name: "no generated columns leaves the list alone", + table: &table{columns: []string{"id", "name"}}, + wantColumns: []string{"id", "name"}, + }, + { + name: "generated columns are dropped", + table: &table{ + columns: []string{"id", "name", "username", "slug"}, + generatedColumns: []string{"username", "slug"}, + }, + wantColumns: []string{"id", "name"}, + }, + { + name: "a generated column that is not selected is ignored", + table: &table{ + columns: []string{"id"}, + generatedColumns: []string{"username"}, + }, + wantColumns: []string{"id"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + original := slices.Clone(tc.table.columns) + got := tc.table.withoutGeneratedColumns() + + require.Equal(t, tc.wantColumns, got.columns) + require.Equal(t, original, tc.table.columns, "the original table must not be modified") + }) + } +} + +func TestBuildCopySQL_excludesGeneratedColumns(t *testing.T) { + t.Parallel() + + testTable := (&table{ + schema: "public", + name: "users", + columns: []string{"id", "name", "username"}, + generatedColumns: []string{"username"}, + }).withoutGeneratedColumns() + + require.Equal(t, + `COPY (SELECT "id", "name" FROM ONLY "public"."users" WHERE ctid BETWEEN '(0,0)' AND '(10,0)') TO STDOUT WITH (FORMAT binary)`, + buildCopyToSQL(testTable, pageRange{start: 0, end: 10})) + require.Equal(t, + `COPY "public"."users" ("id", "name") FROM STDIN WITH (FORMAT binary)`, + buildCopyFromSQL(testTable)) +} + +func TestTable_hasCopyableColumns(t *testing.T) { + t.Parallel() + + require.True(t, (&table{columns: []string{"id"}}).hasCopyableColumns()) + require.True(t, (&table{ + columns: []string{"id", "slug"}, + generatedColumns: []string{"slug"}, + }).hasCopyableColumns()) + + // no COPY carries all-generated rows + require.False(t, (&table{ + columns: []string{"slug"}, + generatedColumns: []string{"slug"}, + }).hasCopyableColumns()) + // a star select the target cannot receive + require.False(t, (&table{generatedColumns: []string{"slug"}}).hasCopyableColumns()) +} + +func TestCopyPassthroughSnapshotter_snapshotRange_fallback(t *testing.T) { + t.Parallel() + + copyable := &table{schema: "public", name: "users", columns: []string{"id"}} + allGenerated := &table{ + schema: "public", + name: "users", + columns: []string{"slug"}, + generatedColumns: []string{"slug"}, + } + + tests := []struct { + name string + table *table + + wantCopied bool + wantFellBack bool + }{ + { + name: "copies the rows COPY can carry", + table: copyable, + wantCopied: true, + }, + { + name: "delegates a table COPY cannot carry", + table: allGenerated, + wantFellBack: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var copied bool + sourceTx := &mocks.Tx{ + CopyToWriterFn: func(_ context.Context, w io.Writer, _ string) (int64, error) { + copied = true + _, err := io.WriteString(w, "1\n") + return 1, err + }, + } + targetConn := &mocks.Querier{ + ExecInTxFn: func(ctx context.Context, fn func(tx pglib.Tx) error) error { + return fn(&mocks.Tx{ + ExecFn: func(context.Context, uint, string, ...any) (pglib.CommandTag, error) { + return pglib.CommandTag{}, nil + }, + CopyFromReaderFn: func(_ context.Context, r io.Reader, _ string) (int64, error) { + if _, err := io.ReadAll(r); err != nil { + return -1, err + } + return 1, nil + }, + }) + }, + } + + fallback := &stubSnapshotter{} + s := newTestCopyPassthrough(targetConn, fallback) + + run := func(ctx context.Context, fn func(tx pglib.Tx) error) error { return fn(sourceTx) } + _, err := s.snapshotRange(t.Context(), run, tc.table, pageRange{start: 0, end: 1}) + require.NoError(t, err) + + require.Equal(t, tc.wantCopied, copied) + require.Equal(t, tc.wantFellBack, fallback.called) + }) + } +} + +func newTestCopyPassthrough(targetConn pglib.Querier, fallback rangeSnapshotter) *copyPassthroughSnapshotter { + return ©PassthroughSnapshotter{ + cfg: &CopyPassthroughConfig{}, + logger: loglib.NewNoopLogger(), + targetConn: targetConn, + budget: synclib.NewWeightedSemaphore(1), + fallback: fallback, + } +} + +type stubSnapshotter struct{ called bool } + +func (s *stubSnapshotter) prepareTable(context.Context, *table) error { return nil } +func (s *stubSnapshotter) close(context.Context) error { return nil } + +func (s *stubSnapshotter) snapshotRange(context.Context, runInSnapshotTx, *table, pageRange) (int64, error) { + s.called = true + return 0, nil +} + +// a retry must re-read the source, not resume a drained pipe +func TestCopyPassthroughSnapshotter_snapshotRange_retries(t *testing.T) { + t.Parallel() + + testTable := &table{schema: "public", name: "users", columns: []string{"id"}} + errRetriable := errors.New("connection reset by peer") + + tests := []struct { + name string + targetErr func(attempt int) error + targetRow func(attempt int) int64 + + wantAttempts int + wantErr error + }{ + { + name: "a retriable failure copies the range again", + targetErr: func(attempt int) error { return map[bool]error{true: errRetriable, false: nil}[attempt == 1] }, + targetRow: func(int) int64 { return 3 }, + wantAttempts: 2, + }, + { + name: "a row count mismatch is not retried", + targetErr: func(int) error { return nil }, + targetRow: func(int) int64 { return 2 }, + wantAttempts: 1, + wantErr: errUnexpectedCopiedRows, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var attempts int + var payloads []string + run := func(ctx context.Context, fn func(tx pglib.Tx) error) error { + attempts++ + return fn(&mocks.Tx{ + CopyToWriterFn: func(_ context.Context, w io.Writer, _ string) (int64, error) { + _, err := io.WriteString(w, "1\n2\n3\n") + return 3, err + }, + }) + } + + s := newTestCopyPassthrough(&mocks.Querier{ + ExecInTxFn: func(ctx context.Context, fn func(tx pglib.Tx) error) error { + return fn(&mocks.Tx{ + ExecFn: func(context.Context, uint, string, ...any) (pglib.CommandTag, error) { + return pglib.CommandTag{}, nil + }, + CopyFromReaderFn: func(_ context.Context, r io.Reader, _ string) (int64, error) { + b, readErr := io.ReadAll(r) + if readErr != nil { + return -1, readErr + } + payloads = append(payloads, string(b)) + if err := tc.targetErr(attempts); err != nil { + return -1, err + } + return tc.targetRow(attempts), nil + }, + }) + }, + }, nil) + s.backoffProvider = backoff.NewProvider(&backoff.Config{ + Constant: &backoff.ConstantConfig{Interval: time.Millisecond, MaxRetries: 3}, + }) + + _, err := s.snapshotRange(t.Context(), run, testTable, pageRange{start: 0, end: 10}) + if tc.wantErr != nil { + require.ErrorIs(t, err, tc.wantErr) + } else { + require.NoError(t, err) + } + + require.Equal(t, tc.wantAttempts, attempts) + // every attempt read the whole range, never a partial pipe + for _, payload := range payloads { + require.Equal(t, "1\n2\n3\n", payload) + } + }) + } +} diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go index 36a6d23e6..577726dcd 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go @@ -42,6 +42,10 @@ type SnapshotGenerator struct { // Function called for processing produced rows. processor processor.Processor tableSnapshotGenerator snapshotTableFn + // decides what becomes of the rows of a page range + rangeSnapshotter rangeSnapshotter + + instrumentation *otel.Instrumentation progressTracking bool progressBars *synclib.Map[string, progress.Bar] @@ -78,6 +82,8 @@ type table struct { rowSize int64 // one list per page range columns []string + // the target computes these itself + generatedColumns []string } type snapshotTableFn func(ctx context.Context, snapshotID string, table *table) error @@ -94,6 +100,31 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor. return nil, err } + sg, err := newSnapshotGenerator(cfg, conn, processor, func(sg *SnapshotGenerator) (rangeSnapshotter, error) { + decoding := newDecodingSnapshotter(sg.adapter, sg.processor) + if cfg.CopyPassthrough == nil { + return decoding, nil + } + return newCopyPassthroughSnapshotter(ctx, cfg.CopyPassthrough, sg.logger, sg.instrumentation, decoding) + }, opts...) + if err != nil { + return nil, errors.Join(err, conn.Close(ctx)) + } + + return sg, nil +} + +// buildSnapshotter is called once the generator holds everything a page range +// reader is built from: the source connection behind the adapter, and the +// logger and instrumentation the options carry. +type buildSnapshotter func(*SnapshotGenerator) (rangeSnapshotter, error) + +// newSnapshotGenerator is the only way to build a generator, so the page range +// reader cannot be left unset. Everything a reader varies on is resolved +// before build is called. +func newSnapshotGenerator(cfg *Config, conn pglib.Querier, processor processor.Processor, + build buildSnapshotter, opts ...Option, +) (*SnapshotGenerator, error) { sg := &SnapshotGenerator{ logger: loglib.NewNoopLogger(), conn: conn, @@ -112,6 +143,12 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor. sg.adapter = newAdapter(pglib.NewMapper(conn), sg.logger) + snapshotter, err := build(sg) + if err != nil { + return nil, err + } + sg.rangeSnapshotter = snapshotter + return sg, nil } @@ -126,6 +163,7 @@ func WithLogger(logger loglib.Logger) Option { func WithInstrumentation(i *otel.Instrumentation) Option { return func(sg *SnapshotGenerator) { var err error + sg.instrumentation = i sg.conn, err = pglibinstrumentation.NewQuerier(sg.conn, i) if err != nil { // this should never happen @@ -196,7 +234,8 @@ func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sn } func (sg *SnapshotGenerator) Close() error { - return sg.conn.Close(context.Background()) + ctx := context.Background() + return errors.Join(sg.conn.Close(ctx), sg.rangeSnapshotter.close(ctx)) } func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTables *schemaTables) error { @@ -354,6 +393,9 @@ func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID strin pglib.UnquoteIdentifier(table.schema), pglib.UnquoteIdentifier(table.name)) } table.columns = columns + if err := sg.rangeSnapshotter.prepareTable(ctx, table); err != nil { + return err + } // If one page range fails, we abort the entire table snapshot. The // snapshot relies on the transaction snapshot id to ensure all workers @@ -407,63 +449,40 @@ func buildPageRangeQuery(t *table, r pageRange) string { } 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, - }) + 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 := buildPageRangeQuery(table, pageRange) - rows, err := tx.Query(ctx, query) - if err != nil { - // something this query names vanished - var relationErr *pglib.ErrRelationDoesNotExist - if errors.As(err, &relationErr) { - return fmt.Errorf("%w: querying table rows: %w", ErrSchemaChangedDuringSnapshot, err) - } - 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) - } + run := func(ctx context.Context, fn func(tx pglib.Tx) error) error { + return sg.execInSnapshotTx(ctx, snapshotID, fn) + } - event := rowAdapter.rowToWalEvent(values) - if event == nil { - continue - } + rowCount, err := sg.rangeSnapshotter.snapshotRange(ctx, run, table, pageRange) + if err != nil { + return err + } - 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(rowCount * table.rowSize) } + } - 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, + }) - sg.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{ - "schema": table.schema, "table": table.name, "snapshotID": snapshotID, - }) + return nil +} - return rows.Err() - }) +// a vanished relation means schema drift +func wrapPageRangeQueryError(err error) error { + var relationErr *pglib.ErrRelationDoesNotExist + if errors.As(err, &relationErr) { + return fmt.Errorf("%w: querying table rows: %w", ErrSchemaChangedDuringSnapshot, err) + } + return fmt.Errorf("querying table rows: %w", err) } func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID string, schemaTables *schemaTables) 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 a4b624339..b471eba00 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go @@ -20,8 +20,6 @@ import ( pgmocks "github.com/xataio/pgstream/internal/postgres/mocks" "github.com/xataio/pgstream/internal/progress" progressmocks "github.com/xataio/pgstream/internal/progress/mocks" - 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" "github.com/xataio/pgstream/pkg/wal/processor" @@ -1382,32 +1380,25 @@ 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 - }, + sg, err := newTestSnapshotGenerator(&Config{ + SnapshotWorkers: 1, + SchemaWorkers: 1, + TableWorkers: 1, + BatchBytes: 1024 * 1024, // 1MB + }, tc.querier, &processormocks.Processor{ + ProcessWALEventFn: func(ctx context.Context, e *wal.Event) error { + eventChan <- e + return nil }, - schemaWorkers: 1, - tableWorkers: 1, - batchBytes: 1024 * 1024, // 1MB - snapshotWorkers: 1, - progressTracking: tc.progressBar != nil, - progressBars: synclib.NewMap[string, progress.Bar](), - progressBarBuilder: func(totalBytes int64, description string) progress.Bar { - return tc.progressBar + CloseFn: func() error { + return tc.processorCloseErr }, + }) + require.NoError(t, err) + sg.progressTracking = tc.progressBar != nil + sg.progressBarBuilder = func(totalBytes int64, description string) progress.Bar { + return tc.progressBar } - sg.tableSnapshotGenerator = sg.snapshotTable if tc.schemaWorkers != 0 { sg.schemaWorkers = tc.schemaWorkers @@ -1418,7 +1409,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) { s = tc.snapshot } - err := sg.CreateSnapshot(context.Background(), s) + err = sg.CreateSnapshot(context.Background(), s) require.Equal(t, tc.wantErr, err) close(eventChan) @@ -2018,32 +2009,25 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { }, } - 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, walEvent *wal.Event) error { - if tc.processor != nil { - if err := tc.processor.ProcessWALEvent(ctx, walEvent); err != nil { - return err - } + sg, err := newTestSnapshotGenerator(&Config{}, 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 { + return err } - eventChan <- walEvent - return nil - }, + } + eventChan <- walEvent + return nil }, - progressTracking: tc.name == "ok - with progress tracking", - progressBars: synclib.NewMap[string, progress.Bar](), - } + }) + require.NoError(t, err) + sg.progressTracking = tc.name == "ok - with progress tracking" if sg.progressTracking { sg.progressBars.Set(tc.table.schema, progressBar) } - err := sg.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange) + err = sg.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange) require.Equal(t, tc.wantErr, err) close(eventChan) @@ -2059,3 +2043,12 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { }) } } + +// newTestSnapshotGenerator goes through the same constructor production does, +// so the page range reader is always set. +func newTestSnapshotGenerator(cfg *Config, conn pglib.Querier, p processor.Processor) (*SnapshotGenerator, error) { + return newSnapshotGenerator(cfg, conn, p, func(sg *SnapshotGenerator) (rangeSnapshotter, error) { + return newDecodingSnapshotter(sg.adapter, sg.processor), nil + }, WithLogger(zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{LogLevel: "debug"}))), + WithProgressTracking()) +} diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_range_snapshotter.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_range_snapshotter.go new file mode 100644 index 000000000..2578a2023 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_range_snapshotter.go @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/pkg/wal/processor" +) + +// runInSnapshotTx runs fn against the source's exported snapshot. +type runInSnapshotTx func(ctx context.Context, fn func(tx pglib.Tx) error) error + +// rangeSnapshotter moves the rows of one page range. The generator owns the +// source connection and decides which ranges exist; a snapshotter decides what +// becomes of the rows. +// +// snapshotRange is handed the transaction runner rather than a transaction, so +// a snapshotter that has to wait for something can wait before the transaction +// opens instead of holding one idle while it does. +type rangeSnapshotter interface { + prepareTable(ctx context.Context, table *table) error + snapshotRange(ctx context.Context, run runInSnapshotTx, table *table, r pageRange) (int64, error) + close(ctx context.Context) error +} + +// decodingSnapshotter reads the rows, decodes them into Go values and emits +// them to the processor as wal events. +type decodingSnapshotter struct { + adapter *adapter + processor processor.Processor +} + +func newDecodingSnapshotter(adapter *adapter, p processor.Processor) *decodingSnapshotter { + return &decodingSnapshotter{adapter: adapter, processor: p} +} + +func (s *decodingSnapshotter) prepareTable(context.Context, *table) error { return nil } + +func (s *decodingSnapshotter) close(context.Context) error { return nil } + +func (s *decodingSnapshotter) snapshotRange(ctx context.Context, run runInSnapshotTx, table *table, r pageRange) (int64, error) { + var rowCount int64 + err := run(ctx, func(tx pglib.Tx) error { + var err error + rowCount, err = s.readRange(ctx, tx, table, r) + return err + }) + return rowCount, err +} + +func (s *decodingSnapshotter) readRange(ctx context.Context, tx pglib.Tx, table *table, r pageRange) (int64, error) { + rows, err := tx.Query(ctx, buildPageRangeQuery(table, r)) + if err != nil { + return 0, wrapPageRangeQueryError(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 := s.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions()) + rowCount := int64(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) + } + } + } + + return rowCount, rows.Err() +} diff --git a/pkg/stream/config.go b/pkg/stream/config.go index cfbe55e49..64afd548e 100644 --- a/pkg/stream/config.go +++ b/pkg/stream/config.go @@ -12,6 +12,7 @@ import ( pglib "github.com/xataio/pgstream/internal/postgres" "github.com/xataio/pgstream/pkg/backoff" "github.com/xataio/pgstream/pkg/kafka" + pgsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/data" kafkacheckpoint "github.com/xataio/pgstream/pkg/wal/checkpointer/kafka" snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder" "github.com/xataio/pgstream/pkg/wal/processor/filter" @@ -248,6 +249,33 @@ func (c *Config) restoreConflictTargetsBeforeData() bool { return !bw.BulkIngestEnabled && strings.EqualFold(bw.OnConflictAction, "update") } +// the chain is the authority on what reads rows +func (c *Config) snapshotCopyPassthroughEligible(chain *processorChain) bool { + if c.Processor.Postgres == nil || !c.Processor.Postgres.BatchWriter.BulkIngestEnabled { + return false + } + if c.Listener.Postgres == nil || c.Listener.Postgres.Snapshot == nil || c.Listener.Postgres.Snapshot.Data == nil { + return false + } + return !chain.hasRowVisibleLayers() +} + +// takes the bypassed writer's settings +func (c *Config) applySnapshotCopyPassthrough(chain *processorChain) bool { + if !c.snapshotCopyPassthroughEligible(chain) { + return false + } + + target := c.Processor.Postgres.BatchWriter + c.Listener.Postgres.Snapshot.Data.CopyPassthrough = &pgsnapshotgenerator.CopyPassthroughConfig{ + TargetURL: target.URL, + DisableTriggers: target.DisableTriggers, + MaxConnections: target.MaxConnections, + RetryPolicy: target.EffectiveRetryPolicy(), + } + return true +} + // applySnapshotRawJSONValues enables raw (text) decoding of json/jsonb values // on the snapshot data generator when the target is postgres. The default pgx // decoding unmarshals json/jsonb into Go values, turning the JSON null value diff --git a/pkg/stream/snapshot_copy_passthrough_test.go b/pkg/stream/snapshot_copy_passthrough_test.go new file mode 100644 index 000000000..4a60441d5 --- /dev/null +++ b/pkg/stream/snapshot_copy_passthrough_test.go @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 + +package stream + +import ( + "testing" + + "github.com/stretchr/testify/require" + loglib "github.com/xataio/pgstream/pkg/log" + pgsnapshotgenerator "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/data" + "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder" + "github.com/xataio/pgstream/pkg/wal/processor/filter" + "github.com/xataio/pgstream/pkg/wal/processor/mocks" + pgwriter "github.com/xataio/pgstream/pkg/wal/processor/postgres" + "github.com/xataio/pgstream/pkg/wal/processor/transformer" +) + +// every row visible layer must block it +func TestConfig_snapshotCopyPassthroughEligible_blockedByEveryModifier(t *testing.T) { + t.Parallel() + + // injector is left out: injector.New dials the source + modifiers := map[string]func(*ProcessorConfig){ + "transformer": func(c *ProcessorConfig) { c.Transformer = &transformer.Config{} }, + "filter": func(c *ProcessorConfig) { c.Filter = &filter.Config{IncludeTables: []string{"*"}} }, + "sanitizer": func(c *ProcessorConfig) { c.Sanitize = &SanitizeConfig{StripNullCharBytes: true} }, + } + + for name, enable := range modifiers { + t.Run(name, func(t *testing.T) { + t.Parallel() + + config := newPassthroughEligibleConfig() + enable(&config.Processor) + + require.False(t, config.snapshotCopyPassthroughEligible(newTestChain(t, config)), + "%s must block the copy passthrough", name) + }) + } +} + +func TestConfig_snapshotCopyPassthroughEligible(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + build func(*Config) + + want bool + }{ + { + name: "postgres target with bulk ingest and no modifiers", + build: func(*Config) {}, + want: true, + }, + { + name: "non postgres target", + build: func(c *Config) { c.Processor.Postgres = nil }, + }, + { + name: "bulk ingest disabled", + build: func(c *Config) { c.Processor.Postgres.BatchWriter.BulkIngestEnabled = false }, + }, + { + name: "no data snapshot configured", + build: func(c *Config) { c.Listener.Postgres.Snapshot.Data = nil }, + }, + { + name: "a sanitizer that strips nothing is never wrapped", + build: func(c *Config) { c.Processor.Sanitize = &SanitizeConfig{} }, + want: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + config := newPassthroughEligibleConfig() + tc.build(config) + + require.Equal(t, tc.want, config.snapshotCopyPassthroughEligible(newTestChain(t, config))) + }) + } +} + +func TestConfig_applySnapshotCopyPassthrough(t *testing.T) { + t.Parallel() + + t.Run("carries the target settings the writer would have applied", func(t *testing.T) { + t.Parallel() + + config := newPassthroughEligibleConfig() + config.Processor.Postgres.BatchWriter.DisableTriggers = true + config.Processor.Postgres.BatchWriter.MaxConnections = 25 + + require.True(t, config.applySnapshotCopyPassthrough(newTestChain(t, config))) + + require.Equal(t, &pgsnapshotgenerator.CopyPassthroughConfig{ + TargetURL: "postgresql://target", + DisableTriggers: true, + MaxConnections: 25, + RetryPolicy: config.Processor.Postgres.BatchWriter.EffectiveRetryPolicy(), + }, config.Listener.Postgres.Snapshot.Data.CopyPassthrough) + }) + + t.Run("leaves the data snapshot alone when not eligible", func(t *testing.T) { + t.Parallel() + + config := newPassthroughEligibleConfig() + config.Processor.Transformer = &transformer.Config{} + + require.False(t, config.applySnapshotCopyPassthrough(newTestChain(t, config))) + require.Nil(t, config.Listener.Postgres.Snapshot.Data.CopyPassthrough) + }) +} + +// assembled the way the pipeline assembles it +func newTestChain(t *testing.T, config *Config) *processorChain { + t.Helper() + + chain, closer, err := addProcessorModifiers(t.Context(), config, loglib.NewNoopLogger(), + &mocks.Processor{}, nil) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, closer()) }) + return chain +} + +func newPassthroughEligibleConfig() *Config { + return &Config{ + Listener: ListenerConfig{ + Postgres: &PostgresListenerConfig{ + // empty: a source url makes the transformer layer dial it + Snapshot: &builder.SnapshotListenerConfig{ + Data: &pgsnapshotgenerator.Config{URL: "postgresql://source"}, + }, + }, + }, + Processor: ProcessorConfig{ + Postgres: &PostgresProcessorConfig{ + BatchWriter: pgwriter.Config{ + URL: "postgresql://target", + BulkIngestEnabled: true, + }, + }, + }, + } +} diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go index ff35964a8..e1f80f516 100644 --- a/pkg/stream/stream_run.go +++ b/pkg/stream/stream_run.go @@ -146,6 +146,7 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, init bool, i } config.applySnapshotRawJSONValues() + config.applySnapshotCopyPassthrough(snapshotChain) snapshotGenerator, err := snapshotbuilder.NewSnapshotGenerator( ctx, config.Listener.Postgres.Snapshot, diff --git a/pkg/stream/stream_snapshot.go b/pkg/stream/stream_snapshot.go index ef20e9a99..5ff681a26 100644 --- a/pkg/stream/stream_snapshot.go +++ b/pkg/stream/stream_snapshot.go @@ -44,6 +44,7 @@ func Snapshot(ctx context.Context, logger loglib.Logger, config *Config, instrum // Listener config.applySnapshotRawJSONValues() + config.applySnapshotCopyPassthrough(chain) snapshotGenerator, err := snapshotbuilder.NewSnapshotGenerator( ctx, config.Listener.Postgres.Snapshot, diff --git a/pkg/wal/processor/postgres/config.go b/pkg/wal/processor/postgres/config.go index bf37b92d6..dd4ab0375 100644 --- a/pkg/wal/processor/postgres/config.go +++ b/pkg/wal/processor/postgres/config.go @@ -50,6 +50,9 @@ func (c *Config) retryPolicy() backoff.Config { } } +// EffectiveRetryPolicy returns the retry policy once the default is applied. +func (c *Config) EffectiveRetryPolicy() backoff.Config { return c.retryPolicy() } + func (c *Config) poolOptions() []pglib.PoolOption { if c.MaxConnections == 0 { return nil diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go index a214ba81f..886a9a607 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go +++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go @@ -27,15 +27,12 @@ type BulkIngestWriter struct { // copyBudget caps the total number of concurrent COPYs across all tables // (and all their send drainers) so they never exhaust the target // connection pool. It is sized from the resolved pool max-connections - // value, minus copyBudgetReserve. + // value, minus synclib.CopyBudgetReserve. copyBudget synclib.WeightedSemaphore } const bulkIngestWriter = "postgres_bulk_ingest_writer" -// batch writer and retrier reset share this pool -const copyBudgetReserve = 5 - var errUnexpectedCopiedRows = errors.New("number of rows copied doesn't match the source rows") // NewBulkIngestWriter returns a postgres processor that batches and writes data @@ -54,7 +51,7 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti biw := &BulkIngestWriter{ Writer: w, batchSenderMap: synclib.NewMap[string, queryBatchSender](), - copyBudget: synclib.NewWeightedSemaphore(copyBudgetSize(w.maxConnections)), + copyBudget: synclib.NewWeightedSemaphore(synclib.CopyBudgetSize(w.maxConnections)), } biw.batchSenderBuilder = func(ctx context.Context, schema, table string) (queryBatchSender, error) { @@ -68,10 +65,6 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti return biw, nil } -func copyBudgetSize(maxConnections int32) int64 { - return max(1, int64(maxConnections)-copyBudgetReserve) -} - // ProcessWALEvent is called on every new message from the wal. It can be called // concurrently. func (w *BulkIngestWriter) ProcessWALEvent(ctx context.Context, walEvent *wal.Event) (err error) { diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go index aa8a7c07f..5e8b1c5dc 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go +++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go @@ -416,7 +416,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) { pgConn: tc.pgConn, disableTriggers: tc.disableTriggers, }, - copyBudget: synclib.NewWeightedSemaphore(pglib.MaxConns - copyBudgetReserve), + copyBudget: synclib.NewWeightedSemaphore(pglib.MaxConns - synclib.CopyBudgetReserve), } err := writer.sendBatch(context.Background(), tc.batch) @@ -500,14 +500,14 @@ func TestCopyBudgetSize(t *testing.T) { }{ {name: "default pool", maxConnections: pglib.MaxConns, expected: 45}, {name: "configured pool", maxConnections: 12, expected: 7}, - {name: "reserve matches pool", maxConnections: copyBudgetReserve, expected: 1}, + {name: "reserve matches pool", maxConnections: synclib.CopyBudgetReserve, expected: 1}, {name: "pool smaller than reserve", maxConnections: 2, expected: 1}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - require.Equal(t, tt.expected, copyBudgetSize(tt.maxConnections)) + require.Equal(t, tt.expected, synclib.CopyBudgetSize(tt.maxConnections)) }) } } @@ -571,7 +571,7 @@ func TestNewBulkIngestWriter_maxConnections(t *testing.T) { require.True(t, ok) require.Equal(t, tt.expectedObserver, observerPool.Config().MaxConns) - budget := copyBudgetSize(tt.expected) + budget := synclib.CopyBudgetSize(tt.expected) for range budget { require.True(t, writer.copyBudget.TryAcquire(1)) }