Don't use SQLBaseStore in state database

2 files changed, 9 insertions, 6 deletions

diff --git a/synapse/storage/databases/state/bg_updates.py b/synapse/storage/databases/state/bg_updates.py
index 1fd333b707..87e14132f5 100644
--- a/synapse/storage/databases/state/bg_updates.py
+++ b/synapse/storage/databases/state/bg_updates.py
@@ -26,17 +26,20 @@ logger = logging.getLogger(__name__)
 MAX_STATE_DELTA_HOPS = 100
 
 
-class StateGroupBackgroundUpdateStore(SQLBaseStore):
+class StateGroupBackgroundUpdateStore:
     """Defines functions related to state groups needed to run the state background
     updates.
     """
 
+    def __init__(self, database: DatabasePool, db_conn, hs):
+        self.db_pool = database
+
     def _count_state_group_hops_txn(self, txn, state_group):
         """Given a state group, count how many hops there are in the tree.
 
         This is used to ensure the delta chains don't get too long.
         """
-        if isinstance(self.database_engine, PostgresEngine):
+        if isinstance(self.db_pool.engine, PostgresEngine):
             sql = """
                 WITH RECURSIVE state(state_group) AS (
                     VALUES(?::bigint)
@@ -84,7 +87,7 @@ class StateGroupBackgroundUpdateStore(SQLBaseStore):
         if where_clause:
             where_clause = " AND (%s)" % (where_clause,)
 
-        if isinstance(self.database_engine, PostgresEngine):
+        if isinstance(self.db_pool.engine, PostgresEngine):
             # Temporarily disable sequential scans in this transaction. This is
             # a temporary hack until we can add the right indices in
             txn.execute("SET LOCAL enable_seqscan=off")
@@ -341,7 +344,7 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore):
     async def _background_index_state(self, progress, batch_size):
         def reindex_txn(conn):
             conn.rollback()
-            if isinstance(self.database_engine, PostgresEngine):
+            if isinstance(self.db_pool.engine, PostgresEngine):
                 # postgres insists on autocommit for the index
                 conn.set_session(autocommit=True)
                 try:
diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py
index e2240703a7..86b06809b3 100644
--- a/synapse/storage/databases/state/store.py
+++ b/synapse/storage/databases/state/store.py
@@ -47,7 +47,7 @@ class _GetStateGroupDelta(
         return len(self.delta_ids) if self.delta_ids else 0
 
 
-class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
+class StateGroupDataStore(StateBackgroundUpdateStore):
     """A data store for fetching/storing state groups."""
 
     def __init__(self, database: DatabasePool, db_conn, hs):
@@ -98,7 +98,7 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
 
         self._state_group_seq_gen = build_sequence_generator(
             db_conn,
-            self.database_engine,
+            self.db_pool.engine,
             get_max_state_group_txn,
             "state_group_id_seq",
             table="state_groups",