db: Use cursor_factory psycopg2 option.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg 2022-07-05 13:38:55 -07:00 committed by Tim Abbott
parent 55915c26a7
commit b4cf9ad777
3 changed files with 8 additions and 40 deletions

View File

@ -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

View File

@ -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,
@ -289,6 +289,7 @@ DATABASES: Dict[str, Dict[str, Any]] = {
"CONN_MAX_AGE": 600,
"OPTIONS": {
"connection_factory": TimeTrackingConnection,
"cursor_factory": TimeTrackingCursor,
},
}
}

View File

@ -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,
},
}