Annotate synapse.storage.util (#10892)

Also mark `synapse.streams` as having has no untyped defs

Co-authored-by: Sean Quah <8349537+squahtx@users.noreply.github.com>

changelog.d/10892.misc

@@ -0,0 +1 @@
+Add further type hints to `synapse.storage.util`.

mypy.ini

@@ -105,6 +105,12 @@ disallow_untyped_defs = True
 [mypy-synapse.state.*]
 disallow_untyped_defs = True
 
+[mypy-synapse.storage.util.*]
+disallow_untyped_defs = True
+
+[mypy-synapse.streams.*]
+disallow_untyped_defs = True
+
 [mypy-synapse.util.batching_queue]
 disallow_untyped_defs = True
 

synapse/replication/slave/storage/_slaved_id_tracker.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 from typing import List, Optional, Tuple
 
-from synapse.storage.types import Connection
+from synapse.storage.database import LoggingDatabaseConnection
 from synapse.storage.util.id_generators import _load_current_id
 
 
 class SlavedIdTracker:
     def __init__(
         self,
-        db_conn: Connection,
+        db_conn: LoggingDatabaseConnection,
         table: str,
         column: str,
         extra_tables: Optional[List[Tuple[str, str]]] = None,

synapse/replication/slave/storage/pushers.py

@@ -15,9 +15,8 @@
 from typing import TYPE_CHECKING
 
 from synapse.replication.tcp.streams import PushersStream
-from synapse.storage.database import DatabasePool
+from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
 from synapse.storage.databases.main.pusher import PusherWorkerStore
-from synapse.storage.types import Connection
 
 from ._base import BaseSlavedStore
 from ._slaved_id_tracker import SlavedIdTracker
@@ -27,7 +26,12 @@
 
 
 class SlavedPusherStore(PusherWorkerStore, BaseSlavedStore):
-    def __init__(self, database: DatabasePool, db_conn: Connection, hs: "HomeServer"):
+    def __init__(
+        self,
+        database: DatabasePool,
+        db_conn: LoggingDatabaseConnection,
+        hs: "HomeServer",
+    ):
         super().__init__(database, db_conn, hs)
         self._pushers_id_gen = SlavedIdTracker(  # type: ignore
             db_conn, "pushers", "id", extra_tables=[("deleted_pushers", "stream_id")]

synapse/storage/databases/main/pusher.py

@@ -18,8 +18,7 @@
 
 from synapse.push import PusherConfig, ThrottleParams
 from synapse.storage._base import SQLBaseStore, db_to_json
-from synapse.storage.database import DatabasePool
-from synapse.storage.types import Connection
+from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
 from synapse.storage.util.id_generators import StreamIdGenerator
 from synapse.types import JsonDict
 from synapse.util import json_encoder
@@ -32,7 +31,12 @@
 
 
 class PusherWorkerStore(SQLBaseStore):
-    def __init__(self, database: DatabasePool, db_conn: Connection, hs: "HomeServer"):
+    def __init__(
+        self,
+        database: DatabasePool,
+        db_conn: LoggingDatabaseConnection,
+        hs: "HomeServer",
+    ):
         super().__init__(database, db_conn, hs)
         self._pushers_id_gen = StreamIdGenerator(
             db_conn, "pushers", "id", extra_tables=[("deleted_pushers", "stream_id")]

synapse/storage/databases/main/registration.py

@@ -26,7 +26,7 @@
 from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
 from synapse.storage.databases.main.cache import CacheInvalidationWorkerStore
 from synapse.storage.databases.main.stats import StatsStore
-from synapse.storage.types import Connection, Cursor
+from synapse.storage.types import Cursor
 from synapse.storage.util.id_generators import IdGenerator
 from synapse.storage.util.sequence import build_sequence_generator
 from synapse.types import UserID, UserInfo
@@ -1775,7 +1775,12 @@ async def is_guest(self, user_id: str) -> bool:
 
 
 class RegistrationStore(StatsStore, RegistrationBackgroundUpdateStore):
-    def __init__(self, database: DatabasePool, db_conn: Connection, hs: "HomeServer"):
+    def __init__(
+        self,
+        database: DatabasePool,
+        db_conn: LoggingDatabaseConnection,
+        hs: "HomeServer",
+    ):
         super().__init__(database, db_conn, hs)
 
         self._ignore_unknown_session_error = (