diff --git a/pyproject.toml b/pyproject.toml index e2262e4acd49f9..9dbbe262a8a0fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -188,7 +188,6 @@ module = [ "sentry.auth.system", "sentry.db.mixin", "sentry.db.postgres.base", - "sentry.db.router", "sentry.discover.endpoints.discover_key_transactions", "sentry.eventstore.models", "sentry.grouping.strategies.legacy", diff --git a/src/sentry/db/router.py b/src/sentry/db/router.py index 7dee83f541a705..5668c49b9ac3da 100644 --- a/src/sentry/db/router.py +++ b/src/sentry/db/router.py @@ -5,9 +5,10 @@ from django.apps import apps from django.db import connections +from django.db.models.base import Model from django.utils.connection import ConnectionDoesNotExist -from sentry.db.models.base import Model, ModelSiloLimit +from sentry.db.models.base import ModelSiloLimit from sentry.silo.base import SiloLimit, SiloMode from sentry.utils.env import in_test_environment @@ -44,9 +45,6 @@ class SiloRouter: SiloMode.CONTROL: "control", } - __table_to_silo = {} - """Memoized results of table : silo pairings""" - __is_simulated = False """Whether or not we're operating in a simulated silo environment""" @@ -82,7 +80,8 @@ class SiloRouter: """ def __init__(self): - self.__table_to_silo = {} + # Memoized results of table : silo pairings + self.__table_to_silo: dict[str, str | None] = {} try: # By accessing the connections Django will raise # Use `assert` to appease linters @@ -122,14 +121,14 @@ def _resolve_silo_connection(self, silo_modes: Iterable[SiloMode], table: str) - else: return None - def _find_model(self, table: str, app_label: str) -> Model | None: + def _find_model(self, table: str) -> type[Model] | None: # Use django's model inventory to find our table and what silo it is on. - for model in apps.get_models(app_label): + for model in apps.get_models(include_auto_created=True): if model._meta.db_table == table: return model return None - def _silo_limit(self, model: Model) -> SiloLimit | None: + def _silo_limit(self, model: type[Model]) -> SiloLimit | None: silo_limit = getattr(model._meta, "silo_limit", None) if silo_limit: return silo_limit @@ -140,24 +139,24 @@ def _silo_limit(self, model: Model) -> SiloLimit | None: # If we didn't find a silo_limit we could be working with __fake__ model # from django, so we need to locate the real class by table. - real_model = self._find_model(db_table, model._meta.app_label) + real_model = self._find_model(db_table) if real_model: return getattr(real_model._meta, "silo_limit", None) return None - def _db_for_model(self, model: Model): + def _db_for_model(self, model: type[Model]) -> str | None: silo_limit = self._silo_limit(model) if not silo_limit: return "default" return self._resolve_silo_connection(silo_limit.modes, table=model._meta.db_table) - def _db_for_table(self, table, app_label): + def _db_for_table(self, table: str, app_label: str) -> str | None: if table in self.__table_to_silo: return self.__table_to_silo[table] - model = self._find_model(table, app_label) + model = self._find_model(table) if model: # Incrementally build up our result cache so we don't # have to scan through models more than once. @@ -171,7 +170,7 @@ def _db_for_table(self, table, app_label): # Default to None for sentry/getsentry app_label as models # in those apps must have silo assignments, and 'default' # for other app_labels that can't have silo assignments. - fallback = "default" + fallback: str | None = "default" if app_label in {"sentry", "getsentry"}: fallback = None self.__table_to_silo[table] = fallback