Method to set statement timeout. github/clokep/psycopg3-driver clokep/psycopg3-driver

1 files changed, 7 insertions, 6 deletions

diff --git a/synapse/storage/background_updates.py b/synapse/storage/background_updates.py
index 62fbd05534..58284223f5 100644
--- a/synapse/storage/background_updates.py
+++ b/synapse/storage/background_updates.py
@@ -755,6 +755,8 @@ class BackgroundUpdater:
             # postgres insists on autocommit for the index
             conn.engine.attempt_to_set_autocommit(conn.conn, True)
 
+            assert isinstance(self.db_pool.engine, PostgresEngine)
+
             try:
                 c = conn.cursor()
 
@@ -768,8 +770,7 @@ class BackgroundUpdater:
 
                 # 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)
+                self.db_pool.engine.set_statement_timeout(c, 0)
 
                 sql = (
                     "CREATE %(unique)s INDEX CONCURRENTLY %(name)s"
@@ -791,11 +792,11 @@ 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)
+                if self.db_pool.engine.statement_timeout is not None:
+                    self.db_pool.engine.set_statement_timeout(
+                        conn.cursor(), self.db_pool.engine.statement_timeout
+                    )
 
                 conn.engine.attempt_to_set_autocommit(conn.conn, False)