package pg import ( "context" "errors" "hash/fnv" "io" "log/slog" "reflect" "strings" "sync" "testing" "testing/fstest" ) func testLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } // testFS is a three-file migration set in deliberately unsorted map order. func testFS() fstest.MapFS { return fstest.MapFS{ "0002_second.sql": {Data: []byte("CREATE TABLE IF NOT EXISTS b ();")}, "0001_first.sql": {Data: []byte("CREATE TABLE IF NOT EXISTS a ();")}, "0003_third.sql": {Data: []byte("CREATE TABLE IF NOT EXISTS c ();")}, "README.md": {Data: []byte("not a migration")}, } } var testNames = []string{"0001_first.sql", "0002_second.sql", "0003_third.sql"} // fakeDB stands in for Postgres: the advisory lock is a mutex, the version // table an in-memory set, and a migration is "run" by recording its version. type fakeDB struct { lock sync.Mutex // the advisory lock: one holder at a time, cluster-wide mu sync.Mutex held bool tableCreated bool applied []string applyCalls []string bodies []string failOn string unlockErr error discards int violations []string } type fakeMigrator struct{ db *fakeDB } func (f *fakeMigrator) Lock(ctx context.Context) error { f.db.lock.Lock() f.db.mu.Lock() defer f.db.mu.Unlock() f.db.held = true return nil } func (f *fakeMigrator) Unlock(ctx context.Context) error { f.db.mu.Lock() f.db.held = false err := f.db.unlockErr f.db.mu.Unlock() f.db.lock.Unlock() return err } func (f *fakeMigrator) Discard(ctx context.Context) { f.db.mu.Lock() defer f.db.mu.Unlock() f.db.discards++ } // requireHeld records any access made without the advisory lock; every schema // read or write must happen inside the locked section. func (f *fakeMigrator) requireHeld(op string) { if !f.db.held { f.db.violations = append(f.db.violations, op) } } func (f *fakeMigrator) EnsureVersionTable(ctx context.Context) error { f.db.mu.Lock() defer f.db.mu.Unlock() f.requireHeld("EnsureVersionTable") f.db.tableCreated = true return nil } func (f *fakeMigrator) AppliedVersions(ctx context.Context) (map[string]bool, error) { f.db.mu.Lock() defer f.db.mu.Unlock() f.requireHeld("AppliedVersions") if !f.db.tableCreated { return nil, errors.New("schema_migrations does not exist") } out := map[string]bool{} for _, v := range f.db.applied { out[v] = true } return out, nil } func (f *fakeMigrator) Apply(ctx context.Context, version, sql string) error { f.db.mu.Lock() defer f.db.mu.Unlock() f.requireHeld("Apply") f.db.applyCalls = append(f.db.applyCalls, version) f.db.bodies = append(f.db.bodies, sql) if sql == "" { return errors.New("empty migration body") } // A failing migration rolls back, so neither the SQL nor the version row lands. if version == f.db.failOn { return errors.New("boom") } f.db.applied = append(f.db.applied, version) return nil } func (f *fakeDB) snapshot() (applied, calls, violations []string) { f.mu.Lock() defer f.mu.Unlock() return append([]string(nil), f.applied...), append([]string(nil), f.applyCalls...), append([]string(nil), f.violations...) } func (f *fakeDB) discardCount() int { f.mu.Lock() defer f.mu.Unlock() return f.discards } // arrproxy pinned its advisory lock key as a literal derived from // FNV-1a/64("arrproxy-migrations"). LockKey must reproduce it exactly, or // adopting this package would silently stop excluding the old deployment. func TestLockKey_MatchesTheEstatesPinnedKey(t *testing.T) { const arrproxyKey int64 = 7816645656172167846 if got := LockKey("arrproxy-migrations"); got != arrproxyKey { t.Fatalf("LockKey(%q) = %d, want %d", "arrproxy-migrations", got, arrproxyKey) } } func TestLockKey_IsFNV1a64(t *testing.T) { for _, name := range []string{"", "encapi-migrations", "a much longer lock name"} { h := fnv.New64a() if _, err := h.Write([]byte(name)); err != nil { t.Fatalf("hash: %v", err) } if got, want := LockKey(name), int64(h.Sum64()); got != want { t.Errorf("LockKey(%q) = %d, want %d", name, got, want) } } } // Distinct names must not collide, or two services sharing a cluster would // serialise against each other by accident. func TestLockKey_DistinctNamesDistinctKeys(t *testing.T) { if LockKey("encapi-migrations") == LockKey("artifactapi-migrations") { t.Fatal("two different lock names produced the same key") } } func TestMigrationNames_SortedSQLOnly(t *testing.T) { got, err := migrationNames(testFS()) if err != nil { t.Fatalf("migrationNames: %v", err) } if !reflect.DeepEqual(got, testNames) { t.Fatalf("got %v, want %v", got, testNames) } } func TestMigrationNames_IgnoresDirectories(t *testing.T) { fsys := fstest.MapFS{ "0001_first.sql": {Data: []byte("SELECT 1;")}, "sub.sql/keep.sql": {Data: []byte("SELECT 1;")}, "0002_second.sql.gz": {Data: []byte("SELECT 1;")}, } got, err := migrationNames(fsys) if err != nil { t.Fatalf("migrationNames: %v", err) } if want := []string{"0001_first.sql"}; !reflect.DeepEqual(got, want) { t.Fatalf("got %v, want %v", got, want) } } func TestMigrationNames_EmptySetIsAnError(t *testing.T) { if _, err := migrationNames(fstest.MapFS{}); err == nil { t.Fatal("expected an error for an empty migration set") } } func TestMigrationNames_UnreadableFSIsAnError(t *testing.T) { if _, err := migrationNames(brokenFS{}); err == nil { t.Fatal("expected an error when the migration directory cannot be read") } } func TestRunMigrations_FreshDatabaseAppliesAllInOrder(t *testing.T) { db := &fakeDB{} if err := runMigrations(context.Background(), &fakeMigrator{db: db}, testFS(), testLogger()); err != nil { t.Fatalf("runMigrations: %v", err) } applied, _, violations := db.snapshot() if !reflect.DeepEqual(applied, testNames) { t.Fatalf("applied %v, want %v", applied, testNames) } if !db.tableCreated { t.Error("schema_migrations was not created") } if len(violations) != 0 { t.Errorf("schema accessed without the advisory lock: %v", violations) } } // The runner must hand Apply the file's bytes, not just its name. func TestRunMigrations_PassesFileBodies(t *testing.T) { db := &fakeDB{} fsys := testFS() if err := runMigrations(context.Background(), &fakeMigrator{db: db}, fsys, testLogger()); err != nil { t.Fatalf("runMigrations: %v", err) } db.mu.Lock() defer db.mu.Unlock() for i, name := range testNames { if want := string(fsys[name].Data); db.bodies[i] != want { t.Errorf("body %d = %q, want %q", i, db.bodies[i], want) } } } func TestRunMigrations_SecondRunIsANoOp(t *testing.T) { db := &fakeDB{} ctx := context.Background() for i := range 2 { if err := runMigrations(ctx, &fakeMigrator{db: db}, testFS(), testLogger()); err != nil { t.Fatalf("run %d: %v", i, err) } } applied, calls, _ := db.snapshot() if !reflect.DeepEqual(applied, testNames) { t.Fatalf("applied %v, want %v", applied, testNames) } if len(calls) != len(testNames) { t.Fatalf("Apply called %d times across two runs, want %d", len(calls), len(testNames)) } } // A file missing from schema_migrations is re-applied even when the rest of the // set is recorded: that is how a schema applied out of band is adopted. func TestRunMigrations_AppliesOnlyTheMissingVersion(t *testing.T) { db := &fakeDB{tableCreated: true, applied: []string{testNames[0], testNames[2]}} if err := runMigrations(context.Background(), &fakeMigrator{db: db}, testFS(), testLogger()); err != nil { t.Fatalf("runMigrations: %v", err) } _, calls, _ := db.snapshot() if want := testNames[1:2]; !reflect.DeepEqual(calls, want) { t.Fatalf("Apply calls %v, want %v", calls, want) } } // Replicas starting together queue on the advisory lock: the first applies the // set, the rest find it already recorded and do nothing. func TestRunMigrations_ConcurrentStartersSerialize(t *testing.T) { db := &fakeDB{} ctx := context.Background() var wg sync.WaitGroup errs := make([]error, 4) for i := range errs { wg.Add(1) go func() { defer wg.Done() errs[i] = runMigrations(ctx, &fakeMigrator{db: db}, testFS(), testLogger()) }() } wg.Wait() for i, err := range errs { if err != nil { t.Fatalf("starter %d: %v", i, err) } } applied, calls, violations := db.snapshot() if !reflect.DeepEqual(applied, testNames) { t.Fatalf("applied %v, want %v", applied, testNames) } if len(calls) != len(testNames) { t.Fatalf("Apply called %d times, want %d (one starter should have applied)", len(calls), len(testNames)) } if len(violations) != 0 { t.Errorf("schema accessed without the advisory lock: %v", violations) } } // A migration that fails mid-set aborts the run, leaves earlier versions // recorded, and never reaches later ones. func TestRunMigrations_FailureStopsAndKeepsEarlierVersions(t *testing.T) { db := &fakeDB{failOn: testNames[1]} err := runMigrations(context.Background(), &fakeMigrator{db: db}, testFS(), testLogger()) if err == nil { t.Fatal("expected the failing migration to abort the run") } if !strings.Contains(err.Error(), testNames[1]) { t.Errorf("error %q does not name the failing migration", err) } applied, calls, _ := db.snapshot() if !reflect.DeepEqual(applied, testNames[:1]) { t.Fatalf("applied %v, want %v", applied, testNames[:1]) } if !reflect.DeepEqual(calls, testNames[:2]) { t.Fatalf("Apply calls %v, want %v (later migrations must not run)", calls, testNames[:2]) } // The lock must be released even on failure, or every later pod deadlocks. if !db.lock.TryLock() { t.Fatal("advisory lock still held after a failed run") } db.lock.Unlock() } // A run that cannot take the lock must not touch the schema. func TestRunMigrations_LockFailureAborts(t *testing.T) { db := &fakeDB{} m := &failingLockMigrator{fakeMigrator{db: db}} if err := runMigrations(context.Background(), m, testFS(), testLogger()); err == nil { t.Fatal("expected a lock failure to abort the run") } applied, calls, _ := db.snapshot() if len(applied) != 0 || len(calls) != 0 { t.Fatalf("migrations ran without the lock: applied %v, calls %v", applied, calls) } if db.tableCreated { t.Error("schema_migrations was created without the lock") } } // A run that cannot see its bookkeeping table fails instead of re-applying blind. func TestRunMigrations_ReadAppliedFailureAborts(t *testing.T) { // tableCreated stays false because EnsureVersionTable is a no-op here, so // the fake's AppliedVersions reports the table as missing. db := &fakeDB{} m := &noTableMigrator{fakeMigrator{db: db}} err := runMigrations(context.Background(), m, testFS(), testLogger()) if err == nil { t.Fatal("expected the run to fail when applied versions cannot be read") } if !strings.Contains(err.Error(), "read applied migrations") { t.Fatalf("error %q does not identify the failing step", err) } _, calls, _ := db.snapshot() if len(calls) != 0 { t.Fatalf("migrations applied without reading the version table: %v", calls) } } func TestRunMigrations_EnsureVersionTableFailureAborts(t *testing.T) { db := &fakeDB{} m := &failingEnsureMigrator{fakeMigrator{db: db}} err := runMigrations(context.Background(), m, testFS(), testLogger()) if err == nil { t.Fatal("expected the run to fail when schema_migrations cannot be created") } if !strings.Contains(err.Error(), "ensure schema_migrations") { t.Fatalf("error %q does not identify the failing step", err) } // Even this early failure must release the lock. if !db.lock.TryLock() { t.Fatal("advisory lock still held") } db.lock.Unlock() } func TestRunMigrations_UnreadableMigrationAborts(t *testing.T) { db := &fakeDB{} // A directory entry that ReadDir reports as a file but ReadFile cannot open. err := runMigrations(context.Background(), &fakeMigrator{db: db}, missingFileFS{}, testLogger()) if err == nil { t.Fatal("expected an unreadable migration to abort the run") } if !strings.Contains(err.Error(), "read migration") { t.Fatalf("error %q does not identify the failing step", err) } } // An unlock that does not land leaves a session that may still hold the lock, so // the connection is discarded rather than returned to the pool. The migrations // themselves already succeeded, so the run still reports success. func TestRunMigrations_UnlockFailureDiscardsConnection(t *testing.T) { db := &fakeDB{unlockErr: errors.New("connection reset")} if err := runMigrations(context.Background(), &fakeMigrator{db: db}, testFS(), testLogger()); err != nil { t.Fatalf("runMigrations: %v", err) } if got := db.discardCount(); got != 1 { t.Fatalf("Discard called %d times after a failed unlock, want 1", got) } applied, _, _ := db.snapshot() if !reflect.DeepEqual(applied, testNames) { t.Fatalf("applied %v, want %v", applied, testNames) } } // A clean unlock keeps the connection: discarding every migration connection // would churn the pool on every start. func TestRunMigrations_CleanUnlockKeepsConnection(t *testing.T) { db := &fakeDB{} if err := runMigrations(context.Background(), &fakeMigrator{db: db}, testFS(), testLogger()); err != nil { t.Fatalf("runMigrations: %v", err) } if got := db.discardCount(); got != 0 { t.Fatalf("Discard called %d times after a clean unlock, want 0", got) } } type failingLockMigrator struct{ fakeMigrator } func (f *failingLockMigrator) Lock(ctx context.Context) error { return errors.New("lock unavailable") } type failingEnsureMigrator struct{ fakeMigrator } func (f *failingEnsureMigrator) EnsureVersionTable(ctx context.Context) error { return errors.New("permission denied") } type noTableMigrator struct{ fakeMigrator } // EnsureVersionTable silently does nothing, so AppliedVersions then fails. func (f *noTableMigrator) EnsureVersionTable(ctx context.Context) error { return nil }