snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)
This commit is contained in:
390
internal/pgschema/pgschema.go
Normal file
390
internal/pgschema/pgschema.go
Normal file
@@ -0,0 +1,390 @@
|
||||
// Package pgschema contains the shared schema-versioning primitives used by
|
||||
// every PostgreSQL-backed broker / controller in this project. Each consumer
|
||||
// stores its own integer schemaVersion plus a map[int]string of upgrade
|
||||
// migrations and wires them together in its own EnsureSchema; the safety
|
||||
// invariants (gap-free migrations, error-discriminated reads, transactional
|
||||
// upgrades, downgrade rejection) all live here so they can't drift apart.
|
||||
package pgschema
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"math/rand/v2"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/jackc/pgerrcode"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// querier is the minimal subset of *pgxpool.Pool / pgx.Tx that the
|
||||
// schema-management helpers need. Accepting an interface keeps the package
|
||||
// free of a hard pgxpool dependency and makes it usable from inside an
|
||||
// outer transaction (e.g. a test fixture that wraps everything in a txn).
|
||||
type querier interface {
|
||||
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
|
||||
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
|
||||
}
|
||||
|
||||
// txBeginner is what ApplyMigrationInTx needs from the connection source —
|
||||
// the ability to start a new transaction. Pools satisfy this directly.
|
||||
type txBeginner interface {
|
||||
BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error)
|
||||
}
|
||||
|
||||
// ValidateMigrationMap panics if migrations doesn't contain exactly one entry
|
||||
// for every version in [2..target] (with nothing outside that range). Call
|
||||
// this from a package-level init() in every broker that ships migrations.
|
||||
//
|
||||
// Rationale: bumping schemaVersion without registering the matching migration
|
||||
// would silently skip it for existing installs — schema_version would jump
|
||||
// forward via the final UPDATE while the DB shape stays at the old version.
|
||||
// Panic at process start rather than corrupt at runtime.
|
||||
func ValidateMigrationMap(label string, target int, migrations map[int]string) {
|
||||
for v := 2; v <= target; v++ {
|
||||
if _, ok := migrations[v]; !ok {
|
||||
panic(fmt.Sprintf("%s: missing schemaMigrations[%d] (schemaVersion=%d) — every version 2..schemaVersion must have a migration registered", label, v, target))
|
||||
}
|
||||
}
|
||||
for v := range migrations {
|
||||
if v < 2 || v > target {
|
||||
panic(fmt.Sprintf("%s: schemaMigrations[%d] is outside the valid range [2..%d]", label, v, target))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateTablePrefix rejects a configured table prefix that can never work as
|
||||
// a PostgreSQL identifier. Every table, index and function name is this prefix
|
||||
// concatenated with a fixed suffix and interpolated into DDL unquoted, so a bad
|
||||
// prefix surfaces as a syntax error deep inside a generated CREATE statement,
|
||||
// naming the table it failed to create rather than the setting that produced
|
||||
// it. Call it once at construction, after defaults are applied.
|
||||
//
|
||||
// The rule is PostgreSQL's own for unquoted identifiers, so it rejects only
|
||||
// what could never have worked: a letter or underscore first, then letters,
|
||||
// digits, underscores or dollar signs. Uppercase stays valid — PostgreSQL
|
||||
// folds it, so a prefix of "MyApp" yields myapp_* tables and every generated
|
||||
// query folds to match — and so do non-ASCII letters.
|
||||
//
|
||||
// It deliberately says nothing about length. PostgreSQL truncates identifiers
|
||||
// at 63 bytes rather than erroring, so a long prefix degrades (partition names
|
||||
// losing their date suffix, for one) instead of failing outright, and rejecting
|
||||
// it here would refuse deployments that are running today.
|
||||
func ValidateTablePrefix(label, prefix string) error {
|
||||
if prefix == "" {
|
||||
return fmt.Errorf("%s: table_prefix is empty", label)
|
||||
}
|
||||
for i, r := range prefix {
|
||||
switch {
|
||||
case unicode.IsLetter(r) || r == '_':
|
||||
// Valid anywhere, including first.
|
||||
case (unicode.IsDigit(r) || r == '$') && i > 0:
|
||||
// Valid only after the first character.
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"%s: table_prefix %q contains %q, which PostgreSQL does not accept in an unquoted identifier — use letters, digits and underscores, and don't start with a digit",
|
||||
label, prefix, r)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReadSchemaVersion reads the schema_version row from `table` and discriminates
|
||||
// the error class. Returns:
|
||||
//
|
||||
// - (version, false, nil) if the row was read successfully
|
||||
// - (0, true, nil) if the table is missing (Postgres 42P01) or the
|
||||
// row is missing (pgx.ErrNoRows). The caller should run the fresh-install
|
||||
// path: DDL creates everything at the latest baseline, then SetSchemaVersion
|
||||
// bumps from 1 to the current schemaVersion.
|
||||
// - (0, false, err) for any other error (transient connection failure,
|
||||
// permission denied, timeout, column missing, …). The caller MUST propagate.
|
||||
//
|
||||
// Why: collapsing every error into "treat as fresh" lets a transient SELECT
|
||||
// failure silently bypass migrations and force schema_version forward via the
|
||||
// final UPDATE — leaving the DB at the old shape while the row claims the new
|
||||
// version (silent corruption).
|
||||
func ReadSchemaVersion(ctx context.Context, q querier, table string) (int, bool, error) {
|
||||
var v int
|
||||
err := q.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %s WHERE id = 1`, table)).Scan(&v)
|
||||
if err == nil {
|
||||
return v, false, nil
|
||||
}
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return 0, true, nil
|
||||
}
|
||||
var pgErr *pgconn.PgError
|
||||
if errors.As(err, &pgErr) && pgErr.Code == pgerrcode.UndefinedTable {
|
||||
return 0, true, nil
|
||||
}
|
||||
return 0, false, fmt.Errorf("read schema_version from %s: %w", table, err)
|
||||
}
|
||||
|
||||
// CheckDowngrade returns a clear error when the DB has been touched by a
|
||||
// newer binary than this one. Silently rewriting schema_version backward (the
|
||||
// pre-pgschema behaviour) would leave the DB structurally newer than the
|
||||
// running code expects.
|
||||
func CheckDowngrade(label string, dbVersion, supportedVersion int) error {
|
||||
if dbVersion > supportedVersion {
|
||||
return fmt.Errorf("%s: DB schema_version is %d but this binary supports only up to %d — downgrade not supported, run a binary >= %d or restore an older DB snapshot", label, dbVersion, supportedVersion, dbVersion)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrationVariant pairs a fully-rendered migration SQL with the version
|
||||
// tracking table it should bump. ApplyMigrationInTx runs every variant in
|
||||
// the SAME transaction — either all variants commit or all roll back.
|
||||
//
|
||||
// The "variant" concept exists because brokers can store both a JSONB-typed
|
||||
// and a BYTEA-typed copy of the schema under different table prefixes; each
|
||||
// variant's SQL is the same migration template rendered for one prefix. A
|
||||
// broker with no variants (e.g. the postgres controller) passes a one-element
|
||||
// slice.
|
||||
type MigrationVariant struct {
|
||||
// SQL is the migration body, with all template placeholders already
|
||||
// substituted by the caller (e.g. via the broker's own render helper).
|
||||
// pgschema is intentionally agnostic about substitution rules so each
|
||||
// broker can use its own renderer (pgstreambroker, for instance, also
|
||||
// substitutes __STREAM_TABLE__).
|
||||
SQL string
|
||||
// VersionTable is the fully-qualified `<prefix>schema_version` table
|
||||
// whose id=1 row this variant bumps. A row with no matching id=1 returns
|
||||
// an error rather than silently no-op'ing.
|
||||
VersionTable string
|
||||
}
|
||||
|
||||
// IsRetryableSchemaExecErr reports whether err is a conflict caused by several
|
||||
// nodes running the same idempotent schema DDL at the same time. That is the
|
||||
// normal situation on a cluster start or a rolling restart: every node runs its
|
||||
// own EnsureSchema against the same database and nothing serialises them.
|
||||
//
|
||||
// - 40P01 deadlock_detected — nodes taking the same catalog locks in
|
||||
// different orders.
|
||||
// - XX000 internal_error, which is how "tuple concurrently updated" arrives
|
||||
// — raised by concurrent CREATE OR REPLACE FUNCTION on the same function.
|
||||
// - 42P07 duplicate_table, 42710 duplicate_object and 23505
|
||||
// unique_violation — CREATE TABLE/INDEX IF NOT EXISTS probes the catalog
|
||||
// before it takes any lock, so two nodes both pass the probe and the loser
|
||||
// gets a hard error instead of the NOTICE a sequential re-run produces.
|
||||
// Which one it gets depends on when the winner committed relative to the
|
||||
// loser's later checks: 42P07 from the relation-name check in
|
||||
// heap_create_with_catalog, 42710 ("type ... already exists") from the
|
||||
// row-type-name check that follows it if the commit landed between the
|
||||
// two, 23505 from the unique index on pg_type/pg_class if the loser reached
|
||||
// its own catalog inserts first.
|
||||
//
|
||||
// The name is deliberately bound to schema execs: 23505 and XX000 are ordinary
|
||||
// application failures anywhere else, and retrying them on a data path would
|
||||
// paper over real bugs. Only classify statements from a schema batch with this.
|
||||
//
|
||||
// Migrations deliberately do not use it: they run under AcquireMigrationLock,
|
||||
// which serialises nodes cluster-wide, so a duplicate object there is a real
|
||||
// bug and must surface rather than be retried away.
|
||||
func IsRetryableSchemaExecErr(err error) bool {
|
||||
var pgErr *pgconn.PgError
|
||||
if !errors.As(err, &pgErr) {
|
||||
return false
|
||||
}
|
||||
switch pgErr.Code {
|
||||
case pgerrcode.DeadlockDetected, pgerrcode.InternalError, pgerrcode.DuplicateTable, pgerrcode.DuplicateObject, pgerrcode.UniqueViolation:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
const (
|
||||
schemaExecAttempts = 5
|
||||
schemaExecBaseBackoff = 100 * time.Millisecond
|
||||
schemaExecMaxBackoff = 800 * time.Millisecond
|
||||
)
|
||||
|
||||
// RetrySchemaExec runs exec, repeating it while it fails with a concurrent-DDL
|
||||
// conflict (see IsRetryableSchemaExecErr) and attempts remain. Any other error
|
||||
// is returned on the first occurrence, as is the last conflict once the
|
||||
// attempts are exhausted; a cancelled ctx aborts the loop and returns
|
||||
// ctx.Err(), so a shutdown mid-retry is reported as cancellation rather than
|
||||
// as a schema failure.
|
||||
//
|
||||
// Callers must make exec idempotent AND self-contained: a retry re-runs the
|
||||
// whole thing, so a schema batch has to be a single multi-statement Exec with
|
||||
// no arguments (pgx sends those over the simple query protocol, so the server
|
||||
// wraps them in one implicit transaction and a failed attempt rolls back
|
||||
// completely). A batch that opens its own BEGIN/COMMIT, or contains a
|
||||
// statement that cannot run inside a transaction block such as CREATE INDEX
|
||||
// CONCURRENTLY, breaks that invariant and must not be retried this way — the
|
||||
// schema-template tests in each consumer package pin it.
|
||||
//
|
||||
// The backoff is exponential with jitter: nodes that started together lose the
|
||||
// same race at the same instant, and retrying them in lockstep just recreates
|
||||
// the collision.
|
||||
func RetrySchemaExec(ctx context.Context, exec func(ctx context.Context) error) error {
|
||||
var err error
|
||||
for attempt := range schemaExecAttempts {
|
||||
err = exec(ctx)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if !IsRetryableSchemaExecErr(err) || attempt == schemaExecAttempts-1 {
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(schemaExecBackoff(attempt)):
|
||||
}
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// schemaExecBackoff returns the delay before the attempt after this one:
|
||||
// exponential from schemaExecBaseBackoff, capped at schemaExecMaxBackoff, and
|
||||
// spread over the upper half of that window so concurrent nodes separate
|
||||
// without any of them waiting the full cap.
|
||||
func schemaExecBackoff(attempt int) time.Duration {
|
||||
d := schemaExecBaseBackoff << attempt
|
||||
if d > schemaExecMaxBackoff {
|
||||
d = schemaExecMaxBackoff
|
||||
}
|
||||
//nolint:gosec // it's a jitter.
|
||||
return d/2 + time.Duration(rand.Int64N(int64(d/2)+1))
|
||||
}
|
||||
|
||||
// ApplyMigrationInTx applies one migration step. For each variant in the
|
||||
// slice:
|
||||
// - The variant's SQL is executed.
|
||||
// - The variant's VersionTable id=1 row is updated to `version`.
|
||||
//
|
||||
// All variants run in one transaction; a failure on any rolls back every
|
||||
// effect, including any partial version bumps. The transactional wrap makes
|
||||
// retries safe even when migration SQL is not perfectly idempotent — the
|
||||
// previous attempt has fully rolled back.
|
||||
//
|
||||
// The body is retried up to migrationRetryAttempts times on transient
|
||||
// conflicts (40P01 deadlock_detected, XX000 "tuple concurrently updated" —
|
||||
// surfaces from concurrent CREATE OR REPLACE in a rolling deploy). ctx
|
||||
// cancellation aborts the retry loop immediately. This is deliberately
|
||||
// narrower than IsRetryableSchemaExecErr — see the note there on why
|
||||
// migrations must not retry away duplicate-object errors.
|
||||
func ApplyMigrationInTx(ctx context.Context, p txBeginner, label string, version int, variants []MigrationVariant) error {
|
||||
if len(variants) == 0 {
|
||||
return fmt.Errorf("%s: ApplyMigrationInTx v%d called with no variants", label, version)
|
||||
}
|
||||
var lastErr error
|
||||
for attempt := 0; attempt < migrationRetryAttempts; attempt++ {
|
||||
err := applyMigrationOnce(ctx, p, label, version, variants)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
var pgErr *pgconn.PgError
|
||||
if errors.As(err, &pgErr) && (pgErr.Code == pgerrcode.DeadlockDetected || pgErr.Code == pgerrcode.InternalError) && attempt < migrationRetryAttempts-1 {
|
||||
lastErr = err
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(migrationRetryBackoff):
|
||||
}
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
const (
|
||||
migrationRetryAttempts = 3
|
||||
migrationRetryBackoff = 200 * time.Millisecond
|
||||
)
|
||||
|
||||
func applyMigrationOnce(ctx context.Context, p txBeginner, label string, version int, variants []MigrationVariant) error {
|
||||
tx, err := p.BeginTx(ctx, pgx.TxOptions{})
|
||||
if err != nil {
|
||||
return fmt.Errorf("%s: begin tx for migration v%d: %w", label, version, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
|
||||
for _, variant := range variants {
|
||||
if _, err := tx.Exec(ctx, variant.SQL); err != nil {
|
||||
return fmt.Errorf("%s: apply migration v%d (target %s): %w", label, version, variant.VersionTable, err)
|
||||
}
|
||||
ct, err := tx.Exec(ctx, fmt.Sprintf(
|
||||
`UPDATE %s SET schema_version = $1 WHERE id = 1`, variant.VersionTable), version)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%s: bump schema_version to v%d (%s): %w", label, version, variant.VersionTable, err)
|
||||
}
|
||||
if ct.RowsAffected() == 0 {
|
||||
return fmt.Errorf("%s: %s has no row id=1 — version tracking is corrupt, manually INSERT (1, %d) and retry", label, variant.VersionTable, version-1)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return fmt.Errorf("%s: commit migration v%d: %w", label, version, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AcquireMigrationLock takes a session-scoped Postgres advisory lock keyed on
|
||||
// `label`. Callers MUST defer the returned release function. While the lock
|
||||
// is held, every concurrent AcquireMigrationLock call against the same
|
||||
// Postgres instance and same label blocks — serialising the migration loop
|
||||
// across nodes during a rolling deploy.
|
||||
//
|
||||
// Implementation: a dedicated connection is checked out of the pool and held
|
||||
// for the lifetime of the lock. Migrations themselves run on different pool
|
||||
// connections; advisory locks are visible globally on the server, so the
|
||||
// lock-holding connection's identity is sufficient to serialise others.
|
||||
//
|
||||
// The release function explicitly executes pg_advisory_unlock BEFORE returning
|
||||
// the connection to the pool, so the next pool user does not inherit the lock.
|
||||
// If unlock fails (e.g. the connection has died), Postgres releases the lock
|
||||
// automatically on session end — no lock is ever stranded across process
|
||||
// crashes.
|
||||
//
|
||||
// The lock ID is derived from a FNV-64 hash of "pgschema/migration:" + label.
|
||||
// Different label strings → different lock IDs → independent serialisation
|
||||
// scopes (e.g. pgmapbroker migrations never block pgstreambroker migrations).
|
||||
func AcquireMigrationLock(ctx context.Context, p *pgxpool.Pool, label string) (release func(), err error) {
|
||||
lockID := migrationLockID(label)
|
||||
conn, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: acquire migration lock connection: %w", label, err)
|
||||
}
|
||||
if _, err := conn.Exec(ctx, `SELECT pg_advisory_lock($1)`, lockID); err != nil {
|
||||
conn.Release()
|
||||
return nil, fmt.Errorf("%s: pg_advisory_lock(%d): %w", label, lockID, err)
|
||||
}
|
||||
return func() {
|
||||
// context.Background() so cleanup runs even when the caller's ctx is
|
||||
// already canceled. Best-effort: a dead connection still gets reaped
|
||||
// server-side, releasing the lock at session end.
|
||||
_, _ = conn.Exec(context.Background(), `SELECT pg_advisory_unlock($1)`, lockID)
|
||||
conn.Release()
|
||||
}, nil
|
||||
}
|
||||
|
||||
// migrationLockID returns a deterministic int64 lock ID for `label`. Different
|
||||
// labels get different IDs so unrelated consumers (e.g. pgmapbroker vs
|
||||
// pgstreambroker) never share a serialisation scope.
|
||||
func migrationLockID(label string) int64 {
|
||||
h := fnv.New64a()
|
||||
_, _ = h.Write([]byte("pgschema/migration:" + label))
|
||||
// Postgres advisory locks take an int64. Cast preserves all 64 bits.
|
||||
return int64(h.Sum64())
|
||||
}
|
||||
|
||||
// SetSchemaVersion updates schema_version=v in every table listed. Returns
|
||||
// an error if any UPDATE fails — used at the end of EnsureSchema to bring
|
||||
// the fresh-install baseline (which the DDL initialised at value 1) up to
|
||||
// the current schemaVersion. Failing softly here would leave a fresh install
|
||||
// stuck at version 1 and re-run the entire migration chain on next start.
|
||||
func SetSchemaVersion(ctx context.Context, q querier, label string, v int, versionTables []string) error {
|
||||
for _, t := range versionTables {
|
||||
if _, err := q.Exec(ctx, fmt.Sprintf(
|
||||
`UPDATE %s SET schema_version = $1 WHERE id = 1`, t), v); err != nil {
|
||||
return fmt.Errorf("%s: set schema_version=%d (%s): %w", label, v, t, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
643
internal/pgschema/pgschema_test.go
Normal file
643
internal/pgschema/pgschema_test.go
Normal file
@@ -0,0 +1,643 @@
|
||||
package pgschema
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgerrcode"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ----- ValidateMigrationMap -----
|
||||
|
||||
func TestValidateMigrationMap_Contiguous_OK(t *testing.T) {
|
||||
require.NotPanics(t, func() {
|
||||
ValidateMigrationMap("test", 3, map[int]string{
|
||||
2: "-- migration v2",
|
||||
3: "-- migration v3",
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateMigrationMap_EmptyAtV1_OK(t *testing.T) {
|
||||
require.NotPanics(t, func() {
|
||||
ValidateMigrationMap("test", 1, map[int]string{})
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateMigrationMap_MissingEntry_Panics(t *testing.T) {
|
||||
require.PanicsWithValue(t,
|
||||
"test: missing schemaMigrations[2] (schemaVersion=3) — every version 2..schemaVersion must have a migration registered",
|
||||
func() {
|
||||
ValidateMigrationMap("test", 3, map[int]string{3: "-- v3 only"})
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateMigrationMap_OutOfRange_Panics(t *testing.T) {
|
||||
// Version 5 is out of range when target is 3.
|
||||
require.Panics(t, func() {
|
||||
ValidateMigrationMap("test", 3, map[int]string{
|
||||
2: "-- v2",
|
||||
3: "-- v3",
|
||||
5: "-- stray v5",
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateMigrationMap_Version1_Panics(t *testing.T) {
|
||||
// Version 1 is the baseline (DDL), not a migration target.
|
||||
require.Panics(t, func() {
|
||||
ValidateMigrationMap("test", 1, map[int]string{1: "-- bad v1"})
|
||||
})
|
||||
}
|
||||
|
||||
// ----- CheckDowngrade -----
|
||||
|
||||
func TestCheckDowngrade_EqualOrLower_OK(t *testing.T) {
|
||||
require.NoError(t, CheckDowngrade("x", 1, 3))
|
||||
require.NoError(t, CheckDowngrade("x", 3, 3))
|
||||
require.NoError(t, CheckDowngrade("x", 0, 3))
|
||||
}
|
||||
|
||||
func TestCheckDowngrade_DBNewer_Rejected(t *testing.T) {
|
||||
err := CheckDowngrade("pgmapbroker", 5, 3)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "schema_version is 5")
|
||||
require.Contains(t, err.Error(), "supports only up to 3")
|
||||
require.Contains(t, err.Error(), "downgrade not supported")
|
||||
}
|
||||
|
||||
// ----- ValidateTablePrefix -----
|
||||
|
||||
func TestValidateTablePrefix(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
prefix string
|
||||
wantErr bool
|
||||
}{
|
||||
{"default", "cf", false},
|
||||
{"underscores", "my_app_v2", false},
|
||||
{"leading underscore", "_cf", false},
|
||||
{"dollar after first", "cf$1", false},
|
||||
// PostgreSQL folds unquoted identifiers, so this yields myapp_* tables
|
||||
// and every generated query folds to match. Deployments run on it.
|
||||
{"uppercase folds, stays valid", "MyApp", false},
|
||||
{"non-ascii letter", "приложение", false},
|
||||
{"empty", "", true},
|
||||
{"leading digit", "1cf", true},
|
||||
{"leading dollar", "$cf", true},
|
||||
{"hyphen", "cf-prod", true},
|
||||
{"space", "cf prod", true},
|
||||
{"dot", "public.cf", true},
|
||||
{"quote", `cf"x`, true},
|
||||
{"semicolon", "cf;DROP TABLE x", true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := ValidateTablePrefix("test", tc.prefix)
|
||||
if tc.wantErr {
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "test:")
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ----- IsRetryableSchemaExecErr / RetrySchemaExec -----
|
||||
|
||||
func TestIsRetryableSchemaExecErr(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
err error
|
||||
want bool
|
||||
}{
|
||||
{"deadlock", &pgconn.PgError{Code: pgerrcode.DeadlockDetected}, true},
|
||||
{"tuple concurrently updated", &pgconn.PgError{Code: pgerrcode.InternalError}, true},
|
||||
{"duplicate table from a concurrent creator", &pgconn.PgError{Code: pgerrcode.DuplicateTable}, true},
|
||||
{"duplicate row type from a concurrent creator", &pgconn.PgError{Code: pgerrcode.DuplicateObject}, true},
|
||||
{"catalog unique violation from a concurrent creator", &pgconn.PgError{Code: pgerrcode.UniqueViolation}, true},
|
||||
{"wrapped", fmt.Errorf("schema exec: %w", &pgconn.PgError{Code: pgerrcode.DuplicateTable}), true},
|
||||
{"permission denied is a real failure", &pgconn.PgError{Code: pgerrcode.InsufficientPrivilege}, false},
|
||||
{"undefined table is a real failure", &pgconn.PgError{Code: pgerrcode.UndefinedTable}, false},
|
||||
{"non-server error", errors.New("connection reset"), false},
|
||||
{"nil", nil, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require.Equal(t, tc.want, IsRetryableSchemaExecErr(tc.err))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrySchemaExec_RetriesConflictThenSucceeds(t *testing.T) {
|
||||
var calls int
|
||||
err := RetrySchemaExec(context.Background(), func(context.Context) error {
|
||||
calls++
|
||||
if calls < 3 {
|
||||
return &pgconn.PgError{Code: pgerrcode.DuplicateTable}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, calls)
|
||||
}
|
||||
|
||||
func TestRetrySchemaExec_RealFailureNotRetried(t *testing.T) {
|
||||
var calls int
|
||||
want := &pgconn.PgError{Code: pgerrcode.InsufficientPrivilege}
|
||||
err := RetrySchemaExec(context.Background(), func(context.Context) error {
|
||||
calls++
|
||||
return want
|
||||
})
|
||||
require.ErrorIs(t, err, want)
|
||||
require.Equal(t, 1, calls, "a permission error must surface on the first attempt")
|
||||
}
|
||||
|
||||
func TestRetrySchemaExec_ExhaustedReturnsLastError(t *testing.T) {
|
||||
var calls int
|
||||
err := RetrySchemaExec(context.Background(), func(context.Context) error {
|
||||
calls++
|
||||
return &pgconn.PgError{Code: pgerrcode.UniqueViolation}
|
||||
})
|
||||
var pgErr *pgconn.PgError
|
||||
require.ErrorAs(t, err, &pgErr)
|
||||
require.Equal(t, pgerrcode.UniqueViolation, pgErr.Code)
|
||||
require.Equal(t, schemaExecAttempts, calls)
|
||||
}
|
||||
|
||||
// A node shutting down mid-retry must report cancellation, not a schema
|
||||
// failure — pgoutbox's Run loop discriminates on it to keep shutdown quiet.
|
||||
func TestRetrySchemaExec_ContextCancelledDuringBackoff(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
var calls int
|
||||
err := RetrySchemaExec(ctx, func(context.Context) error {
|
||||
calls++
|
||||
cancel()
|
||||
return &pgconn.PgError{Code: pgerrcode.DeadlockDetected}
|
||||
})
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
require.Equal(t, 1, calls)
|
||||
}
|
||||
|
||||
func TestSchemaExecBackoff_StaysInsideCap(t *testing.T) {
|
||||
for attempt := range schemaExecAttempts {
|
||||
d := schemaExecBackoff(attempt)
|
||||
require.Positive(t, d)
|
||||
require.LessOrEqual(t, d, schemaExecMaxBackoff)
|
||||
}
|
||||
}
|
||||
|
||||
// ----- Integration helpers -----
|
||||
//
|
||||
// The remaining tests need a real Postgres because the queries hit table-
|
||||
// existence/permission paths that no fake-row implementation can simulate
|
||||
// faithfully. Skip when not configured.
|
||||
|
||||
func getPostgresConnString(tb testing.TB) string {
|
||||
tb.Helper()
|
||||
// Match the env var the broker test suites use so a single local Postgres
|
||||
// runs the whole suite.
|
||||
dsn := os.Getenv("CENTRIFUGE_POSTGRES_URL")
|
||||
if dsn == "" {
|
||||
dsn = "postgres://test:test@localhost:5432/test?sslmode=disable"
|
||||
}
|
||||
// Probe before each test so a missing DB skips, not fails.
|
||||
cfg, err := pgxpool.ParseConfig(dsn)
|
||||
if err != nil {
|
||||
tb.Skipf("invalid PG_TEST_DSN: %v", err)
|
||||
}
|
||||
cfg.MaxConns = 2
|
||||
pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
|
||||
if err != nil {
|
||||
tb.Skipf("postgres not reachable: %v", err)
|
||||
}
|
||||
defer pool.Close()
|
||||
if err := pool.Ping(context.Background()); err != nil {
|
||||
tb.Skipf("postgres ping failed: %v", err)
|
||||
}
|
||||
return dsn
|
||||
}
|
||||
|
||||
// newTestPool gives each test a private prefix to avoid cross-test races.
|
||||
func newTestPool(t *testing.T) (*pgxpool.Pool, string) {
|
||||
t.Helper()
|
||||
dsn := getPostgresConnString(t)
|
||||
pool, err := pgxpool.New(context.Background(), dsn)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(pool.Close)
|
||||
// Use a stable but per-test prefix derived from the test name. Postgres
|
||||
// identifiers max 63 bytes — truncate just in case.
|
||||
prefix := "pgschema_test_" + sanitize(t.Name()) + "_"
|
||||
if len(prefix) > 50 {
|
||||
prefix = prefix[:50] + "_"
|
||||
}
|
||||
cleanup(t, pool, prefix)
|
||||
t.Cleanup(func() { cleanup(t, pool, prefix) })
|
||||
return pool, prefix
|
||||
}
|
||||
|
||||
func sanitize(s string) string {
|
||||
r := strings.NewReplacer("/", "_", " ", "_", "-", "_")
|
||||
return strings.ToLower(r.Replace(s))
|
||||
}
|
||||
|
||||
func cleanup(t *testing.T, pool *pgxpool.Pool, prefix string) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
_, _ = pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %sschema_version CASCADE", prefix))
|
||||
_, _ = pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %sprobe CASCADE", prefix))
|
||||
}
|
||||
|
||||
// createVersionTable creates a `<prefix>schema_version` table seeded at value `v`.
|
||||
func createVersionTable(t *testing.T, pool *pgxpool.Pool, prefix string, v int) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`
|
||||
CREATE TABLE %sschema_version (
|
||||
id INTEGER PRIMARY KEY,
|
||||
schema_version INTEGER NOT NULL
|
||||
)`, prefix))
|
||||
require.NoError(t, err)
|
||||
_, err = pool.Exec(ctx, fmt.Sprintf(
|
||||
`INSERT INTO %sschema_version (id, schema_version) VALUES (1, $1)`, prefix), v)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// ----- ReadSchemaVersion -----
|
||||
|
||||
func TestReadSchemaVersion_RowPresent(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 3)
|
||||
v, isFresh, err := ReadSchemaVersion(context.Background(), pool, prefix+"schema_version")
|
||||
require.NoError(t, err)
|
||||
require.False(t, isFresh)
|
||||
require.Equal(t, 3, v)
|
||||
}
|
||||
|
||||
func TestReadSchemaVersion_TableMissing_TreatedAsFresh(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
// Don't create the table.
|
||||
v, isFresh, err := ReadSchemaVersion(context.Background(), pool, prefix+"schema_version")
|
||||
require.NoError(t, err)
|
||||
require.True(t, isFresh)
|
||||
require.Equal(t, 0, v)
|
||||
}
|
||||
|
||||
func TestReadSchemaVersion_RowMissing_TreatedAsFresh(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
ctx := context.Background()
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`
|
||||
CREATE TABLE %sschema_version (
|
||||
id INTEGER PRIMARY KEY,
|
||||
schema_version INTEGER NOT NULL
|
||||
)`, prefix))
|
||||
require.NoError(t, err)
|
||||
// No INSERT — row missing.
|
||||
v, isFresh, err := ReadSchemaVersion(ctx, pool, prefix+"schema_version")
|
||||
require.NoError(t, err)
|
||||
require.True(t, isFresh)
|
||||
require.Equal(t, 0, v)
|
||||
}
|
||||
|
||||
// TestReadSchemaVersion_ColumnMissing_PropagatesError verifies the critical
|
||||
// safety property: a column-missing error (any error other than 42P01 or
|
||||
// ErrNoRows) must propagate, NOT be silently treated as "fresh install" and
|
||||
// force schema_version forward. Stand-in for the broader class of
|
||||
// transient/permission/timeout failures.
|
||||
func TestReadSchemaVersion_ColumnMissing_PropagatesError(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
ctx := context.Background()
|
||||
// Create the table but with the WRONG column name. Reading `schema_version`
|
||||
// then errors with 42703 (undefined_column) — neither ErrNoRows nor 42P01.
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`
|
||||
CREATE TABLE %sschema_version (
|
||||
id INTEGER PRIMARY KEY,
|
||||
wrong_column INTEGER NOT NULL
|
||||
)`, prefix))
|
||||
require.NoError(t, err)
|
||||
_, err = pool.Exec(ctx, fmt.Sprintf(
|
||||
`INSERT INTO %sschema_version (id, wrong_column) VALUES (1, 3)`, prefix))
|
||||
require.NoError(t, err)
|
||||
|
||||
v, isFresh, err := ReadSchemaVersion(ctx, pool, prefix+"schema_version")
|
||||
require.Error(t, err, "transient/non-fresh errors must propagate")
|
||||
require.False(t, isFresh, "isFresh must be false when an error is returned")
|
||||
require.Equal(t, 0, v)
|
||||
// Verify it surfaces as a pg error (column missing).
|
||||
var pgErr *pgconn.PgError
|
||||
require.True(t, errors.As(err, &pgErr))
|
||||
require.Equal(t, pgerrcode.UndefinedColumn, pgErr.Code)
|
||||
}
|
||||
|
||||
// ----- ApplyMigrationInTx -----
|
||||
|
||||
func TestApplyMigrationInTx_SuccessUpdatesVersion(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 1)
|
||||
ctx := context.Background()
|
||||
|
||||
// Migration adds a column to a probe table.
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`CREATE TABLE %sprobe (id INTEGER)`, prefix))
|
||||
require.NoError(t, err)
|
||||
|
||||
migration := fmt.Sprintf(`ALTER TABLE %sprobe ADD COLUMN extra TEXT`, prefix)
|
||||
|
||||
err = ApplyMigrationInTx(ctx, pool, "test", 2, []MigrationVariant{
|
||||
{SQL: migration, VersionTable: prefix + "schema_version"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify migration applied.
|
||||
var colExists bool
|
||||
err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'extra')`,
|
||||
prefix+"probe").Scan(&colExists)
|
||||
require.NoError(t, err)
|
||||
require.True(t, colExists)
|
||||
|
||||
// Verify schema_version bumped.
|
||||
var v int
|
||||
err = pool.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %sschema_version WHERE id = 1`, prefix)).Scan(&v)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, v)
|
||||
}
|
||||
|
||||
// TestApplyMigrationInTx_FailureRollsBack is the key safety test: a failed
|
||||
// migration must NOT leave schema_version bumped, and must NOT leave partial
|
||||
// schema changes committed. Atomicity is the whole point of running migration
|
||||
// + version bump in one transaction.
|
||||
func TestApplyMigrationInTx_FailureRollsBack(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 1)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`CREATE TABLE %sprobe (id INTEGER)`, prefix))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Migration adds a column AND then runs invalid SQL — the whole tx must roll back.
|
||||
migration := fmt.Sprintf(`
|
||||
ALTER TABLE %sprobe ADD COLUMN extra TEXT;
|
||||
ALTER TABLE %sprobe ADD COLUMN extra TEXT; -- duplicate, fails
|
||||
`, prefix, prefix)
|
||||
|
||||
err = ApplyMigrationInTx(ctx, pool, "test", 2, []MigrationVariant{
|
||||
{SQL: migration, VersionTable: prefix + "schema_version"},
|
||||
})
|
||||
require.Error(t, err)
|
||||
|
||||
// Verify the first ADD COLUMN was rolled back — `extra` must NOT exist.
|
||||
var colExists bool
|
||||
err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'extra')`,
|
||||
prefix+"probe").Scan(&colExists)
|
||||
require.NoError(t, err)
|
||||
require.False(t, colExists, "ADD COLUMN must have rolled back")
|
||||
|
||||
// Verify schema_version unchanged.
|
||||
var v int
|
||||
err = pool.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %sschema_version WHERE id = 1`, prefix)).Scan(&v)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, v, "schema_version must not have advanced after rollback")
|
||||
}
|
||||
|
||||
func TestApplyMigrationInTx_MissingRowFailsLoudly(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
ctx := context.Background()
|
||||
// Create version table but DON'T insert a row.
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`
|
||||
CREATE TABLE %sschema_version (
|
||||
id INTEGER PRIMARY KEY,
|
||||
schema_version INTEGER NOT NULL
|
||||
)`, prefix))
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ApplyMigrationInTx(ctx, pool, "test", 2, []MigrationVariant{
|
||||
{SQL: `SELECT 1`, VersionTable: prefix + "schema_version"},
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "no row id=1")
|
||||
}
|
||||
|
||||
func TestApplyMigrationInTx_MultipleVariants(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 1)
|
||||
prefix2 := prefix + "alt_"
|
||||
t.Cleanup(func() { cleanup(t, pool, prefix2) })
|
||||
createVersionTable(t, pool, prefix2, 1)
|
||||
ctx := context.Background()
|
||||
|
||||
err := ApplyMigrationInTx(ctx, pool, "test", 2, []MigrationVariant{
|
||||
{SQL: `SELECT 1`, VersionTable: prefix + "schema_version"},
|
||||
{SQL: `SELECT 1`, VersionTable: prefix2 + "schema_version"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, p := range []string{prefix, prefix2} {
|
||||
var v int
|
||||
err = pool.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %sschema_version WHERE id = 1`, p)).Scan(&v)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, v, "prefix %s should be at v2", p)
|
||||
}
|
||||
}
|
||||
|
||||
// ----- SetSchemaVersion -----
|
||||
|
||||
func TestSetSchemaVersion_UpdatesAllTables(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 1)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, SetSchemaVersion(ctx, pool, "test", 5, []string{prefix + "schema_version"}))
|
||||
|
||||
var v int
|
||||
err := pool.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %sschema_version WHERE id = 1`, prefix)).Scan(&v)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 5, v)
|
||||
}
|
||||
|
||||
func TestSetSchemaVersion_TableMissing_PropagatesError(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
// Don't create the table.
|
||||
err := SetSchemaVersion(context.Background(), pool, "test", 2, []string{prefix + "schema_version"})
|
||||
require.Error(t, err, "UPDATE against a missing table must propagate (no longer silent)")
|
||||
}
|
||||
|
||||
// TestApplyMigrationInTx_NoVariants returns an error rather than running with
|
||||
// nothing — guards against a programming mistake where a broker forgets to
|
||||
// build the variant slice.
|
||||
func TestApplyMigrationInTx_NoVariants(t *testing.T) {
|
||||
pool, _ := newTestPool(t)
|
||||
err := ApplyMigrationInTx(context.Background(), pool, "test", 2, nil)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "no variants")
|
||||
}
|
||||
|
||||
// TestApplyMigrationInTx_PerVariantSQL verifies the per-variant SQL is run
|
||||
// in the right order against the right tables — distinct migrations for each
|
||||
// variant landed correctly, and the version row of each variant is bumped.
|
||||
func TestApplyMigrationInTx_PerVariantSQL(t *testing.T) {
|
||||
pool, prefix := newTestPool(t)
|
||||
createVersionTable(t, pool, prefix, 1)
|
||||
prefix2 := prefix + "alt_"
|
||||
t.Cleanup(func() { cleanup(t, pool, prefix2) })
|
||||
createVersionTable(t, pool, prefix2, 1)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := pool.Exec(ctx, fmt.Sprintf(`CREATE TABLE %sprobe (id INTEGER)`, prefix))
|
||||
require.NoError(t, err)
|
||||
_, err = pool.Exec(ctx, fmt.Sprintf(`CREATE TABLE %sprobe (id INTEGER)`, prefix2))
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ApplyMigrationInTx(ctx, pool, "test", 2, []MigrationVariant{
|
||||
{
|
||||
SQL: fmt.Sprintf(`ALTER TABLE %sprobe ADD COLUMN col_a TEXT`, prefix),
|
||||
VersionTable: prefix + "schema_version",
|
||||
},
|
||||
{
|
||||
SQL: fmt.Sprintf(`ALTER TABLE %sprobe ADD COLUMN col_b TEXT`, prefix2),
|
||||
VersionTable: prefix2 + "schema_version",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify each variant got its own column.
|
||||
for _, c := range []struct {
|
||||
table, column string
|
||||
}{
|
||||
{prefix + "probe", "col_a"},
|
||||
{prefix2 + "probe", "col_b"},
|
||||
} {
|
||||
var exists bool
|
||||
err = pool.QueryRow(ctx,
|
||||
`SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = $2)`,
|
||||
c.table, c.column,
|
||||
).Scan(&exists)
|
||||
require.NoError(t, err)
|
||||
require.True(t, exists, "column %s should exist on %s", c.column, c.table)
|
||||
}
|
||||
|
||||
// Both version tables bumped.
|
||||
for _, p := range []string{prefix, prefix2} {
|
||||
var v int
|
||||
err = pool.QueryRow(ctx, fmt.Sprintf(`SELECT schema_version FROM %sschema_version WHERE id = 1`, p)).Scan(&v)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, v)
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure we don't accidentally tie the package to pgx-specific row types that
|
||||
// would break the querier abstraction at compile time.
|
||||
var _ = pgx.ErrNoRows
|
||||
|
||||
// ----- migrationLockID -----
|
||||
|
||||
func TestMigrationLockID_Deterministic(t *testing.T) {
|
||||
require.Equal(t, migrationLockID("pgmapbroker"), migrationLockID("pgmapbroker"))
|
||||
}
|
||||
|
||||
func TestMigrationLockID_DifferentLabelsDifferentIDs(t *testing.T) {
|
||||
// Independent broker labels must get independent serialisation scopes.
|
||||
require.NotEqual(t, migrationLockID("pgmapbroker"), migrationLockID("pgstreambroker"))
|
||||
require.NotEqual(t, migrationLockID("pgmapbroker"), migrationLockID("postgres-controller"))
|
||||
}
|
||||
|
||||
// ----- AcquireMigrationLock -----
|
||||
|
||||
// TestAcquireMigrationLock_SerialisesConcurrentCallers proves the core
|
||||
// rolling-deploy guarantee at the pgschema layer: two concurrent
|
||||
// AcquireMigrationLock calls block on each other, only one holds the lock at
|
||||
// a time. Without this serialisation, two nodes upgrading simultaneously
|
||||
// could race through the migration loop in parallel.
|
||||
func TestAcquireMigrationLock_SerialisesConcurrentCallers(t *testing.T) {
|
||||
pool, _ := newTestPool(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// First caller holds the lock.
|
||||
release1, err := AcquireMigrationLock(ctx, pool, "concurrent-test")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Second caller attempts to acquire — must block until release1 is called.
|
||||
acquired := make(chan struct{})
|
||||
go func() {
|
||||
release2, err := AcquireMigrationLock(ctx, pool, "concurrent-test")
|
||||
require.NoError(t, err)
|
||||
close(acquired)
|
||||
release2()
|
||||
}()
|
||||
|
||||
// Give the goroutine a chance to be waiting on pg_advisory_lock.
|
||||
select {
|
||||
case <-acquired:
|
||||
t.Fatal("second caller acquired lock while first still holds it")
|
||||
case <-time.After(150 * time.Millisecond):
|
||||
}
|
||||
|
||||
// Release the first lock; second goroutine must now make progress.
|
||||
release1()
|
||||
select {
|
||||
case <-acquired:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("second caller did not acquire lock after release")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAcquireMigrationLock_DifferentLabelsDoNotBlock verifies the scope of
|
||||
// the lock — two callers with DIFFERENT labels do not block each other (e.g.
|
||||
// pgmapbroker and pgstreambroker migrations can run in parallel).
|
||||
func TestAcquireMigrationLock_DifferentLabelsDoNotBlock(t *testing.T) {
|
||||
pool, _ := newTestPool(t)
|
||||
ctx := context.Background()
|
||||
|
||||
release1, err := AcquireMigrationLock(ctx, pool, "label-a")
|
||||
require.NoError(t, err)
|
||||
defer release1()
|
||||
|
||||
// label-b must acquire immediately — different lock IDs.
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
release2, err := AcquireMigrationLock(ctx, pool, "label-b")
|
||||
require.NoError(t, err)
|
||||
release2()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("label-b should not be blocked by label-a's lock")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAcquireMigrationLock_ReleaseUnlocksPoolConnection verifies the release
|
||||
// function explicitly unlocks BEFORE returning the connection to the pool —
|
||||
// so a subsequent unrelated caller acquiring the same connection doesn't
|
||||
// inherit an active advisory lock.
|
||||
func TestAcquireMigrationLock_ReleaseUnlocksPoolConnection(t *testing.T) {
|
||||
dsn := getPostgresConnString(t)
|
||||
// Single-conn pool so the next Acquire definitely reuses the same conn.
|
||||
cfg, err := pgxpool.ParseConfig(dsn)
|
||||
require.NoError(t, err)
|
||||
cfg.MaxConns = 1
|
||||
pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
|
||||
require.NoError(t, err)
|
||||
defer pool.Close()
|
||||
ctx := context.Background()
|
||||
|
||||
release, err := AcquireMigrationLock(ctx, pool, "release-test")
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
|
||||
// After release, the same conn is back in the pool. A second
|
||||
// AcquireMigrationLock for the same label must not see the lock as held.
|
||||
release2, err := AcquireMigrationLock(ctx, pool, "release-test")
|
||||
require.NoError(t, err)
|
||||
release2()
|
||||
}
|
||||
Reference in New Issue
Block a user