diff --git a/zerver/lib/db.py b/zerver/lib/db.py index 5691682fa8..10231aa9f5 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 @@ -60,41 +48,6 @@ class TimeTrackingConnection(connection): self.queries: List[Dict[str, str]] = [] super().__init__(*args, **kwargs) - @overload - def cursor( - self, - name: Union[str, bytes, None] = ..., - *, - withhold: bool = ..., - scrollable: Optional[bool] = ..., - ) -> TimeTrackingCursor: - ... - - @overload - def cursor( - self, - name: Union[str, bytes, None] = ..., - *, - cursor_factory: Callable[..., CursorT] = ..., - withhold: bool = ..., - scrollable: Optional[bool] = ..., - ) -> CursorT: - ... - - @overload - def cursor( - self, - name: Union[str, bytes, None], - cursor_factory: Callable[..., CursorT] = ..., - withhold: bool = ..., - scrollable: Optional[bool] = ..., - ) -> CursorT: - ... - - def cursor(self, *args: Any, **kwargs: Any) -> cursor: - kwargs["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 edc9d0f55a..54e91428b7 100644 --- a/zproject/computed_settings.py +++ b/zproject/computed_settings.py @@ -7,7 +7,7 @@ from typing import Any, Dict, Final, List, Tuple, Union from urllib.parse import urljoin 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, @@ -281,6 +281,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 0a8f9b7e3f..cbea88b7d6 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 zproject.settings_types import OIDCIdPConfigDict, SAMLIdPConfigDict, SCIMConfigDict from .config import DEPLOY_ROOT, get_from_file_if_exists @@ -36,6 +36,7 @@ DATABASES["default"] = { "TEST_NAME": "django_zulip_tests", "OPTIONS": { "connection_factory": TimeTrackingConnection, + "cursor_factory": TimeTrackingCursor, }, }