chore(hybridcloud) Remove role switching and update guard method

src/sentry/db/postgres/roles.py

@@ -3,15 +3,11 @@
 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,17 +22,11 @@ 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])
         try:
             yield
         finally:
             conn.execute("SET ROLE %s", [cur])
-            conn.execute("SELECT %s", [f"end_role_override_{fence_value}"])

src/sentry/models/counter.py

@@ -13,7 +13,6 @@
     region_silo_only_model,
     sane_repr,
 )
-from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.silo import SiloMode
 
 
@@ -94,6 +93,8 @@ def increment_project_counter(project, delta=1, using="default"):
 # this must be idempotent because it seems to execute twice
 # (at least during test runs)
 def create_counter_function(app_config, using, **kwargs):
+    from sentry.testutils.silo import unguarded_write
+
     if app_config and app_config.name != "sentry":
         return
 
@@ -103,7 +104,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,7 +26,6 @@
     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,
@@ -467,6 +466,7 @@ class OutboxContext(threading.local):
 
 @contextlib.contextmanager
 def outbox_context(inner: Atomic | None = None, flush: bool | None = None) -> ContextManager[None]:
+    from sentry.testutils.silo import unguarded_write
 
     # If we don't specify our flush, use the outer specified override
     if flush is None:
@@ -480,7 +480,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 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.user import RpcUser
 from sentry.services.hybrid_cloud.util import flags_to_bits
+from sentry.testutils.silo import unguarded_write
 from sentry.types.region import find_regions_for_orgs
 
 
@@ -467,7 +467,7 @@ def merge_users(self, *, organization_id: int, from_user_id: int, to_user_id: in
                     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,
@@ -482,7 +482,7 @@ def update_region_user(self, *, user: RpcRegionUser, region_name: str) -> None:
         # 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,13 +1,13 @@
 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,
     RpcOrganizationMapping,
     RpcOrganizationMappingUpdate,
 )
 from sentry.services.hybrid_cloud.organization_mapping.serial import serialize_organization_mapping
+from sentry.testutils.silo import unguarded_write
 
 
 class DatabaseBackedOrganizationMappingService(OrganizationMappingService):
@@ -61,7 +61,7 @@ def get_many(self, *, organization_ids: List[int]) -> List[RpcOrganizationMappin
 
     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 @@ def update(self, organization_id: int, update: RpcOrganizationMappingUpdate) ->
     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 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.serial import (
     serialize_org_member_mapping,
 )
+from sentry.testutils.silo import unguarded_write
 
 
 class DatabaseBackedOrganizationMemberMappingService(OrganizationMemberMappingService):
@@ -101,5 +101,5 @@ def delete(
             organizationmember_id=organizationmember_id,
         )
         if org_member_map:
-            with in_test_psql_role_override("postgres"):
+            with unguarded_write():
                 org_member_map.delete()

src/sentry/tasks/check_auth.py

@@ -4,11 +4,11 @@
 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.base import SiloMode
 from sentry.tasks.base import instrumented_task
+from sentry.testutils.silo import unguarded_write
 from sentry.utils import metrics
 
 logger = logging.getLogger("sentry.auth")
@@ -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,11 +6,11 @@
 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.base import SiloMode
 from sentry.tasks.base import instrumented_task, retry
+from sentry.testutils.silo import unguarded_write
 from sentry.utils import metrics
 from sentry.utils.query import RangeQuerySetWrapper
 
@@ -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/asserts.py

@@ -1,7 +1,5 @@
 from typing import Optional
 
-from sentry.models import AuditLogEntry, CommitFileChange
-
 
 def assert_mock_called_once_with_partial(mock, *args, **kwargs):
     """
@@ -16,10 +14,11 @@ def assert_mock_called_once_with_partial(mock, *args, **kwargs):
         assert m_kwargs[kwarg] == kwargs[kwarg], (m_kwargs[kwarg], kwargs[kwarg])
 
 
-commit_file_type_choices = {c[0] for c in CommitFileChange._meta.get_field("type").choices}
+def assert_commit_shape(commit):
+    from sentry.models import CommitFileChange
 
+    commit_file_type_choices = {c[0] for c in CommitFileChange._meta.get_field("type").choices}
 
-def assert_commit_shape(commit):
     assert commit["id"]
     assert commit["repository"]
     assert commit["author_email"]
@@ -40,6 +39,8 @@ def assert_status_code(response, minimum: int, maximum: Optional[int] = None):
 
 
 def org_audit_log_exists(**kwargs):
+    from sentry.models import AuditLogEntry
+
     assert kwargs
     if "organization" in kwargs:
         kwargs["organization_id"] = kwargs.pop("organization").id
@@ -55,4 +56,6 @@ def assert_org_audit_log_does_not_exist(**kwargs):
 
 
 def delete_all_org_audit_logs():
+    from sentry.models import AuditLogEntry
+
     return AuditLogEntry.objects.all().delete()

src/sentry/testutils/helpers/api_gateway.py

@@ -10,9 +10,9 @@
 
 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.testutils import APITestCase
+from sentry.testutils.silo import unguarded_write
 from sentry.types.region import Region, RegionCategory, clear_global_regions
 from sentry.utils import json
 
@@ -136,7 +136,7 @@ def setUp(self):
             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"
             )