From bcdf9ae274a81f0f54fef592dd39855d1b868354 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Mon, 22 Jun 2026 13:51:50 -0700 Subject: [PATCH 1/7] Switch migrator to orchestrator and manage downstream migrations Signed-off-by: Jeremy Alvis --- go/core/internal/dbtest/dbtest.go | 2 +- go/core/pkg/app/app.go | 27 +- go/core/pkg/migrations/runner.go | 429 ++++++++++++++++++++++---- go/core/pkg/migrations/runner_test.go | 406 ++++++++++++++++++++---- 4 files changed, 711 insertions(+), 153 deletions(-) diff --git a/go/core/internal/dbtest/dbtest.go b/go/core/internal/dbtest/dbtest.go index b4c4a77ecf..0a3e7d7236 100644 --- a/go/core/internal/dbtest/dbtest.go +++ b/go/core/internal/dbtest/dbtest.go @@ -64,7 +64,7 @@ func StartT(ctx context.Context, t *testing.T) string { // If vectorEnabled is true the vector pass is also applied. // Use MigrateT in tests that have a *testing.T; use Migrate in TestMain where no T is available. func Migrate(connStr string, vectorEnabled bool) error { - return migrations.RunUp(connStr, migrations.FS, vectorEnabled) + return migrations.RunUp(context.Background(), connStr, migrations.BuiltinSources(vectorEnabled)) } // MigrateT runs the embedded migrations against connStr and calls t.Fatal on error. diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 3ee2ce8cb5..9e2802ee34 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -309,17 +309,11 @@ type ExtensionConfig struct { type GetExtensionConfig func(bootstrap BootstrapConfig) (*ExtensionConfig, error) -// MigrationRunner applies database migrations given the resolved connection URL. -// vectorEnabled mirrors the --database-vector-enabled flag; custom runners can use it -// to conditionally apply vector-specific migrations. -// Returning a non-nil error causes the app to exit. -// -// Pass nil to Start to use the default migration runner (migrations.RunUp with migrations.FS). -// Provide a custom runner to take over the migration process entirely. -// Custom runners that want to include the built-in migrations can call migrations.RunUp directly. -type MigrationRunner func(ctx context.Context, url string, vectorEnabled bool) error - -func Start(getExtensionConfig GetExtensionConfig, migrationRunner MigrationRunner) { +// Start boots the controller. extraSources registers additional migration +// tracks beyond the built-in sources; they are applied after the built-in +// (core, vector) tracks, in the order given. Pass nil to run only the built-in +// migrations. +func Start(getExtensionConfig GetExtensionConfig, extraSources []migrations.Source) { var tlsOpts []func(*tls.Config) var cfg Config @@ -483,16 +477,11 @@ func Start(getExtensionConfig GetExtensionConfig, migrationRunner MigrationRunne os.Exit(1) } - // Use the built-in migration runner when none is provided. - if migrationRunner == nil { - migrationRunner = func(_ context.Context, url string, vectorEnabled bool) error { - return migrations.RunUp(url, migrations.FS, vectorEnabled) - } - } - // Run migrations before connecting; schema must exist before queries. + // Built-in sources run first, then any downstream-registered extras. setupLog.Info("running database migrations") - if err := migrationRunner(ctx, dbURL, cfg.Database.VectorEnabled); err != nil { + sources := append(migrations.BuiltinSources(cfg.Database.VectorEnabled), extraSources...) + if err := migrations.RunUp(ctx, dbURL, sources); err != nil { setupLog.Error(err, "database migration failed") os.Exit(1) } diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index cf71be7eb1..7fa6603140 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -1,10 +1,14 @@ package migrations import ( + "context" "database/sql" "errors" "fmt" "io/fs" + nurl "net/url" + "regexp" + "slices" "strings" "github.com/golang-migrate/migrate/v4" @@ -16,51 +20,254 @@ import ( var log = ctrl.Log.WithName("migrations") -// RunUp applies all pending migrations for the given FS. -// vectorEnabled controls whether the vector track is also applied. -// Returns an error if any track fails (and attempts rollback of previously applied tracks). -func RunUp(url string, migrationsFS fs.FS, vectorEnabled bool) error { +// Source describes one migration track for the orchestrator to apply. Downstream +// consumers register their own Sources alongside the built-in ones rather than +// owning a runner, so track ordering and failure handling stay centralized. +type Source struct { + // Name labels the source in logs and errors (e.g. "core", "vector"). + Name string + // Schema is the Postgres schema the track lives in. Empty means the + // connection's default schema (resolved via search_path / current_schema), + // which is what the built-in tracks use. A non-empty value scopes the + // tracking table and migration objects to that schema: the orchestrator + // creates it (CREATE SCHEMA IF NOT EXISTS) and sets search_path on the + // connection. Schema is treated as an untrusted identifier and validated. + // + // When Schema is empty (the built-in tracks), the DSN is left untouched: + // the connection keeps the server's default search_path ("$user", public), so + // migration objects and the tracking table land in public exactly as before + // this orchestrator existed. The schema handling below applies only when + // Schema is set. + // + // When Schema is set, search_path is pinned to it alone (pg_catalog is always + // implicitly searched, so built-in types and functions still resolve). public + // is NOT on the path, which keeps a schema-scoped track strictly isolated. A + // migration that needs a shared extension installed in public (e.g. the + // pgvector "vector" type) must therefore either install/relocate the extension + // into this schema or schema-qualify the reference; it cannot rely on public. + // + // Do not register two sources whose schemas resolve to the same value with the + // same TrackingTable — e.g. one source with Schema == "" and another naming the + // connection's current_schema() explicitly. They would share one tracking-table + // row and one advisory lock and corrupt each other's state. The collision unit + // is (resolved schema, TrackingTable). validateSources catches collisions on the + // literal Schema; RunUp additionally resolves "" to current_schema() and rejects + // collisions that only appear after resolution. + Schema string + // TrackingTable is the golang-migrate bookkeeping table for this track. + TrackingTable string + // FS holds the embedded migration files. + FS fs.FS + // Dir is the subdirectory within FS that holds this track's files. + Dir string + // PreCheck, if set, runs before any source is applied. A non-nil error + // aborts the whole run before any migration executes (fail-fast). + PreCheck func(url string) error +} + +// BuiltinSources returns the built-in source set: the core track always, and the +// vector track when vectorEnabled. Downstream consumers append their own Sources +// to this slice before calling RunUp. +func BuiltinSources(vectorEnabled bool) []Source { + sources := []Source{{ + Name: "core", + TrackingTable: "schema_migrations", + FS: FS, + Dir: "core", + }} if vectorEnabled { - if err := checkPgvector(url); err != nil { - return fmt.Errorf("vector migrations require pgvector: %w", err) + sources = append(sources, Source{ + Name: "vector", + TrackingTable: "vector_schema_migrations", + FS: FS, + Dir: "vector", + PreCheck: checkPgvector, + }) + } + return sources +} + +// RunUp applies all pending migrations for each source, in slice order. +// +// All PreChecks run first, before any source is applied, so a failed precheck +// aborts the run before touching the database. Each source is then applied with +// the same per-track safety behavior: it tolerates a database ahead of this +// binary (compatibility mode), refuses a dirty-and-ahead database, and rolls +// itself back if its own Up fails. If a later source fails, previously-applied +// sources are rolled back to their pre-run versions in reverse order. +// +// ctx is honored at source boundaries and during schema setup and prechecks. +// golang-migrate's apply is not context-aware, so an in-flight migration is not +// cancellable. +func RunUp(ctx context.Context, url string, sources []Source) error { + if len(sources) == 0 { + return nil + } + if err := validateSources(sources); err != nil { + return err + } + // Catch collisions that only appear once "" schemas resolve to current_schema(). + if err := checkResolvedSchemaCollisions(ctx, url, sources); err != nil { + return err + } + + // Run every precheck up front so a failure aborts before any source applies + // (e.g. pgvector is verified before the core track runs). + for _, src := range sources { + if src.PreCheck == nil { + continue + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("migrations cancelled before %s precheck: %w", src.Name, err) + } + if err := src.PreCheck(url); err != nil { + return fmt.Errorf("%s precheck: %w", src.Name, err) } } - corePrev, err := applyDir(url, migrationsFS, "core", "schema_migrations") - if err != nil { - return fmt.Errorf("core migrations: %w", err) + type applied struct { + src Source + prev uint } + var done []applied - if vectorEnabled { - if _, err := applyDir(url, migrationsFS, "vector", "vector_schema_migrations"); err != nil { - if corePrev == 0 { - log.Info("vector migration failed; skipping core rollback to version 0 to protect pre-existing data") - } else { - log.Info("rolling back core after vector failure", "targetVersion", corePrev) - rollbackDir(url, migrationsFS, "core", "schema_migrations", corePrev) + for _, src := range sources { + if err := ctx.Err(); err != nil { + return fmt.Errorf("migrations cancelled before %s: %w", src.Name, err) + } + prev, err := applySource(ctx, url, src) + if err != nil { + // Compensating rollback: undo previously-applied sources in reverse + // order, each to its own pre-run version. The failing source has + // already rolled itself back in applySource. + var compErrs []error + for _, a := range slices.Backward(done) { + if a.prev == 0 { + log.Info("skipping compensating rollback to version 0 to protect pre-existing data", "source", a.src.Name) + continue + } + log.Info("rolling back source after later failure", "source", a.src.Name, "targetVersion", a.prev) + if rbErr := rollbackSource(ctx, url, a.src, a.prev); rbErr != nil { + compErrs = append(compErrs, rbErr) + } + } + runErr := fmt.Errorf("%s migrations: %w", src.Name, err) + if len(compErrs) > 0 { + // A compensating rollback failed: the database may be left in a + // partially rolled-back state. Surface it alongside the original + // failure rather than only in the logs. + return errors.Join(append([]error{runErr}, compErrs...)...) } - return fmt.Errorf("vector migrations: %w", err) + return runErr } + done = append(done, applied{src: src, prev: prev}) } return nil } -// applyDir runs Up for dir and rolls back on failure. If prevVersion is 0 -// (no migrations have ever been applied), rollback is skipped to avoid dropping -// pre-existing tables on a GORM-to-golang-migrate upgrade. -// It returns the pre-run version so the caller can roll back this track if a later track fails. -func applyDir(url string, migrationsFS fs.FS, dir, migrationsTable string) (prevVersion uint, err error) { - mg, err := newMigrate(url, migrationsFS, dir, migrationsTable) +// validateSources rejects two sources that share the same (schema, tracking +// table). That pair is the collision unit: it is what determines the +// golang-migrate version row and the advisory-lock id, so two such sources would +// fight over one row and lock regardless of their Dir/FS. Distinct tracking +// tables in the same schema are fine (each gets its own bookkeeping). +// +// This keys on the *literal* Schema string and is intentionally DB-free so it +// stays unit-testable. It therefore cannot see a collision between a source with +// Schema == "" and one naming the connection's current_schema() explicitly, since +// "" resolves to that schema only at connection time. checkResolvedSchemaCollisions +// (called from RunUp, where a connection is available) closes that gap. +func validateSources(sources []Source) error { + seen := make(map[string]string, len(sources)) + for _, s := range sources { + // Validate schema names up front so a bad name on a later source aborts + // the whole run before any earlier source applies, matching the fail-fast + // guarantee RunUp gives prechecks. newMigrate re-validates as a safety net + // for callers that bypass RunUp (e.g. applySource directly in tests). + if s.Schema != "" { + if err := validateSchemaName(s.Schema); err != nil { + return err + } + } + key := s.Schema + "\x00" + s.TrackingTable + if other, ok := seen[key]; ok { + return fmt.Errorf("sources %q and %q share tracking table %q in schema %q", other, s.Name, s.TrackingTable, s.Schema) + } + seen[key] = s.Name + } + return nil +} + +// checkResolvedSchemaCollisions rejects sources that collide only after their +// schemas are resolved — a source with Schema == "" and another naming the +// connection's current_schema() explicitly both land in the same schema, so with +// the same TrackingTable they would share one golang-migrate version row and one +// advisory lock. validateSources cannot see this because it keys on the literal +// Schema; resolving "" requires a connection, which is why this lives here. +// +// The query runs only when the source set mixes empty and explicit schemas; +// all-empty (the built-in core+vector default) and all-explicit sets are already fully +// covered by validateSources' literal-key check, so the common paths pay nothing. +func checkResolvedSchemaCollisions(ctx context.Context, url string, sources []Source) error { + var hasDefault, hasExplicit bool + for _, s := range sources { + if s.Schema == "" { + hasDefault = true + } else { + hasExplicit = true + } + } + if !hasDefault || !hasExplicit { + return nil + } + + db, err := sql.Open("pgx", url) + if err != nil { + return fmt.Errorf("open database to resolve default schema: %w", err) + } + defer db.Close() + + var current sql.NullString + if err := db.QueryRowContext(ctx, "SELECT current_schema()").Scan(¤t); err != nil { + return fmt.Errorf("resolve default schema: %w", err) + } + if !current.Valid { + // search_path resolves to no existing schema; "" sources have no resolved + // schema to collide on yet (their tables would fail to create later anyway). + return nil + } + + seen := make(map[string]string, len(sources)) + for _, s := range sources { + schema := s.Schema + if schema == "" { + schema = current.String + } + key := schema + "\x00" + s.TrackingTable + if other, ok := seen[key]; ok { + return fmt.Errorf("sources %q and %q resolve to the same tracking table %q in schema %q (an empty Schema resolves to current_schema() = %q)", + other, s.Name, s.TrackingTable, schema, current.String) + } + seen[key] = s.Name + } + return nil +} + +// applySource runs Up for one source and rolls it back on failure. If prevVersion +// is 0 (no migrations have ever been applied) rollback is skipped to avoid +// dropping pre-existing tables on a GORM-to-golang-migrate upgrade. It returns +// the pre-run version so the caller can compensate this source if a later one fails. +func applySource(ctx context.Context, url string, src Source) (prevVersion uint, err error) { + mg, err := newMigrate(ctx, url, src) if err != nil { return 0, err } - defer closeMigrate(dir, mg) + defer closeMigrate(src.Name, mg) var dirty bool prevVersion, dirty, err = mg.Version() if err != nil && !errors.Is(err, migrate.ErrNilVersion) { - return 0, fmt.Errorf("get pre-migration version for %s: %w", dir, err) + return 0, fmt.Errorf("get pre-migration version for %s: %w", src.Name, err) } // prevVersion == 0 when ErrNilVersion (no migrations applied yet). @@ -72,8 +279,8 @@ func applyDir(url string, migrationsFS fs.FS, dir, migrationsTable string) (prev // numbers don't align with release versions. // A dirty database is excluded: dirty state means a previous migration // attempt failed and must be resolved, not silently accepted. - if maxVer, scanErr := maxEmbeddedVersion(migrationsFS, dir); scanErr != nil { - log.Error(scanErr, "could not determine max embedded migration version; proceeding with Up", "track", dir) + if maxVer, scanErr := maxEmbeddedVersion(src.FS, src.Dir); scanErr != nil { + log.Error(scanErr, "could not determine max embedded migration version; proceeding with Up", "track", src.Name) } else if prevVersion > maxVer { if dirty { // DB is both dirty and ahead of this binary. Attempting Up/rollback would @@ -81,10 +288,10 @@ func applyDir(url string, migrationsFS fs.FS, dir, migrationsTable string) (prev // and misleading logs. Return a clear error so operators act on the real // problem rather than chasing rollback noise. return prevVersion, fmt.Errorf("database is dirty at version %d and ahead of this binary's max known version %d for track %s: manual operator intervention required: %w", - prevVersion, maxVer, dir, migrate.ErrDirty{Version: int(prevVersion)}) + prevVersion, maxVer, src.Name, migrate.ErrDirty{Version: int(prevVersion)}) } log.Info("database schema is ahead of this binary; running in compatibility mode", - "track", dir, "dbVersion", prevVersion, "binaryMax", maxVer) + "track", src.Name, "dbVersion", prevVersion, "binaryMax", maxVer) return prevVersion, nil } @@ -93,46 +300,50 @@ func applyDir(url string, migrationsFS fs.FS, dir, migrationsTable string) (prev return prevVersion, nil } if prevVersion == 0 { - log.Info("migration failed; skipping rollback to version 0 to protect pre-existing data", "track", dir) + log.Info("migration failed; skipping rollback to version 0 to protect pre-existing data", "track", src.Name) } else { - log.Info("migration failed, attempting rollback", "track", dir, "targetVersion", prevVersion) - if rbErr := rollbackToVersion(mg, dir, prevVersion); rbErr != nil { - log.Error(rbErr, "rollback failed", "track", dir) + log.Info("migration failed, attempting rollback", "track", src.Name, "targetVersion", prevVersion) + if rbErr := rollbackToVersion(mg, src.Name, prevVersion); rbErr != nil { + log.Error(rbErr, "rollback failed", "track", src.Name) } else { - log.Info("rollback complete", "track", dir, "version", prevVersion) + log.Info("rollback complete", "track", src.Name, "version", prevVersion) } } - return prevVersion, fmt.Errorf("run migrations for %s: %w", dir, upErr) + return prevVersion, fmt.Errorf("run migrations for %s: %w", src.Name, upErr) } return prevVersion, nil } -// rollbackDir opens a fresh migrate instance and rolls dir back to targetVersion. -// Used to roll back a previously-succeeded track when a later track fails. -func rollbackDir(url string, migrationsFS fs.FS, dir, migrationsTable string, targetVersion uint) { - mg, err := newMigrate(url, migrationsFS, dir, migrationsTable) +// rollbackSource opens a fresh migrate instance and rolls a source back to +// targetVersion. Used to compensate a previously-succeeded source when a later +// source fails. It returns an error (also logged) when the rollback fails, so +// the orchestrator can surface that the database may be left partially rolled +// back rather than leaving it as a log-only signal. +func rollbackSource(ctx context.Context, url string, src Source, targetVersion uint) error { + mg, err := newMigrate(ctx, url, src) if err != nil { - log.Error(err, "rollback failed (open)", "track", dir) - return + log.Error(err, "rollback failed (open)", "track", src.Name) + return fmt.Errorf("open %s for rollback: %w", src.Name, err) } - defer closeMigrate(dir, mg) - if err := rollbackToVersion(mg, dir, targetVersion); err != nil { - log.Error(err, "rollback failed", "track", dir) - } else { - log.Info("rollback complete", "track", dir, "version", targetVersion) + defer closeMigrate(src.Name, mg) + if err := rollbackToVersion(mg, src.Name, targetVersion); err != nil { + log.Error(err, "rollback failed", "track", src.Name) + return fmt.Errorf("roll back %s to version %d: %w", src.Name, targetVersion, err) } + log.Info("rollback complete", "track", src.Name, "version", targetVersion) + return nil } // rollbackToVersion rolls the migration state back to targetVersion. // It handles the dirty-state cleanup golang-migrate requires after a failed // Up run before down steps can be applied. -func rollbackToVersion(mg *migrate.Migrate, dir string, targetVersion uint) error { +func rollbackToVersion(mg *migrate.Migrate, name string, targetVersion uint) error { currentVersion, dirty, err := mg.Version() if err != nil { if errors.Is(err, migrate.ErrNilVersion) { return nil // nothing was applied; nothing to roll back } - return fmt.Errorf("get version after failure for %s: %w", dir, err) + return fmt.Errorf("get version after failure for %s: %w", name, err) } if dirty { @@ -144,7 +355,7 @@ func rollbackToVersion(mg *migrate.Migrate, dir string, targetVersion uint) erro forceTarget = -1 // negative tells golang-migrate to remove the version record entirely } if err := mg.Force(forceTarget); err != nil { - return fmt.Errorf("clear dirty state for %s: %w", dir, err) + return fmt.Errorf("clear dirty state for %s: %w", name, err) } if forceTarget < 0 { return nil // first migration failed and was cleared; nothing left to roll back @@ -157,7 +368,7 @@ func rollbackToVersion(mg *migrate.Migrate, dir string, targetVersion uint) erro return nil } if err := mg.Steps(-steps); err != nil && !errors.Is(err, migrate.ErrNoChange) { - return fmt.Errorf("roll back %d step(s) for %s: %w", steps, dir, err) + return fmt.Errorf("roll back %d step(s) for %s: %w", steps, name, err) } return nil } @@ -183,35 +394,117 @@ func checkPgvector(url string) error { return nil } -// newMigrate opens a dedicated database connection and constructs a migrate.Migrate -// for the given dir/table. The caller must call closeMigrate when done. -// Uses sql.Open (pgx stdlib shim) — a single dedicated connection — not a pool, -// because the advisory lock is session-level and must not be shared. -func newMigrate(url string, migrationsFS fs.FS, dir, migrationsTable string) (*migrate.Migrate, error) { - db, err := sql.Open("pgx", url) +// newMigrate opens a database handle (sql.Open with the pgx stdlib shim) and +// constructs a migrate.Migrate for the given source. The caller must call +// closeMigrate when done. +// +// sql.Open returns a *sql.DB pool, but the session-scoped advisory lock and the +// migration run are safe regardless: the migratepgx driver checks out a single +// dedicated *sql.Conn and pins all lock/migration work to it. The schema setup +// below (CREATE SCHEMA, and the driver's current_schema() probe) may run on any +// pooled connection, which is fine because the schema name is quoted and +// search_path is set on the DSN — so every pooled connection targets the same +// schema. The orchestrator deliberately does not cap the pool to one connection. +// +// When src.Schema is set, the connection's search_path is pinned to that schema +// (so migration DDL lands there) and the schema is created if missing. The +// tracking table is also scoped to the schema via migratepgx.Config.SchemaName. +func newMigrate(ctx context.Context, dbURL string, src Source) (*migrate.Migrate, error) { + connURL := dbURL + if src.Schema != "" { + if err := validateSchemaName(src.Schema); err != nil { + return nil, err + } + var err error + connURL, err = withSearchPath(dbURL, src.Schema) + if err != nil { + return nil, fmt.Errorf("set search_path for %s: %w", src.Name, err) + } + } + + db, err := sql.Open("pgx", connURL) if err != nil { - return nil, fmt.Errorf("open database for %s: %w", dir, err) + return nil, fmt.Errorf("open database for %s: %w", src.Name, err) } - src, err := iofs.New(migrationsFS, dir) + if src.Schema != "" { + if _, err := db.ExecContext(ctx, "CREATE SCHEMA IF NOT EXISTS "+quoteIdentifier(src.Schema)); err != nil { + _ = db.Close() + return nil, fmt.Errorf("create schema %q for %s: %w", src.Schema, src.Name, err) + } + } + + srcDriver, err := iofs.New(src.FS, src.Dir) if err != nil { - return nil, fmt.Errorf("load migration files from %s: %w", dir, err) + _ = db.Close() + return nil, fmt.Errorf("load migration files from %s: %w", src.Dir, err) } - driver, err := migratepgx.WithInstance(db, &migratepgx.Config{ - MigrationsTable: migrationsTable, - }) + cfg := &migratepgx.Config{MigrationsTable: src.TrackingTable} + if src.Schema != "" { + cfg.SchemaName = src.Schema + } + driver, err := migratepgx.WithInstance(db, cfg) if err != nil { - return nil, fmt.Errorf("create migration driver for %s: %w", dir, err) + _ = db.Close() + return nil, fmt.Errorf("create migration driver for %s: %w", src.Name, err) } - mg, err := migrate.NewWithInstance("iofs", src, "postgres", driver) + mg, err := migrate.NewWithInstance("iofs", srcDriver, "postgres", driver) if err != nil { - return nil, fmt.Errorf("create migrator for %s: %w", dir, err) + _ = srcDriver.Close() + _ = db.Close() + return nil, fmt.Errorf("create migrator for %s: %w", src.Name, err) } return mg, nil } +// withSearchPath returns dbURL with the search_path connection parameter set to +// schema, so every connection in the pool (including the one golang-migrate +// checks out) targets that schema for migration DDL. dbURL must be a URL-form +// DSN (postgres://...); a libpq keyword/value DSN is rejected because net/url +// parses it without error into a meaningless URL, which would silently corrupt +// the connection string rather than fail here. +func withSearchPath(dbURL, schema string) (string, error) { + u, err := nurl.Parse(dbURL) + if err != nil { + return "", fmt.Errorf("parse database url: %w", err) + } + if u.Scheme == "" { + return "", fmt.Errorf("database url must be a URL-form DSN (postgres://...) to scope a schema; got a non-URL DSN") + } + q := u.Query() + q.Set("search_path", schema) + u.RawQuery = q.Encode() + return u.String(), nil +} + +// schemaNameRe constrains a schema name to a lowercase identifier. The name is +// used both quoted (CREATE SCHEMA, the tracking table's SchemaName) and unquoted +// (the search_path connection parameter); Postgres case-folds the unquoted form, +// so a mixed-case name like "MySchema" would create the quoted schema "MySchema" +// while search_path resolved to the folded "myschema" and never matched. Keeping +// the name lowercase makes the two forms identical. +var schemaNameRe = regexp.MustCompile(`^[a-z_][a-z0-9_]*$`) + +// validateSchemaName rejects schema identifiers that are not safe to interpolate +// into DDL. Schema names come from downstream-registered Sources and cannot be +// passed as bind parameters, so they are constrained to a conservative pattern. +func validateSchemaName(schema string) error { + if len(schema) == 0 || len(schema) > 63 { + return fmt.Errorf("invalid schema name %q: must be 1-63 characters", schema) + } + if !schemaNameRe.MatchString(schema) { + return fmt.Errorf("invalid schema name %q: must match %s", schema, schemaNameRe.String()) + } + return nil +} + +// quoteIdentifier double-quotes a SQL identifier, escaping embedded quotes. +func quoteIdentifier(id string) string { + return `"` + strings.ReplaceAll(id, `"`, `""`) + `"` +} + // maxEmbeddedVersion scans dir inside migrationsFS and returns the highest migration // version number found. Only files with a ".up.sql" suffix are considered. Version // numbers are parsed from the leading decimal digits of each filename; the remainder @@ -248,12 +541,12 @@ func maxEmbeddedVersion(migrationsFS fs.FS, dir string) (uint, error) { } // closeMigrate closes mg, logging source and database close errors separately. -func closeMigrate(dir string, mg *migrate.Migrate) { +func closeMigrate(name string, mg *migrate.Migrate) { srcErr, dbErr := mg.Close() if srcErr != nil { - log.Error(srcErr, "closing migration source", "track", dir) + log.Error(srcErr, "closing migration source", "track", name) } if dbErr != nil { - log.Error(dbErr, "closing migration database", "track", dir) + log.Error(dbErr, "closing migration database", "track", name) } } diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 2418ec633d..f9b595fccb 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "maps" + "strings" "testing" "testing/fstest" "time" @@ -70,6 +71,12 @@ var failVectorWithDependencyFS = fstest.MapFS{ "vector/000001_bad_depends_on_core.down.sql": {Data: []byte(`ALTER TABLE shared_data DROP COLUMN IF EXISTS vec_col;`)}, } +// goodVectorFS has a valid vector migration. +var goodVectorFS = fstest.MapFS{ + "vector/000001_create.up.sql": {Data: []byte(`CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE IF NOT EXISTS vec_test (id SERIAL PRIMARY KEY, embedding vector(3));`)}, + "vector/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS vec_test; DROP EXTENSION IF EXISTS vector;`)}, +} + // mergeFS combines multiple MapFS values into one. func mergeFS(fsMaps ...fstest.MapFS) fstest.MapFS { out := fstest.MapFS{} @@ -79,6 +86,17 @@ func mergeFS(fsMaps ...fstest.MapFS) fstest.MapFS { return out } +// coreSource builds a core Source backed by fsys (subdir "core"). +func coreSource(fsys fstest.MapFS) Source { + return Source{Name: "core", TrackingTable: "schema_migrations", FS: fsys, Dir: "core"} +} + +// vectorSource builds a vector Source backed by fsys (subdir "vector"), with the +// pgvector precheck wired in to mirror BuiltinSources. +func vectorSource(fsys fstest.MapFS) Source { + return Source{Name: "vector", TrackingTable: "vector_schema_migrations", FS: fsys, Dir: "vector", PreCheck: checkPgvector} +} + // trackVersion reads the current version from a golang-migrate tracking table. // Returns 0 if the table is empty or does not exist (fully rolled back). func trackVersion(t *testing.T, connStr, table string) uint { @@ -128,12 +146,6 @@ func startTestDB(t *testing.T) string { return connStr } -// goodVectorFS has a valid vector migration. -var goodVectorFS = fstest.MapFS{ - "vector/000001_create.up.sql": {Data: []byte(`CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE IF NOT EXISTS vec_test (id SERIAL PRIMARY KEY, embedding vector(3));`)}, - "vector/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS vec_test; DROP EXTENSION IF EXISTS vector;`)}, -} - // startTestDBWithoutPgvector spins up a plain Postgres container (no pgvector) // and returns its connection string, registering cleanup with t. func startTestDBWithoutPgvector(t *testing.T) string { @@ -167,30 +179,36 @@ func startTestDBWithoutPgvector(t *testing.T) string { // tableExists checks whether a table exists in the public schema. func tableExists(t *testing.T, connStr, table string) bool { + t.Helper() + return tableExistsInSchema(t, connStr, "public", table) +} + +// tableExistsInSchema checks whether a table exists in the given schema. +func tableExistsInSchema(t *testing.T, connStr, schema, table string) bool { t.Helper() db, err := sql.Open("pgx", connStr) if err != nil { - t.Fatalf("tableExists: open db: %v", err) + t.Fatalf("tableExistsInSchema: open db: %v", err) } defer db.Close() var exists bool err = db.QueryRowContext(context.Background(), - "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = $1)", - table).Scan(&exists) + "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2)", + schema, table).Scan(&exists) if err != nil { - t.Fatalf("tableExists: query: %v", err) + t.Fatalf("tableExistsInSchema: query: %v", err) } return exists } -// --- applyDir tests --- +// --- applySource tests --- -func TestApplyDir_HappyPath(t *testing.T) { +func TestApplySource_HappyPath(t *testing.T) { connStr := startTestDB(t) - prev, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations") + prev, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)) if err != nil { - t.Fatalf("applyDir: %v", err) + t.Fatalf("applySource: %v", err) } if prev != 0 { t.Errorf("prevVersion = %d, want 0", prev) @@ -200,13 +218,13 @@ func TestApplyDir_HappyPath(t *testing.T) { } } -func TestApplyDir_NoOpWhenAlreadyAtLatest(t *testing.T) { +func TestApplySource_NoOpWhenAlreadyAtLatest(t *testing.T) { connStr := startTestDB(t) - if _, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)); err != nil { t.Fatalf("first apply: %v", err) } - prev, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations") + prev, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)) if err != nil { t.Fatalf("second apply: %v", err) } @@ -218,10 +236,10 @@ func TestApplyDir_NoOpWhenAlreadyAtLatest(t *testing.T) { } } -func TestApplyDir_NoRollbackWhenFirstMigrationFails(t *testing.T) { +func TestApplySource_NoRollbackWhenFirstMigrationFails(t *testing.T) { connStr := startTestDB(t) - if _, err := applyDir(connStr, failOnFirstCoreFS, "core", "schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, coreSource(failOnFirstCoreFS)); err == nil { t.Fatal("expected error, got nil") } // prevVersion was 0 so rollback is skipped to protect pre-existing data. @@ -231,10 +249,10 @@ func TestApplyDir_NoRollbackWhenFirstMigrationFails(t *testing.T) { } } -func TestApplyDir_NoRollbackWhenLaterMigrationFails(t *testing.T) { +func TestApplySource_NoRollbackWhenLaterMigrationFails(t *testing.T) { connStr := startTestDB(t) - if _, err := applyDir(connStr, failOnSecondCoreFS, "core", "schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, coreSource(failOnSecondCoreFS)); err == nil { t.Fatal("expected error, got nil") } // Migration 1 succeeded, migration 2 failed. Rollback is skipped because @@ -244,16 +262,16 @@ func TestApplyDir_NoRollbackWhenLaterMigrationFails(t *testing.T) { } } -func TestApplyDir_RollsBackToExistingVersion(t *testing.T) { +func TestApplySource_RollsBackToExistingVersion(t *testing.T) { connStr := startTestDB(t) // Establish a baseline at version 1. - if _, err := applyDir(connStr, oneCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(oneCoreFS)); err != nil { t.Fatalf("setup: %v", err) } // Advance to version 2 — should fail and roll back to version 1, not 0. - if _, err := applyDir(connStr, failOnSecondCoreFS, "core", "schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, coreSource(failOnSecondCoreFS)); err == nil { t.Fatal("expected error, got nil") } if got := trackVersion(t, connStr, "schema_migrations"); got != 1 { @@ -261,15 +279,15 @@ func TestApplyDir_RollsBackToExistingVersion(t *testing.T) { } } -// TestApplyDir_RollsBackWithExistingVersion verifies that when migrations have +// TestApplySource_RollsBackWithExistingVersion verifies that when migrations have // previously been applied (prevVersion > 0), rollback always happens on failure. // This ensures the rollback protection only affects the initial migration run // (prevVersion == 0), not subsequent upgrades. -func TestApplyDir_RollsBackWithExistingVersion(t *testing.T) { +func TestApplySource_RollsBackWithExistingVersion(t *testing.T) { connStr := startTestDB(t) // Establish a baseline at version 1. - if _, err := applyDir(connStr, oneCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(oneCoreFS)); err != nil { t.Fatalf("setup: %v", err) } @@ -279,7 +297,7 @@ func TestApplyDir_RollsBackWithExistingVersion(t *testing.T) { } // Advance to version 2 — should roll back because prevVersion > 0. - if _, err := applyDir(connStr, failOnSecondCoreFS, "core", "schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, coreSource(failOnSecondCoreFS)); err == nil { t.Fatal("expected error, got nil") } if got := trackVersion(t, connStr, "schema_migrations"); got != 1 { @@ -357,21 +375,21 @@ func TestMaxEmbeddedVersion(t *testing.T) { } } -// TestApplyDir_SucceedsWhenDBVersionAhead verifies that an older binary starting against +// TestApplySource_SucceedsWhenDBVersionAhead verifies that an older binary starting against // a database that a newer binary has migrated does not crash-loop. It skips Up entirely // and returns success, leaving the schema unchanged. Safe rollback relies on the // expand-then-contract discipline in database-migrations.md and rolling back one release // at a time. -func TestApplyDir_SucceedsWhenDBVersionAhead(t *testing.T) { +func TestApplySource_SucceedsWhenDBVersionAhead(t *testing.T) { connStr := startTestDB(t) // Newer binary applies v1 and v2. - if _, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)); err != nil { t.Fatalf("newer binary apply: %v", err) } // Older binary (max v1) starts against the v2 schema — must not error. - if _, err := applyDir(connStr, oneCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(oneCoreFS)); err != nil { t.Fatalf("older binary apply against newer schema: %v", err) } @@ -381,14 +399,14 @@ func TestApplyDir_SucceedsWhenDBVersionAhead(t *testing.T) { } } -// TestApplyDir_DirtyStateNotMaskedByCompatibilityMode verifies that a dirty database +// TestApplySource_DirtyStateNotMaskedByCompatibilityMode verifies that a dirty database // is not silently accepted by compatibility mode. If the DB is both dirty and ahead of // the binary's max known version, the dirty state must still be surfaced as an error. -func TestApplyDir_DirtyStateNotMaskedByCompatibilityMode(t *testing.T) { +func TestApplySource_DirtyStateNotMaskedByCompatibilityMode(t *testing.T) { connStr := startTestDB(t) // Apply v1 cleanly first so the tracking table exists. - if _, err := applyDir(connStr, oneCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(oneCoreFS)); err != nil { t.Fatalf("setup: %v", err) } @@ -404,7 +422,7 @@ func TestApplyDir_DirtyStateNotMaskedByCompatibilityMode(t *testing.T) { // Older binary (max v1) starts: DB is at v2 dirty. Compatibility mode must NOT // trigger — dirty state must be returned as an error so the operator can act. - _, err = applyDir(connStr, oneCoreFS, "core", "schema_migrations") + _, err = applySource(context.Background(), connStr, coreSource(oneCoreFS)) if err == nil { t.Fatal("expected error for dirty database, got nil") } @@ -414,38 +432,38 @@ func TestApplyDir_DirtyStateNotMaskedByCompatibilityMode(t *testing.T) { } } -// --- rollbackDir tests --- +// --- rollbackSource tests --- -func TestRollbackDir_RollsBackToTarget(t *testing.T) { +func TestRollbackSource_RollsBackToTarget(t *testing.T) { connStr := startTestDB(t) - if _, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)); err != nil { t.Fatalf("setup: %v", err) } - rollbackDir(connStr, goodCoreFS, "core", "schema_migrations", 0) + rollbackSource(context.Background(), connStr, coreSource(goodCoreFS), 0) if got := trackVersion(t, connStr, "schema_migrations"); got != 0 { t.Errorf("version after rollback = %d, want 0", got) } } -func TestRollbackDir_PartialRollback(t *testing.T) { +func TestRollbackSource_PartialRollback(t *testing.T) { connStr := startTestDB(t) - if _, err := applyDir(connStr, goodCoreFS, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(goodCoreFS)); err != nil { t.Fatalf("setup: %v", err) } // Roll back only one step (2 → 1). - rollbackDir(connStr, goodCoreFS, "core", "schema_migrations", 1) + rollbackSource(context.Background(), connStr, coreSource(goodCoreFS), 1) if got := trackVersion(t, connStr, "schema_migrations"); got != 1 { t.Errorf("version after partial rollback = %d, want 1", got) } } -// --- cross-track rollback --- +// --- cross-track rollback (per-source primitives) --- // TestCrossTrackRollback_CoreUnchangedWhenVectorFails covers the case where // core has no new migrations (ErrNoChange) and vector fails. Core should not @@ -456,12 +474,12 @@ func TestCrossTrackRollback_CoreUnchangedWhenVectorFails(t *testing.T) { combined := mergeFS(goodCoreFS, failVectorFS) // Establish core at its latest version before the run. - if _, err := applyDir(connStr, combined, "core", "schema_migrations"); err != nil { + if _, err := applySource(context.Background(), connStr, coreSource(combined)); err != nil { t.Fatalf("setup core: %v", err) } - // Core has no new migrations — applyDir returns ErrNoChange. - corePrev, err := applyDir(connStr, combined, "core", "schema_migrations") + // Core has no new migrations — applySource returns ErrNoChange. + corePrev, err := applySource(context.Background(), connStr, coreSource(combined)) if err != nil { t.Fatalf("core apply (no-op): %v", err) } @@ -470,12 +488,12 @@ func TestCrossTrackRollback_CoreUnchangedWhenVectorFails(t *testing.T) { } // Vector fails and self-rolls-back. - if _, err := applyDir(connStr, combined, "vector", "vector_schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, vectorSource(combined)); err == nil { t.Fatal("expected vector error, got nil") } // Cross-track rollback: core should be untouched since corePrev == current version. - rollbackDir(connStr, combined, "core", "schema_migrations", corePrev) + rollbackSource(context.Background(), connStr, coreSource(combined), corePrev) if got := trackVersion(t, connStr, "schema_migrations"); got != 2 { t.Errorf("core version = %d, want 2 (should not have been downgraded)", got) } @@ -487,7 +505,7 @@ func TestCrossTrackRollback_CoreRolledBackWhenVectorFails(t *testing.T) { combined := mergeFS(goodCoreFS, failVectorFS) // Core succeeds. - corePrev, err := applyDir(connStr, combined, "core", "schema_migrations") + corePrev, err := applySource(context.Background(), connStr, coreSource(combined)) if err != nil { t.Fatalf("core apply: %v", err) } @@ -496,7 +514,7 @@ func TestCrossTrackRollback_CoreRolledBackWhenVectorFails(t *testing.T) { } // Vector fails. Self-rollback is skipped because vector prevVersion is 0. - if _, err := applyDir(connStr, combined, "vector", "vector_schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, vectorSource(combined)); err == nil { t.Fatal("expected vector error, got nil") } if got := trackVersion(t, connStr, "vector_schema_migrations"); got != 1 { @@ -504,7 +522,7 @@ func TestCrossTrackRollback_CoreRolledBackWhenVectorFails(t *testing.T) { } // Cross-track rollback: core should be rolled back to its pre-run version. - rollbackDir(connStr, combined, "core", "schema_migrations", corePrev) + rollbackSource(context.Background(), connStr, coreSource(combined), corePrev) if got := trackVersion(t, connStr, "schema_migrations"); got != corePrev { t.Errorf("core version after cross-track rollback = %d, want %d", got, corePrev) } @@ -520,7 +538,7 @@ func TestCrossTrackRollback_IfExistsGuardsSafeOnVectorFailure(t *testing.T) { combined := mergeFS(expandCoreFS, failVectorWithDependencyFS) // Core succeeds (shared_data created with col_a and col_b). - corePrev, err := applyDir(connStr, combined, "core", "schema_migrations") + corePrev, err := applySource(context.Background(), connStr, coreSource(combined)) if err != nil { t.Fatalf("core apply: %v", err) } @@ -529,7 +547,7 @@ func TestCrossTrackRollback_IfExistsGuardsSafeOnVectorFailure(t *testing.T) { } // Vector fails. Self-rollback is skipped because vector prevVersion is 0. - if _, err := applyDir(connStr, combined, "vector", "vector_schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, vectorSource(combined)); err == nil { t.Fatal("expected vector error, got nil") } if got := trackVersion(t, connStr, "vector_schema_migrations"); got != 1 { @@ -537,7 +555,7 @@ func TestCrossTrackRollback_IfExistsGuardsSafeOnVectorFailure(t *testing.T) { } // Cross-track rollback: core rolls back to its pre-run version. - rollbackDir(connStr, combined, "core", "schema_migrations", corePrev) + rollbackSource(context.Background(), connStr, coreSource(combined), corePrev) if got := trackVersion(t, connStr, "schema_migrations"); got != corePrev { t.Errorf("core version after cross-track rollback = %d, want %d", got, corePrev) } @@ -565,7 +583,7 @@ func TestRunUp_CoreAndVector(t *testing.T) { connStr := startTestDB(t) combined := mergeFS(goodCoreFS, goodVectorFS) - if err := RunUp(connStr, combined, true); err != nil { + if err := RunUp(context.Background(), connStr, []Source{coreSource(combined), vectorSource(combined)}); err != nil { t.Fatalf("RunUp: %v", err) } if got := trackVersion(t, connStr, "schema_migrations"); got != 2 { @@ -580,7 +598,7 @@ func TestRunUp_CoreOnlyWhenVectorDisabled(t *testing.T) { connStr := startTestDB(t) combined := mergeFS(goodCoreFS, goodVectorFS) - if err := RunUp(connStr, combined, false); err != nil { + if err := RunUp(context.Background(), connStr, []Source{coreSource(combined)}); err != nil { t.Fatalf("RunUp: %v", err) } if got := trackVersion(t, connStr, "schema_migrations"); got != 2 { @@ -588,18 +606,19 @@ func TestRunUp_CoreOnlyWhenVectorDisabled(t *testing.T) { } // Vector tracking table should not exist. if tableExists(t, connStr, "vector_schema_migrations") { - t.Error("vector_schema_migrations should not exist when vectorEnabled=false") + t.Error("vector_schema_migrations should not exist when vector source is not registered") } } func TestRunUp_FailsBeforeMigrationsWhenPgvectorMissing(t *testing.T) { connStr := startTestDBWithoutPgvector(t) - err := RunUp(connStr, goodCoreFS, true) + err := RunUp(context.Background(), connStr, []Source{coreSource(goodCoreFS), vectorSource(goodCoreFS)}) if err == nil { t.Fatal("expected error, got nil") } - // Core migrations should NOT have run — no tracking table created. + // Core migrations should NOT have run — no tracking table created. The vector + // precheck runs up front, before any source is applied. if tableExists(t, connStr, "schema_migrations") { t.Error("schema_migrations should not exist — pgvector check should fail before any migrations") } @@ -612,7 +631,7 @@ func TestRunUp_SkipsCoreRollbackWhenVectorFailsOnFirstRun(t *testing.T) { connStr := startTestDB(t) // pgvector available so checkPgvector passes combined := mergeFS(goodCoreFS, failVectorFS) - err := RunUp(connStr, combined, true) + err := RunUp(context.Background(), connStr, []Source{coreSource(combined), vectorSource(combined)}) if err == nil { t.Fatal("expected error, got nil") } @@ -622,18 +641,275 @@ func TestRunUp_SkipsCoreRollbackWhenVectorFailsOnFirstRun(t *testing.T) { } } +// TestRunUp_MultiSourceOrdering verifies sources apply in slice order: source "b" +// references a table created by source "a", so it can only succeed if "a" ran first. +func TestRunUp_MultiSourceOrdering(t *testing.T) { + connStr := startTestDB(t) + + ordFS := fstest.MapFS{ + "a/000001_a.up.sql": {Data: []byte(`CREATE TABLE ord_a (id INT PRIMARY KEY);`)}, + "a/000001_a.down.sql": {Data: []byte(`DROP TABLE IF EXISTS ord_a;`)}, + "b/000001_b.up.sql": {Data: []byte(`CREATE TABLE ord_b (id INT PRIMARY KEY, a_id INT REFERENCES ord_a(id));`)}, + "b/000001_b.down.sql": {Data: []byte(`DROP TABLE IF EXISTS ord_b;`)}, + "c/000001_c.up.sql": {Data: []byte(`CREATE TABLE ord_c (id INT PRIMARY KEY);`)}, + "c/000001_c.down.sql": {Data: []byte(`DROP TABLE IF EXISTS ord_c;`)}, + } + sources := []Source{ + {Name: "a", TrackingTable: "ord_a_migrations", FS: ordFS, Dir: "a"}, + {Name: "b", TrackingTable: "ord_b_migrations", FS: ordFS, Dir: "b"}, + {Name: "c", TrackingTable: "ord_c_migrations", FS: ordFS, Dir: "c"}, + } + + if err := RunUp(context.Background(), connStr, sources); err != nil { + t.Fatalf("RunUp: %v", err) + } + for _, table := range []string{"ord_a", "ord_b", "ord_c"} { + if !tableExists(t, connStr, table) { + t.Errorf("table %s should exist", table) + } + } + for _, tbl := range []string{"ord_a_migrations", "ord_b_migrations", "ord_c_migrations"} { + if got := trackVersion(t, connStr, tbl); got != 1 { + t.Errorf("%s version = %d, want 1", tbl, got) + } + } +} + +// TestRunUp_CompensatingRollbackAcrossThreeSources verifies that when a later +// source fails, previously-applied sources are rolled back to their pre-run +// versions in reverse order — and that the prevVersion==0 guard skips a source +// applied for the first time. Source "a" is fresh (prev 0, not compensated); +// source "b" is pre-seeded at v1 (prev 1, rolled back); source "c" fails. +func TestRunUp_CompensatingRollbackAcrossThreeSources(t *testing.T) { + connStr := startTestDB(t) + + aFS := fstest.MapFS{ + "a/000001_a.up.sql": {Data: []byte(`CREATE TABLE comp_a (id INT PRIMARY KEY);`)}, + "a/000001_a.down.sql": {Data: []byte(`DROP TABLE IF EXISTS comp_a;`)}, + } + bV1FS := fstest.MapFS{ + "b/000001_b.up.sql": {Data: []byte(`CREATE TABLE comp_b (id INT PRIMARY KEY);`)}, + "b/000001_b.down.sql": {Data: []byte(`DROP TABLE IF EXISTS comp_b;`)}, + } + bFullFS := fstest.MapFS{ + "b/000001_b.up.sql": {Data: []byte(`CREATE TABLE comp_b (id INT PRIMARY KEY);`)}, + "b/000001_b.down.sql": {Data: []byte(`DROP TABLE IF EXISTS comp_b;`)}, + "b/000002_b_addcol.up.sql": {Data: []byte(`ALTER TABLE comp_b ADD COLUMN extra TEXT;`)}, + "b/000002_b_addcol.down.sql": {Data: []byte(`ALTER TABLE comp_b DROP COLUMN IF EXISTS extra;`)}, + } + cFS := fstest.MapFS{ + "c/000001_bad.up.sql": {Data: []byte(`ALTER TABLE no_such_table ADD COLUMN x TEXT;`)}, + "c/000001_bad.down.sql": {Data: []byte(`SELECT 1;`)}, + } + + aSrc := Source{Name: "a", TrackingTable: "comp_a_migrations", FS: aFS, Dir: "a"} + bSrcFull := Source{Name: "b", TrackingTable: "comp_b_migrations", FS: bFullFS, Dir: "b"} + cSrc := Source{Name: "c", TrackingTable: "comp_c_migrations", FS: cFS, Dir: "c"} + + // Pre-seed b at v1 so its prevVersion is > 0 during the run below. + if _, err := applySource(context.Background(), connStr, Source{Name: "b", TrackingTable: "comp_b_migrations", FS: bV1FS, Dir: "b"}); err != nil { + t.Fatalf("seed b: %v", err) + } + + err := RunUp(context.Background(), connStr, []Source{aSrc, bSrcFull, cSrc}) + if err == nil { + t.Fatal("expected error from failing source c, got nil") + } + + // a: applied for the first time (prev 0) → compensation skipped → stays at v1. + if got := trackVersion(t, connStr, "comp_a_migrations"); got != 1 { + t.Errorf("comp_a version = %d, want 1 (prev==0 guard skips compensation)", got) + } + // b: prev was 1, advanced to 2, compensated back to 1. + if got := trackVersion(t, connStr, "comp_b_migrations"); got != 1 { + t.Errorf("comp_b version = %d, want 1 (rolled back to pre-run version)", got) + } +} + +// TestRunUp_SchemaScopedSource verifies a Source with a non-empty Schema creates +// the schema and lands both its objects and its tracking table there, not in public. +func TestRunUp_SchemaScopedSource(t *testing.T) { + connStr := startTestDB(t) + + schemaFS := fstest.MapFS{ + "s/000001_create.up.sql": {Data: []byte(`CREATE TABLE scoped_t (id INT PRIMARY KEY);`)}, + "s/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS scoped_t;`)}, + } + src := Source{Name: "scoped", Schema: "myschema", TrackingTable: "schema_migrations", FS: schemaFS, Dir: "s"} + + if err := RunUp(context.Background(), connStr, []Source{src}); err != nil { + t.Fatalf("RunUp: %v", err) + } + + if !tableExistsInSchema(t, connStr, "myschema", "scoped_t") { + t.Error("scoped_t should exist in myschema") + } + if tableExistsInSchema(t, connStr, "public", "scoped_t") { + t.Error("scoped_t should NOT exist in public") + } + if !tableExistsInSchema(t, connStr, "myschema", "schema_migrations") { + t.Error("tracking table should exist in myschema") + } + if tableExistsInSchema(t, connStr, "public", "schema_migrations") { + t.Error("tracking table should NOT exist in public") + } +} + +// TestRunUp_RejectsResolvedSchemaCollision verifies the runtime guard catches a +// collision validateSources cannot: an unscoped source (Schema "") and an explicit +// source naming the connection's current_schema() ("public" here) share one +// tracking table once "" resolves, so RunUp must reject the set before applying +// anything. validateSources alone passes them (distinct literal Schema keys). +func TestRunUp_RejectsResolvedSchemaCollision(t *testing.T) { + connStr := startTestDB(t) + + collide := []Source{ + {Name: "implicit", Schema: "", TrackingTable: "schema_migrations", FS: goodCoreFS, Dir: "core"}, + {Name: "explicit", Schema: "public", TrackingTable: "schema_migrations", FS: goodCoreFS, Dir: "core"}, + } + + // validateSources keys on the literal Schema, so it does NOT catch this. + if err := validateSources(collide); err != nil { + t.Fatalf("validateSources should pass on distinct literal schemas, got %v", err) + } + + // RunUp resolves "" to current_schema() (public) and must reject. + err := RunUp(context.Background(), connStr, collide) + if err == nil { + t.Fatal("expected resolved-collision error, got nil") + } + if !strings.Contains(err.Error(), "resolve to the same tracking table") { + t.Errorf("error %q should describe a resolved tracking-table collision", err) + } + // Guard runs before any source applies — no tracking table created. + if tableExists(t, connStr, "schema_migrations") { + t.Error("schema_migrations should not exist — guard must fire before any source applies") + } +} + +// TestRunUp_PreCheckRunsBeforeAnyApply verifies that a failing PreCheck on a later +// source aborts the run before any earlier source is applied. +func TestRunUp_PreCheckRunsBeforeAnyApply(t *testing.T) { + connStr := startTestDB(t) + + preFS := fstest.MapFS{ + "p/000001_create.up.sql": {Data: []byte(`CREATE TABLE precheck_t (id INT PRIMARY KEY);`)}, + "p/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS precheck_t;`)}, + } + first := Source{Name: "first", TrackingTable: "first_migrations", FS: preFS, Dir: "p"} + second := Source{Name: "second", TrackingTable: "second_migrations", FS: preFS, Dir: "p", + PreCheck: func(string) error { return fmt.Errorf("precheck boom") }} + + err := RunUp(context.Background(), connStr, []Source{first, second}) + if err == nil { + t.Fatal("expected error from failing precheck, got nil") + } + if tableExists(t, connStr, "precheck_t") { + t.Error("precheck_t should not exist — no source should apply when a precheck fails") + } + if tableExists(t, connStr, "first_migrations") { + t.Error("first_migrations tracking table should not exist — first source must not have applied") + } +} + +func TestValidateSchemaName(t *testing.T) { + valid := []string{"a", "myschema", "tenant_1", "_x", "s123"} + for _, s := range valid { + if err := validateSchemaName(s); err != nil { + t.Errorf("validateSchemaName(%q) = %v, want nil", s, err) + } + } + // Uppercase is rejected: the name is used unquoted in search_path (which + // Postgres case-folds) and quoted in CREATE SCHEMA, so a mixed-case name + // would split across two schemas. + invalid := []string{"", "1abc", "has space", "a;b", `a"b`, "a-b", "a.b", "drop table", "ABC", "MySchema", strings.Repeat("x", 64)} + for _, s := range invalid { + if err := validateSchemaName(s); err == nil { + t.Errorf("validateSchemaName(%q) = nil, want error", s) + } + } +} + +// --- container-free unit tests --- + +// TestValidateSources_RejectsDuplicateSchemaAndTable verifies the collision unit +// is (Schema, TrackingTable): two sources sharing both are rejected, while the +// same tracking-table name in different schemas, or different tables in the same +// schema, are allowed. +func TestValidateSources_RejectsDuplicateSchemaAndTable(t *testing.T) { + collide := []Source{ + {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", Dir: "a"}, + {Name: "b", Schema: "s1", TrackingTable: "schema_migrations", Dir: "b"}, + } + if err := validateSources(collide); err == nil { + t.Error("expected collision error for same (schema, tracking table), got nil") + } + + ok := []Source{ + // Same table name, different schema — fine. + {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", Dir: "a"}, + {Name: "b", Schema: "s2", TrackingTable: "schema_migrations", Dir: "b"}, + // Different table, same (default) schema — fine. + {Name: "core", TrackingTable: "schema_migrations", Dir: "core"}, + {Name: "vector", TrackingTable: "vector_schema_migrations", Dir: "vector"}, + } + if err := validateSources(ok); err != nil { + t.Errorf("validateSources(distinct sources) = %v, want nil", err) + } + + // An invalid schema name on any source is rejected up front, before any + // source applies (fail-fast). + badSchema := []Source{ + {Name: "good", TrackingTable: "schema_migrations", Dir: "core"}, + {Name: "bad", Schema: "MixedCase", TrackingTable: "schema_migrations", Dir: "x"}, + } + if err := validateSources(badSchema); err == nil { + t.Error("expected error for invalid schema name on a later source, got nil") + } +} + +// TestRunUp_EmptyAndNilSources verifies the no-op fast paths run no migration +// logic and need no database. +func TestRunUp_EmptyAndNilSources(t *testing.T) { + if err := RunUp(context.Background(), "postgres://invalid", nil); err != nil { + t.Errorf("RunUp(nil sources) = %v, want nil", err) + } + if err := RunUp(context.Background(), "postgres://invalid", []Source{}); err != nil { + t.Errorf("RunUp(empty sources) = %v, want nil", err) + } +} + +// TestWithSearchPath verifies a URL-form DSN gains the search_path param and a +// libpq keyword/value DSN is rejected rather than silently corrupted. +func TestWithSearchPath(t *testing.T) { + got, err := withSearchPath("postgres://u:p@host:5432/db?sslmode=disable", "myschema") + if err != nil { + t.Fatalf("withSearchPath(url DSN) = %v, want nil", err) + } + if !strings.Contains(got, "search_path=myschema") { + t.Errorf("result %q missing search_path=myschema", got) + } + if !strings.Contains(got, "sslmode=disable") { + t.Errorf("result %q dropped existing sslmode param", got) + } + + if _, err := withSearchPath("host=localhost dbname=foo user=bar", "myschema"); err == nil { + t.Error("withSearchPath(keyword/value DSN) = nil, want error") + } +} + // --- dirty state recovery tests --- -// TestApplyDir_DirtyStateRecoveryOnRestart simulates a restart after a failed +// TestApplySource_DirtyStateRecoveryOnRestart simulates a restart after a failed // migration left the database in a dirty state. On the second call, prevVersion // is > 0 (the dirty version), so rollback is enabled. The runner should clear // the dirty state and roll back to the last clean version. -func TestApplyDir_DirtyStateRecoveryOnRestart(t *testing.T) { +func TestApplySource_DirtyStateRecoveryOnRestart(t *testing.T) { connStr := startTestDB(t) // First run: apply version 1, then version 2 fails. prevVersion is 0, so // rollback is skipped. Database left at version 2 dirty. - if _, err := applyDir(connStr, failOnSecondCoreFS, "core", "schema_migrations"); err == nil { + if _, err := applySource(context.Background(), connStr, coreSource(failOnSecondCoreFS)); err == nil { t.Fatal("expected error, got nil") } if got := trackVersion(t, connStr, "schema_migrations"); got != 2 { @@ -643,7 +919,7 @@ func TestApplyDir_DirtyStateRecoveryOnRestart(t *testing.T) { // Second run (simulating restart): prevVersion is now 2 (dirty). The runner // should detect dirty state and attempt to clear it. mg.Up() will fail because // the database is dirty, then rollbackToVersion clears dirty to version 1. - _, err := applyDir(connStr, failOnSecondCoreFS, "core", "schema_migrations") + _, err := applySource(context.Background(), connStr, coreSource(failOnSecondCoreFS)) if err == nil { t.Fatal("expected error on second run, got nil") } From 7f9019febf3bf4fd9c9f4a514d780867073e2770 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Wed, 1 Jul 2026 12:48:24 -0700 Subject: [PATCH 2/7] Fix wording Signed-off-by: Jeremy Alvis --- go/core/pkg/migrations/runner.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index 7fa6603140..bd4eab038e 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -66,8 +66,12 @@ type Source struct { } // BuiltinSources returns the built-in source set: the core track always, and the -// vector track when vectorEnabled. Downstream consumers append their own Sources -// to this slice before calling RunUp. +// vector track when vectorEnabled. app.Start prepends these to any +// downstream-registered extra sources before calling RunUp, so the built-in +// tracks always run first and downstream consumers only supply their own extras +// (never assembling this slice themselves). A caller that invokes RunUp directly +// (e.g. a migration CLI) composes the list the same way: BuiltinSources first, +// then extras. func BuiltinSources(vectorEnabled bool) []Source { sources := []Source{{ Name: "core", From 820e97bc2b1907c959af80776b5f3f48f7ca4535 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Thu, 2 Jul 2026 14:14:01 -0700 Subject: [PATCH 3/7] Fix compilation issue from merge Signed-off-by: Jeremy Alvis --- go/core/test/upgrade/roundtrip_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go/core/test/upgrade/roundtrip_test.go b/go/core/test/upgrade/roundtrip_test.go index 73f878ec49..8af9cd8da2 100644 --- a/go/core/test/upgrade/roundtrip_test.go +++ b/go/core/test/upgrade/roundtrip_test.go @@ -126,7 +126,7 @@ func buildCleanInstallSchema(t *testing.T, env upgradeEnv, dbName string, vector defer stop() url := fmt.Sprintf("postgres://kagent:kagent@127.0.0.1:%d/%s?sslmode=disable", localPort, dbName) - require.NoError(t, migrations.RunUp(url, migrations.FS, vectorEnabled), + require.NoError(t, migrations.RunUp(t.Context(), url, migrations.BuiltinSources(vectorEnabled)), "apply embedded migrations to clean reference database %s", dbName) return pgSchemaDump(t, env, dbName) From f08ea36105224b34c375bbacf6ef614df27e1e1b Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Thu, 2 Jul 2026 15:27:36 -0700 Subject: [PATCH 4/7] Updates based on review Signed-off-by: Jeremy Alvis --- go/core/pkg/migrations/runner.go | 71 ++++++++++++++++------- go/core/pkg/migrations/runner_test.go | 82 ++++++++++++++++++++------- 2 files changed, 115 insertions(+), 38 deletions(-) diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index bd4eab038e..2b697430b2 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -144,6 +144,15 @@ func RunUp(ctx context.Context, url string, sources []Source) error { // Compensating rollback: undo previously-applied sources in reverse // order, each to its own pre-run version. The failing source has // already rolled itself back in applySource. + // + // Run the rollback under a context detached from cancellation. If the + // caller's ctx is already canceled (e.g. SIGTERM arriving mid-startup + // as a source fails), using it here would abort the rollback at schema + // setup and leave the database mid-migration. WithoutCancel keeps any + // request-scoped values but lets best-effort cleanup finish; + // golang-migrate's own steps aren't ctx-aware regardless, so only the + // schema-setup ExecContext consults it. + rbCtx := context.WithoutCancel(ctx) var compErrs []error for _, a := range slices.Backward(done) { if a.prev == 0 { @@ -151,7 +160,7 @@ func RunUp(ctx context.Context, url string, sources []Source) error { continue } log.Info("rolling back source after later failure", "source", a.src.Name, "targetVersion", a.prev) - if rbErr := rollbackSource(ctx, url, a.src, a.prev); rbErr != nil { + if rbErr := rollbackSource(rbCtx, url, a.src, a.prev); rbErr != nil { compErrs = append(compErrs, rbErr) } } @@ -170,27 +179,50 @@ func RunUp(ctx context.Context, url string, sources []Source) error { return nil } -// validateSources rejects two sources that share the same (schema, tracking -// table). That pair is the collision unit: it is what determines the -// golang-migrate version row and the advisory-lock id, so two such sources would -// fight over one row and lock regardless of their Dir/FS. Distinct tracking -// tables in the same schema are fine (each gets its own bookkeeping). +// validateSources rejects a source set that cannot be run safely. It checks two +// things. // -// This keys on the *literal* Schema string and is intentionally DB-free so it -// stays unit-testable. It therefore cannot see a collision between a source with -// Schema == "" and one naming the connection's current_schema() explicitly, since -// "" resolves to that schema only at connection time. checkResolvedSchemaCollisions -// (called from RunUp, where a connection is available) closes that gap. +// First, required fields. Source is a public extension point, so a source with a +// missing field is a caller mistake that should fail fast with an actionable +// error rather than surface later as a vague golang-migrate failure. Name, FS, +// Dir, and TrackingTable are required; Schema is intentionally optional +// ("" = the connection's default schema) and PreCheck is optional (nil = none). +// TrackingTable has no safe default: sources can share a schema (core and vector +// both live in the connection default), so an empty TrackingTable would silently +// fall back to golang-migrate's "schema_migrations" and collide. +// +// Second, collisions on the (schema, tracking table) pair. That pair is the +// collision unit: it determines the golang-migrate version row and the +// advisory-lock id, so two such sources would fight over one row and lock +// regardless of their Dir/FS. Distinct tracking tables in the same schema are +// fine (each gets its own bookkeeping). +// +// The collision check keys on the *literal* Schema string and is intentionally +// DB-free so it stays unit-testable. It therefore cannot see a collision between +// a source with Schema == "" and one naming the connection's current_schema() +// explicitly, since "" resolves to that schema only at connection time. +// checkResolvedSchemaCollisions (called from RunUp, where a connection is +// available) closes that gap. func validateSources(sources []Source) error { seen := make(map[string]string, len(sources)) for _, s := range sources { + switch { + case s.Name == "": + return fmt.Errorf("migration source has empty Name") + case s.TrackingTable == "": + return fmt.Errorf("migration source %q has empty TrackingTable", s.Name) + case s.FS == nil: + return fmt.Errorf("migration source %q has nil FS", s.Name) + case s.Dir == "": + return fmt.Errorf("migration source %q has empty Dir", s.Name) + } // Validate schema names up front so a bad name on a later source aborts // the whole run before any earlier source applies, matching the fail-fast // guarantee RunUp gives prechecks. newMigrate re-validates as a safety net // for callers that bypass RunUp (e.g. applySource directly in tests). if s.Schema != "" { if err := validateSchemaName(s.Schema); err != nil { - return err + return fmt.Errorf("source %q: %w", s.Name, err) } } key := s.Schema + "\x00" + s.TrackingTable @@ -417,7 +449,7 @@ func newMigrate(ctx context.Context, dbURL string, src Source) (*migrate.Migrate connURL := dbURL if src.Schema != "" { if err := validateSchemaName(src.Schema); err != nil { - return nil, err + return nil, fmt.Errorf("source %q: %w", src.Name, err) } var err error connURL, err = withSearchPath(dbURL, src.Schema) @@ -465,17 +497,18 @@ func newMigrate(ctx context.Context, dbURL string, src Source) (*migrate.Migrate // withSearchPath returns dbURL with the search_path connection parameter set to // schema, so every connection in the pool (including the one golang-migrate -// checks out) targets that schema for migration DDL. dbURL must be a URL-form -// DSN (postgres://...); a libpq keyword/value DSN is rejected because net/url -// parses it without error into a meaningless URL, which would silently corrupt -// the connection string rather than fail here. +// checks out) targets that schema for migration DDL. dbURL must be a postgres:// +// or postgresql:// URL. Other inputs are rejected: net/url parses a libpq +// keyword/value DSN or a bare "host:port/db" without error (the latter as scheme +// "host"), so requiring a known Postgres scheme is what makes this fail fast +// rather than silently rewrite a meaningless URL. func withSearchPath(dbURL, schema string) (string, error) { u, err := nurl.Parse(dbURL) if err != nil { return "", fmt.Errorf("parse database url: %w", err) } - if u.Scheme == "" { - return "", fmt.Errorf("database url must be a URL-form DSN (postgres://...) to scope a schema; got a non-URL DSN") + if u.Scheme != "postgres" && u.Scheme != "postgresql" { + return "", fmt.Errorf("database url must be a postgres:// or postgresql:// DSN to scope a schema; got scheme %q", u.Scheme) } q := u.Query() q.Set("search_path", schema) diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index f9b595fccb..81b36e48e5 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -837,9 +837,11 @@ func TestValidateSchemaName(t *testing.T) { // same tracking-table name in different schemas, or different tables in the same // schema, are allowed. func TestValidateSources_RejectsDuplicateSchemaAndTable(t *testing.T) { + dummy := fstest.MapFS{} // non-nil so the required-field checks pass + collide := []Source{ - {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", Dir: "a"}, - {Name: "b", Schema: "s1", TrackingTable: "schema_migrations", Dir: "b"}, + {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", FS: dummy, Dir: "a"}, + {Name: "b", Schema: "s1", TrackingTable: "schema_migrations", FS: dummy, Dir: "b"}, } if err := validateSources(collide); err == nil { t.Error("expected collision error for same (schema, tracking table), got nil") @@ -847,11 +849,11 @@ func TestValidateSources_RejectsDuplicateSchemaAndTable(t *testing.T) { ok := []Source{ // Same table name, different schema — fine. - {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", Dir: "a"}, - {Name: "b", Schema: "s2", TrackingTable: "schema_migrations", Dir: "b"}, + {Name: "a", Schema: "s1", TrackingTable: "schema_migrations", FS: dummy, Dir: "a"}, + {Name: "b", Schema: "s2", TrackingTable: "schema_migrations", FS: dummy, Dir: "b"}, // Different table, same (default) schema — fine. - {Name: "core", TrackingTable: "schema_migrations", Dir: "core"}, - {Name: "vector", TrackingTable: "vector_schema_migrations", Dir: "vector"}, + {Name: "core", TrackingTable: "schema_migrations", FS: dummy, Dir: "core"}, + {Name: "vector", TrackingTable: "vector_schema_migrations", FS: dummy, Dir: "vector"}, } if err := validateSources(ok); err != nil { t.Errorf("validateSources(distinct sources) = %v, want nil", err) @@ -860,14 +862,40 @@ func TestValidateSources_RejectsDuplicateSchemaAndTable(t *testing.T) { // An invalid schema name on any source is rejected up front, before any // source applies (fail-fast). badSchema := []Source{ - {Name: "good", TrackingTable: "schema_migrations", Dir: "core"}, - {Name: "bad", Schema: "MixedCase", TrackingTable: "schema_migrations", Dir: "x"}, + {Name: "good", TrackingTable: "schema_migrations", FS: dummy, Dir: "core"}, + {Name: "bad", Schema: "MixedCase", TrackingTable: "schema_migrations", FS: dummy, Dir: "x"}, } if err := validateSources(badSchema); err == nil { t.Error("expected error for invalid schema name on a later source, got nil") } } +// TestValidateSources_RejectsMissingRequiredFields verifies Source's required +// fields (Name, TrackingTable, FS, Dir) are enforced up front, while Schema +// stays optional ("" = connection default). +func TestValidateSources_RejectsMissingRequiredFields(t *testing.T) { + base := Source{Name: "x", TrackingTable: "x_migrations", FS: fstest.MapFS{}, Dir: "x"} + + // The fully-populated base (with empty Schema) is valid. + if err := validateSources([]Source{base}); err != nil { + t.Fatalf("validateSources(valid source) = %v, want nil", err) + } + + mutations := map[string]func(Source) Source{ + "empty Name": func(s Source) Source { s.Name = ""; return s }, + "empty TrackingTable": func(s Source) Source { s.TrackingTable = ""; return s }, + "nil FS": func(s Source) Source { s.FS = nil; return s }, + "empty Dir": func(s Source) Source { s.Dir = ""; return s }, + } + for name, mut := range mutations { + t.Run(name, func(t *testing.T) { + if err := validateSources([]Source{mut(base)}); err == nil { + t.Errorf("validateSources with %s = nil, want error", name) + } + }) + } +} + // TestRunUp_EmptyAndNilSources verifies the no-op fast paths run no migration // logic and need no database. func TestRunUp_EmptyAndNilSources(t *testing.T) { @@ -879,22 +907,38 @@ func TestRunUp_EmptyAndNilSources(t *testing.T) { } } -// TestWithSearchPath verifies a URL-form DSN gains the search_path param and a -// libpq keyword/value DSN is rejected rather than silently corrupted. +// TestWithSearchPath verifies a postgres:// or postgresql:// DSN gains the +// search_path param (preserving existing params) and that anything else is +// rejected rather than silently rewritten. func TestWithSearchPath(t *testing.T) { - got, err := withSearchPath("postgres://u:p@host:5432/db?sslmode=disable", "myschema") - if err != nil { - t.Fatalf("withSearchPath(url DSN) = %v, want nil", err) - } - if !strings.Contains(got, "search_path=myschema") { - t.Errorf("result %q missing search_path=myschema", got) + for _, dsn := range []string{ + "postgres://u:p@host:5432/db?sslmode=disable", + "postgresql://u:p@host:5432/db", + } { + got, err := withSearchPath(dsn, "myschema") + if err != nil { + t.Fatalf("withSearchPath(%q) = %v, want nil", dsn, err) + } + if !strings.Contains(got, "search_path=myschema") { + t.Errorf("withSearchPath(%q) = %q, missing search_path=myschema", dsn, got) + } } - if !strings.Contains(got, "sslmode=disable") { + + // Existing query params are preserved. + if got, _ := withSearchPath("postgres://u:p@host:5432/db?sslmode=disable", "myschema"); !strings.Contains(got, "sslmode=disable") { t.Errorf("result %q dropped existing sslmode param", got) } - if _, err := withSearchPath("host=localhost dbname=foo user=bar", "myschema"); err == nil { - t.Error("withSearchPath(keyword/value DSN) = nil, want error") + // Non-postgres inputs fail fast rather than being silently rewritten. + for _, bad := range []string{ + "host=localhost dbname=foo user=bar", // libpq keyword/value DSN (scheme "") + "localhost:5432/db", // parses as scheme "localhost" + "mysql://u:p@host/db", // wrong scheme + "", // empty + } { + if _, err := withSearchPath(bad, "myschema"); err == nil { + t.Errorf("withSearchPath(%q) = nil, want error", bad) + } } } From 03bf28304c12c890b01f00c74463123dfb9b1227 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Mon, 6 Jul 2026 07:46:13 -0700 Subject: [PATCH 5/7] add kagent db migrate CLI for out-of-band migrations Signed-off-by: Jeremy Alvis --- go/core/cli/cmd/kagent/main.go | 22 +- go/core/pkg/cli/db/db.go | 43 ++ go/core/pkg/cli/db/migrate/migrate.go | 729 +++++++++++++++++++++ go/core/pkg/cli/db/migrate/migrate_test.go | 314 +++++++++ go/core/pkg/migrations/runner.go | 17 + 5 files changed, 1124 insertions(+), 1 deletion(-) create mode 100644 go/core/pkg/cli/db/db.go create mode 100644 go/core/pkg/cli/db/migrate/migrate.go create mode 100644 go/core/pkg/cli/db/migrate/migrate_test.go diff --git a/go/core/cli/cmd/kagent/main.go b/go/core/cli/cmd/kagent/main.go index 78b6c062e5..765a5cb0f3 100644 --- a/go/core/cli/cmd/kagent/main.go +++ b/go/core/cli/cmd/kagent/main.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "os/signal" + "strconv" "syscall" cli "github.com/kagent-dev/kagent/go/core/cli/internal/cli/agent" @@ -13,6 +14,8 @@ import ( "github.com/kagent-dev/kagent/go/core/cli/internal/config" "github.com/kagent-dev/kagent/go/core/cli/internal/profiles" "github.com/kagent-dev/kagent/go/core/cli/internal/tui" + dbcli "github.com/kagent-dev/kagent/go/core/pkg/cli/db" + "github.com/kagent-dev/kagent/go/core/pkg/migrations" "github.com/spf13/cobra" "sigs.k8s.io/controller-runtime/pkg/client" ) @@ -449,11 +452,28 @@ Examples: runCmd.Flags().StringVar(&runCfg.ProjectDir, "project-dir", "", "Project directory (default: current directory)") runCmd.Flags().BoolVar(&runCfg.Build, "build", false, "Rebuild the Docker image before running") - rootCmd.AddCommand(installCmd, uninstallCmd, invokeCmd, bugReportCmd, versionCmd, dashboardCmd, getCmd, initCmd, buildCmd, deployCmd, addMcpCmd, runCmd, mcp.NewMCPCmd(), envdoc.NewEnvCmd()) + rootCmd.AddCommand(installCmd, uninstallCmd, invokeCmd, bugReportCmd, versionCmd, dashboardCmd, getCmd, initCmd, buildCmd, deployCmd, addMcpCmd, runCmd, mcp.NewMCPCmd(), envdoc.NewEnvCmd(), dbcli.NewCommand(builtinMigrationSources()...)) return rootCmd } +// builtinMigrationSources returns the built-in migration tracks for `kagent +// db migrate`. The vector track is gated on the same DATABASE_VECTOR_ENABLED +// env var the controller reads for --database-vector-enabled, with the same +// default (enabled), so the CLI operates on the tracks the server migrates. +func builtinMigrationSources() []migrations.Source { + vectorEnabled := true + if v := os.Getenv("DATABASE_VECTOR_ENABLED"); v != "" { + b, err := strconv.ParseBool(v) + if err != nil { + fmt.Fprintf(os.Stderr, "warning: invalid DATABASE_VECTOR_ENABLED=%q; assuming true\n", v) + } else { + vectorEnabled = b + } + } + return migrations.BuiltinSources(vectorEnabled) +} + func runInteractive(cmd *cobra.Command, args []string, cfg *config.Config) { client := cfg.Client() diff --git a/go/core/pkg/cli/db/db.go b/go/core/pkg/cli/db/db.go new file mode 100644 index 0000000000..56299bab6b --- /dev/null +++ b/go/core/pkg/cli/db/db.go @@ -0,0 +1,43 @@ +// Package db hosts the `kagent db` parent command and its subcommands. +// Currently only `migrate` is wired; future siblings attach here. +package db + +import ( + "github.com/spf13/cobra" + "github.com/spf13/pflag" + + "github.com/kagent-dev/kagent/go/core/pkg/cli/db/migrate" + "github.com/kagent-dev/kagent/go/core/pkg/migrations" +) + +// NewCommand returns the `db` parent command with `migrate` attached. The +// given sources define the migration tracks the subcommands operate on, in +// orchestrator registration order. +func NewCommand(sources ...migrations.Source) *cobra.Command { + cmd := &cobra.Command{ + Use: "db", + Short: "Database operations (migrations, inspection)", + } + cmd.AddCommand(migrate.NewCommand(sources...)) + + // Hide the root's API-oriented persistent flags from help across the + // entire `db` subtree. They target the kagent server / Kubernetes, + // but db commands talk to Postgres directly via --db-url. + // + // Hidden is a property of the flag itself (shared across the whole + // tree), so we can't flip it permanently. The HelpFunc override + // toggles it for the duration of the help render and restores it + // after. Children of `db` that don't set their own HelpFunc walk the + // parent chain and pick this one up. + cmd.SetHelpFunc(func(c *cobra.Command, args []string) { + for _, name := range []string{"config", "kagent-url", "namespace", "output-format", "timeout", "verbose"} { + if f := c.InheritedFlags().Lookup(name); f != nil { + f.Hidden = true + defer func(f *pflag.Flag) { f.Hidden = false }(f) + } + } + c.Root().HelpFunc()(c, args) + }) + + return cmd +} diff --git a/go/core/pkg/cli/db/migrate/migrate.go b/go/core/pkg/cli/db/migrate/migrate.go new file mode 100644 index 0000000000..9a00b75511 --- /dev/null +++ b/go/core/pkg/cli/db/migrate/migrate.go @@ -0,0 +1,729 @@ +// Package migrate exposes the `kagent db migrate` subcommand: out-of-band +// application, rollback, and inspection of the database migration tracks. +// +// The command operates on the same migrations.Source values app.Start +// registers and opens migrators through migrations.WithMigrator, so every +// operation shares the orchestrator's schema handling, tracking tables, and +// advisory-lock identity — a CLI invocation racing a booting server +// serializes instead of corrupting state. +// +// Unlike the in-app startup path, the CLI never auto-recovers a dirty +// tracking table: up, down, and goto refuse a dirty source and report the +// `force` invocation that clears it. Deliberate recovery is the operator's +// call. +package migrate + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "regexp" + "slices" + "strconv" + "strings" + + "github.com/golang-migrate/migrate/v4" + "github.com/spf13/cobra" + + "github.com/kagent-dev/kagent/go/core/pkg/migrations" +) + +const ( + // dbURLEnv is the controller's env var for --postgres-database-url + // (see app.LoadFromEnv); the CLI falls back to it so an operator with + // the controller's environment needs no extra flags. + dbURLEnv = "POSTGRES_DATABASE_URL" + sourceFlag = "source" +) + +// sourceNameRE constrains Source.Name to lowercase identifiers with +// underscores or hyphens (e.g. "core", "vector", and hyphenated names from +// downstream-registered sources). Name flows into `--source ` handling +// and operator-facing messages; the regex keeps those strings predictable. +var sourceNameRE = regexp.MustCompile(`^[a-z][a-z0-9_-]*$`) + +type commandState struct { + dbURL string + source string + sources []migrations.Source +} + +// NewCommand returns the `migrate` parent command with all subcommands +// attached, operating on the given sources in orchestrator registration +// order. Panics on an invalid or duplicate source Name so misconfiguration +// fails at wiring time, not mid-operation. +func NewCommand(sources ...migrations.Source) *cobra.Command { + seen := map[string]bool{} + for _, source := range sources { + if !sourceNameRE.MatchString(source.Name) { + panic(fmt.Sprintf("migrate.NewCommand: source Name=%q must match %s", source.Name, sourceNameRE.String())) + } + if seen[source.Name] { + panic(fmt.Sprintf("migrate.NewCommand: source %q already configured; each source must have a unique Name", source.Name)) + } + seen[source.Name] = true + } + state := &commandState{sources: append([]migrations.Source(nil), sources...)} + + cmd := &cobra.Command{ + Use: "migrate", + Short: "Apply, roll back, and inspect database migrations", + Long: `Apply, roll back, and inspect database migrations independently +of server startup. Reads ` + dbURLEnv + ` from the environment when +--db-url is omitted.`, + } + cmd.PersistentFlags().StringVar(&state.dbURL, "db-url", "", + "PostgreSQL connection URL (defaults to value of "+dbURLEnv+" env var)") + cmd.PersistentFlags().StringVar(&state.source, sourceFlag, "", + "Migration source name for per-source ops (down/goto/force/version); inferred when only one source is registered. Not applicable to up or status — those aggregate across every registered source.") + + cmd.AddCommand(newUpCmd(state)) + cmd.AddCommand(newDownCmd(state)) + cmd.AddCommand(newStatusCmd(state)) + cmd.AddCommand(newVersionCmd(state)) + cmd.AddCommand(newGotoCmd(state)) + cmd.AddCommand(newForceCmd(state)) + return cmd +} + +func (s *commandState) resolveDSN() (string, error) { + dsn := strings.TrimSpace(s.dbURL) + if dsn == "" { + dsn = os.Getenv(dbURLEnv) + } + if dsn == "" { + return "", fmt.Errorf("database URL not set; pass --db-url or set %s", dbURLEnv) + } + return dsn, nil +} + +// resolveSource picks the source for a per-source operation. With one source +// registered it's returned directly; with more than one the operator must +// pass --source and we report the registered set when they don't. +func (s *commandState) resolveSource() (migrations.Source, error) { + srcs := s.sources + if len(srcs) == 0 { + return migrations.Source{}, errors.New("no migration sources registered") + } + if len(srcs) == 1 { + if s.source != "" && s.source != srcs[0].Name { + return migrations.Source{}, fmt.Errorf("--source %q not registered; registered source: %s", s.source, srcs[0].Name) + } + return srcs[0], nil + } + if s.source == "" { + return migrations.Source{}, fmt.Errorf("registered sources: %s; pass --source", sourceNames(srcs)) + } + for _, src := range srcs { + if src.Name == s.source { + return src, nil + } + } + return migrations.Source{}, fmt.Errorf("--source %q not registered; registered sources: %s", s.source, sourceNames(srcs)) +} + +func sourceNames(srcs []migrations.Source) string { + names := make([]string, len(srcs)) + for i, s := range srcs { + names[i] = s.Name + } + return strings.Join(names, ", ") +} + +// readVersion returns mg's highest applied version and whether the tracking +// row is dirty (mid-failed-migration). ErrNilVersion (nothing applied) is +// normalized to (0, false, nil). +func readVersion(mg *migrate.Migrate) (uint, bool, error) { + v, dirty, err := mg.Version() + if err != nil { + if errors.Is(err, migrate.ErrNilVersion) { + return 0, false, nil + } + return 0, false, fmt.Errorf("read version: %w", err) + } + return v, dirty, nil +} + +// ensureClean returns an actionable error when src's tracking table is dirty. +// The CLI refuses to operate on a dirty source rather than auto-recovering +// (the in-app startup path is the automatic tier; the CLI is the manual one). +func ensureClean(src migrations.Source, mg *migrate.Migrate) error { + v, dirty, err := readVersion(mg) + if err != nil { + return err + } + if dirty { + return fmt.Errorf("source %q is dirty at version %d (a previous migration attempt failed and must be resolved manually); inspect the schema, then clear the flag with \"kagent db migrate force --source %s\" where is the version the schema actually reflects", + src.Name, v, src.Name) + } + return nil +} + +// sourceFileVersions returns the (ascending-sorted) NNN versions parsed from +// every NNN_name.up.sql file in src.FS/src.Dir. +// +// The set isn't required to be contiguous — gaps (e.g. a deleted 005) are +// real and the missing numbers are treated as not-applicable-to-this-binary. +// status/desync math goes via this list so the count of files and the highest +// version stay distinct. +func sourceFileVersions(src migrations.Source) ([]int, error) { + entries, err := fs.ReadDir(src.FS, src.Dir) + if err != nil { + return nil, fmt.Errorf("read migration dir %s: %w", src.Dir, err) + } + var versions []int + for _, e := range entries { + if e.IsDir() { + continue + } + name := e.Name() + if !strings.HasSuffix(name, ".up.sql") { + continue + } + parts := strings.SplitN(name, "_", 2) + if len(parts) != 2 { + continue + } + v, err := strconv.Atoi(parts[0]) + if err != nil { + continue + } + versions = append(versions, v) + } + slices.Sort(versions) + return versions, nil +} + +// lineRow carries per-source status data through the status command's text +// and JSON output paths. +type lineRow struct { + src migrations.Source + applied int + pending int + dbVersion int // raw DB version for desync reporting + downgraded bool // dbVersion > highest shipped version + dirty bool // mid-failed-migration; surfaced as a (dirty) annotation +} + +func newUpCmd(state *commandState) *cobra.Command { + return &cobra.Command{ + Use: "up", + Short: "Apply all pending migrations across every registered source", + Long: `Applies pending migrations for every registered source in +registration order, through the same orchestrator the server runs at +startup: per-source advisory locking, pre-run version snapshots, and +compensating rollback of earlier sources when a later one fails. + +Refuses to run while any source's tracking table is dirty; clear it +with 'force' first. + +The --source flag is intentionally not applicable to up; pass it only +on the per-source subcommands (down/goto/force).`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + if state.source != "" { + return errors.New("up aggregates across all registered sources; --source is not applicable") + } + dsn, err := state.resolveDSN() + if err != nil { + return err + } + srcs := state.sources + if len(srcs) == 0 { + return errors.New("no migration sources registered") + } + + ctx := cmd.Context() + // One pre-pass over the sources: refuse dirty state, and snapshot + // pending counts so we can report "applied N migration(s)" after + // the orchestrator succeeds. + prePending := 0 + for _, src := range srcs { + p, err := pendingCount(ctx, src, dsn) + if err != nil { + return err + } + prePending += p + } + + if err := migrations.RunUp(ctx, dsn, srcs); err != nil { + return err + } + + if prePending == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "no pending migrations; schema is up to date") + return nil + } + fmt.Fprintf(cmd.OutOrStdout(), "applied %d migration(s); schema is up to date\n", prePending) + return nil + }, + } +} + +// pendingCount counts NNN_*.up.sql files whose version is greater than the +// source's current applied version, refusing a dirty source. Uses the same +// `sourceFileVersions` primitive as `status` so the two paths can't drift. +func pendingCount(ctx context.Context, src migrations.Source, dsn string) (int, error) { + versions, err := sourceFileVersions(src) + if err != nil { + return 0, err + } + var pending int + err = migrations.WithMigrator(ctx, dsn, src, func(mg *migrate.Migrate) error { + if err := ensureClean(src, mg); err != nil { + return err + } + v, _, verr := readVersion(mg) + if verr != nil { + return verr + } + for _, fv := range versions { + if uint(fv) > v { + pending++ + } + } + return nil + }) + if err != nil { + return 0, err + } + return pending, nil +} + +func newDownCmd(state *commandState) *cobra.Command { + return &cobra.Command{ + Use: "down N", + Short: "Roll back the N most-recent applied migrations for the selected source", + Long: `Roll back the N most-recent applied migrations for the selected source. + +Down migrations can lose data by design — a rolled-back column loses +its contents. Refuses to run while the source's tracking table is +dirty; clear it with 'force' first.`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + n, err := strconv.Atoi(args[0]) + if err != nil || n < 1 { + return fmt.Errorf("expected a positive integer for N, got %q", args[0]) + } + dsn, err := state.resolveDSN() + if err != nil { + return err + } + src, err := state.resolveSource() + if err != nil { + return err + } + return migrations.WithMigrator(cmd.Context(), dsn, src, func(mg *migrate.Migrate) error { + if err := ensureClean(src, mg); err != nil { + return err + } + preV, _, err := readVersion(mg) + if err != nil { + return err + } + // Guarded here rather than relying on golang-migrate: Steps + // on an empty track reports a confusing "file does not + // exist" instead of ErrNoChange. + if preV == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "no migrations to roll back") + return nil + } + if err := mg.Steps(-n); err != nil { + if errors.Is(err, migrate.ErrNoChange) { + fmt.Fprintln(cmd.OutOrStdout(), "no migrations to roll back") + return nil + } + return err + } + postV, _, verr := readVersion(mg) + if verr != nil { + return fmt.Errorf("read version after rollback: %w", verr) + } + rolled := countVersionsBetween(src, postV, preV) + fmt.Fprintf(cmd.OutOrStdout(), "rolled back %d migration(s)\n", rolled) + return nil + }) + }, + } +} + +func newStatusCmd(state *commandState) *cobra.Command { + var output string + cmd := &cobra.Command{ + Use: "status", + Short: "Show how many migrations are applied vs pending across all sources", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + if state.source != "" { + return errors.New("status aggregates across all registered sources; --source is not applicable") + } + if output != "text" && output != "json" { + return fmt.Errorf("invalid --output %q; supported: text, json", output) + } + dsn, err := state.resolveDSN() + if err != nil { + return err + } + srcs := state.sources + if len(srcs) == 0 { + return errors.New("no migration sources registered") + } + + lines := make([]lineRow, 0, len(srcs)) + appliedTotal, pendingTotal := 0, 0 + // Single-source builds print no per-source breakdown, so the + // stderr desync warning below is gated to them as their only + // signal; multi-source builds carry it in the stdout breakdown. + multiSource := len(srcs) > 1 + for _, src := range srcs { + versions, err := sourceFileVersions(src) + if err != nil { + return err + } + maxFileVersion := 0 + if len(versions) > 0 { + maxFileVersion = versions[len(versions)-1] + } + var applied, dbVersion int + var downgraded, dirty bool + if rerr := migrations.WithMigrator(cmd.Context(), dsn, src, func(mg *migrate.Migrate) error { + v, d, err := readVersion(mg) + if err != nil { + return err + } + dbVersion = int(v) + dirty = d + // Count files at/below the DB version as applied + // (counting by version, not file count, survives gaps + // like a deleted v5). + for _, fv := range versions { + if fv <= dbVersion { + applied++ + } + } + if dbVersion > maxFileVersion { + // Older binary against a DB migrated by a newer + // build. Warn, don't fail. + downgraded = true + if !multiSource { + fmt.Fprintf(cmd.ErrOrStderr(), + "warning: %s reports version %d but this binary's highest shipped migration is v%d (older binary against newer DB?)\n", + src.Name, v, maxFileVersion) + } + } + return nil + }); rerr != nil { + return rerr + } + pending := len(versions) - applied + lines = append(lines, lineRow{src: src, applied: applied, pending: pending, dbVersion: dbVersion, downgraded: downgraded, dirty: dirty}) + appliedTotal += applied + pendingTotal += pending + } + + out := cmd.OutOrStdout() + if output == "json" { + return writeStatusJSON(out, lines, appliedTotal, pendingTotal) + } + if multiSource { + fmt.Fprintf(out, "%d migration(s) applied, %d pending\n", appliedTotal, pendingTotal) + // Reuses the same `multiSource` gate as the stderr desync + // warning above on purpose: a per-source skip branch must + // not let the two diverge (a desync gets a source warned + // twice or not at all). + for _, l := range lines { + if l.downgraded { + fmt.Fprintf(out, " %s: %d applied, %d pending (db reports v%d%s — binary out of date)\n", + l.src.Name, l.applied, l.pending, l.dbVersion, dirtyTag(l.dirty)) + } else { + fmt.Fprintf(out, " %s: %d applied (at v%d%s), %d pending\n", + l.src.Name, l.applied, l.dbVersion, dirtyTag(l.dirty), l.pending) + } + } + } else { + // Single-source: fold the version into the headline so + // operators needn't run `version` separately. dbVersion is + // the raw tracking-table value (matches `force V`). + l := lines[0] + fmt.Fprintf(out, "%d migration(s) applied (at v%d%s), %d pending\n", + l.applied, l.dbVersion, dirtyTag(l.dirty), l.pending) + } + return nil + }, + } + // No -o shorthand: the kagent root command already owns -o + // (--output-format) as a persistent flag, and cobra panics on a + // shorthand redefinition. + cmd.Flags().StringVar(&output, "output", "text", + `Output format: "text" (default) or "json"`) + return cmd +} + +// statusJSON is the wire format for `kagent db migrate status -o json`. +// Operators consume it via `jq`, so the field names and types are a frozen +// contract; TestStatusJSONShape locks them and fails CI on a rename or +// retype. +type statusJSON struct { + Applied int `json:"applied"` + Pending int `json:"pending"` + Sources []statusSourceJSON `json:"sources"` +} + +// statusSourceJSON is the per-source object inside statusJSON; same +// frozen-shape contract, also locked by TestStatusJSONShape. +type statusSourceJSON struct { + Name string `json:"name"` + Applied int `json:"applied"` + Pending int `json:"pending"` + Version int `json:"version"` + Downgraded bool `json:"downgraded"` + Dirty bool `json:"dirty"` +} + +func writeStatusJSON(out io.Writer, lines []lineRow, appliedTotal, pendingTotal int) error { + payload := statusJSON{ + Applied: appliedTotal, + Pending: pendingTotal, + Sources: make([]statusSourceJSON, 0, len(lines)), + } + for _, l := range lines { + payload.Sources = append(payload.Sources, statusSourceJSON{ + Name: l.src.Name, + Applied: l.applied, + Pending: l.pending, + Version: l.dbVersion, + Downgraded: l.downgraded, + Dirty: l.dirty, + }) + } + enc := json.NewEncoder(out) + enc.SetIndent("", " ") + return enc.Encode(payload) +} + +// dirtyTag returns " (dirty)" when the source is mid-failed-migration, "" +// otherwise. Used to annotate the version in status/version text output +// without adding a separate line. +func dirtyTag(dirty bool) string { + if dirty { + return " (dirty)" + } + return "" +} + +// versionAnnotation renders the trailing annotation for `version` output. +// Disambiguates an unapplied-migrations state (v=0, !dirty) from a versioned +// state by tagging the former; dirty wins over the "no migrations applied" +// tag because it's the more actionable signal. +func versionAnnotation(v uint, dirty bool) string { + if dirty { + return " (dirty)" + } + if v == 0 { + return " (no migrations applied)" + } + return "" +} + +func newVersionCmd(state *commandState) *cobra.Command { + return &cobra.Command{ + Use: "version", + Short: "Print the highest applied migration version", + Long: `Print the highest applied migration version. +For a single registered source the value is on one line; multi-source +binaries print one line per source. When multiple sources are +registered, --source filters to a single track.`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + dsn, err := state.resolveDSN() + if err != nil { + return err + } + srcs := state.sources + if len(srcs) == 0 { + return errors.New("no migration sources registered") + } + // --source filters the output even though version is otherwise + // an aggregate op. Empty flag = print all. + if state.source != "" { + picked := -1 + for i, s := range srcs { + if s.Name == state.source { + picked = i + break + } + } + if picked < 0 { + return fmt.Errorf("--source %q not registered; registered sources: %s", state.source, sourceNames(srcs)) + } + srcs = []migrations.Source{srcs[picked]} + } + out := cmd.OutOrStdout() + if len(srcs) == 1 { + return migrations.WithMigrator(cmd.Context(), dsn, srcs[0], func(mg *migrate.Migrate) error { + v, dirty, err := readVersion(mg) + if err != nil { + return err + } + fmt.Fprintf(out, "%d%s\n", v, versionAnnotation(v, dirty)) + return nil + }) + } + for _, src := range srcs { + if err := migrations.WithMigrator(cmd.Context(), dsn, src, func(mg *migrate.Migrate) error { + v, dirty, err := readVersion(mg) + if err != nil { + return err + } + fmt.Fprintf(out, "%s: %d%s\n", src.Name, v, versionAnnotation(v, dirty)) + return nil + }); err != nil { + return err + } + } + return nil + }, + } +} + +func newGotoCmd(state *commandState) *cobra.Command { + return &cobra.Command{ + Use: "goto V", + Short: "Move the selected source's schema to version V", + Long: `Move the selected source's schema to version V (forward or backward). +V=0 is the special "empty schema" target: every applied migration in +the source is rolled back. + +Refuses to run while the source's tracking table is dirty; clear it +with 'force' first.`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + v, err := strconv.Atoi(args[0]) + if err != nil || v < 0 { + return fmt.Errorf("expected a non-negative integer for V, got %q", args[0]) + } + dsn, err := state.resolveDSN() + if err != nil { + return err + } + src, err := state.resolveSource() + if err != nil { + return err + } + return migrations.WithMigrator(cmd.Context(), dsn, src, func(mg *migrate.Migrate) error { + if err := ensureClean(src, mg); err != nil { + return err + } + if v == 0 { + if err := mg.Down(); err != nil && !errors.Is(err, migrate.ErrNoChange) { + return err + } + fmt.Fprintln(cmd.OutOrStdout(), "schema is at version 0 (empty)") + return nil + } + if err := mg.Migrate(uint(v)); err != nil && !errors.Is(err, migrate.ErrNoChange) { + return err + } + actual, dirty, aerr := readVersion(mg) + if aerr != nil { + return aerr + } + fmt.Fprintf(cmd.OutOrStdout(), "schema is at version %d%s\n", actual, versionAnnotation(actual, dirty)) + return nil + }) + }, + } +} + +func newForceCmd(state *commandState) *cobra.Command { + return &cobra.Command{ + Use: "force V", + Short: "Mark version V as applied without running its SQL", + Long: `Used to reconcile the selected source's tracking table after manual +remediation, e.g. to clear a dirty flag left by a failed migration. +V=0 clears the version record entirely (the "no migrations applied" +state). Any other V must correspond to a shipped migration file in +the selected source — otherwise the tracking row would point at a +version the binary cannot apply or roll back to, wedging the DB.`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + v, err := strconv.Atoi(args[0]) + if err != nil || v < 0 { + return fmt.Errorf("expected a non-negative integer for V, got %q", args[0]) + } + dsn, err := state.resolveDSN() + if err != nil { + return err + } + src, err := state.resolveSource() + if err != nil { + return err + } + if v > 0 { + versions, err := sourceFileVersions(src) + if err != nil { + return err + } + if !slices.Contains(versions, v) { + return fmt.Errorf( + "version %d is not a shipped migration for source %q; valid versions are %s", + v, src.Name, formatVersionList(versions)) + } + } + return migrations.WithMigrator(cmd.Context(), dsn, src, func(mg *migrate.Migrate) error { + if v == 0 { + // golang-migrate's Force takes -1 to delete the version + // record; 0 is not a valid stored version. + if err := mg.Force(-1); err != nil { + return err + } + fmt.Fprintln(cmd.OutOrStdout(), "version record cleared (no migrations applied)") + return nil + } + if err := mg.Force(v); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "version %d marked as applied\n", v) + return nil + }) + }, + } +} + +// countVersionsBetween returns the count of shipped source migrations in the +// half-open interval (low, high]. Used by `down N` to report the actual +// number of migrations rolled back regardless of whether the user-supplied N +// exceeded the applied count. Returns 0 on a source-enumeration error; the +// call site already surfaces success/failure via Steps(-N) and the count is +// operator-facing display only. +func countVersionsBetween(src migrations.Source, low, high uint) int { + versions, err := sourceFileVersions(src) + if err != nil { + return 0 + } + count := 0 + for _, v := range versions { + uv := uint(v) + if uv > low && uv <= high { + count++ + } + } + return count +} + +// formatVersionList renders a small []int as a human-readable list for error +// messages: "1, 2, 5" or "(none)" if empty. +func formatVersionList(versions []int) string { + if len(versions) == 0 { + return "(none)" + } + parts := make([]string, len(versions)) + for i, v := range versions { + parts[i] = strconv.Itoa(v) + } + return strings.Join(parts, ", ") +} diff --git a/go/core/pkg/cli/db/migrate/migrate_test.go b/go/core/pkg/cli/db/migrate/migrate_test.go new file mode 100644 index 0000000000..86660acdd1 --- /dev/null +++ b/go/core/pkg/cli/db/migrate/migrate_test.go @@ -0,0 +1,314 @@ +package migrate + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "strings" + "testing" + "testing/fstest" + + _ "github.com/jackc/pgx/v5/stdlib" + + "github.com/kagent-dev/kagent/go/core/internal/dbtest" + "github.com/kagent-dev/kagent/go/core/pkg/migrations" +) + +// --- fixtures --- + +var alphaFS = fstest.MapFS{ + "alpha/000001_create.up.sql": {Data: []byte(`CREATE TABLE IF NOT EXISTS cli_alpha (id SERIAL PRIMARY KEY);`)}, + "alpha/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS cli_alpha;`)}, + "alpha/000002_alter.up.sql": {Data: []byte(`ALTER TABLE cli_alpha ADD COLUMN IF NOT EXISTS name TEXT;`)}, + "alpha/000002_alter.down.sql": {Data: []byte(`ALTER TABLE cli_alpha DROP COLUMN IF EXISTS name;`)}, +} + +var betaFS = fstest.MapFS{ + "beta/000001_create.up.sql": {Data: []byte(`CREATE TABLE IF NOT EXISTS cli_beta (id SERIAL PRIMARY KEY);`)}, + "beta/000001_create.down.sql": {Data: []byte(`DROP TABLE IF EXISTS cli_beta;`)}, +} + +func testSources() []migrations.Source { + return []migrations.Source{ + {Name: "alpha", TrackingTable: "alpha_schema_migrations", FS: alphaFS, Dir: "alpha"}, + {Name: "beta", TrackingTable: "beta_schema_migrations", FS: betaFS, Dir: "beta"}, + } +} + +// runCLI executes `db migrate ` against a fresh command tree and +// returns combined stdout and the error, plus stderr separately. +func runCLI(t *testing.T, sources []migrations.Source, args ...string) (string, string, error) { + t.Helper() + cmd := NewCommand(sources...) + var out, errOut bytes.Buffer + cmd.SetOut(&out) + cmd.SetErr(&errOut) + cmd.SetArgs(args) + err := cmd.ExecuteContext(context.Background()) + return out.String(), errOut.String(), err +} + +// --- unit tests (no database) --- + +func TestNewCommandRejectsBadSources(t *testing.T) { + tests := []struct { + name string + sources []migrations.Source + }{ + {name: "invalid name", sources: []migrations.Source{{Name: "Bad Name", TrackingTable: "t", FS: alphaFS, Dir: "alpha"}}}, + {name: "duplicate name", sources: []migrations.Source{ + {Name: "core", TrackingTable: "a", FS: alphaFS, Dir: "alpha"}, + {Name: "core", TrackingTable: "b", FS: betaFS, Dir: "beta"}, + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatal("expected NewCommand to panic") + } + }() + NewCommand(tt.sources...) + }) + } +} + +func TestNewCommandAcceptsHyphenatedNames(t *testing.T) { + // Downstream-registered sources may use hyphenated names. + NewCommand(migrations.Source{Name: "extra-track", TrackingTable: "t", FS: alphaFS, Dir: "alpha"}) +} + +func TestResolveDSN(t *testing.T) { + tests := []struct { + name string + flag string + env string + want string + wantErr bool + }{ + {name: "flag wins over env", flag: "postgres://flag", env: "postgres://env", want: "postgres://flag"}, + {name: "env fallback", env: "postgres://env", want: "postgres://env"}, + {name: "neither set", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv(dbURLEnv, tt.env) + s := &commandState{dbURL: tt.flag} + got, err := s.resolveDSN() + if (err != nil) != tt.wantErr { + t.Fatalf("resolveDSN() error = %v, wantErr %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("resolveDSN() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestResolveSource(t *testing.T) { + multi := testSources() + single := multi[:1] + tests := []struct { + name string + sources []migrations.Source + flag string + want string + wantErr string + }{ + {name: "single source inferred", sources: single, want: "alpha"}, + {name: "single source explicit match", sources: single, flag: "alpha", want: "alpha"}, + {name: "single source mismatch", sources: single, flag: "beta", wantErr: "not registered"}, + {name: "multi requires flag", sources: multi, wantErr: "pass --source"}, + {name: "multi explicit", sources: multi, flag: "beta", want: "beta"}, + {name: "multi unknown", sources: multi, flag: "nope", wantErr: "not registered"}, + {name: "none registered", sources: nil, wantErr: "no migration sources"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := &commandState{sources: tt.sources, source: tt.flag} + got, err := s.resolveSource() + if tt.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("resolveSource() error = %v, want containing %q", err, tt.wantErr) + } + return + } + if err != nil { + t.Fatalf("resolveSource() unexpected error: %v", err) + } + if got.Name != tt.want { + t.Errorf("resolveSource() = %q, want %q", got.Name, tt.want) + } + }) + } +} + +// TestArgValidation covers rejections that fire before any database access. +func TestArgValidation(t *testing.T) { + tests := []struct { + name string + args []string + wantErr string + }{ + {name: "down non-integer", args: []string{"down", "abc"}, wantErr: "positive integer"}, + {name: "down zero", args: []string{"down", "0"}, wantErr: "positive integer"}, + {name: "goto non-integer", args: []string{"goto", "abc"}, wantErr: "non-negative integer"}, + {name: "goto negative", args: []string{"goto", "--", "-1"}, wantErr: "non-negative integer"}, + {name: "force non-integer", args: []string{"force", "abc"}, wantErr: "non-negative integer"}, + {name: "up rejects source flag", args: []string{"up", "--source", "alpha"}, wantErr: "--source is not applicable"}, + {name: "status rejects source flag", args: []string{"status", "--source", "alpha"}, wantErr: "--source is not applicable"}, + {name: "status invalid output", args: []string{"status", "--output", "yaml"}, wantErr: "invalid --output"}, + {name: "no dsn", args: []string{"version"}, wantErr: "database URL not set"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv(dbURLEnv, "") + _, _, err := runCLI(t, testSources(), tt.args...) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %v, want containing %q", err, tt.wantErr) + } + }) + } +} + +// TestStatusJSONShape freezes the `status -o json` wire format. Operators +// consume it via jq, so a rename or retype is a breaking change; update this +// test only with a deliberate contract change. +func TestStatusJSONShape(t *testing.T) { + payload := statusJSON{ + Applied: 3, + Pending: 1, + Sources: []statusSourceJSON{{Name: "alpha", Applied: 2, Pending: 1, Version: 2, Downgraded: false, Dirty: true}}, + } + got, err := json.Marshal(payload) + if err != nil { + t.Fatal(err) + } + want := `{"applied":3,"pending":1,"sources":[{"name":"alpha","applied":2,"pending":1,"version":2,"downgraded":false,"dirty":true}]}` + if string(got) != want { + t.Errorf("status JSON shape changed:\n got: %s\nwant: %s", got, want) + } +} + +// --- database-backed tests --- + +// TestCLIAgainstPostgres walks the operator surface end to end against a real +// Postgres: up, status, version, down, goto, dirty refusal, and force. The +// subtests share one container and run in order — each builds on the schema +// state the previous one left. +func TestCLIAgainstPostgres(t *testing.T) { + if testing.Short() { + t.Skip("skipping container-backed test in -short mode") + } + ctx := context.Background() + dsn := dbtest.StartT(ctx, t) + sources := testSources() + + mustContain := func(t *testing.T, out, want string) { + t.Helper() + if !strings.Contains(out, want) { + t.Fatalf("output %q does not contain %q", out, want) + } + } + mustRun := func(t *testing.T, args ...string) string { + t.Helper() + out, _, err := runCLI(t, sources, append(args, "--db-url", dsn)...) + if err != nil { + t.Fatalf("%v: %v", args, err) + } + return out + } + + t.Run("up applies all sources", func(t *testing.T) { + mustContain(t, mustRun(t, "up"), "applied 3 migration(s)") + }) + + t.Run("up is idempotent", func(t *testing.T) { + mustContain(t, mustRun(t, "up"), "no pending migrations") + }) + + t.Run("status text", func(t *testing.T) { + out := mustRun(t, "status") + mustContain(t, out, "3 migration(s) applied, 0 pending") + mustContain(t, out, "alpha: 2 applied (at v2), 0 pending") + mustContain(t, out, "beta: 1 applied (at v1), 0 pending") + }) + + t.Run("status json", func(t *testing.T) { + var got statusJSON + if err := json.Unmarshal([]byte(mustRun(t, "status", "--output", "json")), &got); err != nil { + t.Fatal(err) + } + if got.Applied != 3 || got.Pending != 0 || len(got.Sources) != 2 { + t.Fatalf("unexpected status: %+v", got) + } + }) + + t.Run("version per source and filtered", func(t *testing.T) { + out := mustRun(t, "version") + mustContain(t, out, "alpha: 2") + mustContain(t, out, "beta: 1") + mustContain(t, mustRun(t, "version", "--source", "alpha"), "2") + }) + + t.Run("down rolls back one step", func(t *testing.T) { + mustContain(t, mustRun(t, "down", "1", "--source", "alpha"), "rolled back 1 migration(s)") + mustContain(t, mustRun(t, "version", "--source", "alpha"), "1") + }) + + t.Run("goto moves forward", func(t *testing.T) { + mustContain(t, mustRun(t, "goto", "2", "--source", "alpha"), "schema is at version 2") + }) + + t.Run("goto zero empties the source", func(t *testing.T) { + mustContain(t, mustRun(t, "goto", "0", "--source", "beta"), "version 0 (empty)") + mustContain(t, mustRun(t, "version", "--source", "beta"), "no migrations applied") + }) + + t.Run("down with nothing to roll back", func(t *testing.T) { + mustContain(t, mustRun(t, "down", "1", "--source", "beta"), "no migrations to roll back") + }) + + t.Run("dirty source is refused and force recovers", func(t *testing.T) { + markDirty(t, dsn, "alpha_schema_migrations") + for _, args := range [][]string{{"up"}, {"down", "1", "--source", "alpha"}, {"goto", "1", "--source", "alpha"}} { + _, _, err := runCLI(t, sources, append(args, "--db-url", dsn)...) + if err == nil || !strings.Contains(err.Error(), "dirty") { + t.Fatalf("%v: error = %v, want dirty refusal", args, err) + } + } + // status still reports rather than refusing, and annotates the row. + mustContain(t, mustRun(t, "status"), "(dirty)") + + mustContain(t, mustRun(t, "force", "2", "--source", "alpha"), "version 2 marked as applied") + mustContain(t, mustRun(t, "up"), "applied 1 migration(s)") // beta was left at 0 by the goto above + }) + + t.Run("force rejects unshipped version", func(t *testing.T) { + _, _, err := runCLI(t, sources, "force", "99", "--source", "alpha", "--db-url", dsn) + if err == nil || !strings.Contains(err.Error(), "not a shipped migration") { + t.Fatalf("error = %v, want unshipped-version rejection", err) + } + }) + + t.Run("force zero clears the version record", func(t *testing.T) { + mustContain(t, mustRun(t, "force", "0", "--source", "beta"), "version record cleared") + mustContain(t, mustRun(t, "version", "--source", "beta"), "no migrations applied") + mustContain(t, mustRun(t, "up"), "applied 1 migration(s)") + }) +} + +// markDirty flips the dirty flag on a tracking table, simulating a process +// that died mid-migration. +func markDirty(t *testing.T, dsn, table string) { + t.Helper() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + if _, err := db.Exec("UPDATE " + table + " SET dirty = true"); err != nil { + t.Fatalf("mark %s dirty: %v", table, err) + } +} diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index 2b697430b2..a90dd946ac 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -350,6 +350,23 @@ func applySource(ctx context.Context, url string, src Source) (prevVersion uint, return prevVersion, nil } +// WithMigrator opens a migrator for src against url, runs fn against it, and +// closes it. The migrator carries the same schema handling, tracking-table +// configuration, and advisory-lock identity as the orchestrator's own runs, so +// out-of-band tooling (the `kagent db migrate` CLI) built on this serializes +// correctly against a concurrently booting server and cannot drift from the +// startup path. fn's migration operations (Up/Down/Steps/Migrate/Force) each +// take golang-migrate's per-(database, schema) advisory lock; Version reads do +// not. +func WithMigrator(ctx context.Context, url string, src Source, fn func(*migrate.Migrate) error) error { + mg, err := newMigrate(ctx, url, src) + if err != nil { + return err + } + defer closeMigrate(src.Name, mg) + return fn(mg) +} + // rollbackSource opens a fresh migrate instance and rolls a source back to // targetVersion. Used to compensate a previously-succeeded source when a later // source fails. It returns an error (also logged) when the rollback fails, so From c75eee187669a44d494159071c6401e380330517 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Mon, 6 Jul 2026 09:59:56 -0700 Subject: [PATCH 6/7] add SKIP_MIGRATIONS to run migrations out-of-band Signed-off-by: Jeremy Alvis --- go/core/pkg/app/app.go | 30 ++++-- go/core/pkg/migrations/runner.go | 71 ++++++++++++++ go/core/pkg/migrations/runner_test.go | 94 +++++++++++++++++++ helm/kagent/templates/NOTES.txt | 10 ++ .../templates/controller-configmap.yaml | 1 + helm/kagent/values.yaml | 4 + 6 files changed, 202 insertions(+), 8 deletions(-) diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 6cadbdd6d8..1cf46122c0 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -137,9 +137,10 @@ type Config struct { // that originates TLS upstream. Off by default; MCPEgressPlaintext bool Database struct { - Url string - UrlFile string - VectorEnabled bool + Url string + UrlFile string + VectorEnabled bool + SkipMigrations bool } Substrate struct { AteAPIEndpoint string @@ -181,6 +182,7 @@ func (cfg *Config) SetFlags(commandLine *flag.FlagSet) { commandLine.StringVar(&cfg.Database.Url, "postgres-database-url", "postgres://postgres:kagent@kagent-postgresql.kagent.svc.cluster.local:5432/postgres", "The URL of the PostgreSQL database.") commandLine.StringVar(&cfg.Database.UrlFile, "postgres-database-url-file", "", "Path to a file containing the PostgreSQL database URL. Takes precedence over --postgres-database-url.") commandLine.BoolVar(&cfg.Database.VectorEnabled, "database-vector-enabled", true, "Enable pgvector extension and memory table. Requires pgvector to be installed on the PostgreSQL server.") + commandLine.BoolVar(&cfg.Database.SkipMigrations, "skip-migrations", false, "Do not run database migrations at startup; instead verify the database is already migrated and fail if it is not. Migrations must be applied out-of-band (e.g. from a pipeline or pre-upgrade hook). Settable via the SKIP_MIGRATIONS env var.") commandLine.StringVar(&cfg.WatchNamespaces, "watch-namespaces", "", "The namespaces to watch for .") @@ -470,13 +472,25 @@ func Start(getExtensionConfig GetExtensionConfig, extraSources []migrations.Sour // Run migrations before connecting; schema must exist before queries. // Built-in sources run first, then any downstream-registered extras. - setupLog.Info("running database migrations") + // With --skip-migrations (SKIP_MIGRATIONS) the server applies nothing and + // instead verifies the database is already migrated, so migrations can run + // out-of-band and this connection needs no DDL privileges. sources := append(migrations.BuiltinSources(cfg.Database.VectorEnabled), extraSources...) - if err := migrations.RunUp(ctx, dbURL, sources); err != nil { - setupLog.Error(err, "database migration failed") - os.Exit(1) + if cfg.Database.SkipMigrations { + setupLog.Info("skipping database migrations; verifying schema is migrated") + if err := migrations.VerifyMigrated(ctx, dbURL, sources); err != nil { + setupLog.Error(err, "database migration verification failed") + os.Exit(1) + } + setupLog.Info("database schema verified") + } else { + setupLog.Info("running database migrations") + if err := migrations.RunUp(ctx, dbURL, sources); err != nil { + setupLog.Error(err, "database migration failed") + os.Exit(1) + } + setupLog.Info("database migrations complete") } - setupLog.Info("database migrations complete") // Connect to database db, err := database.Connect(ctx, &database.PostgresConfig{ diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index a90dd946ac..0687c16fe7 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -179,6 +179,77 @@ func RunUp(ctx context.Context, url string, sources []Source) error { return nil } +// VerifyMigrated checks, without applying or reverting anything, that every +// source's migrations have been applied to the database. It is the boot-time +// guard for the SKIP_MIGRATIONS deployment mode, where migrations run +// out-of-band (a pipeline or pre-upgrade hook) and the server must refuse to +// serve a wrong-shaped schema. It issues only SELECTs — never golang-migrate, +// which creates the tracking table on open — so it is safe on a connection +// whose role has no DDL privileges. +// +// Per source: a missing tracking table or a version behind this binary's +// embedded max is an error; a dirty tracking table is an error; a database +// ahead of the binary is tolerated (compatibility mode), matching RunUp. +func VerifyMigrated(ctx context.Context, url string, sources []Source) error { + if len(sources) == 0 { + return nil + } + if err := validateSources(sources); err != nil { + return err + } + + db, err := sql.Open("pgx", url) + if err != nil { + return fmt.Errorf("open database to verify migrations: %w", err) + } + defer db.Close() + + for _, src := range sources { + if err := ctx.Err(); err != nil { + return fmt.Errorf("migration verification cancelled before %s: %w", src.Name, err) + } + maxVer, err := maxEmbeddedVersion(src.FS, src.Dir) + if err != nil { + return fmt.Errorf("determine max embedded version for %s: %w", src.Name, err) + } + + // For Schema == "" the unqualified name resolves via the connection's + // search_path — the same place RunUp put the table. + table := quoteIdentifier(src.TrackingTable) + if src.Schema != "" { + table = quoteIdentifier(src.Schema) + "." + table + } + + var exists bool + if err := db.QueryRowContext(ctx, "SELECT to_regclass($1) IS NOT NULL", table).Scan(&exists); err != nil { + return fmt.Errorf("check tracking table for %s: %w", src.Name, err) + } + if !exists { + return fmt.Errorf("source %s: tracking table %s does not exist — the database has not been migrated; apply migrations out-of-band or unset SKIP_MIGRATIONS", src.Name, table) + } + + var version int64 + var dirty bool + err = db.QueryRowContext(ctx, "SELECT version, dirty FROM "+table+" LIMIT 1").Scan(&version, &dirty) + if errors.Is(err, sql.ErrNoRows) { + version, dirty = 0, false // table exists but nothing applied yet + } else if err != nil { + return fmt.Errorf("read tracking table for %s: %w", src.Name, err) + } + + switch { + case dirty: + return fmt.Errorf("source %s is dirty at version %d: a previous migration attempt failed and must be resolved before starting with SKIP_MIGRATIONS", src.Name, version) + case version < int64(maxVer): + return fmt.Errorf("source %s is at version %d but this binary requires version %d: apply migrations out-of-band or unset SKIP_MIGRATIONS", src.Name, version, maxVer) + case version > int64(maxVer): + log.Info("database schema is ahead of this binary; running in compatibility mode", + "track", src.Name, "dbVersion", version, "binaryMax", maxVer) + } + } + return nil +} + // validateSources rejects a source set that cannot be run safely. It checks two // things. // diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 81b36e48e5..9019aef821 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "maps" + "net/url" "strings" "testing" "testing/fstest" @@ -972,3 +973,96 @@ func TestApplySource_DirtyStateRecoveryOnRestart(t *testing.T) { t.Errorf("after restart: version = %d, want 1 (dirty cleared, rolled back)", got) } } + +// TestVerifyMigrated covers the SKIP_MIGRATIONS boot guard: refuse an +// un-migrated or behind or dirty database, tolerate exact-match and ahead +// (compatibility mode), and work on a connection whose role has no DDL +// privileges — the deployment mode the guard exists for. +func TestVerifyMigrated(t *testing.T) { + ctx := context.Background() + connStr := startTestDB(t) + full := []Source{coreSource(goodCoreFS)} // max embedded version 2 + + t.Run("unmigrated database is refused", func(t *testing.T) { + err := VerifyMigrated(ctx, connStr, full) + if err == nil || !strings.Contains(err.Error(), "has not been migrated") { + t.Fatalf("error = %v, want missing-tracking-table refusal", err) + } + }) + + t.Run("pending migrations are refused", func(t *testing.T) { + if _, err := applySource(ctx, connStr, coreSource(oneCoreFS)); err != nil { + t.Fatalf("apply v1: %v", err) + } + err := VerifyMigrated(ctx, connStr, full) + if err == nil || !strings.Contains(err.Error(), "requires version 2") { + t.Fatalf("error = %v, want behind-binary refusal", err) + } + }) + + t.Run("fully migrated database passes", func(t *testing.T) { + if _, err := applySource(ctx, connStr, coreSource(goodCoreFS)); err != nil { + t.Fatalf("apply v2: %v", err) + } + if err := VerifyMigrated(ctx, connStr, full); err != nil { + t.Fatalf("VerifyMigrated() = %v, want nil", err) + } + }) + + t.Run("works without DDL privileges", func(t *testing.T) { + db, err := sql.Open("pgx", connStr) + if err != nil { + t.Fatal(err) + } + defer db.Close() + for _, q := range []string{ + `CREATE ROLE readonly LOGIN PASSWORD 'ro'`, + `GRANT USAGE ON SCHEMA public TO readonly`, + `GRANT SELECT ON ALL TABLES IN SCHEMA public TO readonly`, + } { + if _, err := db.ExecContext(ctx, q); err != nil { + t.Fatalf("%s: %v", q, err) + } + } + u, err := url.Parse(connStr) + if err != nil { + t.Fatal(err) + } + u.User = url.UserPassword("readonly", "ro") + if err := VerifyMigrated(ctx, u.String(), full); err != nil { + t.Fatalf("VerifyMigrated() as readonly = %v, want nil", err) + } + }) + + t.Run("database ahead of binary passes", func(t *testing.T) { + if err := VerifyMigrated(ctx, connStr, []Source{coreSource(oneCoreFS)}); err != nil { + t.Fatalf("VerifyMigrated() with older binary = %v, want nil (compatibility mode)", err) + } + }) + + t.Run("dirty tracking table is refused", func(t *testing.T) { + db, err := sql.Open("pgx", connStr) + if err != nil { + t.Fatal(err) + } + defer db.Close() + if _, err := db.ExecContext(ctx, `UPDATE schema_migrations SET dirty = true`); err != nil { + t.Fatal(err) + } + defer func() { + if _, err := db.ExecContext(ctx, `UPDATE schema_migrations SET dirty = false`); err != nil { + t.Fatal(err) + } + }() + err = VerifyMigrated(ctx, connStr, full) + if err == nil || !strings.Contains(err.Error(), "dirty") { + t.Fatalf("error = %v, want dirty refusal", err) + } + }) + + t.Run("no sources is a no-op", func(t *testing.T) { + if err := VerifyMigrated(ctx, "postgres://unused", nil); err != nil { + t.Fatalf("VerifyMigrated() with no sources = %v, want nil", err) + } + }) +} diff --git a/helm/kagent/templates/NOTES.txt b/helm/kagent/templates/NOTES.txt index 1bffe87b15..96f4c75ab4 100644 --- a/helm/kagent/templates/NOTES.txt +++ b/helm/kagent/templates/NOTES.txt @@ -92,6 +92,16 @@ DOCUMENTATION: database.postgres.bundled.enabled=false {{- end }} {{- end }} +{{- if .Values.database.postgres.skipMigrations }} +################################################################################ +# NOTE: STARTUP MIGRATIONS ARE DISABLED # +################################################################################ + database.postgres.skipMigrations is set: the controller will not run database + migrations at startup. It verifies the schema is already migrated and fails + to start if it is not. + + Ensure migrations are applied out-of-band before installing or upgrading. +{{- end }} {{- if .Values.substrate.enabled }} ################################################################################ # WARNING: SUBSTRATE IS EXPERIMENTAL, USE AT OWN RISK # diff --git a/helm/kagent/templates/controller-configmap.yaml b/helm/kagent/templates/controller-configmap.yaml index 0436ac15fd..cadb2df438 100644 --- a/helm/kagent/templates/controller-configmap.yaml +++ b/helm/kagent/templates/controller-configmap.yaml @@ -55,6 +55,7 @@ data: PROXY_URL: {{ .Values.proxy.url | quote }} {{- end }} DATABASE_VECTOR_ENABLED: {{ .Values.database.postgres.vectorEnabled | quote }} + SKIP_MIGRATIONS: {{ .Values.database.postgres.skipMigrations | default false | quote }} WATCH_NAMESPACES: {{ include "kagent.watchNamespaces" . | quote }} MCP_EGRESS_PLAINTEXT: {{ .Values.controller.mcpEgressPlaintext | default false | quote }} {{- if .Values.controller.a2aClientTimeout }} diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 5fbe1bd0e7..6117696ebd 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -76,6 +76,10 @@ database: # Required to use features that depend on database vector capability. (e.g. long-term memory) # Set to true when using an external PostgreSQL that has the pgvector extension installed. vectorEnabled: false + # -- Skip running database migrations at controller startup. + # The controller instead verifies the database is already migrated and fails if it is not. + # Migrations must be applied out-of-band (e.g. from a CI/CD pipeline) before install/upgrade. + skipMigrations: false # -- Bundled PostgreSQL instance — for development and evaluation only. # Not suitable for production. Deployed when enabled is true and url/urlFile are not set. bundled: From 78757252648e93399f7f6d0c629efa598e679bb1 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Mon, 6 Jul 2026 10:41:31 -0700 Subject: [PATCH 7/7] Fix text Signed-off-by: Jeremy Alvis --- go/core/pkg/migrations/runner.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go/core/pkg/migrations/runner.go b/go/core/pkg/migrations/runner.go index 0687c16fe7..06762e75ca 100644 --- a/go/core/pkg/migrations/runner.go +++ b/go/core/pkg/migrations/runner.go @@ -225,7 +225,7 @@ func VerifyMigrated(ctx context.Context, url string, sources []Source) error { return fmt.Errorf("check tracking table for %s: %w", src.Name, err) } if !exists { - return fmt.Errorf("source %s: tracking table %s does not exist — the database has not been migrated; apply migrations out-of-band or unset SKIP_MIGRATIONS", src.Name, table) + return fmt.Errorf("source %s: tracking table %s does not exist - the database has not been migrated; apply migrations out-of-band or unset SKIP_MIGRATIONS", src.Name, table) } var version int64