diff --git a/README.md b/README.md index 2d25cbd..c86eb40 100644 --- a/README.md +++ b/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 ` | 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 diff --git a/internal/app/app.go b/internal/app/app.go index 87542d3..aab359e 100644 --- a/internal/app/app.go +++ b/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 } diff --git a/internal/app/replident.go b/internal/app/replident.go index 524f071..e38f56b 100644 --- a/internal/app/replident.go +++ b/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 } diff --git a/internal/app/tuning.go b/internal/app/tuning.go index ea0b3e0..972594a 100644 --- a/internal/app/tuning.go +++ b/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 } diff --git a/internal/app/vacuum.go b/internal/app/vacuum.go index 89e1a9d..2b23a59 100644 --- a/internal/app/vacuum.go +++ b/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 } diff --git a/internal/app/verify.go b/internal/app/verify.go index 2c4bd27..c6a3bc1 100644 --- a/internal/app/verify.go +++ b/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 } diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 7e98c00..aba1c9b 100644 --- a/internal/cdc/applier.go +++ b/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) } diff --git a/internal/copy/copy.go b/internal/copy/copy.go index 8ecec7f..845cc7f 100644 --- a/internal/copy/copy.go +++ b/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) diff --git a/internal/copy/copy_integration_test.go b/internal/copy/copy_integration_test.go index 75e485f..b556687 100644 --- a/internal/copy/copy_integration_test.go +++ b/internal/copy/copy_integration_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/GetStream/pgmigrate/internal/pgtest" + "github.com/GetStream/pgmigrate/internal/postgres" "github.com/GetStream/pgmigrate/internal/state" "github.com/jackc/pgx/v5" ) @@ -435,6 +436,70 @@ func TestPG17PartsCopyIntoTheirOwnTable(t *testing.T) { } } +func TestPG17CopySurvivesInheritedStatementTimeout(t *testing.T) { + source := pgtest.Start(t, 17) + target := pgtest.Start(t, 17) + ctx := context.Background() + src := source.Connect(t) + dst := target.Connect(t) + if _, err := src.Exec(ctx, ` + CREATE TABLE timeout_rows (id bigint PRIMARY KEY, value text); + INSERT INTO timeout_rows SELECT i, repeat('x', 100) FROM generate_series(1, 5000) i`); err != nil { + t.Fatal(err) + } + if _, err := dst.Exec(ctx, "CREATE TABLE timeout_rows (id bigint PRIMARY KEY, value text)"); err != nil { + t.Fatal(err) + } + exporter, err := pgx.Connect(ctx, source.URI) + if err != nil { + t.Fatal(err) + } + defer exporter.Close(ctx) + tx, err := exporter.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly}) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(ctx) + var snapshot string + if err := tx.QueryRow(ctx, "SELECT pg_export_snapshot()").Scan(&snapshot); err != nil { + t.Fatal(err) + } + if _, err := src.Exec(ctx, "ALTER DATABASE pgmigrate SET statement_timeout = '1ms'"); err != nil { + t.Fatal(err) + } + if _, err := dst.Exec(ctx, "ALTER DATABASE pgmigrate SET statement_timeout = '1ms'"); err != nil { + t.Fatal(err) + } + tables, err := InventorySnapshot(ctx, func(ctx context.Context) (*pgx.Conn, error) { + return postgres.Connect(ctx, source.URI) + }, snapshot, func(_, name string) bool { return name == "timeout_rows" }) + if err != nil { + t.Fatal(err) + } + if len(tables) != 1 { + t.Fatalf("inventory tables=%d, want 1", len(tables)) + } + store, err := state.Open(ctx, t.TempDir(), state.Fingerprints{Source: "source", Filter: "timeout"}) + if err != nil { + t.Fatal(err) + } + defer store.Close() + runner := Runner{ + Source: func(ctx context.Context) (*pgx.Conn, error) { return postgres.Connect(ctx, source.URI) }, + Target: func(ctx context.Context) (*pgx.Conn, error) { return postgres.Connect(ctx, target.URI) }, + Snapshot: snapshot, + Workers: 1, + State: store, + } + if err := runner.Run(ctx, Plan(tables[0], 0, 1, Binary)); err != nil { + t.Fatal(err) + } + var count int + if err := dst.QueryRow(ctx, "SELECT count(*) FROM timeout_rows").Scan(&count); err != nil || count != 5000 { + t.Fatalf("copied rows=%d err=%v", count, err) + } +} + func connectWithDefaults(ctx context.Context, uri string, defaults map[string]string) (*pgx.Conn, error) { config, err := pgx.ParseConfig(uri) if err != nil { diff --git a/internal/copy/copy_test.go b/internal/copy/copy_test.go index 704e97d..2944cfc 100644 --- a/internal/copy/copy_test.go +++ b/internal/copy/copy_test.go @@ -65,6 +65,24 @@ func TestRetryClassification(t *testing.T) { if retryableConnectionError(&pgconn.PgError{Code: "22000"}) { t.Fatal("data exception was retryable") } + if retryableConnectionError(&pgconn.PgError{Code: "57014"}) { + t.Fatal("query_canceled was retryable") + } +} + +func TestAnnotateCopyErrorNamesExternalCancel(t *testing.T) { + err := annotateCopyError("copy out of source", &pgconn.PgError{Code: "57014", Message: "canceling statement due to statement timeout"}) + if err == nil || !strings.Contains(err.Error(), "SQLSTATE 57014") || + !strings.Contains(err.Error(), "external cancel") { + t.Fatalf("annotated 57014 = %v", err) + } + plain := annotateCopyError("copy into target", errors.New("pipe closed")) + if plain == nil || !strings.Contains(plain.Error(), "copy into target: pipe closed") { + t.Fatalf("annotated plain error = %v", plain) + } + if got := annotateCopyError("copy out of source", nil); got != nil { + t.Fatalf("nil error = %v", got) + } } type retryError struct{} diff --git a/internal/postgres/connect.go b/internal/postgres/connect.go new file mode 100644 index 0000000..9ce8db8 --- /dev/null +++ b/internal/postgres/connect.go @@ -0,0 +1,81 @@ +package postgres + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5" +) + +// SessionTimeoutNames are the GUCs that cancel a long COPY or restore. A copy +// part is one statement and holds a transaction for its whole duration, so any +// of these inherited from a role or parameter group will kill it. +var SessionTimeoutNames = []string{ + "statement_timeout", + "lock_timeout", + "idle_in_transaction_session_timeout", + "idle_session_timeout", +} + +// SessionTimeoutPGOPTIONS disables those GUCs for libpq subprocesses +// (pg_dump, pg_restore) that never see AfterConnect. +const SessionTimeoutPGOPTIONS = "-c statement_timeout=0 -c lock_timeout=0 -c idle_in_transaction_session_timeout=0 -c idle_session_timeout=0" + +// SessionTimeout is one inherited timeout, in milliseconds. Zero means disabled. +type SessionTimeout struct { + Name string + Milliseconds int64 +} + +// Connect opens a SQL session and disables the timeouts that would cancel bulk +// work. Replication-protocol connections must not use this: SET is not a +// replication command. +func Connect(ctx context.Context, dsn string) (*pgx.Conn, error) { + conn, err := pgx.Connect(ctx, dsn) + if err != nil { + return nil, err + } + if err := DisableSessionTimeouts(ctx, conn); err != nil { + conn.Close(context.Background()) + return nil, err + } + return conn, nil +} + +// DisableSessionTimeouts sets the bulk-work timeouts to 0 for the rest of this +// session. It overrides ALTER ROLE and parameter-group defaults; it does not +// persist. +func DisableSessionTimeouts(ctx context.Context, conn *pgx.Conn) error { + if _, err := conn.Exec(ctx, ` + SELECT set_config('statement_timeout','0',false), + set_config('lock_timeout','0',false), + set_config('idle_in_transaction_session_timeout','0',false), + set_config('idle_session_timeout','0',false)`); err != nil { + return fmt.Errorf("disable session timeouts: %w", err) + } + return nil +} + +// InheritedSessionTimeouts reports the values RESET would restore, which is +// what the session inherited from role, database, and parameter group, even +// after DisableSessionTimeouts. +func InheritedSessionTimeouts(ctx context.Context, conn *pgx.Conn) ([]SessionTimeout, error) { + rows, err := conn.Query(ctx, ` + SELECT name, reset_val::bigint + FROM pg_catalog.pg_settings + WHERE name = ANY($1) + ORDER BY name`, SessionTimeoutNames) + if err != nil { + return nil, err + } + defer rows.Close() + var timeouts []SessionTimeout + for rows.Next() { + var timeout SessionTimeout + if err := rows.Scan(&timeout.Name, &timeout.Milliseconds); err != nil { + return nil, err + } + timeouts = append(timeouts, timeout) + } + return timeouts, rows.Err() +} diff --git a/internal/postgres/connect_integration_test.go b/internal/postgres/connect_integration_test.go new file mode 100644 index 0000000..219afed --- /dev/null +++ b/internal/postgres/connect_integration_test.go @@ -0,0 +1,73 @@ +//go:build integration + +package postgres_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/GetStream/pgmigrate/internal/pgtest" + "github.com/GetStream/pgmigrate/internal/postgres" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +func TestPG17ConnectDisablesInheritedStatementTimeout(t *testing.T) { + instance := pgtest.Start(t, 17) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + admin := instance.Connect(t) + if _, err := admin.Exec(ctx, "ALTER DATABASE pgmigrate SET statement_timeout = '1s'"); err != nil { + t.Fatal(err) + } + + raw, err := pgx.Connect(ctx, instance.URI) + if err != nil { + t.Fatal(err) + } + defer raw.Close(context.Background()) + var inherited string + if err := raw.QueryRow(ctx, "SHOW statement_timeout").Scan(&inherited); err != nil { + t.Fatal(err) + } + if inherited != "1s" && inherited != "1000ms" { + t.Fatalf("inherited statement_timeout = %q", inherited) + } + + conn, err := postgres.Connect(ctx, instance.URI) + if err != nil { + t.Fatal(err) + } + defer conn.Close(context.Background()) + var current string + if err := conn.QueryRow(ctx, "SHOW statement_timeout").Scan(¤t); err != nil { + t.Fatal(err) + } + if current != "0" { + t.Fatalf("session statement_timeout = %q, want 0", current) + } + timeouts, err := postgres.InheritedSessionTimeouts(ctx, conn) + if err != nil { + t.Fatal(err) + } + var statement int64 = -1 + for _, timeout := range timeouts { + if timeout.Name == "statement_timeout" { + statement = timeout.Milliseconds + } + } + if statement != 1000 { + t.Fatalf("reset_val statement_timeout = %d, want 1000", statement) + } + if _, err := conn.Exec(ctx, "SELECT pg_sleep(1.2)"); err != nil { + t.Fatalf("pg_sleep on a disabled-timeout session: %v", err) + } + _, err = raw.Exec(ctx, "SELECT pg_sleep(1.2)") + var pgErr *pgconn.PgError + if err == nil || !errors.As(err, &pgErr) || pgErr.Code != "57014" { + t.Fatalf("raw pg_sleep error = %v, want SQLSTATE 57014", err) + } +} diff --git a/internal/preflight/preflight.go b/internal/preflight/preflight.go index 0b20670..f21e463 100644 --- a/internal/preflight/preflight.go +++ b/internal/preflight/preflight.go @@ -97,12 +97,12 @@ func Run(ctx context.Context, cfg Config) (Result, error) { if strings.TrimSpace(cfg.SourceDSN) == "" || strings.TrimSpace(cfg.TargetDSN) == "" { return Result{}, errors.New("source and target DSNs are required") } - source, err := pgx.Connect(ctx, cfg.SourceDSN) + source, err := postgres.Connect(ctx, cfg.SourceDSN) if err != nil { return Result{}, fmt.Errorf("connect source: %w", err) } defer source.Close(context.Background()) - target, err := pgx.Connect(ctx, cfg.TargetDSN) + target, err := postgres.Connect(ctx, cfg.TargetDSN) if err != nil { return Result{}, fmt.Errorf("connect target: %w", err) } @@ -148,6 +148,12 @@ func RunConnections(ctx context.Context, source, target *pgx.Conn, cfg Config) ( return decide(result, cfg), nil } + runCheck(add, "source-timeouts", func() ([]Finding, error) { + return checkSessionTimeouts(ctx, source, "source") + }) + runCheck(add, "target-timeouts", func() ([]Finding, error) { + return checkSessionTimeouts(ctx, target, "target") + }) runCheck(add, "source-replication", func() ([]Finding, error) { return checkSourceReplication(ctx, source, cfg.WALSampleDuration, cfg.WALRetentionDuration) }) @@ -380,6 +386,43 @@ func walRetentionFinding( return nil } +func checkSessionTimeouts(ctx context.Context, conn *pgx.Conn, side string) ([]Finding, error) { + timeouts, err := postgres.InheritedSessionTimeouts(ctx, conn) + if err != nil { + return nil, err + } + return sessionTimeoutFindings(side, timeouts), nil +} + +// sessionTimeoutFindings warns when a role or parameter group would cancel a +// COPY part. pgmigrate sets these to 0 on its own sessions, so the finding is +// acknowledgeable rather than blocking. +func sessionTimeoutFindings(side string, timeouts []postgres.SessionTimeout) []Finding { + var findings []Finding + for _, timeout := range timeouts { + if timeout.Milliseconds == 0 { + continue + } + findings = append(findings, Finding{ + ID: side + "-" + strings.ReplaceAll(timeout.Name, "_", "-"), + Kind: "timeout", + Severity: SeverityWarning, + Message: fmt.Sprintf( + "%s %s is %s; a COPY part is one statement and holds its transaction until it finishes, so pgmigrate sets this to 0 on its own SQL sessions", + side, timeout.Name, formatMilliseconds(timeout.Milliseconds), + ), + }) + } + return findings +} + +func formatMilliseconds(ms int64) string { + if ms%1000 == 0 { + return fmt.Sprintf("%ds", ms/1000) + } + return fmt.Sprintf("%dms", ms) +} + func checkReplicaIdentity(ctx context.Context, conn *pgx.Conn, tables []Table) ([]Finding, error) { var findings []Finding for _, table := range tables { diff --git a/internal/preflight/preflight_integration_test.go b/internal/preflight/preflight_integration_test.go index 6236c22..b1809c3 100644 --- a/internal/preflight/preflight_integration_test.go +++ b/internal/preflight/preflight_integration_test.go @@ -91,6 +91,50 @@ func roleDSN(t testing.TB, uri, user, password string) string { return parsed.String() } +func TestPG17SessionTimeouts(t *testing.T) { + source := pgtest.Start(t, 17) + target := pgtest.Start(t, 17) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + sourceConn := source.Connect(t) + if _, err := sourceConn.Exec(ctx, "CREATE TABLE selected (id bigint PRIMARY KEY)"); err != nil { + t.Fatal(err) + } + if _, err := sourceConn.Exec(ctx, "ALTER DATABASE pgmigrate SET statement_timeout = '60s'"); err != nil { + t.Fatal(err) + } + var oid uint32 + if err := sourceConn.QueryRow(ctx, "SELECT 'selected'::regclass::oid").Scan(&oid); err != nil { + t.Fatal(err) + } + tool := filepath.Join(t.TempDir(), "pg-tool") + if err := os.WriteFile(tool, []byte("#!/bin/sh\necho 'pg_dump (PostgreSQL) 17.1'\n"), 0o700); err != nil { + t.Fatal(err) + } + result, err := preflight.Run(ctx, preflight.Config{ + SourceDSN: source.URI, TargetDSN: target.URI, + Tables: []preflight.Table{{OID: oid, Schema: "public", Name: "selected"}}, + PGDumpPath: tool, PGRestorePath: tool, WALSampleDuration: 10 * time.Millisecond, + AcknowledgeWarnings: true, + }) + if err != nil { + t.Fatal(err) + } + var found preflight.Finding + for _, finding := range result.Findings { + if finding.ID == "source-statement-timeout" { + found = finding + } + } + if found.Severity != preflight.SeverityWarning || !strings.Contains(found.Message, "60s") { + t.Fatalf("source-statement-timeout = %+v findings=%+v", found, result.Findings) + } + if !result.Allowed { + t.Fatalf("acknowledgeable timeout blocked: %+v", result.Findings) + } +} + // TestPG17SequenceHeadroom checks the headroom report against real sequences: a // sequence with less room left than --sequence-offset stops the migration because // the cutover's setval would be rejected, one merely running low is diff --git a/internal/preflight/preflight_test.go b/internal/preflight/preflight_test.go index 19f1a61..9c97b5b 100644 --- a/internal/preflight/preflight_test.go +++ b/internal/preflight/preflight_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/GetStream/pgmigrate/internal/collation" + "github.com/GetStream/pgmigrate/internal/postgres" "github.com/GetStream/pgmigrate/internal/replident" "github.com/GetStream/pgmigrate/internal/tuning" ) @@ -120,6 +121,37 @@ func TestCollationVersionRemainsAcknowledgeable(t *testing.T) { } } +func TestSessionTimeoutFindingsWarnWhenInheritedAndStaySilentWhenDisabled(t *testing.T) { + findings := sessionTimeoutFindings("source", []postgres.SessionTimeout{ + {Name: "statement_timeout", Milliseconds: 60000}, + {Name: "lock_timeout", Milliseconds: 0}, + {Name: "idle_in_transaction_session_timeout", Milliseconds: 1500}, + {Name: "idle_session_timeout", Milliseconds: 0}, + }) + if len(findings) != 2 { + t.Fatalf("findings = %+v, want one per non-zero timeout", findings) + } + byID := map[string]Finding{} + for _, finding := range findings { + byID[finding.ID] = finding + if finding.Severity != SeverityWarning || finding.Kind != "timeout" { + t.Errorf("finding = %+v, want an acknowledgeable timeout warning", finding) + } + } + if msg := byID["source-statement-timeout"].Message; !strings.Contains(msg, "60s") || + !strings.Contains(msg, "one statement") { + t.Errorf("statement_timeout message = %q", msg) + } + if msg := byID["source-idle-in-transaction-session-timeout"].Message; !strings.Contains(msg, "1500ms") { + t.Errorf("idle_in_transaction message = %q", msg) + } + if extra := sessionTimeoutFindings("target", []postgres.SessionTimeout{ + {Name: "statement_timeout", Milliseconds: 0}, + }); len(extra) != 0 { + t.Errorf("disabled timeouts = %+v", extra) + } +} + // TestCollationFindingRefusesStructuralRisksOutright covers the case the flag // must not unblock, where collation decides which rows are equal and which // partition a row belongs to rather than only the order rows come back in. diff --git a/internal/schema/schema.go b/internal/schema/schema.go index 3348d4f..5b1721f 100644 --- a/internal/schema/schema.go +++ b/internal/schema/schema.go @@ -335,7 +335,7 @@ func (s Service) RestoreDeferred(ctx context.Context, targetURI, archive, useLis if s.DeferredMarkers == nil || s.InspectDeferred == nil { return errors.New("deferred markers and exact inspector are required") } - conn, err := pgx.Connect(ctx, targetURI) + conn, err := postgres.Connect(ctx, targetURI) if err != nil { return fmt.Errorf("connect deferred restore target: %w", err) } @@ -489,7 +489,7 @@ func (s Service) Clean(ctx context.Context, targetURI, archive string) error { if err != nil { return err } - conn, err := pgx.Connect(ctx, targetURI) + conn, err := postgres.Connect(ctx, targetURI) if err != nil { return fmt.Errorf("connect target archive cleanup: %w", err) } @@ -522,6 +522,7 @@ func run(ctx context.Context, name string, args ...string) error { func output(ctx context.Context, name string, args ...string) ([]byte, error) { cmd := exec.CommandContext(ctx, name, args...) + cmd.Env = withSessionTimeoutPGOPTIONS(os.Environ()) var combined bytes.Buffer cmd.Stdout, cmd.Stderr = &combined, &combined err := cmd.Run() @@ -530,3 +531,14 @@ func output(ctx context.Context, name string, args ...string) ([]byte, error) { } return combined.Bytes(), nil } + +func withSessionTimeoutPGOPTIONS(env []string) []string { + out := append([]string(nil), env...) + for i, kv := range out { + if strings.HasPrefix(kv, "PGOPTIONS=") { + out[i] = kv + " " + postgres.SessionTimeoutPGOPTIONS + return out + } + } + return append(out, "PGOPTIONS="+postgres.SessionTimeoutPGOPTIONS) +} diff --git a/internal/schema/schema_test.go b/internal/schema/schema_test.go index f34b2ff..d0daec9 100644 --- a/internal/schema/schema_test.go +++ b/internal/schema/schema_test.go @@ -9,6 +9,8 @@ import ( "slices" "strings" "testing" + + "github.com/GetStream/pgmigrate/internal/postgres" ) func TestParseTOCAndUseList(t *testing.T) { @@ -335,3 +337,25 @@ func TestFilterTOCPreservesArchiveOrder(t *testing.T) { t.Errorf("filtered order = %v, want %v", got, want) } } + +func TestWithSessionTimeoutPGOPTIONSAppendsOrSets(t *testing.T) { + got := withSessionTimeoutPGOPTIONS([]string{"HOME=/tmp", "PGOPTIONS=-c extra_float_digits=3"}) + if !slices.Contains(got, "HOME=/tmp") { + t.Fatalf("lost unrelated env: %v", got) + } + var options string + for _, kv := range got { + if strings.HasPrefix(kv, "PGOPTIONS=") { + options = kv + } + } + if !strings.Contains(options, "-c extra_float_digits=3") || + !strings.Contains(options, "statement_timeout=0") || + !strings.Contains(options, "idle_in_transaction_session_timeout=0") { + t.Fatalf("PGOPTIONS = %q", options) + } + fresh := withSessionTimeoutPGOPTIONS([]string{"HOME=/tmp"}) + if !slices.Contains(fresh, "PGOPTIONS="+postgres.SessionTimeoutPGOPTIONS) { + t.Fatalf("missing PGOPTIONS: %v", fresh) + } +} diff --git a/internal/setup/setup.go b/internal/setup/setup.go index bc14cad..50063c5 100644 --- a/internal/setup/setup.go +++ b/internal/setup/setup.go @@ -150,12 +150,12 @@ func Run(ctx context.Context, cfg Config, state SnapshotState) (_ *Holder, err e return nil, errors.New("source DSN, target DSN, directory, and selected tables are required") } - source, err := pgx.Connect(ctx, cfg.SourceDSN) + source, err := postgres.Connect(ctx, cfg.SourceDSN) if err != nil { return nil, fmt.Errorf("connect source setup: %w", err) } defer source.Close(context.Background()) - target, err := pgx.Connect(ctx, cfg.TargetDSN) + target, err := postgres.Connect(ctx, cfg.TargetDSN) if err != nil { return nil, fmt.Errorf("connect target setup: %w", err) } @@ -232,7 +232,7 @@ func Run(ctx context.Context, cfg Config, state SnapshotState) (_ *Holder, err e if err := postgres.EnsureProgressTable(ctx, target); err != nil { return nil, fmt.Errorf("create target progress table: %w", err) } - monitor, err := pgx.Connect(ctx, cfg.SourceDSN) + monitor, err := postgres.Connect(ctx, cfg.SourceDSN) if err != nil { return nil, fmt.Errorf("connect snapshot monitor: %w", err) } @@ -352,7 +352,7 @@ func RecoverStale(ctx context.Context, cfg Config, confirmation ResumeConfirmati } publication, slot := Names(SourceFingerprint(system.SystemID, system.DBName), cfg.MigrationID) - conn, err := pgx.Connect(ctx, cfg.SourceDSN) + conn, err := postgres.Connect(ctx, cfg.SourceDSN) if err != nil { return fmt.Errorf("connect source recovery: %w", err) } @@ -512,7 +512,7 @@ func CleanupOwned( tables []Table, expectFailover bool, ) error { - conn, err := pgx.Connect(ctx, sourceDSN) + conn, err := postgres.Connect(ctx, sourceDSN) if err != nil { return fmt.Errorf("connect validated source cleanup: %w", err) }