From 3d060eae6c836f4153a3150c6970cb9b10516da6 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <clokep@users.noreply.github.com>
Date: Thu, 9 Mar 2023 07:10:09 -0500
Subject: [PATCH] Add missing type hints to `synapse.storage.database`.
 (#15230)

---
 changelog.d/15230.misc      |  1 +
 mypy.ini                    |  3 ---
 synapse/storage/database.py | 21 ++++++++++++++++-----
 3 files changed, 17 insertions(+), 8 deletions(-)
 create mode 100644 changelog.d/15230.misc

diff --git a/changelog.d/15230.misc b/changelog.d/15230.misc
new file mode 100644
index 000000000..93ceaeafc
--- /dev/null
+++ b/changelog.d/15230.misc
@@ -0,0 +1 @@
+Improve type hints.
diff --git a/mypy.ini b/mypy.ini
index cad371638..945f7925c 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -48,9 +48,6 @@ warn_unused_ignores = False
 [mypy-synapse.util.caches.treecache]
 disallow_untyped_defs = False
 
-[mypy-synapse.storage.database]
-disallow_untyped_defs = False
-
 [mypy-tests.util.caches.test_descriptors]
 disallow_untyped_defs = False
 
diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 5efe31aa1..fec4ae5b9 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -34,6 +34,7 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
+    Union,
     cast,
     overload,
 )
@@ -100,6 +101,15 @@ UNIQUE_INDEX_BACKGROUND_UPDATES = {
 }
 
 
+class _PoolConnection(Connection):
+    """
+    A Connection from twisted.enterprise.adbapi.Connection.
+    """
+
+    def reconnect(self) -> None:
+        ...
+
+
 def make_pool(
     reactor: IReactorCore,
     db_config: DatabaseConnectionConfig,
@@ -856,7 +866,8 @@ class DatabasePool:
             try:
                 with opentracing.start_active_span(f"db.{desc}"):
                     result = await self.runWithConnection(
-                        self.new_transaction,
+                        # mypy seems to have an issue with this, maybe a bug?
+                        self.new_transaction,  # type: ignore[arg-type]
                         desc,
                         after_callbacks,
                         async_after_callbacks,
@@ -892,7 +903,7 @@ class DatabasePool:
 
     async def runWithConnection(
         self,
-        func: Callable[..., R],
+        func: Callable[Concatenate[LoggingDatabaseConnection, P], R],
         *args: Any,
         db_autocommit: bool = False,
         isolation_level: Optional[int] = None,
@@ -926,7 +937,7 @@ class DatabasePool:
 
         start_time = monotonic_time()
 
-        def inner_func(conn, *args, **kwargs):
+        def inner_func(conn: _PoolConnection, *args: P.args, **kwargs: P.kwargs) -> R:
             # We shouldn't be in a transaction. If we are then something
             # somewhere hasn't committed after doing work. (This is likely only
             # possible during startup, as `run*` will ensure changes are
@@ -1019,7 +1030,7 @@ class DatabasePool:
         decoder: Optional[Callable[[Cursor], R]],
         query: str,
         *args: Any,
-    ) -> R:
+    ) -> Union[List[Tuple[Any, ...]], R]:
         """Runs a single query for a result set.
 
         Args:
@@ -1032,7 +1043,7 @@ class DatabasePool:
             The result of decoder(results)
         """
 
-        def interaction(txn):
+        def interaction(txn: LoggingTransaction) -> Union[List[Tuple[Any, ...]], R]:
             txn.execute(query, args)
             if decoder:
                 return decoder(txn)