ref: fix typing for sentry.db.postgres.base

pyproject.toml

@@ -112,7 +112,6 @@ module = [
     "sentry.api.endpoints.organization_events_spans_performance",
     "sentry.api.endpoints.organization_releases",
     "sentry.api.paginator",
-    "sentry.db.postgres.base",
     "sentry.eventstore.models",
     "sentry.identity.gitlab.provider",
     "sentry.identity.oauth2",

src/sentry/db/postgres/base.py

@@ -1,7 +1,7 @@
-import psycopg2 as Database
+from types import TracebackType
+from typing import Self
 
-# Some of these imports are unused, but they are inherited from other engines
-# and should be available as part of the backend ``base.py`` namespace.
+import psycopg2
 from django.db.backends.postgresql.base import DatabaseWrapper as DjangoDatabaseWrapper
 from django.db.backends.postgresql.operations import DatabaseOperations
 
@@ -61,10 +61,7 @@ def clean_bad_params(params):
 
 
 class CursorWrapper:
-    """
-    A wrapper around the postgresql_psycopg2 backend which handles various events
-    from cursors, such as auto reconnects and lazy time zone evaluation.
-    """
+    """A wrapper around the postgresql_psycopg2 backend which handles auto reconnects"""
 
     def __init__(self, db, cursor):
         self.db = db
@@ -76,6 +73,18 @@ def __getattr__(self, attr):
     def __iter__(self):
         return iter(self.cursor)
 
+    def __enter__(self) -> Self:
+        self.cursor.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        type: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        return self.cursor.__exit__(type, value, traceback)
+
     @capture_transaction_exceptions
     @auto_reconnect_cursor
     @more_better_error_messages
@@ -92,29 +101,16 @@ def executemany(self, sql, paramlist=()):
 
 
 class DatabaseWrapper(DjangoDatabaseWrapper):
-    SchemaEditorClass = DatabaseSchemaEditorProxy
+    SchemaEditorClass = DatabaseSchemaEditorProxy  # type: ignore[assignment]  # a proxy class isn't exactly the original type
     queries_limit = 15000
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.ops = DatabaseOperations(self)
 
     @auto_reconnect_connection
-    def _set_isolation_level(self, level):
-        return super()._set_isolation_level(level)
-
-    @auto_reconnect_connection
-    def _cursor(self, *args, **kwargs):
-        return super()._cursor()
-
-    # We're overriding this internal method that's present in Django 1.11+, because
-    # things were shuffled around since 1.10 resulting in not constructing a django CursorWrapper
-    # with our CursorWrapper. We need to be passing our wrapped cursor to their wrapped cursor,
-    # not the other way around since then we'll lose things like __enter__ due to the way this
-    # wrapper is working (getattr on self.cursor).
-    def _prepare_cursor(self, cursor):
-        cursor = super()._prepare_cursor(CursorWrapper(self, cursor))
-        return cursor
+    def cursor(self) -> CursorWrapper:
+        return CursorWrapper(self, super().cursor())
 
     def close(self, reconnect=False):
         """
@@ -124,7 +120,7 @@ def close(self, reconnect=False):
             if not self.connection.closed:
                 try:
                     self.connection.close()
-                except Database.InterfaceError:
+                except psycopg2.InterfaceError:
                     # connection was already closed by something
                     # like pgbouncer idle timeout.
                     pass

src/sentry/db/postgres/decorators.py

@@ -21,7 +21,7 @@ def inner(self, *args, **kwargs):
                 raise
 
             self.db.close(reconnect=True)
-            self.cursor = self.db._cursor()
+            self.cursor = self.db.cursor()
 
             return func(self, *args, **kwargs)