From e4e42797ee787209ee5ddf83795dd627d72bf9e7 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 11 Nov 2022 13:02:32 +0000
Subject: [PATCH] Fix SQLite migration

---
 .../deltas/2022110411000000_server_names.go   | 56 ++++++++++++++-----
 1 file changed, 43 insertions(+), 13 deletions(-)

diff --git a/userapi/storage/sqlite3/deltas/2022110411000000_server_names.go b/userapi/storage/sqlite3/deltas/2022110411000000_server_names.go
index 74e1db811..4426d1589 100644
--- a/userapi/storage/sqlite3/deltas/2022110411000000_server_names.go
+++ b/userapi/storage/sqlite3/deltas/2022110411000000_server_names.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 
 	"github.com/lib/pq"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/sirupsen/logrus"
 )
 
 var serverNamesTables = []string{
@@ -28,6 +30,13 @@ var serverNamesDropPK = []string{
 	"userapi_profiles",
 }
 
+// These indices are out of date so let's drop them. They will get recreated
+// automatically.
+var serverNamesDropIndex = []string{
+	"userapi_pusher_localpart_idx",
+	"userapi_pusher_app_id_pushkey_localpart_idx",
+}
+
 // I know what you're thinking: you're wondering "why doesn't this use $1
 // and pass variadic parameters to ExecContext?" — the answer is because
 // PostgreSQL doesn't expect the table name to be specified as a substituted
@@ -36,37 +45,58 @@ var serverNamesDropPK = []string{
 func UpServerNames(ctx context.Context, tx *sql.Tx, serverName gomatrixserverlib.ServerName) error {
 	for _, table := range serverNamesTables {
 		q := fmt.Sprintf(
-			"SELECT COUNT(name) FROM sqlite_schema WHERE type='table' AND name=%s;",
+			"SELECT COUNT(*) FROM pragma_table_info(%s) WHERE name='server_name'",
 			pq.QuoteIdentifier(table),
 		)
 		var c int
-		if err := tx.QueryRowContext(ctx, q).Scan(&c); err != nil || c == 0 {
+		if err := tx.QueryRowContext(ctx, q).Scan(&c); err != nil || c == 1 {
+			logrus.Infof("Table %s already has column, skipping", table)
 			continue
 		}
-		q = fmt.Sprintf(
-			"ALTER TABLE %s ADD COLUMN IF NOT EXISTS server_name TEXT NOT NULL DEFAULT '';",
-			pq.QuoteIdentifier(table),
-		)
-		if _, err := tx.ExecContext(ctx, q); err != nil {
-			return fmt.Errorf("add server name to %q error: %w", table, err)
+		logrus.Infof("Table %s add column", table)
+		if c == 0 {
+			q = fmt.Sprintf(
+				"ALTER TABLE %s ADD COLUMN server_name TEXT NOT NULL DEFAULT '';",
+				pq.QuoteIdentifier(table),
+			)
+			if _, err := tx.ExecContext(ctx, q); err != nil {
+				return fmt.Errorf("add server name to %q error: %w", table, err)
+			}
 		}
 	}
 	for _, table := range serverNamesDropPK {
 		q := fmt.Sprintf(
-			"SELECT COUNT(name) FROM sqlite_schema WHERE type='table' AND name=%s;",
+			"SELECT COUNT(name), sql FROM sqlite_schema WHERE type='table' AND name=%s;",
 			pq.QuoteIdentifier(table),
 		)
 		var c int
-		if err := tx.QueryRowContext(ctx, q).Scan(&c); err != nil || c == 0 {
+		var sql string
+		if err := tx.QueryRowContext(ctx, q).Scan(&c, &sql); err != nil || c == 0 {
 			continue
 		}
-		q = fmt.Sprintf(
-			"ALTER TABLE %s DROP PRIMARY KEY;",
-			pq.QuoteIdentifier(table),
+		q = fmt.Sprintf(`
+			%s; -- create temporary table
+			INSERT INTO %s SELECT * FROM %s; -- copy data
+			DROP TABLE %s; -- drop original table
+			ALTER TABLE %s RENAME TO %s; -- rename new table
+		`,
+			strings.Replace(sql, table, table+"_tmp", 1), // create temporary table
+			table+"_tmp", table, // copy data
+			table,               // drop original table
+			table+"_tmp", table, // rename new table
 		)
 		if _, err := tx.ExecContext(ctx, q); err != nil {
 			return fmt.Errorf("drop PK from %q error: %w", table, err)
 		}
 	}
+	for _, index := range serverNamesDropIndex {
+		q := fmt.Sprintf(
+			"DROP INDEX IF EXISTS %s;",
+			pq.QuoteIdentifier(index),
+		)
+		if _, err := tx.ExecContext(ctx, q); err != nil {
+			return fmt.Errorf("drop index %q error: %w", index, err)
+		}
+	}
 	return nil
 }