paulmataruso/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2025-11-02 13:03:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

db: Fix types.

Browse Source 
Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg
2022-03-24 17:48:52 -07:00
committed by Tim Abbott
parent e3600900c0
commit 83c90c53df
1 changed files with 9 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
zerver/lib/db.py
View File

@@ -17,18 +17,18 @@ from psycopg2.extensions import connection, cursor
from psycopg2.sql import Composable from psycopg2.sql import Composable
CursorObj = TypeVar("CursorObj", bound=cursor) CursorObj = TypeVar("CursorObj", bound=cursor)
Query = Union[str, Composable] Query = Union[str, bytes, Composable]
Params = Union[Sequence[object], Mapping[str, object]] Params = Union[Sequence[object], Mapping[str, object], None]
ParamsT = TypeVar("ParamsT") ParamsT = TypeVar("ParamsT")
# Similar to the tracking done in Django's CursorDebugWrapper, but done at the # Similar to the tracking done in Django's CursorDebugWrapper, but done at the
# psycopg2 cursor level so it works with SQLAlchemy. # psycopg2 cursor level so it works with SQLAlchemy.
def wrapper_execute( def wrapper_execute(
self: CursorObj, action: Callable[[Query, ParamsT], CursorObj], sql: Query, params: ParamsT self: CursorObj, action: Callable[[Query, ParamsT], None], sql: Query, params: ParamsT
) -> CursorObj: ) -> None:
start = time.time() start = time.time()
try: try:
return action(sql, params) action(sql, params)
finally: finally:
stop = time.time() stop = time.time()
duration = stop - start duration = stop - start
@@ -42,13 +42,11 @@ def wrapper_execute(
class TimeTrackingCursor(cursor): class TimeTrackingCursor(cursor):
"""A psycopg2 cursor class that tracks the time spent executing queries.""" """A psycopg2 cursor class that tracks the time spent executing queries."""
def execute(self, query: Query, vars: Optional[Params] = None) -> "TimeTrackingCursor": def execute(self, query: Query, vars: Params = None) -> None:
return wrapper_execute(self, super().execute, query, vars) wrapper_execute(self, super().execute, query, vars)
def executemany( def executemany(self, query: Query, vars: Iterable[Params]) -> None: # nocoverage
self, query: Query, vars: Iterable[Params] wrapper_execute(self, super().executemany, query, vars)
) -> "TimeTrackingCursor": # nocoverage
return wrapper_execute(self, super().executemany, query, vars)
CursorT = TypeVar("CursorT", bound=cursor) CursorT = TypeVar("CursorT", bound=cursor)