chore(hybridcloud) Remove role switching and update guard method v2 (#52802)

Mulligan on https://github.com/getsentry/sentry/pull/52673 as using
testutils code in application code blows up docker image generation.

Because of that I've moved unguarded_write to sentry.silo instead.

src/sentry/db/postgres/roles.py

@@ -3,15 +3,11 @@ from __future__ import annotations
 import contextlib
 import os
 import sys
-from collections import defaultdict
-from typing import MutableMapping
 
 from django.db.transaction import get_connection
 
 from sentry.silo.patches.silo_aware_transaction_patch import determine_using_by_silo_mode
 
-_fencing_counters: MutableMapping[str, int] = defaultdict(int)
-
 
 @contextlib.contextmanager
 def in_test_psql_role_override(role_name: str, using: str | None = None):
@@ -26,12 +22,7 @@ def in_test_psql_role_override(role_name: str, using: str | None = None):
 
     using = determine_using_by_silo_mode(using)
 
-    # TODO(mark) Move this closer to other silo code.
-    _fencing_counters[using] += 1
-
     with get_connection(using).cursor() as conn:
-        fence_value = _fencing_counters[using]
-        conn.execute("SELECT %s", [f"start_role_override_{fence_value}"])
         conn.execute("SELECT user")
         (cur,) = conn.fetchone()
         conn.execute("SET ROLE %s", [role_name])
@@ -39,4 +30,3 @@ def in_test_psql_role_override(role_name: str, using: str | None = None):
             yield
         finally:
             conn.execute("SET ROLE %s", [cur])
-            conn.execute("SELECT %s", [f"end_role_override_{fence_value}"])

src/sentry/models/counter.py

@@ -13,8 +13,7 @@ from sentry.db.models import (
     region_silo_only_model,
     sane_repr,
 )
-from sentry.db.postgres.roles import in_test_psql_role_override
-from sentry.silo import SiloMode
+from sentry.silo import SiloMode, unguarded_write
 
 
 @region_silo_only_model
@@ -103,7 +102,7 @@ def create_counter_function(app_config, using, **kwargs):
     if SiloMode.get_current_mode() == SiloMode.CONTROL:
         return
 
-    with in_test_psql_role_override("postgres", using), connections[using].cursor() as cursor:
+    with unguarded_write(using), connections[using].cursor() as cursor:
         cursor.execute(
             """
             create or replace function sentry_increment_project_counter(

src/sentry/models/outbox.py

@@ -26,13 +26,12 @@ from sentry.db.models import (
     region_silo_only_model,
     sane_repr,
 )
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.db.postgres.transactions import (
     django_test_transaction_water_mark,
     in_test_assert_no_transaction,
 )
 from sentry.services.hybrid_cloud import REGION_NAME_LENGTH
-from sentry.silo import SiloMode
+from sentry.silo import SiloMode, unguarded_write
 from sentry.utils import metrics
 
 THE_PAST = datetime.datetime(2016, 8, 1, 0, 0, 0, 0, tzinfo=timezone.utc)
@@ -468,7 +467,6 @@ _outbox_context = OutboxContext()
 
 @contextlib.contextmanager
 def outbox_context(inner: Atomic | None = None, flush: bool | None = None) -> ContextManager[None]:
-
     # If we don't specify our flush, use the outer specified override
     if flush is None:
         flush = _outbox_context.flushing_enabled
@@ -481,7 +479,7 @@ def outbox_context(inner: Atomic | None = None, flush: bool | None = None) -> Co
     original = _outbox_context.flushing_enabled
 
     if inner:
-        with in_test_psql_role_override("postgres", using=inner.using), inner:
+        with unguarded_write(using=inner.using), inner:
             _outbox_context.flushing_enabled = flush
             try:
                 yield

src/sentry/services/hybrid_cloud/organization/impl.py

@@ -7,7 +7,6 @@ from django.dispatch import Signal
 
 from sentry import roles
 from sentry.api.serializers import serialize
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models import (
     Activity,
     ControlOutbox,
@@ -48,6 +47,7 @@ from sentry.services.hybrid_cloud.organization.serial import (
 )
 from sentry.services.hybrid_cloud.user import RpcUser
 from sentry.services.hybrid_cloud.util import flags_to_bits
+from sentry.silo import unguarded_write
 from sentry.types.region import find_regions_for_orgs
 
 
@@ -470,7 +470,7 @@ class DatabaseBackedOrganizationService(OrganizationService):
                     pass
 
     def reset_idp_flags(self, *, organization_id: int) -> None:
-        with in_test_psql_role_override("postgres"):
+        with unguarded_write():
             # Flags are not replicated -- these updates are safe without outbox application.
             OrganizationMember.objects.filter(
                 organization_id=organization_id,
@@ -485,7 +485,7 @@ class DatabaseBackedOrganizationService(OrganizationService):
         # Normally, calling update on a QS for organization member fails because we need to ensure that updates to
         # OrganizationMember objects produces outboxes.  In this case, it is safe to do the update directly because
         # the attribute we are changing never needs to produce an outbox.
-        with in_test_psql_role_override("postgres"):
+        with unguarded_write():
             OrganizationMember.objects.filter(user_id=user.id).update(
                 user_is_active=user.is_active, user_email=user.email
             )

src/sentry/services/hybrid_cloud/organization_mapping/impl.py

@@ -1,6 +1,5 @@
 from typing import List, Optional
 
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models.organizationmapping import OrganizationMapping
 from sentry.services.hybrid_cloud.organization_mapping import (
     OrganizationMappingService,
@@ -8,6 +7,7 @@ from sentry.services.hybrid_cloud.organization_mapping import (
     RpcOrganizationMappingUpdate,
 )
 from sentry.services.hybrid_cloud.organization_mapping.serial import serialize_organization_mapping
+from sentry.silo import unguarded_write
 
 
 class DatabaseBackedOrganizationMappingService(OrganizationMappingService):
@@ -61,7 +61,7 @@ class DatabaseBackedOrganizationMappingService(OrganizationMappingService):
 
     def update(self, organization_id: int, update: RpcOrganizationMappingUpdate) -> None:
         # TODO: REMOVE FROM GETSENTRY!
-        with in_test_psql_role_override("postgres"):
+        with unguarded_write():
             try:
                 OrganizationMapping.objects.get(organization_id=organization_id).update(**update)
             except OrganizationMapping.DoesNotExist:
@@ -70,7 +70,7 @@ class DatabaseBackedOrganizationMappingService(OrganizationMappingService):
     def upsert(
         self, organization_id: int, update: RpcOrganizationMappingUpdate
     ) -> RpcOrganizationMapping:
-        with in_test_psql_role_override("postgres"):
+        with unguarded_write():
             org_mapping, _created = OrganizationMapping.objects.update_or_create(
                 organization_id=organization_id, defaults=update
             )

src/sentry/services/hybrid_cloud/organizationmember_mapping/impl.py

@@ -7,7 +7,6 @@ from typing import Optional
 
 from django.db import IntegrityError, transaction
 
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models import outbox_context
 from sentry.models.organizationmembermapping import OrganizationMemberMapping
 from sentry.models.user import User
@@ -19,6 +18,7 @@ from sentry.services.hybrid_cloud.organizationmember_mapping import (
 from sentry.services.hybrid_cloud.organizationmember_mapping.serial import (
     serialize_org_member_mapping,
 )
+from sentry.silo import unguarded_write
 
 
 class DatabaseBackedOrganizationMemberMappingService(OrganizationMemberMappingService):
@@ -101,5 +101,5 @@ class DatabaseBackedOrganizationMemberMappingService(OrganizationMemberMappingSe
             organizationmember_id=organizationmember_id,
         )
         if org_member_map:
-            with in_test_psql_role_override("postgres"):
+            with unguarded_write():
                 org_member_map.delete()

src/sentry/silo/safety.py

@@ -23,10 +23,13 @@ def unguarded_write(using: str | None = None, *args: Any, **kwargs: Any):
     """
     Used to indicate that the wrapped block is safe to do
     mutations on outbox backed records.
+
     In production this context manager has no effect, but
     in tests it emits 'fencing' queries that are audited at the
     end of each test run by:
+
     sentry.testutils.silo.validate_protected_queries
+
     This code can't be co-located with the auditing logic because
     the testutils module cannot be used in production code.
     """

src/sentry/tasks/check_auth.py

@@ -4,9 +4,9 @@ from datetime import timedelta
 from django.utils import timezone
 
 from sentry.auth.exceptions import IdentityNotValid
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models import AuthIdentity
 from sentry.services.hybrid_cloud.organization import RpcOrganizationMember, organization_service
+from sentry.silo import unguarded_write
 from sentry.silo.base import SiloMode
 from sentry.tasks.base import instrumented_task
 from sentry.utils import metrics
@@ -95,7 +95,7 @@ def check_auth_identity(auth_identity_id, **kwargs):
         is_valid = True
 
     if getattr(om.flags, "sso:linked") != is_linked:
-        with in_test_psql_role_override("postgres"):
+        with unguarded_write():
             # flags are not replicated, so it's ok not to create outboxes here.
             setattr(om.flags, "sso:linked", is_linked)
             setattr(om.flags, "sso:invalid", not is_valid)

src/sentry/tasks/organization_mapping.py

@@ -6,9 +6,9 @@ from django.db import router
 from django.db.models import Count
 from django.utils import timezone
 
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models.organizationmapping import OrganizationMapping
 from sentry.services.hybrid_cloud.organization import organization_service
+from sentry.silo import unguarded_write
 from sentry.silo.base import SiloMode
 from sentry.tasks.base import instrumented_task, retry
 from sentry.utils import metrics
@@ -42,7 +42,7 @@ def _verify_mappings(expiration_threshold_time: datetime) -> None:
         org = organization_service.get_organization_by_id(
             id=mapping.organization_id, slug=mapping.slug
         )
-        with in_test_psql_role_override("postgres", using=router.db_for_write(OrganizationMapping)):
+        with unguarded_write(using=router.db_for_write(OrganizationMapping)):
             if org is None and mapping.date_created <= expiration_threshold_time:
                 mapping.delete()
             elif org is not None:
@@ -65,7 +65,7 @@ def _remove_duplicate_mappings(expiration_threshold_time: datetime) -> None:
         )
         organization_id = dupe["organization_id"]
 
-        with in_test_psql_role_override("postgres", using=router.db_for_write(OrganizationMapping)):
+        with unguarded_write(using=router.db_for_write(OrganizationMapping)):
             # Organization exists in the region silo
             found_org = organization_service.get_organization_by_id(id=organization_id)
             if found_org is None:

src/sentry/testutils/helpers/api_gateway.py

@@ -10,8 +10,8 @@ from rest_framework.response import Response
 
 from sentry.api.base import control_silo_endpoint, region_silo_endpoint
 from sentry.api.bases.organization import OrganizationEndpoint
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models.organizationmapping import OrganizationMapping
+from sentry.silo import unguarded_write
 from sentry.testutils import APITestCase
 from sentry.types.region import Region, RegionCategory, clear_global_regions
 from sentry.utils import json
@@ -136,7 +136,7 @@ class ApiGatewayTestCase(APITestCase):
             adding_headers={"test": "header"},
         )
 
-        with in_test_psql_role_override("postgres", using=router.db_for_write(OrganizationMapping)):
+        with unguarded_write(using=router.db_for_write(OrganizationMapping)):
             OrganizationMapping.objects.get(organization_id=self.organization.id).update(
                 region_name="region1"
             )