postgres: disable the session timeouts that cancel COPY

README.md

@@ -236,6 +236,13 @@ warning: what remains has to cover both databases until traffic moves. One with
 less room than the offset is an error, because `setval` refuses a value past the
 bound and the cutover would fail at its sequence step.
 
+Inherited `statement_timeout`, `lock_timeout`, `idle_in_transaction_session_timeout`,
+and `idle_session_timeout` on the source or target are warnings. pgmigrate sets
+those GUCs to 0 on every SQL session it opens, including `pg_dump` and
+`pg_restore`, so a COPY that would otherwise die at 60s can finish. The warning
+is so the inherited values are visible before a long run, not a request to
+`ALTER ROLE` for pgmigrate.
+
 | flag | default | what it does |
 |---|---|---|
 | `--dir <path>` | required | migration state directory, created if absent |
@@ -689,6 +696,10 @@ Re-run `pgmigrate run` with the same DSNs, filter, and directory.
 - Copy parts retry classified connection failures up to five times. CDC
   transport failures reconnect automatically; corruption, protocol errors,
   divergence, and prolonged handoff backpressure stop the run for diagnosis.
+  SQLSTATE 57014 is not a connection failure: after pgmigrate has set
+  `statement_timeout`, `lock_timeout`, `idle_in_transaction_session_timeout`,
+  and `idle_session_timeout` to 0 on its SQL sessions, a cancel is external
+  (`pg_cancel_backend`, a proxy idle timeout) and stops the part.
 - Cutover records each successful step and resumes at the first incomplete one,
   reusing the end position the first attempt recorded. It never moves that
   boundary: the target has already been drained to it, and a fresh one would
@@ -718,6 +729,10 @@ the same reason rather than resuming the loop. Resolve the cause and pass
 `--retry-base-copy` once; the record clears by itself as soon as the run reaches
 `indexes`, after which restarts resume instead of discarding work. A process
 killed outright, or stopped with a signal, is not a failed attempt.
+`--retry-base-copy` is not the fix for SQLSTATE 57014: pgmigrate already
+disables the session timeouts that cancel long COPY, dump, and restore. If a
+cancel still happens, it is external, and retrying the whole base copy will
+hit it again.
 
 ## Security
 

internal/app/app.go

