fix(hybrid-cloud): Fixes tests for split DB mode app installs using transaction routing and silo annotations (#52932)

Co-authored-by: Zachary Collins <zachary.collins@sentry.io>

src/sentry/db/postgres/transactions.py

@@ -5,7 +5,7 @@ import sys
 import threading
 
 from django.conf import settings
-from django.db import transaction
+from django.db import connections, transaction
 
 
 @contextlib.contextmanager
@@ -87,7 +87,7 @@ def in_test_assert_no_transaction(msg: str):
 
     from sentry.testutils import hybrid_cloud
 
-    for using in settings.DATABASES:  # type: ignore
+    for conn in connections.all():
         assert not hybrid_cloud.simulated_transaction_watermarks.connection_above_watermark(
-            using
+            connection=conn
         ), msg

src/sentry/models/organization.py

@@ -5,7 +5,7 @@ from enum import IntEnum
 from typing import Collection, FrozenSet, Optional, Sequence
 
 from django.conf import settings
-from django.db import models, transaction
+from django.db import models, router, transaction
 from django.db.models import QuerySet
 from django.urls import NoReverseMatch, reverse
 from django.utils import timezone
@@ -255,7 +255,7 @@ class Organization(Model, OptionMixin, OrganizationAbsoluteUrlMixin, SnowflakeId
                 lambda: self.save_with_update_outbox(*args, **kwargs),
             )
         else:
-            with outbox_context(transaction.atomic()):
+            with outbox_context(transaction.atomic(using=router.db_for_write(Organization))):
                 self.save_with_update_outbox(*args, **kwargs)
 
     # Override for the default update method to ensure that most atomic updates

src/sentry/models/organizationmember.py

@@ -9,7 +9,7 @@ from urllib.parse import urlencode
 from uuid import uuid4
 
 from django.conf import settings
-from django.db import models, transaction
+from django.db import models, router, transaction
 from django.db.models import Q, QuerySet
 from django.urls import reverse
 from django.utils import timezone
@@ -260,7 +260,7 @@ class OrganizationMember(Model):
             self.user_id and self.email is None
         ), "Must set either user or email"
 
-        with outbox_context(transaction.atomic()):
+        with outbox_context(transaction.atomic(using=router.db_for_write(OrganizationMember))):
             if self.token and not self.token_expires_at:
                 self.refresh_expires_at()
             super().save(*args, **kwargs)

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

@@ -5,7 +5,7 @@
 
 from typing import Optional
 
-from django.db import IntegrityError, transaction
+from django.db import IntegrityError, router, transaction
 
 from sentry.models import outbox_context
 from sentry.models.organizationmembermapping import OrganizationMemberMapping
@@ -48,7 +48,9 @@ class DatabaseBackedOrganizationMemberMappingService(OrganizationMemberMappingSe
                     outbox.save()
 
         try:
-            with outbox_context(transaction.atomic()):
+            with outbox_context(
+                transaction.atomic(using=router.db_for_write(OrganizationMemberMapping))
+            ):
                 existing = self._find_organization_member(
                     organization_id=organization_id,
                     organizationmember_id=organizationmember_id,
@@ -76,7 +78,9 @@ class DatabaseBackedOrganizationMemberMappingService(OrganizationMemberMappingSe
             if existing is None:
                 raise e
 
-            with outbox_context(transaction.atomic()):
+            with outbox_context(
+                transaction.atomic(using=router.db_for_write(OrganizationMemberMapping))
+            ):
                 apply_update(existing)
 
         return serialize_org_member_mapping(existing)

src/sentry/utils/snowflake.py

@@ -3,7 +3,7 @@ from datetime import datetime, timedelta
 from typing import Tuple
 
 from django.conf import settings
-from django.db import IntegrityError, transaction
+from django.db import IntegrityError, router, transaction
 from rest_framework import status
 from rest_framework.exceptions import APIException
 
@@ -24,7 +24,9 @@ class SnowflakeIdMixin:
             if not self.id:
                 self.id = generate_snowflake_id(snowflake_redis_key)
             try:
-                with outbox_context(transaction.atomic()):
+                # We need to route to the correct database in a split DB mode
+                #  so we use the class being mixed in to do this.
+                with outbox_context(transaction.atomic(using=router.db_for_write(type(self)))):
                     save_callback()
                 return
             except IntegrityError: