Copy passthrough during snapshot mode if no transformers are configured

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{}

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)
+}

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))
+}

pkg/snapshot/generator/postgres/data/chunk_mover.go

@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/otel"
+)
+
+// runInTx runs fn against a connection that can read the chunk. A reading
+// strategy supplies its own: the ctid reader imports the exported transaction
+// snapshot, a keyset reader need not open one at all.
+type runInTx func(ctx context.Context, fn func(tx pglib.Tx) error) error
+
+// chunkMover moves the rows a reading strategy has selected. The strategy
+// decides which rows to read and says so as a complete query; the mover decides
+// what becomes of them, so every strategy gets every mover.
+//
+// It takes the query rather than the rows because the copy passthrough must
+// never run it as an ordinary query: it wraps it in COPY. It takes a way to run
+// rather than an open transaction because the passthrough retries a whole chunk
+// on a fresh pipe, which needs a fresh transaction.
+//
+// The query must be complete SQL. COPY accepts no bind parameters, so a
+// strategy that phrases its chunk with placeholders can only ever be decoded.
+type chunkMover interface {
+	// prepareTable resolves whatever the mover needs before the table's chunks
+	// are moved, once its columns are known.
+	prepareTable(ctx context.Context, table *table) error
+	// move reads the chunk the query selects and returns the rows moved.
+	move(ctx context.Context, run runInTx, table *table, query string) (uint, error)
+	close(ctx context.Context) error
+}
+
+// newChunkMover picks how a chunk is moved. Selecting it here, next to the
+// sink and before any reader exists, is what keeps it independent of the
+// reading strategy.
+func newChunkMover(ctx context.Context, cfg *Config, logger loglib.Logger,
+	instrumentation *otel.Instrumentation, progress progressTracker, decoding chunkMover,
+) (chunkMover, error) {
+	if cfg.CopyPassthrough == nil {
+		return decoding, nil
+	}
+	return newCopyPassthroughMover(ctx, cfg.CopyPassthrough, logger, instrumentation, progress, decoding)
+}
+
+// decodingMover reads the rows and hands them to the sink, which adapts them
+// into wal events for the processor.
+type decodingMover struct {
+	sink rowSink
+}
+
+func newDecodingMover(sink rowSink) decodingMover {
+	return decodingMover{sink: sink}
+}
+
+func (m decodingMover) prepareTable(context.Context, *table) error { return nil }
+
+func (m decodingMover) close(context.Context) error { return nil }
+
+func (m decodingMover) move(ctx context.Context, run runInTx, table *table, query string) (uint, error) {
+	var rowCount uint
+	err := run(ctx, func(tx pglib.Tx) error {
+		rows, err := tx.Query(ctx, query)
+		if err != nil {
+			return wrapChunkQueryError(err)
+		}
+		defer rows.Close()
+
+		rowCount, err = m.sink.emit(ctx, table, rows)
+		return err
+	})
+	return rowCount, err
+}
+
+// a vanished relation means schema drift
+func wrapChunkQueryError(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)
+}

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,8 @@ 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, not users
+	CopyPassthrough *CopyPassthroughConfig
 }
 
 const (
@@ -40,6 +47,26 @@ const (
 	defaultMaxConnections  = 50
 )
 
+// the generator writes the target
+type CopyPassthroughConfig struct {
+	TargetURL       string
+	DisableTriggers bool
+	// 0 defers to the url
+	MaxConnections uint
+	RetryPolicy    backoff.Config
+}
+
+// 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

pkg/snapshot/generator/postgres/data/ctid_table_reader.go

@@ -4,7 +4,6 @@ package postgres
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"strings"
 
@@ -23,7 +22,7 @@ import (
 type ctidReader struct {
 	conn         pglib.Querier
 	logger       loglib.Logger
-	sink         rowSink
+	mover        chunkMover
 	tableWorkers uint
 	batchBytes   uint64
 }
@@ -96,6 +95,10 @@ func (s *ctidSession) readTable(ctx context.Context, table *table) error {
 	}
 	table.columns = columns
 
+	if err := s.reader.mover.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
 	// have the same table view, which allows us to use the ctid to
@@ -148,34 +151,20 @@ func buildPageRangeQuery(t *table, r pageRange) string {
 }
 
 func (s *ctidSession) snapshotTableRange(ctx context.Context, table *table, pageRange pageRange) error {
-	return s.execInTx(ctx, func(tx pglib.Tx) error {
-		s.reader.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
-		})
-
-		query := buildPageRangeQuery(table, pageRange)
-		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()
-
-		rowCount, err := s.reader.sink.emit(ctx, table, rows)
-		if err != nil {
-			return err
-		}
+	s.reader.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
+		"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
+	})
 
-		s.reader.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
-			"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
-		})
+	rowCount, err := s.reader.mover.move(ctx, s.execInTx, table, buildPageRangeQuery(table, pageRange))
+	if err != nil {
+		return err
+	}
 
-		return nil
+	s.reader.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
+		"schema": table.schema, "table": table.name, "snapshotID": s.snapshotID,
 	})
+
+	return nil
 }
 
 // tableInfoQuery shares the capture rule.