@@ -60,7 +60,7 @@ func (a App) progressOutput() io.Writer {
 }
 
 func connector(dsn string) func(context.Context) (*pgx.Conn, error) {
-	return func(ctx context.Context) (*pgx.Conn, error) { return pgx.Connect(ctx, dsn) }
+	return func(ctx context.Context) (*pgx.Conn, error) { return postgres.Connect(ctx, dsn) }
 }
 
 func loadFilter(path string) (config.Filter, error) {
@@ -89,7 +89,7 @@ func sourceFingerprint(ctx context.Context, dsn string) (string, error) {
 }
 
 func inventory(ctx context.Context, cfg config.Config, filter config.Filter) ([]pgcopy.Table, error) {
-	conn, err := pgx.Connect(ctx, cfg.Source)
+	conn, err := postgres.Connect(ctx, cfg.Source)
 	if err != nil {
 		return nil, fmt.Errorf("connect source inventory: %w", err)
 	}
@@ -260,7 +260,7 @@ func loadCDCBinaryMode(ctx context.Context, store *state.Store) (bool, error) {
 }
 
 func initializeTargetProgress(ctx context.Context, targetDSN, streamID, generation string) error {
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}
@@ -286,7 +286,7 @@ func initializeTargetProgress(ctx context.Context, targetDSN, streamID, generati
 }
 
 func validateTargetProgress(ctx context.Context, targetDSN, streamID, generation string) error {
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}
@@ -338,7 +338,7 @@ func finalizeTargetCleanup(ctx context.Context, targetDSN string, store *state.S
 	if err := validateTargetOnly(ctx, targetDSN, migration); err != nil {
 		return err
 	}
-	target, err := pgx.Connect(ctx, targetDSN)
+	target, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}
@@ -572,7 +572,7 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) {
 		if err := pauseForCrashTest(groupCtx, state.PhaseIndexes); err != nil {
 			return err
 		}
-		source, err := pgx.Connect(groupCtx, cfg.Source)
+		source, err := postgres.Connect(groupCtx, cfg.Source)
 		if err != nil {
 			return err
 		}
@@ -760,7 +760,7 @@ func resumeIndexes(ctx context.Context, cfg config.Config, store *state.Store) e
 	for _, table := range tables {
 		selected[table.OID] = true
 	}
-	source, err := pgx.Connect(ctx, cfg.Source)
+	source, err := postgres.Connect(ctx, cfg.Source)
 	if err != nil {
 		return err
 	}
@@ -940,7 +940,7 @@ func dumpSelection(ctx context.Context, sourceDSN, snapshot string, tables []pgc
 		selection.Tables[i] = schema.QualifiedName{Schema: table.Schema, Name: table.Name}
 		oids[i] = table.OID
 	}
-	conn, err := pgx.Connect(ctx, sourceDSN)
+	conn, err := postgres.Connect(ctx, sourceDSN)
 	if err != nil {
 		return selection, err
 	}
@@ -1125,7 +1125,7 @@ func dumpSelection(ctx context.Context, sourceDSN, snapshot string, tables []pgc
 
 func inspectDeferred(dir, sourceDSN string) schema.DeferredInspector {
 	return func(ctx context.Context, target *pgx.Conn, entry schema.TOCEntry) (schema.DeferredStatus, error) {
-		source, err := pgx.Connect(ctx, sourceDSN)
+		source, err := postgres.Connect(ctx, sourceDSN)
 		if err != nil {
 			return schema.DeferredStatus{}, err
 		}
@@ -1559,7 +1559,7 @@ func followChecks(ctx context.Context, cfg config.Config, store *state.Store, sl
 		if migration.Phase != state.PhaseFollow {
 			continue
 		}
-		conn, err := pgx.Connect(ctx, cfg.Source)
+		conn, err := postgres.Connect(ctx, cfg.Source)
 		if err != nil {
 			_ = store.UpsertFinding(ctx, state.Finding{ID: "follow-source-health", Kind: "health", Severity: "error", Message: err.Error()})
 			continue
@@ -1714,7 +1714,7 @@ func monitorProgress(ctx context.Context, store *state.Store, targetDSN, streamI
 	defer ticker.Stop()
 	nextLog := time.Now()
 	for {
-		conn, err := pgx.Connect(ctx, targetDSN)
+		conn, err := postgres.Connect(ctx, targetDSN)
 		if err != nil {
 			return err
 		}
@@ -1763,7 +1763,7 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta
 	if err != nil && !errors.Is(err, os.ErrNotExist) {
 		return fmt.Errorf("read prior snapshot metadata: %w", err)
 	}
-	target, err := pgx.Connect(ctx, cfg.Target)
+	target, err := postgres.Connect(ctx, cfg.Target)
 	if err != nil {
 		return err
 	}
@@ -1912,7 +1912,7 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta
 }
 
 func recordTargetIdentity(ctx context.Context, targetDSN, sourceFingerprint, filterFingerprint, streamID, generation string) error {
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}
@@ -1975,7 +1975,7 @@ func validateTargetIdentity(ctx context.Context, cfg config.Config, store *state
 }
 
 func validateTargetOnly(ctx context.Context, targetDSN string, migration state.Migration) error {
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}
@@ -2085,7 +2085,7 @@ func verificationInventory(ctx context.Context, cfg config.Config, tables []stat
 	for _, table := range tables {
 		selected[table.Schema+"\x00"+table.Name] = true
 	}
-	source, err := pgx.Connect(ctx, cfg.Source)
+	source, err := postgres.Connect(ctx, cfg.Source)
 	if err != nil {
 		return nil, err
 	}
@@ -2118,7 +2118,7 @@ func waitTargetProgress(ctx context.Context, targetDSN, streamID, wanted string)
 	var reached pglogrepl.LSN
 	advancedAt := time.Now()
 	for {
-		conn, err := pgx.Connect(ctx, targetDSN)
+		conn, err := postgres.Connect(ctx, targetDSN)
 		if err != nil {
 			return err
 		}
@@ -2297,7 +2297,7 @@ func (a App) Cutover(ctx context.Context, cfg config.Config) error {
 		Sequences:      selectedSequences,
 		SequenceOffset: cfg.SequenceOffset,
 		EmitBoundary: func(ctx context.Context) (string, error) {
-			conn, err := pgx.Connect(ctx, cfg.Source)
+			conn, err := postgres.Connect(ctx, cfg.Source)
 			if err != nil {
 				return "", err
 			}
@@ -2360,7 +2360,7 @@ func cleanupAfterCutover(ctx context.Context, cfg config.Config, store *state.St
 
 func waitSlotInactive(ctx context.Context, sourceDSN, slot string) error {
 	for {
-		conn, err := pgx.Connect(ctx, sourceDSN)
+		conn, err := postgres.Connect(ctx, sourceDSN)
 		if err != nil {
 			return err
 		}

internal/app/replident.go

@@ -7,10 +7,9 @@ import (
 	"io"
 	"strings"
 
-	"github.com/jackc/pgx/v5"
-
 	"github.com/GetStream/pgmigrate/internal/config"
 	pgcopy "github.com/GetStream/pgmigrate/internal/copy"
+	"github.com/GetStream/pgmigrate/internal/postgres"
 	"github.com/GetStream/pgmigrate/internal/replident"
 	"github.com/GetStream/pgmigrate/internal/state"
 )
@@ -60,7 +59,7 @@ func (r replidentRecorder) Record(ctx context.Context, record replident.Record)
 func (a App) applyReplicaIdentityFallback(
 	ctx context.Context, cfg config.Config, store *state.Store, tables []pgcopy.Table,
 ) error {
-	conn, err := pgx.Connect(ctx, cfg.Source)
+	conn, err := postgres.Connect(ctx, cfg.Source)
 	if err != nil {
 		return err
 	}
@@ -157,7 +156,7 @@ func restoreReplicaIdentities(ctx context.Context, sourceDSN string, store *stat
 	if len(records) == 0 {
 		return nil
 	}
-	conn, err := pgx.Connect(ctx, sourceDSN)
+	conn, err := postgres.Connect(ctx, sourceDSN)
 	if err != nil {
 		return err
 	}
@@ -191,7 +190,7 @@ func restoreTargetReplicaIdentities(ctx context.Context, targetDSN string, store
 	if len(records) == 0 {
 		return nil
 	}
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}

internal/app/tuning.go

@@ -8,9 +8,8 @@ import (
 	"slices"
 	"strings"
 
-	"github.com/jackc/pgx/v5"
-
 	"github.com/GetStream/pgmigrate/internal/config"
+	"github.com/GetStream/pgmigrate/internal/postgres"
 	"github.com/GetStream/pgmigrate/internal/preflight"
 	"github.com/GetStream/pgmigrate/internal/state"
 	"github.com/GetStream/pgmigrate/internal/tuning"
@@ -69,7 +68,7 @@ func tuneTarget(ctx context.Context, cfg config.Config, store *state.Store) (map
 	if err != nil {
 		return nil, err
 	}
-	conn, err := pgx.Connect(ctx, cfg.Target)
+	conn, err := postgres.Connect(ctx, cfg.Target)
 	if err != nil {
 		return nil, err
 	}
@@ -203,7 +202,7 @@ func revertTargetTuning(ctx context.Context, targetDSN string, store *state.Stor
 	if len(changes) == 0 {
 		return nil
 	}
-	conn, err := pgx.Connect(ctx, targetDSN)
+	conn, err := postgres.Connect(ctx, targetDSN)
 	if err != nil {
 		return err
 	}

internal/app/vacuum.go

@@ -129,7 +129,7 @@ func vacuumOne(
 // settings, which is what lets one vacuum use the memory and parallel workers the
 // target was sized for.
 func vacuumSession(ctx context.Context, cfg config.Config, sessionGUCs map[string]string) (*pgx.Conn, error) {
-	conn, err := pgx.Connect(ctx, cfg.Target)
+	conn, err := postgres.Connect(ctx, cfg.Target)
 	if err != nil {
 		return nil, err
 	}

internal/app/verify.go

@@ -78,7 +78,7 @@ func (m *marker) flushWAL(ctx context.Context) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	if m.nudge == nil {
-		conn, err := pgx.Connect(ctx, m.dsn)
+		conn, err := postgres.Connect(ctx, m.dsn)
 		if err != nil {
 			return fmt.Errorf("connect to flush the verification marker: %w", err)
 		}
@@ -114,7 +114,7 @@ func (m *marker) close() {
 
 // sourceCapabilities reads what the source's release supports.
 func sourceCapabilities(ctx context.Context, dsn string) (postgres.Capabilities, error) {
-	conn, err := pgx.Connect(ctx, dsn)
+	conn, err := postgres.Connect(ctx, dsn)
 	if err != nil {
 		return postgres.Capabilities{}, err
 	}

internal/cdc/applier.go

@@ -119,7 +119,7 @@ func (a *Applier) WaitUntil(ctx context.Context, boundary LSN) error {
 	effectiveBoundary := boundary
 	resolved := false
 	for {
-		conn, err := pgx.Connect(ctx, a.config.ConnString)
+		conn, err := postgres.Connect(ctx, a.config.ConnString)
 		if err != nil {
 			return fmt.Errorf("cdc: connect catch-up observer: %w", err)
 		}
@@ -152,7 +152,7 @@ func (a *Applier) WaitUntil(ctx context.Context, boundary LSN) error {
 }
 
 func (a *Applier) runConnection(ctx context.Context) error {
-	conn, err := pgx.Connect(ctx, a.config.ConnString)
+	conn, err := postgres.Connect(ctx, a.config.ConnString)
 	if err != nil {
 		return fmt.Errorf("cdc: connect applier: %w", err)
 	}

internal/copy/copy.go

@@ -483,6 +483,17 @@ func retryableConnectionError(err error) bool {
 	return pgconn.SafeToRetry(err)
 }
 
+func annotateCopyError(side string, err error) error {
+	if err == nil {
+		return nil
+	}
+	var pgErr *pgconn.PgError
+	if errors.As(err, &pgErr) && pgErr.Code == "57014" {
+		return fmt.Errorf("%s: query canceled (SQLSTATE 57014); pgmigrate already set statement_timeout, lock_timeout, idle_in_transaction_session_timeout, and idle_session_timeout to 0 on this session, so this is an external cancel: %w", side, err)
+	}
+	return fmt.Errorf("%s: %w", side, err)
+}
+
 // LargestFirst returns a stable, independently owned worker schedule.
 func LargestFirst(parts []Part) []Part {
 	result := append([]Part(nil), parts...)
@@ -593,13 +604,7 @@ func (r Runner) copyPart(ctx context.Context, p Part) error {
 	_ = pr.CloseWithError(targetErr)
 	src := <-ch
 	if src.err != nil || targetErr != nil {
-		if src.err != nil {
-			src.err = fmt.Errorf("copy out of source: %w", src.err)
-		}
-		if targetErr != nil {
-			targetErr = fmt.Errorf("copy into target: %w", targetErr)
-		}
-		return errors.Join(src.err, targetErr)
+		return errors.Join(annotateCopyError("copy out of source", src.err), annotateCopyError("copy into target", targetErr))
 	}
 	if _, err := ttx.Exec(ctx, `
 		INSERT INTO pgmigrate_internal.copy_parts(table_oid, part_id, rows_copied, bytes_copied)