diff --git a/zerver/lib/db.py b/zerver/lib/db.py index ac7dd1a721..c59fb2621e 100644 --- a/zerver/lib/db.py +++ b/zerver/lib/db.py @@ -1,17 +1,5 @@ import time -from typing import ( - Any, - Callable, - Dict, - Iterable, - List, - Mapping, - Optional, - Sequence, - TypeVar, - Union, - overload, -) +from typing import Any, Callable, Dict, Iterable, List, Mapping, Sequence, TypeVar, Union from psycopg2.extensions import connection, cursor from psycopg2.sql import Composable @@ -59,30 +47,6 @@ class TimeTrackingConnection(connection): self.queries: List[Dict[str, str]] = [] super().__init__(*args, **kwargs) - @overload - def cursor( - self, - name: str = ..., - *, - scrollable: Optional[bool] = ..., - withhold: bool = ..., - ) -> TimeTrackingCursor: - ... - - @overload - def cursor( - self, - name: str = ..., - cursor_factory: Callable[..., CursorT] = ..., - scrollable: Optional[bool] = ..., - withhold: bool = ..., - ) -> CursorT: - ... - - def cursor(self, *args: object, **kwargs: object) -> cursor: - kwargs.setdefault("cursor_factory", TimeTrackingCursor) - return super().cursor(*args, **kwargs) - def reset_queries() -> None: from django.db import connections diff --git a/zproject/computed_settings.py b/zproject/computed_settings.py index c5c89b6cd3..77edbf843e 100644 --- a/zproject/computed_settings.py +++ b/zproject/computed_settings.py @@ -9,7 +9,7 @@ from django.template.loaders import app_directories import zerver.lib.logging_util from scripts.lib.zulip_tools import get_tornado_ports -from zerver.lib.db import TimeTrackingConnection +from zerver.lib.db import TimeTrackingConnection, TimeTrackingCursor from .config import ( DEPLOY_ROOT, @@ -287,6 +287,7 @@ DATABASES: Dict[str, Dict[str, Any]] = { "CONN_MAX_AGE": 600, "OPTIONS": { "connection_factory": TimeTrackingConnection, + "cursor_factory": TimeTrackingCursor, }, } } diff --git a/zproject/test_extra_settings.py b/zproject/test_extra_settings.py index 9452d51e00..23b5964eec 100644 --- a/zproject/test_extra_settings.py +++ b/zproject/test_extra_settings.py @@ -4,7 +4,7 @@ from typing import Dict, List, Optional, Tuple import ldap from django_auth_ldap.config import LDAPSearch -from zerver.lib.db import TimeTrackingConnection +from zerver.lib.db import TimeTrackingConnection, TimeTrackingCursor from zerver.lib.types import SAMLIdPConfigDict from .config import DEPLOY_ROOT, get_from_file_if_exists @@ -37,7 +37,10 @@ DATABASES["default"] = { "SCHEMA": "zulip", "ENGINE": "django.db.backends.postgresql", "TEST_NAME": "django_zulip_tests", - "OPTIONS": {"connection_factory": TimeTrackingConnection}, + "OPTIONS": { + "connection_factory": TimeTrackingConnection, + "cursor_factory": TimeTrackingCursor, + }, }