Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
Use psycopg2 type stubs (#12269)
Browse files Browse the repository at this point in the history
  • Loading branch information
David Robertson authored Mar 23, 2022
1 parent c577678 commit f4c5e58
Show file tree
Hide file tree
Showing 4 changed files with 14 additions and 4 deletions.
1 change: 1 addition & 0 deletions changelog.d/12269.misc
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Use type stubs for `psycopg2`.
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ def exec_file(path_segments):
"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",
Expand Down
14 changes: 11 additions & 3 deletions synapse/storage/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ def execute_batch(self, sql: str, args: Iterable[Iterable[Any]]) -> None:
"""

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:
Expand All @@ -302,10 +302,18 @@ def execute_values(self, sql: str, *args: Any, fetch: bool = True) -> List[Tuple
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:
Expand Down
2 changes: 1 addition & 1 deletion synapse/storage/engines/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit f4c5e58

Please sign in to comment.