From 778b981879f2fd718c69a822d720200e09affd7f Mon Sep 17 00:00:00 2001
From: Till Faelligen <tfaelligen@gmail.com>
Date: Thu, 28 Apr 2022 14:25:51 +0200
Subject: [PATCH] Fix upgrade

---
 .../storage/postgres/queue_edus_table.go          | 15 +++++++--------
 federationapi/storage/postgres/storage.go         |  3 +++
 federationapi/storage/sqlite3/queue_edus_table.go | 11 ++++++-----
 federationapi/storage/sqlite3/storage.go          |  3 +++
 federationapi/storage/tables/interface.go         |  1 +
 5 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/federationapi/storage/postgres/queue_edus_table.go b/federationapi/storage/postgres/queue_edus_table.go
index 8f68792e1..a7ca328ef 100644
--- a/federationapi/storage/postgres/queue_edus_table.go
+++ b/federationapi/storage/postgres/queue_edus_table.go
@@ -86,21 +86,20 @@ func NewPostgresQueueEDUsTable(db *sql.DB) (s *queueEDUsStatements, err error) {
 		db: db,
 	}
 	_, err = s.db.Exec(queueEDUsSchema)
-	if err != nil {
-		return
-	}
-	if s.deleteQueueEDUStmt, err = s.db.Prepare(deleteQueueEDUSQL); err != nil {
-		return
-	}
-	return s, sqlutil.StatementList{
+	return s, err
+}
+
+func (s *queueEDUsStatements) Prepare() error {
+	return sqlutil.StatementList{
 		{&s.insertQueueEDUStmt, insertQueueEDUSQL},
+		{&s.deleteQueueEDUStmt, deleteQueueEDUSQL},
 		{&s.selectQueueEDUStmt, selectQueueEDUSQL},
 		{&s.selectQueueEDUReferenceJSONCountStmt, selectQueueEDUReferenceJSONCountSQL},
 		{&s.selectQueueEDUCountStmt, selectQueueEDUCountSQL},
 		{&s.selectQueueEDUServerNamesStmt, selectQueueServerNamesSQL},
 		{&s.selectExpiredEDUsStmt, selectExpiredEDUsSQL},
 		{&s.deleteExpiredEDUsStmt, deleteExpiredEDUsSQL},
-	}.Prepare(db)
+	}.Prepare(s.db)
 }
 
 func (s *queueEDUsStatements) InsertQueueEDU(
diff --git a/federationapi/storage/postgres/storage.go b/federationapi/storage/postgres/storage.go
index e087af3d5..993bf7945 100644
--- a/federationapi/storage/postgres/storage.go
+++ b/federationapi/storage/postgres/storage.go
@@ -88,6 +88,9 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationC
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
 		return nil, err
 	}
+	if err = queueEDUs.Prepare(); err != nil {
+		return nil, err
+	}
 	d.Database = shared.Database{
 		DB:                       d.db,
 		ServerName:               serverName,
diff --git a/federationapi/storage/sqlite3/queue_edus_table.go b/federationapi/storage/sqlite3/queue_edus_table.go
index 4935cf262..a2a60b29c 100644
--- a/federationapi/storage/sqlite3/queue_edus_table.go
+++ b/federationapi/storage/sqlite3/queue_edus_table.go
@@ -87,10 +87,11 @@ func NewSQLiteQueueEDUsTable(db *sql.DB) (s *queueEDUsStatements, err error) {
 		db: db,
 	}
 	_, err = db.Exec(queueEDUsSchema)
-	if err != nil {
-		return
-	}
-	return s, sqlutil.StatementList{
+	return s, err
+}
+
+func (s *queueEDUsStatements) Prepare() error {
+	return sqlutil.StatementList{
 		{&s.insertQueueEDUStmt, insertQueueEDUSQL},
 		{&s.selectQueueEDUStmt, selectQueueEDUSQL},
 		{&s.selectQueueEDUReferenceJSONCountStmt, selectQueueEDUReferenceJSONCountSQL},
@@ -98,7 +99,7 @@ func NewSQLiteQueueEDUsTable(db *sql.DB) (s *queueEDUsStatements, err error) {
 		{&s.selectQueueEDUServerNamesStmt, selectQueueServerNamesSQL},
 		{&s.selectExpiredEDUsStmt, selectExpiredEDUsSQL},
 		{&s.deleteExpiredEDUsStmt, deleteExpiredEDUsSQL},
-	}.Prepare(db)
+	}.Prepare(s.db)
 }
 
 func (s *queueEDUsStatements) InsertQueueEDU(
diff --git a/federationapi/storage/sqlite3/storage.go b/federationapi/storage/sqlite3/storage.go
index 939f466c4..63cebfaa1 100644
--- a/federationapi/storage/sqlite3/storage.go
+++ b/federationapi/storage/sqlite3/storage.go
@@ -87,6 +87,9 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationC
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
 		return nil, err
 	}
+	if err = queueEDUs.Prepare(); err != nil {
+		return nil, err
+	}
 	d.Database = shared.Database{
 		DB:                       d.db,
 		ServerName:               serverName,
diff --git a/federationapi/storage/tables/interface.go b/federationapi/storage/tables/interface.go
index 64c66da90..ffb70a47e 100644
--- a/federationapi/storage/tables/interface.go
+++ b/federationapi/storage/tables/interface.go
@@ -42,6 +42,7 @@ type FederationQueueEDUs interface {
 	SelectQueueEDUServerNames(ctx context.Context, txn *sql.Tx) ([]gomatrixserverlib.ServerName, error)
 	SelectExpiredEDUs(ctx context.Context, txn *sql.Tx, expiredBefore gomatrixserverlib.Timestamp) ([]int64, error)
 	DeleteExpiredEDUs(ctx context.Context, txn *sql.Tx, expiredBefore gomatrixserverlib.Timestamp) error
+	Prepare() error
 }
 
 type FederationQueueJSON interface {