From 637119be92338e7637f7645d61bda1d0963d104e Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Tue, 23 Aug 2022 09:09:35 +0200
Subject: [PATCH] Use sqlite_master to determine if migration was already
 executed

---
 .../storage/sqlite3/key_changes_table.go      | 37 ++++++-------------
 roomserver/storage/sqlite3/storage.go         | 37 ++++++-------------
 2 files changed, 24 insertions(+), 50 deletions(-)

diff --git a/keyserver/storage/sqlite3/key_changes_table.go b/keyserver/storage/sqlite3/key_changes_table.go
index f3d743cac..9e0cb9ef5 100644
--- a/keyserver/storage/sqlite3/key_changes_table.go
+++ b/keyserver/storage/sqlite3/key_changes_table.go
@@ -17,9 +17,8 @@ package sqlite3
 import (
 	"context"
 	"database/sql"
-	"strings"
+	"errors"
 
-	"github.com/mattn/go-sqlite3"
 	"github.com/sirupsen/logrus"
 
 	"github.com/matrix-org/dendrite/internal"
@@ -85,35 +84,23 @@ func executeMigration(ctx context.Context, db *sql.DB) error {
 		return nil
 	}
 
-	var partition int
-	err = db.QueryRowContext(ctx, "SELECT partition FROM keyserver_key_changes LIMIT 1;").Scan(&partition)
-	if err == nil {
-		m := sqlutil.NewMigrator(db)
-		m.AddMigrations(sqlutil.Migration{
-			Version: migrationName,
-			Up:      deltas.UpRefactorKeyChanges,
-		})
-		if err = m.Up(ctx); err != nil {
-			return err
-		}
-	} else {
-		switch e := err.(type) {
-		case *sqlite3.Error:
-			// ignore "no such column" errors, as this is expected at this point
-			if !strings.Contains(e.Error(), "no such column") {
-				return err
-			}
-
-			// reset the error to nil, so the deferred function will insert the migration
+	err = db.QueryRowContext(ctx, `SELECT p.name FROM sqlite_master AS m JOIN pragma_table_info(m.name) AS p WHERE m.name = 'keyserver_key_changes' AND p.name = 'partition'`).Err()
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) { // migration was already executed, as the column was removed
 			if err = sqlutil.InsertMigration(ctx, db, migrationName); err != nil {
 				// not a fatal error, log and continue
 				logrus.WithError(err).Warnf("unable to manually insert migration '%s'", migrationName)
 			}
-		default:
-			return err
+			return nil
 		}
+		return err
 	}
-	return nil
+	m := sqlutil.NewMigrator(db)
+	m.AddMigrations(sqlutil.Migration{
+		Version: migrationName,
+		Up:      deltas.UpRefactorKeyChanges,
+	})
+	return m.Up(ctx)
 }
 
 func (s *keyChangesStatements) Prepare() (err error) {
diff --git a/roomserver/storage/sqlite3/storage.go b/roomserver/storage/sqlite3/storage.go
index 3c7996d85..c1b3b41de 100644
--- a/roomserver/storage/sqlite3/storage.go
+++ b/roomserver/storage/sqlite3/storage.go
@@ -18,11 +18,10 @@ package sqlite3
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/mattn/go-sqlite3"
 	"github.com/sirupsen/logrus"
 
 	"github.com/matrix-org/dendrite/internal/caching"
@@ -92,35 +91,23 @@ func executeMigration(ctx context.Context, db *sql.DB) error {
 		return nil
 	}
 
-	var eventNID int
-	err = db.QueryRowContext(ctx, "SELECT event_nid FROM roomserver_state_block LIMIT 1;").Scan(&eventNID)
-	if err == nil {
-		m := sqlutil.NewMigrator(db)
-		m.AddMigrations(sqlutil.Migration{
-			Version: migrationName,
-			Up:      deltas.UpStateBlocksRefactor,
-		})
-		if err = m.Up(ctx); err != nil {
-			return err
-		}
-	} else {
-		switch e := err.(type) {
-		case *sqlite3.Error:
-			// ignore "no such column" errors, as this is expected at this point
-			if !strings.Contains(e.Error(), "no such column") {
-				return err
-			}
-
-			// reset the error to nil, so the deferred function will insert the migration
+	err = db.QueryRowContext(ctx, `SELECT p.name FROM sqlite_master AS m JOIN pragma_table_info(m.name) AS p WHERE m.name = 'roomserver_state_block' AND p.name = 'event_nid'`).Err()
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) { // migration was already executed, as the column was removed
 			if err = sqlutil.InsertMigration(ctx, db, migrationName); err != nil {
 				// not a fatal error, log and continue
 				logrus.WithError(err).Warnf("unable to manually insert migration '%s'", migrationName)
 			}
-		default:
-			return err
+			return nil
 		}
+		return err
 	}
-	return nil
+	m := sqlutil.NewMigrator(db)
+	m.AddMigrations(sqlutil.Migration{
+		Version: migrationName,
+		Up:      deltas.UpStateBlocksRefactor,
+	})
+	return m.Up(ctx)
 }
 
 func (d *Database) create(db *sql.DB) error {