db: Use cursor_factory psycopg2 option.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
(cherry picked from commit b4cf9ad777)
This commit is contained in:
Anders Kaseorg
2022-07-05 13:38:55 -07:00
committed by Tim Abbott
parent 0dbc3917ef
commit 4a618ed973
3 changed files with 8 additions and 40 deletions

View File

@@ -1,17 +1,5 @@
import time import time
from typing import ( from typing import Any, Callable, Dict, Iterable, List, Mapping, Sequence, TypeVar, Union
Any,
Callable,
Dict,
Iterable,
List,
Mapping,
Optional,
Sequence,
TypeVar,
Union,
overload,
)
from psycopg2.extensions import connection, cursor from psycopg2.extensions import connection, cursor
from psycopg2.sql import Composable from psycopg2.sql import Composable
@@ -59,30 +47,6 @@ class TimeTrackingConnection(connection):
self.queries: List[Dict[str, str]] = [] self.queries: List[Dict[str, str]] = []
super().__init__(*args, **kwargs) 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: def reset_queries() -> None:
from django.db import connections from django.db import connections

View File

@@ -9,7 +9,7 @@ from django.template.loaders import app_directories
import zerver.lib.logging_util import zerver.lib.logging_util
from scripts.lib.zulip_tools import get_tornado_ports 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 ( from .config import (
DEPLOY_ROOT, DEPLOY_ROOT,
@@ -287,6 +287,7 @@ DATABASES: Dict[str, Dict[str, Any]] = {
"CONN_MAX_AGE": 600, "CONN_MAX_AGE": 600,
"OPTIONS": { "OPTIONS": {
"connection_factory": TimeTrackingConnection, "connection_factory": TimeTrackingConnection,
"cursor_factory": TimeTrackingCursor,
}, },
} }
} }

View File

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional, Tuple
import ldap import ldap
from django_auth_ldap.config import LDAPSearch 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 zerver.lib.types import SAMLIdPConfigDict
from .config import DEPLOY_ROOT, get_from_file_if_exists from .config import DEPLOY_ROOT, get_from_file_if_exists
@@ -37,7 +37,10 @@ DATABASES["default"] = {
"SCHEMA": "zulip", "SCHEMA": "zulip",
"ENGINE": "django.db.backends.postgresql", "ENGINE": "django.db.backends.postgresql",
"TEST_NAME": "django_zulip_tests", "TEST_NAME": "django_zulip_tests",
"OPTIONS": {"connection_factory": TimeTrackingConnection}, "OPTIONS": {
"connection_factory": TimeTrackingConnection,
"cursor_factory": TimeTrackingCursor,
},
} }