use in_test_hide_transaction_boundary

Author: kcons

src/sentry/incidents/endpoints/organization_alert_rule_index.py

@@ -1,10 +1,9 @@
 import logging
 from collections.abc import Sequence
-from contextlib import contextmanager
 from copy import deepcopy
 from datetime import UTC, datetime
 
-from django.db import connections, router
+from django.db import connections, router, transaction
 from django.db.models import (
     Case,
     DateTimeField,
@@ -45,6 +44,7 @@
 from sentry.apidocs.utils import inline_sentry_response_serializer
 from sentry.constants import ObjectStatus
 from sentry.db.models.manager.base_query_set import BaseQuerySet
+from sentry.db.postgres.transactions import in_test_hide_transaction_boundary
 from sentry.exceptions import InvalidParams
 from sentry.incidents.endpoints.bases import OrganizationAlertRuleBaseEndpoint
 from sentry.incidents.endpoints.serializers.alert_rule import (
@@ -115,25 +115,6 @@
 logger = logging.getLogger(__name__)
 
 
-@contextmanager
-def _postgres_jit_disabled(using: str):
-    """
-    Disable PostgreSQL JIT compilation on the given database connection for the duration of
-    the block.
-
-    Uses session-level SET/RESET rather than SET LOCAL (which requires a transaction) so that
-    callers don't need to open a transaction — important when the calling code may issue queries
-    against multiple databases.
-    """
-    with connections[using].cursor() as cursor:
-        cursor.execute("SET jit = off")
-    try:
-        yield
-    finally:
-        with connections[using].cursor() as cursor:
-            cursor.execute("RESET jit")
-
-
 # Sentinel values for incident_status annotation when sorting combined rules
 # Used to ensure proper sort order for rules without active incidents
 INCIDENT_STATUS_NONE = -1  # Metric alerts with no active incident
@@ -549,11 +530,14 @@ def has_type(rule_type: str) -> bool:
         # Disabling it makes this endpoint considerably faster.
         # The risk of other regression here should be low; our API endpoint isn't generally doing the sort of bulk
         # work that benefits from JIT.
-        #
-        # Note: Monitor lives on a different database, so we can't use a transaction here
-        # (SET LOCAL would require one). Session-level SET + RESET is fine for a request-scoped
-        # connection.
-        with _postgres_jit_disabled(using=router.db_for_write(Detector)):
+        # in_test_hide_transaction_boundary is safe here: this transaction is only
+        # used to scope SET LOCAL, not to guard data mutations. No writes happen
+        # inside this block, so there's no cross-db atomicity concern to enforce.
+        db = router.db_for_write(Detector)
+        with in_test_hide_transaction_boundary(), transaction.atomic(using=db):
+            with connections[db].cursor() as cursor:
+                cursor.execute("SET LOCAL jit = off")
+
             intermediaries: list[CombinedQuerysetIntermediary] = []
             if has_type("alert_rule"):
                 intermediaries.append(CombinedQuerysetIntermediary(metric_detectors, sort_key))