fix(hc): Prevent exception from siloed_atomic (#52326)

An exception rises from `router.db_for_write` if the database doesn't
exist in the environment, as observed with `SENTRY_USE_SILOS` set to
true. Introduce a custom exception type to mark this condition, and
catch it when setting the atomic impl. This should be equivalent
behavior, as we only care about whether "missing" database is equal to
`using` (which a null value won't be).

Author: RyanSkonnord

src/sentry/db/router.py

@@ -12,6 +12,10 @@
 logger = logging.getLogger(__name__)
 
 
+class SiloConnectionUnavailableError(ValueError):
+    pass
+
+
 class SiloRouter:
     """
     Django database router for multi-region deployments.
@@ -75,7 +79,7 @@ def use_simulated(self, value: bool):
             raise ValueError("Cannot mutate simulation mode outside of tests")
         self.__is_simulated = value
 
-    def _resolve_silo_connection(self, silo_modes: Iterable[SiloMode], table: str):
+    def _resolve_silo_connection(self, silo_modes: Iterable[SiloMode], table: str) -> str:
         # XXX This method has an override in getsentry for region silo primary splits.
         active_mode = SiloMode.get_current_mode()
 
@@ -92,7 +96,7 @@ def _resolve_silo_connection(self, silo_modes: Iterable[SiloMode], table: str):
             # If we're in tests raise an error, otherwise return 'no decision'
             # so that django skips migration operations that won't work.
             if "pytest" in sys.modules:
-                raise ValueError(
+                raise SiloConnectionUnavailableError(
                     f"Cannot resolve table {table} in {silo_mode}. "
                     f"Application silo mode is {active_mode} and simulated silos are not enabled."
                 )

src/sentry/silo/patches/silo_aware_transaction_patch.py

@@ -1,16 +1,28 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Optional, Type
 
 from django import get_version
 from django.db import router, transaction
 from django.db.transaction import Atomic
 
+if TYPE_CHECKING:
+    from sentry.db.models import Model
+
 _default_atomic_impl = transaction.atomic
 
 
 class MismatchedSiloTransactionError(Exception):
     pass
 
 
+def _get_db_for_model_if_available(model: Type["Model"]) -> Optional[str]:
+    from sentry.db.router import SiloConnectionUnavailableError
+
+    try:
+        return router.db_for_write(model)
+    except SiloConnectionUnavailableError:
+        return None
+
+
 def siloed_atomic(using: Optional[str] = None, savepoint: bool = True) -> Atomic:
     from sentry.models import ControlOutbox, RegionOutbox
     from sentry.silo import SiloMode
@@ -20,24 +32,17 @@ def siloed_atomic(using: Optional[str] = None, savepoint: bool = True) -> Atomic
         model_to_route_to = RegionOutbox if current_silo_mode == SiloMode.REGION else ControlOutbox
         using = router.db_for_write(model_to_route_to)
 
-    both_silos_route_to_same_db = router.db_for_write(ControlOutbox) == router.db_for_write(
-        RegionOutbox
-    )
+    control_db = _get_db_for_model_if_available(ControlOutbox)
+    region_db = _get_db_for_model_if_available(RegionOutbox)
+    both_silos_route_to_same_db = control_db == region_db
 
     if both_silos_route_to_same_db or current_silo_mode == SiloMode.MONOLITH:
         pass
-    elif (
-        using == router.db_for_write(ControlOutbox)
-        and SiloMode.get_current_mode() != SiloMode.CONTROL
-    ):
+    elif using == control_db and SiloMode.get_current_mode() != SiloMode.CONTROL:
         raise MismatchedSiloTransactionError(
             f"Cannot use transaction.atomic({using}) in Control Mode"
         )
-
-    elif (
-        using == router.db_for_write(RegionOutbox)
-        and SiloMode.get_current_mode() != SiloMode.REGION
-    ):
+    elif using == region_db and SiloMode.get_current_mode() != SiloMode.REGION:
         raise MismatchedSiloTransactionError(
             f"Cannot use transaction.atomic({using}) in Region Mode"
         )