From f4c5e5864cdc04aa61ad13d6f6ba870df811a881 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Wed, 23 Mar 2022 14:03:24 +0000
Subject: [PATCH] Use psycopg2 type stubs (#12269)

---
 changelog.d/12269.misc              |  1 +
 setup.py                            |  1 +
 synapse/storage/database.py         | 14 +++++++++++---
 synapse/storage/engines/__init__.py |  2 +-
 4 files changed, 14 insertions(+), 4 deletions(-)
 create mode 100644 changelog.d/12269.misc

diff --git a/changelog.d/12269.misc b/changelog.d/12269.misc
new file mode 100644
index 000000000..ed79cbb52
--- /dev/null
+++ b/changelog.d/12269.misc
@@ -0,0 +1 @@
+Use type stubs for `psycopg2`.
diff --git a/setup.py b/setup.py
index 439ed75d7..63da71ad7 100755
--- a/setup.py
+++ b/setup.py
@@ -108,6 +108,7 @@ CONDITIONAL_REQUIREMENTS["mypy"] = [
     "types-jsonschema>=3.2.0",
     "types-opentracing>=2.4.2",
     "types-Pillow>=8.3.4",
+    "types-psycopg2>=2.9.9",
     "types-pyOpenSSL>=20.0.7",
     "types-PyYAML>=5.4.10",
     "types-requests>=2.26.0",
diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 9749f0c06..367709a1a 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -288,7 +288,7 @@ class LoggingTransaction:
         """
 
         if isinstance(self.database_engine, PostgresEngine):
-            from psycopg2.extras import execute_batch  # type: ignore
+            from psycopg2.extras import execute_batch
 
             self._do_execute(lambda *x: execute_batch(self.txn, *x), sql, args)
         else:
@@ -302,10 +302,18 @@ class LoggingTransaction:
         rows (e.g. INSERTs).
         """
         assert isinstance(self.database_engine, PostgresEngine)
-        from psycopg2.extras import execute_values  # type: ignore
+        from psycopg2.extras import execute_values
 
         return self._do_execute(
-            lambda *x: execute_values(self.txn, *x, fetch=fetch), sql, *args
+            # Type ignore: mypy is unhappy because if `x` is a 5-tuple, then there will
+            # be two values for `fetch`: one given positionally, and another given
+            # as a keyword argument. We might be able to fix this by
+            # - propagating the signature of psycopg2.extras.execute_values to this
+            #   function, or
+            # - changing `*args: Any` to `values: T` for some appropriate T.
+            lambda *x: execute_values(self.txn, *x, fetch=fetch),  # type: ignore[misc]
+            sql,
+            *args,
         )
 
     def execute(self, sql: str, *args: Any) -> None:
diff --git a/synapse/storage/engines/__init__.py b/synapse/storage/engines/__init__.py
index 9abc02046..afb7d5054 100644
--- a/synapse/storage/engines/__init__.py
+++ b/synapse/storage/engines/__init__.py
@@ -27,7 +27,7 @@ def create_engine(database_config) -> BaseDatabaseEngine:
 
     if name == "psycopg2":
         # Note that psycopg2cffi-compat provides the psycopg2 module on pypy.
-        import psycopg2  # type: ignore
+        import psycopg2
 
         return PostgresEngine(psycopg2, database_config)