From 40bc0db036cadd036a5cb3a27dac8077170db24d Mon Sep 17 00:00:00 2001 From: ~wispem-wantex Date: Sat, 19 Sep 2026 21:36:31 -0700 Subject: [PATCH] refactor: make migration process more intelligent, add test coverage --- pkg/db/connect.go | 45 +++---- pkg/db/connect_test.go | 6 +- pkg/db/errors.go | 2 + pkg/db/migration.go | 38 ++++++ pkg/db/migration_test.go | 260 ++++++++++++++++++++++++++++++++++++++ pkg/db/pragma_fk_check.go | 42 ++++++ 6 files changed, 361 insertions(+), 32 deletions(-) create mode 100644 pkg/db/migration_test.go create mode 100644 pkg/db/pragma_fk_check.go diff --git a/pkg/db/connect.go b/pkg/db/connect.go index 86550a6..b154b28 100644 --- a/pkg/db/connect.go +++ b/pkg/db/connect.go @@ -89,40 +89,27 @@ func (c DBConfig) CheckAndUpdateVersion(db *sqlx.DB) error { fmt.Printf("Database version is out of date. Upgrading database from version %d to version %d!\n", version, c.version_number) fmt.Print(textutils.ColorReset) - c.UpgradeFromXToY(db, version, c.version_number) + + for i := version; i < c.version_number; i++ { + m := (*c.migrations)[i] + if err := m.Apply(db); err != nil { + return fmt.Errorf("migration %d: %w", m.ID, err) + } + + fmt.Print(textutils.ColorYellow) + fmt.Printf("Now at database schema version %d.\n", m.ID) + fmt.Print(textutils.ColorReset) + } + + fmt.Print(textutils.ColorGreen) + fmt.Printf("================================================\n") + fmt.Printf("Database version has been upgraded to version %d.\n", c.version_number) + fmt.Print(textutils.ColorReset) } return nil } -// UpgradeFromXToY runs all the migrations from version X to version Y, and update the `database_version` table's `version_number` -func (c DBConfig) UpgradeFromXToY(db *sqlx.DB, x uint, y uint) { - for i := x; i < y; i++ { - fmt.Print(textutils.ColorCyan) - fmt.Println((*c.migrations)[i].SQL) - fmt.Print(textutils.ColorReset) - - // Execute the migration in a transaction - tx, err := db.Beginx() - if err != nil { - panic(err) - } - tx.MustExec((*c.migrations)[i].SQL) - tx.MustExec("update db_version set version = ?", i+1) - if err := tx.Commit(); err != nil { - panic(err) - } - - fmt.Print(textutils.ColorYellow) - fmt.Printf("Now at database schema version %d.\n", i+1) - fmt.Print(textutils.ColorReset) - } - fmt.Print(textutils.ColorGreen) - fmt.Printf("================================================\n") - fmt.Printf("Database version has been upgraded to version %d.\n", y) - fmt.Print(textutils.ColorReset) -} - type VersionMismatchError struct { EngineVersion uint DatabaseVersion uint diff --git a/pkg/db/connect_test.go b/pkg/db/connect_test.go index 83a6e4d..f7cb288 100644 --- a/pkg/db/connect_test.go +++ b/pkg/db/connect_test.go @@ -62,8 +62,8 @@ func TestVersionUpgrade(t *testing.T) { // Create a migration to add a new Item migrations = append(migrations, db.Migration{ID: 1, SQL: "insert into items (rowid) values (1)"}) - db.Init(&initial_schema, &migrations) // Reinitialize with the new migration - config.UpgradeFromXToY(connection, uint(len(migrations)-1), uint(len(migrations))) + config = db.Init(&initial_schema, &migrations) // Reinitialize with the new migration + require.NoError(config.CheckAndUpdateVersion(connection)) var items2 []int require.NoError(connection.Select(&items2, "select * from items")) @@ -73,7 +73,7 @@ func TestVersionUpgrade(t *testing.T) { // Create a migration to add a new Item migrations = append(migrations, db.Migration{ID: 2, SQL: `alter table items add column name string default 'asdf'`}) config = db.Init(&initial_schema, &migrations) // Reinitialize with the new migration - config.UpgradeFromXToY(connection, uint(len(migrations)-1), uint(len(migrations))) + require.NoError(config.CheckAndUpdateVersion(connection)) var items3 []struct { ID uint64 `db:"rowid"` diff --git a/pkg/db/errors.go b/pkg/db/errors.go index ae70854..1f4e1d4 100644 --- a/pkg/db/errors.go +++ b/pkg/db/errors.go @@ -12,6 +12,8 @@ var ( ErrItemIsDeleted = errors.New("item is deleted") ErrForeignKeyViolation = errors.New("foreign key constraint failed") ErrDatabaseAlreadyExists = errors.New("target already exists") + + ErrMigrationOrder = errors.New("migration in wrong order") ) type ForeignKey interface { diff --git a/pkg/db/migration.go b/pkg/db/migration.go index 8e111ce..05a69ae 100644 --- a/pkg/db/migration.go +++ b/pkg/db/migration.go @@ -1,5 +1,14 @@ package db +import ( + "fmt" + + "github.com/jmoiron/sqlx" + + "git.offline-twitter.com/offline-labs/gas-stack/pkg/must" + "git.offline-twitter.com/offline-labs/gas-stack/pkg/textutils" +) + type Migration struct { // ID is the version number that the schema will be at, after applying the migration. ID int @@ -10,3 +19,32 @@ type Migration struct { // SQL is the body of the migration that will be executed. SQL string } + +// Apply runs the migration in a transaction. Foreign key checking is delayed until right before committing the transaction. +func (m Migration) Apply(db *sqlx.DB) error { + fmt.Print(textutils.ColorCyan) + fmt.Println(m.SQL) + fmt.Print(textutils.ColorReset) + + // Temporarily suspend foreign keys until the migration is finished + db.MustExec("pragma foreign_keys = off") + defer db.MustExec("pragma foreign_keys = on") + + tx := db.MustBegin() + defer func() { _ = tx.Rollback() }() + + // Sanity: should initially be on the previous version + var prevVersion int + must.Do(tx.Get(&prevVersion, `select version from db_version`)) + if prevVersion != m.ID-1 { + return fmt.Errorf("%w: tried to apply migration %d on DB version %d", ErrMigrationOrder, m.ID, prevVersion) + } + + // Execute the migration + tx.MustExec(m.SQL) + tx.MustExec("update db_version set version = ?", m.ID) + must.Do(PragmaForeignKeyCheck(tx)) + must.Do(tx.Commit()) + + return nil +} diff --git a/pkg/db/migration_test.go b/pkg/db/migration_test.go new file mode 100644 index 0000000..bdbd69a --- /dev/null +++ b/pkg/db/migration_test.go @@ -0,0 +1,260 @@ +package db_test + +import ( + "testing" + + "github.com/jmoiron/sqlx" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "git.offline-twitter.com/offline-labs/gas-stack/pkg/db" + "git.offline-twitter.com/offline-labs/gas-stack/pkg/must" +) + +// migrationTestSchema has a foreign key (`children.parent_id` => `parents.rowid`) along with one +// valid parent/child pair, so that migrations can be written which either preserve or break +// referential integrity. +var migrationTestSchema = ` + create table db_version (version integer primary key) strict, without rowid; + insert into db_version values (0); + + create table parents ( + rowid integer primary key, + name text not null + ); + insert into parents (rowid, name) values (1, 'p1'); + + create table children ( + rowid integer primary key, + parent_id integer not null references parents(rowid) + ); + insert into children (rowid, parent_id) values (1, 1); +` + +func newMigrationTestDB() *sqlx.DB { + return must.Get(db.Init(&migrationTestSchema, &[]db.Migration{}).Create(":memory:")) +} + +func assertForeignKeysAreGood(t *testing.T, conn *sqlx.DB) { + t.Helper() + + // FKs should be enabled and intact at all times + var fkEnabled bool + require.NoError(t, conn.Get(&fkEnabled, `pragma foreign_keys`)) + assert.True(t, fkEnabled) + assert.NoError(t, db.PragmaForeignKeyCheck(conn)) + + // ...and enforcement should be in effect again, not merely reported as on + _, err := conn.Exec(`insert into children (rowid, parent_id) values (99, 40235235234)`) + assert.True(t, db.IsSqliteFkError(err), "expected a foreign key error, got: %v", err) +} + +func TestMigrationApply(t *testing.T) { + conn := newMigrationTestDB() + + testcases := []struct { + Migration db.Migration + ExpectedParents int + ExpectedChildren int + }{ + { + Migration: db.Migration{ + ID: 1, + Title: "add a second parent", + SQL: `insert into parents (rowid, name) values (2, 'p2')`, + }, + ExpectedParents: 2, + ExpectedChildren: 1, + }, + { + Migration: db.Migration{ID: 2, SQL: ` + insert into children (rowid, parent_id) values (2, 2); + insert into parents (rowid, name) values (3, 'p3'); + `}, + ExpectedParents: 3, + ExpectedChildren: 2, + }, + } + + for _, tc := range testcases { + t.Run(tc.Migration.Title, func(t *testing.T) { + require.NoError(t, tc.Migration.Apply(conn)) + + // Check that it applied + var parents int + require.NoError(t, conn.Get(&parents, `select count(*) from parents`)) + assert.Equal(t, tc.ExpectedParents, parents) + var children int + require.NoError(t, conn.Get(&children, `select count(*) from children`)) + assert.Equal(t, tc.ExpectedChildren, children) + + // Should bump the DB version + var version int + require.NoError(t, conn.Get(&version, `select version from db_version`)) + assert.Equal(t, tc.Migration.ID, version, "db version should be bumped to the migration's ID") + + assertForeignKeysAreGood(t, conn) + }) + } +} + +// A migration whose SQL fails must leave no trace, including the statements which had already +// succeeded before the failing one. Every migration below therefore starts by creating +// `new_table`, which must never survive. +func TestMigrationFailureRollsBack(t *testing.T) { + conn := newMigrationTestDB() + + nothingHappens := func(t *testing.T) { + t.Helper() + // Everything should be rolled back, including the statements which ran before the failing one + var newTables int + require.NoError(t, conn.Get(&newTables, `select count(*) from sqlite_master where type = 'table' and name = 'new_table'`)) + assert.Equal(t, 0, newTables, "the migration's earlier statements should be rolled back too") + + var version int + require.NoError(t, conn.Get(&version, `select version from db_version`)) + assert.Equal(t, 0, version) + + var parents int + require.NoError(t, conn.Get(&parents, `select count(*) from parents`)) + assert.Equal(t, 1, parents) + + var children int + require.NoError(t, conn.Get(&children, `select count(*) from children`)) + assert.Equal(t, 1, children) + } + + testcases2 := []db.Migration{ + { + ID: 1, + Title: "syntax error", + SQL: ` + create table new_table (rowid integer primary key); + this is not valid sql; + `, + }, + { + ID: 1, + Title: "unknown table", + SQL: ` + create table new_table (rowid integer primary key); + insert into no_such_table (rowid) values (1); + `, + }, + { + ID: 1, + Title: "constraint violation", + SQL: ` + create table new_table (rowid integer primary key); + insert into parents (rowid, name) values (1, 'duplicate rowid'); + `, + }, + { + ID: 1, + Title: "failure after a successful write", + SQL: ` + create table new_table (rowid integer primary key); + insert into parents (rowid, name) values (2, 'p2'); + insert into parents (rowid, name) values (1, 'duplicate rowid'); + `, + }, + } + for _, tc := range testcases2 { + t.Run(tc.Title, func(t *testing.T) { + assert.Panics(t, func() { _ = tc.Apply(conn) }) + nothingHappens(t) + assertForeignKeysAreGood(t, conn) + }) + } + + // --------- + // Foreign key failures should trigger a rollback too + // --------- + + testcases := []db.Migration{ + { + ID: 1, + Title: "child pointing at a nonexistent parent", + SQL: ` + create table new_table (rowid integer primary key); + insert into children (rowid, parent_id) values (2, 404); + `, + }, + { + ID: 1, + Title: "parent deleted out from under an existing child", + SQL: ` + create table new_table (rowid integer primary key); + delete from parents where rowid = 1; + `, + }, + } + for _, tc := range testcases { + t.Run(tc.Title, func(t *testing.T) { + // Should notify with the correct detail about the error + func() { + defer func() { + err := recover() + require.NotNil(t, err) + switch err := err.(type) { + case error: + var violations db.PragmaFKViolations + require.ErrorAs(t, err, &violations) + require.Len(t, violations, 1) + assert.Equal(t, "children", violations[0].SrcTable) + assert.Equal(t, "parents", violations[0].DstTable) + default: + panic(err) // Something weird happened + } + }() + _ = tc.Apply(conn) + }() + + // ...as well as the usual + nothingHappens(t) + assertForeignKeysAreGood(t, conn) + }) + } +} + +// Migrations are sequential: a migration may only be applied to a DB sitting at the immediately +// preceding version. +func TestMigrationApplyRejectsOutOfOrderMigrations(t *testing.T) { + const sql = `insert into parents (name) values ('p2')` + + t.Run("migration from the future", func(t *testing.T) { + conn := newMigrationTestDB() + + err := db.Migration{ID: 2, SQL: sql}.Apply(conn) + require.ErrorIs(t, err, db.ErrMigrationOrder) + + var version int + require.NoError(t, conn.Get(&version, `select version from db_version`)) + assert.Equal(t, 0, version) + + var parents int + require.NoError(t, conn.Get(&parents, `select count(*) from parents`)) + assert.Equal(t, 1, parents, "the migration's SQL should not have run") + + assertForeignKeysAreGood(t, conn) + }) + + t.Run("migration applied twice", func(t *testing.T) { + conn := newMigrationTestDB() + m := db.Migration{ID: 1, SQL: sql} + require.NoError(t, m.Apply(conn)) + + err := m.Apply(conn) + require.ErrorIs(t, err, db.ErrMigrationOrder) + + var version int + require.NoError(t, conn.Get(&version, `select version from db_version`)) + assert.Equal(t, 1, version) + + var parents int + require.NoError(t, conn.Get(&parents, `select count(*) from parents`)) + assert.Equal(t, 2, parents, "the migration's SQL should not have run again") + + assertForeignKeysAreGood(t, conn) + }) +} diff --git a/pkg/db/pragma_fk_check.go b/pkg/db/pragma_fk_check.go new file mode 100644 index 0000000..089ed03 --- /dev/null +++ b/pkg/db/pragma_fk_check.go @@ -0,0 +1,42 @@ +package db + +import ( + "fmt" + "strings" + + "github.com/jmoiron/sqlx" + + "git.offline-twitter.com/offline-labs/gas-stack/pkg/must" +) + +// PragmaFKViolation is a result row of SQLite's `pragma foreign_key_check`. +type PragmaFKViolation struct { + SrcTable string `db:"table"` + SrcRowid int `db:"rowid"` + DstTable string `db:"parent"` + ForeignKeyIdx int `db:"fkid"` +} + +// PragmaFKViolations is a slice of FK violations. It also implements the `error` interface. +type PragmaFKViolations []PragmaFKViolation + +func (v PragmaFKViolations) Error() string { + var b strings.Builder + for i, fkv := range v { + if i > 0 { + b.WriteString("\n") + } + fmt.Fprintf(&b, "Item #%d in table %q has invalid reference to table %q", fkv.SrcRowid, fkv.SrcTable, fkv.DstTable) + } + return b.String() +} + +// PragmaForeignKeyCheck uses `pragma foreign_key_check` to do a full database check for any FK violations. +func PragmaForeignKeyCheck(q sqlx.Queryer) error { + var ret PragmaFKViolations + must.Do(sqlx.Select(q, &ret, `pragma foreign_key_check`)) + if len(ret) == 0 { + return nil + } + return ret +}