Update database migrations, remove goose (#2264)

* Add new db migration

* Update migrations
Remove goose

* Add possibility to test direct upgrades

* Try to fix WASM test

* Add checks for specific migrations

* Remove AddMigration
Use WithTransaction
Add Dendrite version to table

* Fix linter issues

* Update tests

* Update comments, outdent if

* Namespace migrations

* Add direct upgrade tests, skipping over one version

* Split migrations

* Update go version in CI

* Fix copy&paste mistake

* Use contexts in migrations

Co-authored-by: kegsay <kegan@matrix.org>
Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

8 files changed, 129 insertions, 127 deletions

diff --git a/roomserver/storage/postgres/deltas/20201028212440_add_forgotten_column.go b/roomserver/storage/postgres/deltas/20201028212440_add_forgotten_column.go
index f3bd8632..61d4dba8 100644
--- a/roomserver/storage/postgres/deltas/20201028212440_add_forgotten_column.go
+++ b/roomserver/storage/postgres/deltas/20201028212440_add_forgotten_column.go
@@ -15,32 +15,21 @@
 package deltas
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
-
-	"github.com/matrix-org/dendrite/internal/sqlutil"
-	"github.com/pressly/goose"
 )
 
-func LoadFromGoose() {
-	goose.AddMigration(UpAddForgottenColumn, DownAddForgottenColumn)
-	goose.AddMigration(UpStateBlocksRefactor, DownStateBlocksRefactor)
-}
-
-func LoadAddForgottenColumn(m *sqlutil.Migrations) {
-	m.AddMigration(UpAddForgottenColumn, DownAddForgottenColumn)
-}
-
-func UpAddForgottenColumn(tx *sql.Tx) error {
-	_, err := tx.Exec(`ALTER TABLE roomserver_membership ADD COLUMN IF NOT EXISTS forgotten BOOLEAN NOT NULL DEFAULT false;`)
+func UpAddForgottenColumn(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_membership ADD COLUMN IF NOT EXISTS forgotten BOOLEAN NOT NULL DEFAULT false;`)
 	if err != nil {
 		return fmt.Errorf("failed to execute upgrade: %w", err)
 	}
 	return nil
 }
 
-func DownAddForgottenColumn(tx *sql.Tx) error {
-	_, err := tx.Exec(`ALTER TABLE roomserver_membership DROP COLUMN IF EXISTS forgotten;`)
+func DownAddForgottenColumn(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_membership DROP COLUMN IF EXISTS forgotten;`)
 	if err != nil {
 		return fmt.Errorf("failed to execute downgrade: %w", err)
 	}
diff --git a/roomserver/storage/postgres/deltas/2021041615092700_state_blocks_refactor.go b/roomserver/storage/postgres/deltas/2021041615092700_state_blocks_refactor.go
index 06442a4c..355c49b1 100644
--- a/roomserver/storage/postgres/deltas/2021041615092700_state_blocks_refactor.go
+++ b/roomserver/storage/postgres/deltas/2021041615092700_state_blocks_refactor.go
@@ -15,11 +15,11 @@
 package deltas
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 
 	"github.com/lib/pq"
-	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/util"
 	"github.com/sirupsen/logrus"
@@ -36,48 +36,44 @@ type stateBlockData struct {
 	EventNIDs     types.EventNIDs
 }
 
-func LoadStateBlocksRefactor(m *sqlutil.Migrations) {
-	m.AddMigration(UpStateBlocksRefactor, DownStateBlocksRefactor)
-}
-
 // nolint:gocyclo
-func UpStateBlocksRefactor(tx *sql.Tx) error {
+func UpStateBlocksRefactor(ctx context.Context, tx *sql.Tx) error {
 	logrus.Warn("Performing state storage upgrade. Please wait, this may take some time!")
 	defer logrus.Warn("State storage upgrade complete")
 
 	var snapshotcount int
 	var maxsnapshotid int
 	var maxblockid int
-	if err := tx.QueryRow(`SELECT COUNT(DISTINCT state_snapshot_nid) FROM roomserver_state_snapshots;`).Scan(&snapshotcount); err != nil {
-		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	if err := tx.QueryRowContext(ctx, `SELECT COUNT(DISTINCT state_snapshot_nid) FROM roomserver_state_snapshots;`).Scan(&snapshotcount); err != nil {
+		return fmt.Errorf("tx.QueryRowContext.Scan (count snapshots): %w", err)
 	}
-	if err := tx.QueryRow(`SELECT COALESCE(MAX(state_snapshot_nid),0) FROM roomserver_state_snapshots;`).Scan(&maxsnapshotid); err != nil {
-		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	if err := tx.QueryRowContext(ctx, `SELECT COALESCE(MAX(state_snapshot_nid),0) FROM roomserver_state_snapshots;`).Scan(&maxsnapshotid); err != nil {
+		return fmt.Errorf("tx.QueryRowContext.Scan (count snapshots): %w", err)
 	}
-	if err := tx.QueryRow(`SELECT COALESCE(MAX(state_block_nid),0) FROM roomserver_state_block;`).Scan(&maxblockid); err != nil {
-		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	if err := tx.QueryRowContext(ctx, `SELECT COALESCE(MAX(state_block_nid),0) FROM roomserver_state_block;`).Scan(&maxblockid); err != nil {
+		return fmt.Errorf("tx.QueryRowContext.Scan (count snapshots): %w", err)
 	}
 	maxsnapshotid++
 	maxblockid++
 
-	if _, err := tx.Exec(`ALTER TABLE roomserver_state_block RENAME TO _roomserver_state_block;`); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_state_block RENAME TO _roomserver_state_block;`); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	if _, err := tx.Exec(`ALTER TABLE roomserver_state_snapshots RENAME TO _roomserver_state_snapshots;`); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_state_snapshots RENAME TO _roomserver_state_snapshots;`); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
 	// We create new sequences starting with the maximum state snapshot and block NIDs.
 	// This means that all newly created snapshots and blocks by the migration will have
 	// NIDs higher than these values, so that when we come to update the references to
 	// these NIDs using UPDATE statements, we can guarantee we are only ever updating old
 	// values and not accidentally overwriting new ones.
-	if _, err := tx.Exec(fmt.Sprintf(`CREATE SEQUENCE roomserver_state_block_nid_sequence START WITH %d;`, maxblockid)); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, fmt.Sprintf(`CREATE SEQUENCE roomserver_state_block_nid_sequence START WITH %d;`, maxblockid)); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	if _, err := tx.Exec(fmt.Sprintf(`CREATE SEQUENCE roomserver_state_snapshot_nid_sequence START WITH %d;`, maxsnapshotid)); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, fmt.Sprintf(`CREATE SEQUENCE roomserver_state_snapshot_nid_sequence START WITH %d;`, maxsnapshotid)); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	_, err := tx.Exec(`
+	_, err := tx.ExecContext(ctx, `
 		CREATE TABLE IF NOT EXISTS roomserver_state_block (
 			state_block_nid bigint PRIMARY KEY DEFAULT nextval('roomserver_state_block_nid_sequence'),
 			state_block_hash BYTEA UNIQUE,
@@ -87,7 +83,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	if err != nil {
 		return fmt.Errorf("tx.Exec (create blocks table): %w", err)
 	}
-	_, err = tx.Exec(`
+	_, err = tx.ExecContext(ctx, `
 		CREATE TABLE IF NOT EXISTS roomserver_state_snapshots (
 			state_snapshot_nid bigint PRIMARY KEY DEFAULT nextval('roomserver_state_snapshot_nid_sequence'),
 			state_snapshot_hash BYTEA UNIQUE,
@@ -104,7 +100,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	// in question a state snapshot NID of 0 to indicate 'no snapshot'.
 	// If we don't do this, we'll fail the assertions later on which try to ensure we didn't forget
 	// any snapshots.
-	_, err = tx.Exec(
+	_, err = tx.ExecContext(ctx,
 		`UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE event_type_nid = $1 AND event_state_key_nid = $2`,
 		types.MRoomCreateNID, types.EmptyStateKeyNID,
 	)
@@ -115,7 +111,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	batchsize := 100
 	for batchoffset := 0; batchoffset < snapshotcount; batchoffset += batchsize {
 		var snapshotrows *sql.Rows
-		snapshotrows, err = tx.Query(`
+		snapshotrows, err = tx.QueryContext(ctx, `
 			SELECT
 				state_snapshot_nid,
 				room_nid,
@@ -146,7 +142,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 				state_block_nid;
 		`, batchsize, batchoffset)
 		if err != nil {
-			return fmt.Errorf("tx.Query: %w", err)
+			return fmt.Errorf("tx.QueryContext: %w", err)
 		}
 
 		logrus.Warnf("Rewriting snapshots %d-%d of %d...", batchoffset, batchoffset+batchsize, snapshotcount)
@@ -183,7 +179,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 		// fill in bad create snapshots
 		for _, s := range badCreateSnapshots {
 			var createEventNID types.EventNID
-			err = tx.QueryRow(
+			err = tx.QueryRowContext(ctx,
 				`SELECT event_nid FROM roomserver_events WHERE state_snapshot_nid = $1 AND event_type_nid = 1`, s.StateSnapshotNID,
 			).Scan(&createEventNID)
 			if err == sql.ErrNoRows {
@@ -208,7 +204,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			}
 
 			var blocknid types.StateBlockNID
-			err = tx.QueryRow(`
+			err = tx.QueryRowContext(ctx, `
 				INSERT INTO roomserver_state_block (state_block_hash, event_nids)
 					VALUES ($1, $2)
 					ON CONFLICT (state_block_hash) DO UPDATE SET event_nids=$2
@@ -227,7 +223,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			}
 
 			var newNID types.StateSnapshotNID
-			err = tx.QueryRow(`
+			err = tx.QueryRowContext(ctx, `
 					INSERT INTO roomserver_state_snapshots (state_snapshot_hash, room_nid, state_block_nids)
 						VALUES ($1, $2, $3)
 						ON CONFLICT (state_snapshot_hash) DO UPDATE SET room_nid=$2
@@ -237,12 +233,12 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 				return fmt.Errorf("tx.QueryRow.Scan (insert new snapshot): %w", err)
 			}
 
-			if _, err = tx.Exec(`UPDATE roomserver_events SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newNID, snapshotdata.StateSnapshotNID, maxsnapshotid); err != nil {
-				return fmt.Errorf("tx.Exec (update events): %w", err)
+			if _, err = tx.ExecContext(ctx, `UPDATE roomserver_events SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newNID, snapshotdata.StateSnapshotNID, maxsnapshotid); err != nil {
+				return fmt.Errorf("tx.ExecContext (update events): %w", err)
 			}
 
-			if _, err = tx.Exec(`UPDATE roomserver_rooms SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newNID, snapshotdata.StateSnapshotNID, maxsnapshotid); err != nil {
-				return fmt.Errorf("tx.Exec (update rooms): %w", err)
+			if _, err = tx.ExecContext(ctx, `UPDATE roomserver_rooms SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newNID, snapshotdata.StateSnapshotNID, maxsnapshotid); err != nil {
+				return fmt.Errorf("tx.ExecContext (update rooms): %w", err)
 			}
 		}
 	}
@@ -252,13 +248,13 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	// in roomserver_state_snapshots
 	var count int64
 
-	if err = tx.QueryRow(`SELECT COUNT(*) FROM roomserver_events WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid).Scan(&count); err != nil {
+	if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM roomserver_events WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid).Scan(&count); err != nil {
 		return fmt.Errorf("assertion query failed: %s", err)
 	}
 	if count > 0 {
 		var res sql.Result
 		var c int64
-		res, err = tx.Exec(`UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid)
+		res, err = tx.ExecContext(ctx, `UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid)
 		if err != nil && err != sql.ErrNoRows {
 			return fmt.Errorf("failed to reset invalid state snapshots: %w", err)
 		}
@@ -268,13 +264,13 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			return fmt.Errorf("expected to reset %d event(s) but only updated %d event(s)", count, c)
 		}
 	}
-	if err = tx.QueryRow(`SELECT COUNT(*) FROM roomserver_rooms WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid).Scan(&count); err != nil {
+	if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM roomserver_rooms WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid).Scan(&count); err != nil {
 		return fmt.Errorf("assertion query failed: %s", err)
 	}
 	if count > 0 {
 		var debugRoomID string
 		var debugSnapNID, debugLastEventNID int64
-		err = tx.QueryRow(
+		err = tx.QueryRowContext(ctx,
 			`SELECT room_id, state_snapshot_nid, last_event_sent_nid FROM roomserver_rooms WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, maxsnapshotid,
 		).Scan(&debugRoomID, &debugSnapNID, &debugLastEventNID)
 		if err != nil {
@@ -291,13 +287,13 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 		return fmt.Errorf("%d rooms exist in roomserver_rooms which have not been converted to a new state_snapshot_nid; this is a bug, please report", count)
 	}
 
-	if _, err = tx.Exec(`
+	if _, err = tx.ExecContext(ctx, `
 		DROP TABLE _roomserver_state_snapshots;
 		DROP SEQUENCE roomserver_state_snapshot_nid_seq;
 	`); err != nil {
 		return fmt.Errorf("tx.Exec (delete old snapshot table): %w", err)
 	}
-	if _, err = tx.Exec(`
+	if _, err = tx.ExecContext(ctx, `
 		DROP TABLE _roomserver_state_block;
 		DROP SEQUENCE roomserver_state_block_nid_seq;
 	`); err != nil {
@@ -307,6 +303,6 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	return nil
 }
 
-func DownStateBlocksRefactor(tx *sql.Tx) error {
+func DownStateBlocksRefactor(ctx context.Context, tx *sql.Tx) error {
 	panic("Downgrading state storage is not supported")
 }
diff --git a/roomserver/storage/postgres/membership_table.go b/roomserver/storage/postgres/membership_table.go
index eaaea22a..bd3fd559 100644
--- a/roomserver/storage/postgres/membership_table.go
+++ b/roomserver/storage/postgres/membership_table.go
@@ -23,6 +23,7 @@ import (
 	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/postgres/deltas"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -173,7 +174,15 @@ type membershipStatements struct {
 
 func CreateMembershipTable(db *sql.DB) error {
 	_, err := db.Exec(membershipSchema)
-	return err
+	if err != nil {
+		return err
+	}
+	m := sqlutil.NewMigrator(db)
+	m.AddMigrations(sqlutil.Migration{
+		Version: "roomserver: add forgotten column",
+		Up:      deltas.UpAddForgottenColumn,
+	})
+	return m.Up(context.Background())
 }
 
 func PrepareMembershipTable(db *sql.DB) (tables.Membership, error) {
diff --git a/roomserver/storage/postgres/storage.go b/roomserver/storage/postgres/storage.go
index 70ea4d8b..4c271ea9 100644
--- a/roomserver/storage/postgres/storage.go
+++ b/roomserver/storage/postgres/storage.go
@@ -21,7 +21,6 @@ import (
 
 	// Import the postgres database driver.
 	_ "github.com/lib/pq"
-
 	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/postgres/deltas"
@@ -45,17 +44,25 @@ func Open(base *base.BaseDendrite, dbProperties *config.DatabaseOptions, cache c
 	}
 
 	// Create the tables.
-	if err := d.create(db); err != nil {
+	if err = d.create(db); err != nil {
 		return nil, err
 	}
 
-	// Then execute the migrations. By this point the tables are created with the latest
-	// schemas.
-	m := sqlutil.NewMigrations()
-	deltas.LoadAddForgottenColumn(m)
-	deltas.LoadStateBlocksRefactor(m)
-	if err := m.RunDeltas(db, dbProperties); err != nil {
-		return nil, err
+	// Special case, since this migration uses several tables, so it needs to
+	// be sure that all tables are created first.
+	// TODO: Remove when we are sure we are not having goose artefacts in the db
+	// This forces an error, which indicates the migration is already applied, since the
+	// column event_nid was removed from the table
+	err = db.QueryRow("SELECT event_nid FROM roomserver_state_block LIMIT 1;").Scan()
+	if err == nil {
+		m := sqlutil.NewMigrator(db)
+		m.AddMigrations(sqlutil.Migration{
+			Version: "roomserver: state blocks refactor",
+			Up:      deltas.UpStateBlocksRefactor,
+		})
+		if err := m.Up(base.Context()); err != nil {
+			return nil, err
+		}
 	}
 
 	// Then prepare the statements. Now that the migrations have run, any columns referred
diff --git a/roomserver/storage/sqlite3/deltas/20201028212440_add_forgotten_column.go b/roomserver/storage/sqlite3/deltas/20201028212440_add_forgotten_column.go
index d08ab02d..4c002e33 100644
--- a/roomserver/storage/sqlite3/deltas/20201028212440_add_forgotten_column.go
+++ b/roomserver/storage/sqlite3/deltas/20201028212440_add_forgotten_column.go
@@ -15,24 +15,13 @@
 package deltas
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
-
-	"github.com/matrix-org/dendrite/internal/sqlutil"
-	"github.com/pressly/goose"
 )
 
-func LoadFromGoose() {
-	goose.AddMigration(UpAddForgottenColumn, DownAddForgottenColumn)
-	goose.AddMigration(UpStateBlocksRefactor, DownStateBlocksRefactor)
-}
-
-func LoadAddForgottenColumn(m *sqlutil.Migrations) {
-	m.AddMigration(UpAddForgottenColumn, DownAddForgottenColumn)
-}
-
-func UpAddForgottenColumn(tx *sql.Tx) error {
-	_, err := tx.Exec(`	ALTER TABLE roomserver_membership RENAME TO roomserver_membership_tmp;
+func UpAddForgottenColumn(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, `	ALTER TABLE roomserver_membership RENAME TO roomserver_membership_tmp;
 CREATE TABLE IF NOT EXISTS roomserver_membership (
 		room_nid INTEGER NOT NULL,
 		target_nid INTEGER NOT NULL,
@@ -57,8 +46,8 @@ DROP TABLE roomserver_membership_tmp;`)
 	return nil
 }
 
-func DownAddForgottenColumn(tx *sql.Tx) error {
-	_, err := tx.Exec(`	ALTER TABLE roomserver_membership RENAME TO roomserver_membership_tmp;
+func DownAddForgottenColumn(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, `	ALTER TABLE roomserver_membership RENAME TO roomserver_membership_tmp;
 CREATE TABLE IF NOT EXISTS roomserver_membership (
 		room_nid INTEGER NOT NULL,
 		target_nid INTEGER NOT NULL,
diff --git a/roomserver/storage/sqlite3/deltas/2021041615092700_state_blocks_refactor.go b/roomserver/storage/sqlite3/deltas/2021041615092700_state_blocks_refactor.go
index 8f5ab8fc..00978121 100644
--- a/roomserver/storage/sqlite3/deltas/2021041615092700_state_blocks_refactor.go
+++ b/roomserver/storage/sqlite3/deltas/2021041615092700_state_blocks_refactor.go
@@ -21,40 +21,35 @@ import (
 	"fmt"
 
 	"github.com/matrix-org/dendrite/internal"
-	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/util"
 	"github.com/sirupsen/logrus"
 )
 
-func LoadStateBlocksRefactor(m *sqlutil.Migrations) {
-	m.AddMigration(UpStateBlocksRefactor, DownStateBlocksRefactor)
-}
-
 // nolint:gocyclo
-func UpStateBlocksRefactor(tx *sql.Tx) error {
+func UpStateBlocksRefactor(ctx context.Context, tx *sql.Tx) error {
 	logrus.Warn("Performing state storage upgrade. Please wait, this may take some time!")
 	defer logrus.Warn("State storage upgrade complete")
 
 	var maxsnapshotid int
 	var maxblockid int
-	if err := tx.QueryRow(`SELECT IFNULL(MAX(state_snapshot_nid),0) FROM roomserver_state_snapshots;`).Scan(&maxsnapshotid); err != nil {
-		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	if err := tx.QueryRowContext(ctx, `SELECT IFNULL(MAX(state_snapshot_nid),0) FROM roomserver_state_snapshots;`).Scan(&maxsnapshotid); err != nil {
+		return fmt.Errorf("tx.QueryRowContext.Scan (count snapshots): %w", err)
 	}
-	if err := tx.QueryRow(`SELECT IFNULL(MAX(state_block_nid),0) FROM roomserver_state_block;`).Scan(&maxblockid); err != nil {
-		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	if err := tx.QueryRowContext(ctx, `SELECT IFNULL(MAX(state_block_nid),0) FROM roomserver_state_block;`).Scan(&maxblockid); err != nil {
+		return fmt.Errorf("tx.QueryRowContext.Scan (count snapshots): %w", err)
 	}
 	maxsnapshotid++
 	maxblockid++
 	oldMaxSnapshotID := maxsnapshotid
 
-	if _, err := tx.Exec(`ALTER TABLE roomserver_state_block RENAME TO _roomserver_state_block;`); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_state_block RENAME TO _roomserver_state_block;`); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	if _, err := tx.Exec(`ALTER TABLE roomserver_state_snapshots RENAME TO _roomserver_state_snapshots;`); err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+	if _, err := tx.ExecContext(ctx, `ALTER TABLE roomserver_state_snapshots RENAME TO _roomserver_state_snapshots;`); err != nil {
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	_, err := tx.Exec(`
+	_, err := tx.ExecContext(ctx, `
 		CREATE TABLE IF NOT EXISTS roomserver_state_block (
 			state_block_nid INTEGER PRIMARY KEY AUTOINCREMENT,
 			state_block_hash BLOB UNIQUE,
@@ -62,9 +57,9 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 		);
 	`)
 	if err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	_, err = tx.Exec(`
+	_, err = tx.ExecContext(ctx, `
 		CREATE TABLE IF NOT EXISTS roomserver_state_snapshots (
 			state_snapshot_nid INTEGER PRIMARY KEY AUTOINCREMENT,
 			state_snapshot_hash BLOB UNIQUE,
@@ -73,11 +68,11 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	  	);
 	`)
 	if err != nil {
-		return fmt.Errorf("tx.Exec: %w", err)
+		return fmt.Errorf("tx.ExecContext: %w", err)
 	}
-	snapshotrows, err := tx.Query(`SELECT state_snapshot_nid, room_nid, state_block_nids FROM _roomserver_state_snapshots;`)
+	snapshotrows, err := tx.QueryContext(ctx, `SELECT state_snapshot_nid, room_nid, state_block_nids FROM _roomserver_state_snapshots;`)
 	if err != nil {
-		return fmt.Errorf("tx.Query: %w", err)
+		return fmt.Errorf("tx.QueryContext: %w", err)
 	}
 	defer internal.CloseAndLogIfError(context.TODO(), snapshotrows, "rows.close() failed")
 	for snapshotrows.Next() {
@@ -99,7 +94,7 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			// in question a state snapshot NID of 0 to indicate 'no snapshot'.
 			// If we don't do this, we'll fail the assertions later on which try to ensure we didn't forget
 			// any snapshots.
-			_, err = tx.Exec(
+			_, err = tx.ExecContext(ctx,
 				`UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE event_type_nid = $1 AND event_state_key_nid = $2 AND state_snapshot_nid = $3`,
 				types.MRoomCreateNID, types.EmptyStateKeyNID, snapshot,
 			)
@@ -109,9 +104,9 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 		}
 		for _, block := range blocks {
 			if err = func() error {
-				blockrows, berr := tx.Query(`SELECT event_nid FROM _roomserver_state_block WHERE state_block_nid = $1`, block)
+				blockrows, berr := tx.QueryContext(ctx, `SELECT event_nid FROM _roomserver_state_block WHERE state_block_nid = $1`, block)
 				if berr != nil {
-					return fmt.Errorf("tx.Query (event nids from old block): %w", berr)
+					return fmt.Errorf("tx.QueryContext (event nids from old block): %w", berr)
 				}
 				defer internal.CloseAndLogIfError(context.TODO(), blockrows, "rows.close() failed")
 				events := types.EventNIDs{}
@@ -129,14 +124,14 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 				}
 
 				var blocknid types.StateBlockNID
-				err = tx.QueryRow(`
+				err = tx.QueryRowContext(ctx, `
 					INSERT INTO roomserver_state_block (state_block_nid, state_block_hash, event_nids)
 						VALUES ($1, $2, $3)
 						ON CONFLICT (state_block_hash) DO UPDATE SET event_nids=$3
 						RETURNING state_block_nid
 				`, maxblockid, events.Hash(), eventjson).Scan(&blocknid)
 				if err != nil {
-					return fmt.Errorf("tx.QueryRow.Scan (insert new block): %w", err)
+					return fmt.Errorf("tx.QueryRowContext.Scan (insert new block): %w", err)
 				}
 				maxblockid++
 				newblocks = append(newblocks, blocknid)
@@ -151,22 +146,22 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			}
 
 			var newsnapshot types.StateSnapshotNID
-			err = tx.QueryRow(`
+			err = tx.QueryRowContext(ctx, `
 				INSERT INTO roomserver_state_snapshots (state_snapshot_nid, state_snapshot_hash, room_nid, state_block_nids)
 					VALUES ($1, $2, $3, $4)
 					ON CONFLICT (state_snapshot_hash) DO UPDATE SET room_nid=$3
 					RETURNING state_snapshot_nid
 			`, maxsnapshotid, newblocks.Hash(), room, newblocksjson).Scan(&newsnapshot)
 			if err != nil {
-				return fmt.Errorf("tx.QueryRow.Scan (insert new snapshot): %w", err)
+				return fmt.Errorf("tx.QueryRowContext.Scan (insert new snapshot): %w", err)
 			}
 			maxsnapshotid++
-			_, err = tx.Exec(`UPDATE roomserver_events SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newsnapshot, snapshot, maxsnapshotid)
+			_, err = tx.ExecContext(ctx, `UPDATE roomserver_events SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newsnapshot, snapshot, maxsnapshotid)
 			if err != nil {
-				return fmt.Errorf("tx.Exec (update events): %w", err)
+				return fmt.Errorf("tx.ExecContext (update events): %w", err)
 			}
-			if _, err = tx.Exec(`UPDATE roomserver_rooms SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newsnapshot, snapshot, maxsnapshotid); err != nil {
-				return fmt.Errorf("tx.Exec (update rooms): %w", err)
+			if _, err = tx.ExecContext(ctx, `UPDATE roomserver_rooms SET state_snapshot_nid=$1 WHERE state_snapshot_nid=$2 AND state_snapshot_nid<$3`, newsnapshot, snapshot, maxsnapshotid); err != nil {
+				return fmt.Errorf("tx.ExecContext (update rooms): %w", err)
 			}
 		}
 	}
@@ -175,13 +170,13 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 	// If we do, this is a problem if Dendrite tries to load the snapshot as it will not exist
 	// in roomserver_state_snapshots
 	var count int64
-	if err = tx.QueryRow(`SELECT COUNT(*) FROM roomserver_events WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID).Scan(&count); err != nil {
+	if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM roomserver_events WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID).Scan(&count); err != nil {
 		return fmt.Errorf("assertion query failed: %s", err)
 	}
 	if count > 0 {
 		var res sql.Result
 		var c int64
-		res, err = tx.Exec(`UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID)
+		res, err = tx.ExecContext(ctx, `UPDATE roomserver_events SET state_snapshot_nid = 0 WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID)
 		if err != nil && err != sql.ErrNoRows {
 			return fmt.Errorf("failed to reset invalid state snapshots: %w", err)
 		}
@@ -191,23 +186,23 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 			return fmt.Errorf("expected to reset %d event(s) but only updated %d event(s)", count, c)
 		}
 	}
-	if err = tx.QueryRow(`SELECT COUNT(*) FROM roomserver_rooms WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID).Scan(&count); err != nil {
+	if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM roomserver_rooms WHERE state_snapshot_nid < $1 AND state_snapshot_nid != 0`, oldMaxSnapshotID).Scan(&count); err != nil {
 		return fmt.Errorf("assertion query failed: %s", err)
 	}
 	if count > 0 {
 		return fmt.Errorf("%d rooms exist in roomserver_rooms which have not been converted to a new state_snapshot_nid; this is a bug, please report", count)
 	}
 
-	if _, err = tx.Exec(`DROP TABLE _roomserver_state_snapshots;`); err != nil {
+	if _, err = tx.ExecContext(ctx, `DROP TABLE _roomserver_state_snapshots;`); err != nil {
 		return fmt.Errorf("tx.Exec (delete old snapshot table): %w", err)
 	}
-	if _, err = tx.Exec(`DROP TABLE _roomserver_state_block;`); err != nil {
+	if _, err = tx.ExecContext(ctx, `DROP TABLE _roomserver_state_block;`); err != nil {
 		return fmt.Errorf("tx.Exec (delete old block table): %w", err)
 	}
 
 	return nil
 }
 
-func DownStateBlocksRefactor(tx *sql.Tx) error {
+func DownStateBlocksRefactor(ctx context.Context, tx *sql.Tx) error {
 	panic("Downgrading state storage is not supported")
 }
diff --git a/roomserver/storage/sqlite3/membership_table.go b/roomserver/storage/sqlite3/membership_table.go
index 9dda82fb..f3303eb0 100644
--- a/roomserver/storage/sqlite3/membership_table.go
+++ b/roomserver/storage/sqlite3/membership_table.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3/deltas"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -148,7 +149,15 @@ type membershipStatements struct {
 
 func CreateMembershipTable(db *sql.DB) error {
 	_, err := db.Exec(membershipSchema)
-	return err
+	if err != nil {
+		return err
+	}
+	m := sqlutil.NewMigrator(db)
+	m.AddMigrations(sqlutil.Migration{
+		Version: "roomserver: add forgotten column",
+		Up:      deltas.UpAddForgottenColumn,
+	})
+	return m.Up(context.Background())
 }
 
 func PrepareMembershipTable(db *sql.DB) (tables.Membership, error) {
diff --git a/roomserver/storage/sqlite3/storage.go b/roomserver/storage/sqlite3/storage.go
index 8325fdad..bb9c15b5 100644
--- a/roomserver/storage/sqlite3/storage.go
+++ b/roomserver/storage/sqlite3/storage.go
@@ -54,17 +54,25 @@ func Open(base *base.BaseDendrite, dbProperties *config.DatabaseOptions, cache c
 	// db.SetMaxOpenConns(20)
 
 	// Create the tables.
-	if err := d.create(db); err != nil {
+	if err = d.create(db); err != nil {
 		return nil, err
 	}
 
-	// Then execute the migrations. By this point the tables are created with the latest
-	// schemas.
-	m := sqlutil.NewMigrations()
-	deltas.LoadAddForgottenColumn(m)
-	deltas.LoadStateBlocksRefactor(m)
-	if err := m.RunDeltas(db, dbProperties); err != nil {
-		return nil, err
+	// Special case, since this migration uses several tables, so it needs to
+	// be sure that all tables are created first.
+	// TODO: Remove when we are sure we are not having goose artefacts in the db
+	// This forces an error, which indicates the migration is already applied, since the
+	// column event_nid was removed from the table
+	err = db.QueryRow("SELECT event_nid FROM roomserver_state_block LIMIT 1;").Scan()
+	if err == nil {
+		m := sqlutil.NewMigrator(db)
+		m.AddMigrations(sqlutil.Migration{
+			Version: "roomserver: state blocks refactor",
+			Up:      deltas.UpStateBlocksRefactor,
+		})
+		if err := m.Up(base.Context()); err != nil {
+			return nil, err
+		}
 	}
 
 	// Then prepare the statements. Now that the migrations have run, any columns referred