From 205369b05f71520d188d1f6c1936fab686de352c Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Thu, 9 Nov 2023 09:57:12 -0500 Subject: [PATCH 1/4] Use attempt_to_set_autocommit everywhere. --- synapse/storage/background_updates.py | 18 ++++++++++++------ synapse/storage/databases/main/search.py | 8 ++++---- synapse/storage/databases/state/bg_updates.py | 4 ++-- tests/server.py | 15 +++++++-------- 4 files changed, 25 insertions(+), 20 deletions(-) diff --git a/synapse/storage/background_updates.py b/synapse/storage/background_updates.py index 7426dbcad611..62fbd055348a 100644 --- a/synapse/storage/background_updates.py +++ b/synapse/storage/background_updates.py @@ -49,7 +49,11 @@ if TYPE_CHECKING: from synapse.server import HomeServer - from synapse.storage.database import DatabasePool, LoggingTransaction + from synapse.storage.database import ( + DatabasePool, + LoggingDatabaseConnection, + LoggingTransaction, + ) logger = logging.getLogger(__name__) @@ -746,10 +750,10 @@ async def create_index_in_background( The named index will be dropped upon completion of the new index. """ - def create_index_psql(conn: Connection) -> None: + def create_index_psql(conn: "LoggingDatabaseConnection") -> None: conn.rollback() # postgres insists on autocommit for the index - conn.set_session(autocommit=True) # type: ignore + conn.engine.attempt_to_set_autocommit(conn.conn, True) try: c = conn.cursor() @@ -793,9 +797,9 @@ def create_index_psql(conn: Connection) -> None: undo_timeout_sql = f"SET statement_timeout = {default_timeout}" conn.cursor().execute(undo_timeout_sql) - conn.set_session(autocommit=False) # type: ignore + conn.engine.attempt_to_set_autocommit(conn.conn, False) - def create_index_sqlite(conn: Connection) -> None: + def create_index_sqlite(conn: "LoggingDatabaseConnection") -> None: # Sqlite doesn't support concurrent creation of indexes. # # We assume that sqlite doesn't give us invalid indices; however @@ -825,7 +829,9 @@ def create_index_sqlite(conn: Connection) -> None: c.execute(sql) if isinstance(self.db_pool.engine, engines.PostgresEngine): - runner: Optional[Callable[[Connection], None]] = create_index_psql + runner: Optional[ + Callable[[LoggingDatabaseConnection], None] + ] = create_index_psql elif psql_only: runner = None else: diff --git a/synapse/storage/databases/main/search.py b/synapse/storage/databases/main/search.py index f4bef4c99b35..e25d86818b37 100644 --- a/synapse/storage/databases/main/search.py +++ b/synapse/storage/databases/main/search.py @@ -275,7 +275,7 @@ def create_index(conn: LoggingDatabaseConnection) -> None: # we have to set autocommit, because postgres refuses to # CREATE INDEX CONCURRENTLY without it. - conn.set_session(autocommit=True) + conn.engine.attempt_to_set_autocommit(conn.conn, True) try: c = conn.cursor() @@ -301,7 +301,7 @@ def create_index(conn: LoggingDatabaseConnection) -> None: # we should now be able to delete the GIST index. c.execute("DROP INDEX IF EXISTS event_search_fts_idx_gist") finally: - conn.set_session(autocommit=False) + conn.engine.attempt_to_set_autocommit(conn.conn, False) if isinstance(self.database_engine, PostgresEngine): await self.db_pool.runWithConnection(create_index) @@ -323,7 +323,7 @@ async def _background_reindex_search_order( def create_index(conn: LoggingDatabaseConnection) -> None: conn.rollback() - conn.set_session(autocommit=True) + conn.engine.attempt_to_set_autocommit(conn.conn, True) c = conn.cursor() # We create with NULLS FIRST so that when we search *backwards* @@ -340,7 +340,7 @@ def create_index(conn: LoggingDatabaseConnection) -> None: ON event_search(origin_server_ts NULLS FIRST, stream_ordering NULLS FIRST) """ ) - conn.set_session(autocommit=False) + conn.engine.attempt_to_set_autocommit(conn.conn, False) await self.db_pool.runWithConnection(create_index) diff --git a/synapse/storage/databases/state/bg_updates.py b/synapse/storage/databases/state/bg_updates.py index 0f9c550b27e4..2c3151526db8 100644 --- a/synapse/storage/databases/state/bg_updates.py +++ b/synapse/storage/databases/state/bg_updates.py @@ -492,7 +492,7 @@ def reindex_txn(conn: LoggingDatabaseConnection) -> None: conn.rollback() if isinstance(self.database_engine, PostgresEngine): # postgres insists on autocommit for the index - conn.set_session(autocommit=True) + conn.engine.attempt_to_set_autocommit(conn.conn, True) try: txn = conn.cursor() txn.execute( @@ -501,7 +501,7 @@ def reindex_txn(conn: LoggingDatabaseConnection) -> None: ) txn.execute("DROP INDEX IF EXISTS state_groups_state_id") finally: - conn.set_session(autocommit=False) + conn.engine.attempt_to_set_autocommit(conn.conn, False) else: txn = conn.cursor() txn.execute( diff --git a/tests/server.py b/tests/server.py index c8342db399ab..22880b7dc10e 100644 --- a/tests/server.py +++ b/tests/server.py @@ -88,7 +88,7 @@ from synapse.server import HomeServer from synapse.storage import DataStore from synapse.storage.database import LoggingDatabaseConnection -from synapse.storage.engines import PostgresEngine, create_engine +from synapse.storage.engines import PostgresEngine, create_engine, BaseDatabaseEngine from synapse.storage.prepare_database import prepare_database from synapse.types import ISynapseReactor, JsonDict from synapse.util import Clock @@ -1030,8 +1030,6 @@ def setup_test_homeserver( # Create the database before we actually try and connect to it, based off # the template database we generate in setupdb() if isinstance(db_engine, PostgresEngine): - import psycopg2.extensions - db_conn = db_engine.module.connect( database=POSTGRES_BASE_DB, user=POSTGRES_USER, @@ -1039,8 +1037,9 @@ def setup_test_homeserver( port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - assert isinstance(db_conn, psycopg2.extensions.connection) - db_conn.autocommit = True + # Trick mypy into treating this generically so it doesn't complain about + # the connection type. + cast(BaseDatabaseEngine, db_engine).attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,)) cur.execute( @@ -1071,7 +1070,6 @@ def setup_test_homeserver( # We need to do cleanup on PostgreSQL def cleanup() -> None: import psycopg2 - import psycopg2.extensions # Close all the db pools database_pool._db_pool.close() @@ -1086,8 +1084,9 @@ def cleanup() -> None: port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - assert isinstance(db_conn, psycopg2.extensions.connection) - db_conn.autocommit = True + # Trick mypy into treating this generically so it doesn't complain about + # the connection type. + cast(BaseDatabaseEngine, db_engine).attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() # Try a few times to drop the DB. Some things may hold on to the From dd38626b7e6531265a1dd37f83bc7daf4bd481f4 Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Thu, 9 Nov 2023 10:29:15 -0500 Subject: [PATCH 2/4] Newsfragment --- changelog.d/16615.misc | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/16615.misc diff --git a/changelog.d/16615.misc b/changelog.d/16615.misc new file mode 100644 index 000000000000..37ab711dc657 --- /dev/null +++ b/changelog.d/16615.misc @@ -0,0 +1 @@ +Use more generic database methods. From 4eb0b5092385d732e8d2446bec6a69131885f9e9 Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Thu, 9 Nov 2023 10:33:24 -0500 Subject: [PATCH 3/4] Lint --- tests/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server.py b/tests/server.py index 22880b7dc10e..955fa27b7bac 100644 --- a/tests/server.py +++ b/tests/server.py @@ -88,7 +88,7 @@ from synapse.server import HomeServer from synapse.storage import DataStore from synapse.storage.database import LoggingDatabaseConnection -from synapse.storage.engines import PostgresEngine, create_engine, BaseDatabaseEngine +from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, create_engine from synapse.storage.prepare_database import prepare_database from synapse.types import ISynapseReactor, JsonDict from synapse.util import Clock From 1f31fdeac64527e78a2354c0bbcb09abb96a837f Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Thu, 9 Nov 2023 14:07:56 -0500 Subject: [PATCH 4/4] Remove the cast call. --- tests/server.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/tests/server.py b/tests/server.py index 955fa27b7bac..82da23907541 100644 --- a/tests/server.py +++ b/tests/server.py @@ -88,7 +88,7 @@ from synapse.server import HomeServer from synapse.storage import DataStore from synapse.storage.database import LoggingDatabaseConnection -from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, create_engine +from synapse.storage.engines import create_engine from synapse.storage.prepare_database import prepare_database from synapse.types import ISynapseReactor, JsonDict from synapse.util import Clock @@ -1029,7 +1029,7 @@ def setup_test_homeserver( # Create the database before we actually try and connect to it, based off # the template database we generate in setupdb() - if isinstance(db_engine, PostgresEngine): + if USE_POSTGRES_FOR_TESTS: db_conn = db_engine.module.connect( database=POSTGRES_BASE_DB, user=POSTGRES_USER, @@ -1037,9 +1037,7 @@ def setup_test_homeserver( port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - # Trick mypy into treating this generically so it doesn't complain about - # the connection type. - cast(BaseDatabaseEngine, db_engine).attempt_to_set_autocommit(db_conn, True) + db_engine.attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,)) cur.execute( @@ -1064,7 +1062,7 @@ def setup_test_homeserver( hs.setup() - if isinstance(db_engine, PostgresEngine): + if USE_POSTGRES_FOR_TESTS: database_pool = hs.get_datastores().databases[0] # We need to do cleanup on PostgreSQL @@ -1084,9 +1082,7 @@ def cleanup() -> None: port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - # Trick mypy into treating this generically so it doesn't complain about - # the connection type. - cast(BaseDatabaseEngine, db_engine).attempt_to_set_autocommit(db_conn, True) + db_engine.attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() # Try a few times to drop the DB. Some things may hold on to the