From de2a7393a53b485ca6a79cf038273d6f2c109d05 Mon Sep 17 00:00:00 2001
From: anthony sottile <anthony.sottile@sentry.io>
Date: Tue, 29 Jul 2025 15:23:26 -0400
Subject: [PATCH] ref: fix typing for sentry.db.postgres.base

_set_isolation_level was removed in django/django@64a899dc815f1a070dc7a7c22276e8bb41e46ea6 released in django 1.9
---
 pyproject.toml                       |  1 -
 src/sentry/db/postgres/base.py       | 44 +++++++++++++---------------
 src/sentry/db/postgres/decorators.py |  2 +-
 3 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 25c1c772dc2611..43a09775d4ca67 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/src/sentry/db/postgres/base.py b/src/sentry/db/postgres/base.py
index b6207b582c94df..eee2cf2e0e8069 100644
--- a/src/sentry/db/postgres/base.py
+++ b/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,7 +101,7 @@ 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):
@@ -100,21 +109,8 @@ def __init__(self, *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
diff --git a/src/sentry/db/postgres/decorators.py b/src/sentry/db/postgres/decorators.py
index 39e465b0b6d433..aef7a08c2401fb 100644
--- a/src/sentry/db/postgres/decorators.py
+++ b/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)