Override global statement timeout when creating indexes in Postgres (#16085)

1 files changed, 12 insertions, 0 deletions

diff --git a/synapse/storage/background_updates.py b/synapse/storage/background_updates.py
index 2d5ddc3e7b..ddca0af1da 100644
--- a/synapse/storage/background_updates.py
+++ b/synapse/storage/background_updates.py
@@ -238,6 +238,7 @@ class BackgroundUpdater:
     def __init__(self, hs: "HomeServer", database: "DatabasePool"):
         self._clock = hs.get_clock()
         self.db_pool = database
+        self.hs = hs
 
         self._database_name = database.name()
 
@@ -758,6 +759,11 @@ class BackgroundUpdater:
                 logger.debug("[SQL] %s", sql)
                 c.execute(sql)
 
+                # override the global statement timeout to avoid accidentally squashing
+                # a long-running index creation process
+                timeout_sql = "SET SESSION statement_timeout = 0"
+                c.execute(timeout_sql)
+
                 sql = (
                     "CREATE %(unique)s INDEX CONCURRENTLY %(name)s"
                     " ON %(table)s"
@@ -778,6 +784,12 @@ class BackgroundUpdater:
                     logger.debug("[SQL] %s", sql)
                     c.execute(sql)
             finally:
+                # mypy ignore - `statement_timeout` is defined on PostgresEngine
+                # reset the global timeout to the default
+                default_timeout = self.db_pool.engine.statement_timeout  # type: ignore[attr-defined]
+                undo_timeout_sql = f"SET statement_timeout = {default_timeout}"
+                conn.cursor().execute(undo_timeout_sql)
+
                 conn.set_session(autocommit=False)  # type: ignore
 
         def create_index_sqlite(conn: Connection) -> None: