Move base store

47 files changed, 162 insertions, 121 deletions

diff --git a/mypy.ini b/mypy.ini
index 64ed45dac2..45007932fc 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -45,6 +45,7 @@ files =
   synapse/storage/__init__.py,
   synapse/storage/_base.py,
   synapse/storage/background_updates.py,
+  synapse/storage/databases/main/_base.py,
   synapse/storage/databases/main/appservice.py,
   synapse/storage/databases/main/events.py,
   synapse/storage/databases/main/keys.py,
diff --git a/scripts-dev/hash_history.py b/scripts-dev/hash_history.py
index 8d6c3d24db..3149258461 100644
--- a/scripts-dev/hash_history.py
+++ b/scripts-dev/hash_history.py
@@ -8,7 +8,7 @@ from synapse.crypto.event_signing import (
     compute_pdu_event_reference_hash,
 )
 from synapse.federation.units import Pdu
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.pdu import PduStore
 from synapse.storage.signatures import SignatureStore
 
diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py
index a25c4093bc..ad33a39bcc 100644
--- a/synapse/storage/_base.py
+++ b/synapse/storage/_base.py
@@ -15,90 +15,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-import random
-from abc import ABCMeta
-from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
+from typing import Any, Union
 
 from synapse.storage.database import LoggingTransaction  # noqa: F401
 from synapse.storage.database import make_in_list_sql_clause  # noqa: F401
-from synapse.storage.database import DatabasePool
-from synapse.storage.types import Connection
-from synapse.types import Collection, StreamToken, get_domain_from_id
 from synapse.util import json_decoder
 
-if TYPE_CHECKING:
-    from synapse.app.homeserver import HomeServer
-
 logger = logging.getLogger(__name__)
 
 
-# some of our subclasses have abstract methods, so we use the ABCMeta metaclass.
-class SQLBaseStore(metaclass=ABCMeta):
-    """Base class for data stores that holds helper functions.
-
-    Note that multiple instances of this class will exist as there will be one
-    per data store (and not one per physical database).
-    """
-
-    def __init__(self, database: DatabasePool, db_conn: Connection, hs: "HomeServer"):
-        self.hs = hs
-        self._clock = hs.get_clock()
-        self.database_engine = database.engine
-        self.db_pool = database
-        self.rand = random.SystemRandom()
-
-    def process_replication_rows(
-        self,
-        stream_name: str,
-        instance_name: str,
-        token: StreamToken,
-        rows: Iterable[Any],
-    ) -> None:
-        pass
-
-    def _invalidate_state_caches(
-        self, room_id: str, members_changed: Iterable[str]
-    ) -> None:
-        """Invalidates caches that are based on the current state, but does
-        not stream invalidations down replication.
-
-        Args:
-            room_id: Room where state changed
-            members_changed: The user_ids of members that have changed
-        """
-        for host in {get_domain_from_id(u) for u in members_changed}:
-            self._attempt_to_invalidate_cache("is_host_joined", (room_id, host))
-
-        self._attempt_to_invalidate_cache("get_users_in_room", (room_id,))
-        self._attempt_to_invalidate_cache("get_room_summary", (room_id,))
-        self._attempt_to_invalidate_cache("get_current_state_ids", (room_id,))
-
-    def _attempt_to_invalidate_cache(
-        self, cache_name: str, key: Optional[Collection[Any]]
-    ) -> None:
-        """Attempts to invalidate the cache of the given name, ignoring if the
-        cache doesn't exist. Mainly used for invalidating caches on workers,
-        where they may not have the cache.
-
-        Args:
-            cache_name
-            key: Entry to invalidate. If None then invalidates the entire
-                cache.
-        """
-
-        try:
-            cache = getattr(self, cache_name)
-        except AttributeError:
-            # We probably haven't pulled in the cache in this worker,
-            # which is fine.
-            return
-
-        if key is None:
-            cache.invalidate_all()
-        else:
-            cache.invalidate(tuple(key))
-
-
 def db_to_json(db_content: Union[memoryview, bytes, bytearray, str]) -> Any:
     """
     Take some data from a database row and return a JSON-decoded object.
@@ -122,5 +47,5 @@ def db_to_json(db_content: Union[memoryview, bytes, bytearray, str]) -> Any:
     try:
         return json_decoder.decode(db_content)
     except Exception:
-        logging.warning("Tried to decode '%r' as JSON and failed", db_content)
+        logger.warning("Tried to decode '%r' as JSON and failed", db_content)
         raise
diff --git a/synapse/storage/databases/main/_base.py b/synapse/storage/databases/main/_base.py
new file mode 100644
index 0000000000..0cb2f446b8
--- /dev/null
+++ b/synapse/storage/databases/main/_base.py
@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+# Copyright 2014-2016 OpenMarket Ltd
+# Copyright 2017-2018 New Vector Ltd
+# Copyright 2019 The Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import random
+from abc import ABCMeta
+from typing import TYPE_CHECKING, Any, Iterable, Optional
+
+from synapse.storage.database import DatabasePool
+from synapse.storage.types import Connection
+from synapse.types import Collection, StreamToken, get_domain_from_id
+
+if TYPE_CHECKING:
+    from synapse.app.homeserver import HomeServer
+
+logger = logging.getLogger(__name__)
+
+
+# some of our subclasses have abstract methods, so we use the ABCMeta metaclass.
+class SQLBaseStore(metaclass=ABCMeta):
+    """Base class for data stores that holds helper functions.
+
+    Note that multiple instances of this class will exist as there will be one
+    per data store (and not one per physical database).
+    """
+
+    def __init__(self, database: DatabasePool, db_conn: Connection, hs: "HomeServer"):
+        self.hs = hs
+        self._clock = hs.get_clock()
+        self.database_engine = database.engine
+        self.db_pool = database
+        self.rand = random.SystemRandom()
+
+    def process_replication_rows(
+        self,
+        stream_name: str,
+        instance_name: str,
+        token: StreamToken,
+        rows: Iterable[Any],
+    ) -> None:
+        pass
+
+    def _invalidate_state_caches(
+        self, room_id: str, members_changed: Iterable[str]
+    ) -> None:
+        """Invalidates caches that are based on the current state, but does
+        not stream invalidations down replication.
+
+        Args:
+            room_id: Room where state changed
+            members_changed: The user_ids of members that have changed
+        """
+        for host in {get_domain_from_id(u) for u in members_changed}:
+            self._attempt_to_invalidate_cache("is_host_joined", (room_id, host))
+
+        self._attempt_to_invalidate_cache("get_users_in_room", (room_id,))
+        self._attempt_to_invalidate_cache("get_room_summary", (room_id,))
+        self._attempt_to_invalidate_cache("get_current_state_ids", (room_id,))
+
+    def _attempt_to_invalidate_cache(
+        self, cache_name: str, key: Optional[Collection[Any]]
+    ) -> None:
+        """Attempts to invalidate the cache of the given name, ignoring if the
+        cache doesn't exist. Mainly used for invalidating caches on workers,
+        where they may not have the cache.
+
+        Args:
+            cache_name
+            key: Entry to invalidate. If None then invalidates the entire
+                cache.
+        """
+
+        try:
+            cache = getattr(self, cache_name)
+        except AttributeError:
+            # We probably haven't pulled in the cache in this worker,
+            # which is fine.
+            return
+
+        if key is None:
+            cache.invalidate_all()
+        else:
+            cache.invalidate(tuple(key))
diff --git a/synapse/storage/databases/main/account_data.py b/synapse/storage/databases/main/account_data.py
index a277a1ef13..b1496674a5 100644
--- a/synapse/storage/databases/main/account_data.py
+++ b/synapse/storage/databases/main/account_data.py
@@ -20,8 +20,9 @@ from typing import Dict, List, Optional, Set, Tuple
 from synapse.api.constants import AccountDataTypes
 from synapse.replication.slave.storage._slaved_id_tracker import SlavedIdTracker
 from synapse.replication.tcp.streams import AccountDataStream, TagAccountDataStream
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator, StreamIdGenerator
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/appservice.py b/synapse/storage/databases/main/appservice.py
index 03a38422a1..2fbdec247e 100644
--- a/synapse/storage/databases/main/appservice.py
+++ b/synapse/storage/databases/main/appservice.py
@@ -24,8 +24,9 @@ from synapse.appservice import (
 )
 from synapse.config.appservice import load_appservices
 from synapse.events import EventBase
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.types import Connection
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/cache.py b/synapse/storage/databases/main/cache.py
index 1e7637a6f5..8bd0204d5e 100644
--- a/synapse/storage/databases/main/cache.py
+++ b/synapse/storage/databases/main/cache.py
@@ -25,8 +25,8 @@ from synapse.replication.tcp.streams.events import (
     EventsStreamCurrentStateRow,
     EventsStreamEventRow,
 )
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.util.iterutils import batch_iter
 
diff --git a/synapse/storage/databases/main/censor_events.py b/synapse/storage/databases/main/censor_events.py
index 3e26d5ba87..e1bf55edc2 100644
--- a/synapse/storage/databases/main/censor_events.py
+++ b/synapse/storage/databases/main/censor_events.py
@@ -18,8 +18,8 @@ from typing import TYPE_CHECKING
 
 from synapse.events.utils import prune_event_dict
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.cache import CacheInvalidationWorkerStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.util import json_encoder
diff --git a/synapse/storage/databases/main/client_ips.py b/synapse/storage/databases/main/client_ips.py
index 6d18e692b0..f7b4b58e96 100644
--- a/synapse/storage/databases/main/client_ips.py
+++ b/synapse/storage/databases/main/client_ips.py
@@ -17,8 +17,8 @@ import logging
 from typing import Dict, List, Optional, Tuple, Union
 
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool, make_tuple_comparison_clause
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import UserID
 from synapse.util.caches.lrucache import LruCache
 
diff --git a/synapse/storage/databases/main/deviceinbox.py b/synapse/storage/databases/main/deviceinbox.py
index 45ca6620a8..cf1eb1337c 100644
--- a/synapse/storage/databases/main/deviceinbox.py
+++ b/synapse/storage/databases/main/deviceinbox.py
@@ -18,8 +18,9 @@ from typing import List, Tuple
 
 from synapse.logging.opentracing import log_kv, set_tag, trace
 from synapse.replication.tcp.streams import ToDeviceStream
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator, StreamIdGenerator
 from synapse.util import json_encoder
diff --git a/synapse/storage/databases/main/devices.py b/synapse/storage/databases/main/devices.py
index d327e9aa0b..5f8e17d995 100644
--- a/synapse/storage/databases/main/devices.py
+++ b/synapse/storage/databases/main/devices.py
@@ -26,12 +26,13 @@ from synapse.logging.opentracing import (
     whitelisted_homeserver,
 )
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import (
     DatabasePool,
     LoggingTransaction,
     make_tuple_comparison_clause,
 )
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import Collection, JsonDict, get_verify_key_from_cross_signing_key
 from synapse.util import json_decoder, json_encoder
 from synapse.util.caches.descriptors import cached, cachedList
diff --git a/synapse/storage/databases/main/directory.py b/synapse/storage/databases/main/directory.py
index 267b948397..1dc0bc7300 100644
--- a/synapse/storage/databases/main/directory.py
+++ b/synapse/storage/databases/main/directory.py
@@ -17,7 +17,7 @@ from collections import namedtuple
 from typing import Iterable, List, Optional
 
 from synapse.api.errors import SynapseError
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import RoomAlias
 from synapse.util.caches.descriptors import cached
 
diff --git a/synapse/storage/databases/main/e2e_room_keys.py b/synapse/storage/databases/main/e2e_room_keys.py
index 12cecceec2..c97172adb5 100644
--- a/synapse/storage/databases/main/e2e_room_keys.py
+++ b/synapse/storage/databases/main/e2e_room_keys.py
@@ -18,7 +18,8 @@ from typing import Optional
 
 from synapse.api.errors import StoreError
 from synapse.logging.opentracing import log_kv, trace
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.util import json_encoder
 
 
diff --git a/synapse/storage/databases/main/end_to_end_keys.py b/synapse/storage/databases/main/end_to_end_keys.py
index f1e7859d26..b03ffaf969 100644
--- a/synapse/storage/databases/main/end_to_end_keys.py
+++ b/synapse/storage/databases/main/end_to_end_keys.py
@@ -23,8 +23,9 @@ from canonicaljson import encode_canonical_json
 from twisted.enterprise.adbapi import Connection
 
 from synapse.logging.opentracing import log_kv, set_tag, trace
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool, make_in_list_sql_clause
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.types import Cursor
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/event_federation.py b/synapse/storage/databases/main/event_federation.py
index 18ddb92fcc..cade7984ff 100644
--- a/synapse/storage/databases/main/event_federation.py
+++ b/synapse/storage/databases/main/event_federation.py
@@ -20,8 +20,9 @@ from typing import Dict, Iterable, List, Set, Tuple
 from synapse.api.errors import StoreError
 from synapse.events import EventBase
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore, make_in_list_sql_clause
+from synapse.storage._base import make_in_list_sql_clause
 from synapse.storage.database import DatabasePool, LoggingTransaction
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.databases.main.signatures import SignatureWorkerStore
 from synapse.storage.engines import PostgresEngine
diff --git a/synapse/storage/databases/main/event_push_actions.py b/synapse/storage/databases/main/event_push_actions.py
index 78245ad5bd..3861537497 100644
--- a/synapse/storage/databases/main/event_push_actions.py
+++ b/synapse/storage/databases/main/event_push_actions.py
@@ -19,8 +19,9 @@ from typing import Dict, List, Optional, Tuple, Union
 import attr
 
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool, LoggingTransaction
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.util import json_encoder
 from synapse.util.caches.descriptors import cached
 
diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index c1626ccf28..fa1c068730 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -21,8 +21,9 @@ import attr
 from synapse.api.constants import EventContentFields
 from synapse.api.room_versions import KNOWN_ROOM_VERSIONS
 from synapse.events import make_event_from_dict
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import DatabasePool, make_tuple_comparison_clause
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events import PersistEventsStore
 from synapse.storage.types import Cursor
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/events_forward_extremities.py b/synapse/storage/databases/main/events_forward_extremities.py
index b3703ae161..ac013f8ddf 100644
--- a/synapse/storage/databases/main/events_forward_extremities.py
+++ b/synapse/storage/databases/main/events_forward_extremities.py
@@ -17,7 +17,7 @@ import logging
 from typing import Dict, List
 
 from synapse.api.errors import SynapseError
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 
 logger = logging.getLogger(__name__)
 
diff --git a/synapse/storage/databases/main/events_worker.py b/synapse/storage/databases/main/events_worker.py
index edbe42f2bf..ec4c4f617a 100644
--- a/synapse/storage/databases/main/events_worker.py
+++ b/synapse/storage/databases/main/events_worker.py
@@ -41,8 +41,9 @@ from synapse.metrics.background_process_metrics import (
 from synapse.replication.slave.storage._slaved_id_tracker import SlavedIdTracker
 from synapse.replication.tcp.streams import BackfillStream
 from synapse.replication.tcp.streams.events import EventsStream
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator, StreamIdGenerator
 from synapse.storage.util.sequence import build_sequence_generator
diff --git a/synapse/storage/databases/main/filtering.py b/synapse/storage/databases/main/filtering.py
index d2f5b9a502..e0af3b0a25 100644
--- a/synapse/storage/databases/main/filtering.py
+++ b/synapse/storage/databases/main/filtering.py
@@ -16,7 +16,8 @@
 from canonicaljson import encode_canonical_json
 
 from synapse.api.errors import Codes, SynapseError
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import JsonDict
 from synapse.util.caches.descriptors import cached
 
diff --git a/synapse/storage/databases/main/group_server.py b/synapse/storage/databases/main/group_server.py
index ac07e0197b..23366d5af7 100644
--- a/synapse/storage/databases/main/group_server.py
+++ b/synapse/storage/databases/main/group_server.py
@@ -19,7 +19,8 @@ from typing import Any, Dict, List, Optional, Tuple
 from typing_extensions import TypedDict
 
 from synapse.api.errors import SynapseError
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import JsonDict
 from synapse.util import json_encoder
 
diff --git a/synapse/storage/databases/main/keys.py b/synapse/storage/databases/main/keys.py
index d504323b03..56890ed2c2 100644
--- a/synapse/storage/databases/main/keys.py
+++ b/synapse/storage/databases/main/keys.py
@@ -20,7 +20,7 @@ from typing import Dict, Iterable, List, Optional, Tuple
 
 from signedjson.key import decode_verify_key_bytes
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.keys import FetchKeyResult
 from synapse.storage.types import Cursor
 from synapse.util.caches.descriptors import cached, cachedList
diff --git a/synapse/storage/databases/main/media_repository.py b/synapse/storage/databases/main/media_repository.py
index 274f8de595..47ed6f8494 100644
--- a/synapse/storage/databases/main/media_repository.py
+++ b/synapse/storage/databases/main/media_repository.py
@@ -16,8 +16,8 @@
 from enum import Enum
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 
 BG_UPDATE_REMOVE_MEDIA_REPO_INDEX_WITHOUT_METHOD = (
     "media_repository_drop_index_wo_method"
diff --git a/synapse/storage/databases/main/metrics.py b/synapse/storage/databases/main/metrics.py
index 614a418a15..d813b7157a 100644
--- a/synapse/storage/databases/main/metrics.py
+++ b/synapse/storage/databases/main/metrics.py
@@ -19,8 +19,8 @@ from typing import Dict
 
 from synapse.metrics import GaugeBucketCollector
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.event_push_actions import (
     EventPushActionsWorkerStore,
 )
diff --git a/synapse/storage/databases/main/monthly_active_users.py b/synapse/storage/databases/main/monthly_active_users.py
index d788dc0fc6..929183b020 100644
--- a/synapse/storage/databases/main/monthly_active_users.py
+++ b/synapse/storage/databases/main/monthly_active_users.py
@@ -16,8 +16,8 @@ import logging
 from typing import Dict, List
 
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool, make_in_list_sql_clause
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.util.caches.descriptors import cached
 
 logger = logging.getLogger(__name__)
diff --git a/synapse/storage/databases/main/openid.py b/synapse/storage/databases/main/openid.py
index 2aac64901b..96dcba1dd6 100644
--- a/synapse/storage/databases/main/openid.py
+++ b/synapse/storage/databases/main/openid.py
@@ -1,6 +1,6 @@
 from typing import Optional
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 
 
 class OpenIdStore(SQLBaseStore):
diff --git a/synapse/storage/databases/main/presence.py b/synapse/storage/databases/main/presence.py
index 29edab34d4..ab9d08054c 100644
--- a/synapse/storage/databases/main/presence.py
+++ b/synapse/storage/databases/main/presence.py
@@ -16,7 +16,8 @@
 from typing import List, Tuple
 
 from synapse.api.presence import UserPresenceState
-from synapse.storage._base import SQLBaseStore, make_in_list_sql_clause
+from synapse.storage._base import make_in_list_sql_clause
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.util.caches.descriptors import cached, cachedList
 from synapse.util.iterutils import batch_iter
 
diff --git a/synapse/storage/databases/main/profile.py b/synapse/storage/databases/main/profile.py
index ba01d3108a..f606b7e77e 100644
--- a/synapse/storage/databases/main/profile.py
+++ b/synapse/storage/databases/main/profile.py
@@ -15,7 +15,7 @@
 from typing import Any, Dict, List, Optional
 
 from synapse.api.errors import StoreError
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.roommember import ProfileInfo
 
 
diff --git a/synapse/storage/databases/main/purge_events.py b/synapse/storage/databases/main/purge_events.py
index ecfc9f20b1..f58b72b2f0 100644
--- a/synapse/storage/databases/main/purge_events.py
+++ b/synapse/storage/databases/main/purge_events.py
@@ -17,7 +17,7 @@ import logging
 from typing import Any, List, Set, Tuple
 
 from synapse.api.errors import SynapseError
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.state import StateGroupWorkerStore
 from synapse.types import RoomStreamToken
 
diff --git a/synapse/storage/databases/main/push_rule.py b/synapse/storage/databases/main/push_rule.py
index 9e58dc0e6a..37f9b3dd7d 100644
--- a/synapse/storage/databases/main/push_rule.py
+++ b/synapse/storage/databases/main/push_rule.py
@@ -20,8 +20,9 @@ from typing import List, Tuple, Union
 from synapse.api.errors import NotFoundError, StoreError
 from synapse.push.baserules import list_with_base_rules
 from synapse.replication.slave.storage._slaved_id_tracker import SlavedIdTracker
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.appservice import ApplicationServiceWorkerStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.databases.main.pusher import PusherWorkerStore
diff --git a/synapse/storage/databases/main/pusher.py b/synapse/storage/databases/main/pusher.py
index 7cb69dd6bd..d1c63af7c5 100644
--- a/synapse/storage/databases/main/pusher.py
+++ b/synapse/storage/databases/main/pusher.py
@@ -18,8 +18,9 @@ import logging
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Tuple
 
 from synapse.push import PusherConfig, ThrottleParams
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.types import Connection
 from synapse.storage.util.id_generators import StreamIdGenerator
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/receipts.py b/synapse/storage/databases/main/receipts.py
index 43c852c96c..0c7b821233 100644
--- a/synapse/storage/databases/main/receipts.py
+++ b/synapse/storage/databases/main/receipts.py
@@ -21,8 +21,9 @@ from twisted.internet import defer
 
 from synapse.replication.slave.storage._slaved_id_tracker import SlavedIdTracker
 from synapse.replication.tcp.streams import ReceiptsStream
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator, StreamIdGenerator
 from synapse.types import JsonDict
diff --git a/synapse/storage/databases/main/rejections.py b/synapse/storage/databases/main/rejections.py
index 1e361aaa9a..17bc19e14d 100644
--- a/synapse/storage/databases/main/rejections.py
+++ b/synapse/storage/databases/main/rejections.py
@@ -16,7 +16,7 @@
 import logging
 from typing import Optional
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 
 logger = logging.getLogger(__name__)
 
diff --git a/synapse/storage/databases/main/relations.py b/synapse/storage/databases/main/relations.py
index 5cd61547f7..de15edaf67 100644
--- a/synapse/storage/databases/main/relations.py
+++ b/synapse/storage/databases/main/relations.py
@@ -20,7 +20,7 @@ import attr
 
 from synapse.api.constants import RelationTypes
 from synapse.events import EventBase
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.stream import generate_pagination_where_clause
 from synapse.storage.relations import (
     AggregationPaginationToken,
diff --git a/synapse/storage/databases/main/room.py b/synapse/storage/databases/main/room.py
index 9cbcd53026..f6ed69e14e 100644
--- a/synapse/storage/databases/main/room.py
+++ b/synapse/storage/databases/main/room.py
@@ -23,8 +23,9 @@ from typing import Any, Dict, List, Optional, Tuple
 from synapse.api.constants import EventTypes
 from synapse.api.errors import StoreError
 from synapse.api.room_versions import RoomVersion, RoomVersions
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool, LoggingTransaction
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.search import SearchStore
 from synapse.types import JsonDict, ThirdPartyInstanceID
 from synapse.util import json_encoder
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index a9216ca9ae..ffb49e40bd 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -24,8 +24,9 @@ from synapse.metrics.background_process_metrics import (
     run_as_background_process,
     wrap_as_background_process,
 )
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.engines import Sqlite3Engine
 from synapse.storage.roommember import (
diff --git a/synapse/storage/databases/main/search.py b/synapse/storage/databases/main/search.py
index f5e7d9ef98..50390466b1 100644
--- a/synapse/storage/databases/main/search.py
+++ b/synapse/storage/databases/main/search.py
@@ -20,8 +20,9 @@ from typing import List, Optional, Set
 
 from synapse.api.errors import SynapseError
 from synapse.events import EventBase
-from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
+from synapse.storage._base import db_to_json, make_in_list_sql_clause
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.storage.engines import PostgresEngine, Sqlite3Engine
 from synapse.types import Collection
diff --git a/synapse/storage/databases/main/signatures.py b/synapse/storage/databases/main/signatures.py
index c8c67953e4..0a66972456 100644
--- a/synapse/storage/databases/main/signatures.py
+++ b/synapse/storage/databases/main/signatures.py
@@ -17,7 +17,7 @@ from typing import Dict, Iterable, List, Tuple
 
 from unpaddedbase64 import encode_base64
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.types import Cursor
 from synapse.util.caches.descriptors import cached, cachedList
 
diff --git a/synapse/storage/databases/main/state.py b/synapse/storage/databases/main/state.py
index a7f371732f..47c4894659 100644
--- a/synapse/storage/databases/main/state.py
+++ b/synapse/storage/databases/main/state.py
@@ -22,8 +22,8 @@ from synapse.api.constants import EventTypes, Membership
 from synapse.api.errors import NotFoundError, UnsupportedRoomVersionError
 from synapse.api.room_versions import KNOWN_ROOM_VERSIONS, RoomVersion
 from synapse.events import EventBase
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.databases.main.roommember import RoomMemberWorkerStore
 from synapse.storage.state import StateFilter
diff --git a/synapse/storage/databases/main/state_deltas.py b/synapse/storage/databases/main/state_deltas.py
index 0dbb501f16..f4956e87a3 100644
--- a/synapse/storage/databases/main/state_deltas.py
+++ b/synapse/storage/databases/main/state_deltas.py
@@ -16,7 +16,7 @@
 import logging
 from typing import Any, Dict, List, Tuple
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 
 logger = logging.getLogger(__name__)
 
diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 91f8abb67d..edbc3014af 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -45,12 +45,12 @@ from twisted.internet import defer
 from synapse.api.filtering import Filter
 from synapse.events import EventBase
 from synapse.logging.context import make_deferred_yieldable, run_in_background
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import (
     DatabasePool,
     LoggingTransaction,
     make_in_list_sql_clause,
 )
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator
diff --git a/synapse/storage/databases/main/transactions.py b/synapse/storage/databases/main/transactions.py
index b921d63d30..8319a3a08b 100644
--- a/synapse/storage/databases/main/transactions.py
+++ b/synapse/storage/databases/main/transactions.py
@@ -20,8 +20,9 @@ from typing import Iterable, List, Optional, Tuple
 from canonicaljson import encode_canonical_json
 
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import DatabasePool, LoggingTransaction
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import PostgresEngine, Sqlite3Engine
 from synapse.types import JsonDict
 from synapse.util.caches.expiringcache import ExpiringCache
diff --git a/synapse/storage/databases/main/ui_auth.py b/synapse/storage/databases/main/ui_auth.py
index 5473ec1485..b84463a9af 100644
--- a/synapse/storage/databases/main/ui_auth.py
+++ b/synapse/storage/databases/main/ui_auth.py
@@ -17,8 +17,9 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import attr
 
 from synapse.api.errors import StoreError
-from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.storage._base import db_to_json
 from synapse.storage.database import LoggingTransaction
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.types import JsonDict
 from synapse.util import json_encoder, stringutils
 
diff --git a/synapse/storage/databases/main/user_erasure_store.py b/synapse/storage/databases/main/user_erasure_store.py
index f9575b1f1f..f2b43893bf 100644
--- a/synapse/storage/databases/main/user_erasure_store.py
+++ b/synapse/storage/databases/main/user_erasure_store.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from synapse.storage._base import SQLBaseStore
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.util.caches.descriptors import cached, cachedList
 
 
diff --git a/synapse/storage/databases/state/bg_updates.py b/synapse/storage/databases/state/bg_updates.py
index 87e14132f5..b170e6d801 100644
--- a/synapse/storage/databases/state/bg_updates.py
+++ b/synapse/storage/databases/state/bg_updates.py
@@ -15,7 +15,6 @@
 
 import logging
 
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
 from synapse.storage.engines import PostgresEngine
 from synapse.storage.state import StateFilter
diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py
index 86b06809b3..d75ffa866c 100644
--- a/synapse/storage/databases/state/store.py
+++ b/synapse/storage/databases/state/store.py
@@ -18,7 +18,6 @@ from collections import namedtuple
 from typing import Dict, Iterable, List, Set, Tuple
 
 from synapse.api.constants import EventTypes
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
 from synapse.storage.databases.state.bg_updates import StateBackgroundUpdateStore
 from synapse.storage.state import StateFilter
diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index eac7e4dcd2..fdb3cacb99 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -20,8 +20,8 @@ from mock import Mock
 
 from twisted.internet import defer
 
-from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import DatabasePool
+from synapse.storage.databases.main._base import SQLBaseStore
 from synapse.storage.engines import create_engine
 
 from tests import unittest