refactor: make migration process more intelligent, add test coverage

pkg/db/connect.go

@@ -89,38 +89,25 @@ 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)
-	}
 
-	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)
+		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", i+1)
+			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", y)
+		fmt.Printf("Database version has been upgraded to version %d.\n", c.version_number)
 		fmt.Print(textutils.ColorReset)
+	}
+
+	return nil
 }
 
 type VersionMismatchError struct {

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"`

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 {

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
+}

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)
+	})
+}

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
+}