Browse Source

fix(hybrid): Add silo mode to "model exists" conditions (#38836)

In several places where the existence of a model is checked via
`app_config.get_model`, replace with a new utility function that also
checks whether the model is available in the current silo mode.

Promote the `__silo_limit` meta-attributes to publicly visible
attributes, dropping the leading underscores.
Ryan Skonnord 2 years ago
parent
commit
729b8112eb

+ 1 - 1
src/sentry/api/base.py

@@ -503,7 +503,7 @@ class EndpointSiloLimit(SiloLimit):
             (decorated_class,),
             {
                 "dispatch": dispatch_override,
-                "__silo_limit": self,  # For internal tooling only
+                "silo_limit": self,
             },
         )
         new_class.__module__ = decorated_class.__module__

+ 21 - 4
src/sentry/db/models/base.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Iterable, Mapping, Tuple, cast
+from typing import Any, Callable, Iterable, Mapping, Tuple, Type, cast
 
+from django.apps.config import AppConfig
 from django.db import models
 from django.db.models import signals
 from django.utils import timezone
@@ -17,6 +18,7 @@ __all__ = (
     "Model",
     "DefaultFieldsModel",
     "sane_repr",
+    "get_model_if_available",
     "control_silo_model",
     "region_silo_model",
 )
@@ -142,6 +144,23 @@ signals.post_save.connect(__model_post_save)
 signals.class_prepared.connect(__model_class_prepared)
 
 
+def get_model_if_available(app_config: AppConfig, model_name: str) -> Type[models.Model] | None:
+    """Get a named model class if it exists and is available in this silo mode."""
+    try:
+        model = app_config.get_model(model_name)
+    except LookupError:
+        return None
+    assert isinstance(model, type) and issubclass(model, models.Model)
+
+    silo_limit = getattr(model._meta, "silo_limit", None)  # type: ignore
+    if silo_limit is not None:
+        assert isinstance(silo_limit, ModelSiloLimit)
+        if not silo_limit.is_available():
+            return None
+
+    return model
+
+
 class ModelSiloLimit(SiloLimit):
     def __init__(
         self,
@@ -199,9 +218,7 @@ class ModelSiloLimit(SiloLimit):
                 # trigger hooks in Django's ModelBase metaclass a second time.
                 setattr(model_class, model_attr_name, override)
 
-        # For internal tooling only. Having any production logic depend on this is
-        # strongly discouraged.
-        model_class._meta.__silo_limit = self
+        model_class._meta.silo_limit = self
 
         return model_class
 

+ 2 - 3
src/sentry/models/counter.py

@@ -9,6 +9,7 @@ from sentry.db.models import (
     BoundedBigIntegerField,
     FlexibleForeignKey,
     Model,
+    get_model_if_available,
     region_silo_model,
     sane_repr,
 )
@@ -94,9 +95,7 @@ def create_counter_function(app_config, using, **kwargs):
     if app_config and app_config.name != "sentry":
         return
 
-    try:
-        app_config.get_model("Counter")
-    except LookupError:
+    if not get_model_if_available(app_config, "Counter"):
         return
 
     cursor = connections[using].cursor()

+ 8 - 4
src/sentry/models/userrole.py

@@ -4,7 +4,13 @@ from django.conf import settings
 from django.db import models
 from django.db.models.signals import post_migrate
 
-from sentry.db.models import ArrayField, DefaultFieldsModel, region_silo_model, sane_repr
+from sentry.db.models import (
+    ArrayField,
+    DefaultFieldsModel,
+    get_model_if_available,
+    region_silo_model,
+    sane_repr,
+)
 from sentry.db.models.fields.foreignkey import FlexibleForeignKey
 
 
@@ -57,9 +63,7 @@ def manage_default_super_admin_role(app_config, using, **kwargs):
     if app_config and app_config.name != "sentry":
         return
 
-    try:
-        app_config.get_model("UserRole")
-    except LookupError:
+    if not get_model_if_available(app_config, "UserRole"):
         return
 
     role, _ = UserRole.objects.get_or_create(

+ 2 - 3
src/sentry/receivers/core.py

@@ -10,6 +10,7 @@ from django.db.utils import OperationalError, ProgrammingError
 from packaging.version import parse as parse_version
 
 from sentry import options
+from sentry.db.models import get_model_if_available
 from sentry.models import Organization, OrganizationMember, Project, ProjectKey, Team, User
 from sentry.signals import project_created
 
@@ -38,9 +39,7 @@ def create_default_projects(app_config, verbosity=2, **kwargs):
     if app_config and app_config.name != "sentry":
         return
 
-    try:
-        app_config.get_model("Project")
-    except LookupError:
+    if not get_model_if_available(app_config, "Project"):
         return
 
     create_default_project(

+ 4 - 3
src/sentry/receivers/users.py

@@ -1,14 +1,15 @@
 from django.db import router
 from django.db.models.signals import post_migrate
 
+from sentry.db.models import get_model_if_available
+
 
 def create_first_user(app_config, using, interactive, **kwargs):
     if app_config and app_config.name != "sentry":
         return
 
-    try:
-        User = app_config.get_model("User")
-    except LookupError:
+    User = get_model_if_available(app_config, "User")
+    if not User:
         return
 
     if User.objects.filter(is_superuser=True).exists():

+ 12 - 5
src/sentry/silo.py

@@ -65,6 +65,14 @@ class SiloLimit(abc.ABC):
         """
         raise NotImplementedError
 
+    def is_available(self, extra_modes: Iterable[SiloMode] = ()) -> bool:
+        current_mode = SiloMode.get_current_mode()
+        return (
+            current_mode == SiloMode.MONOLITH
+            or current_mode in self.modes
+            or current_mode in extra_modes
+        )
+
     def create_override(
         self,
         original_method: Callable[..., Any],
@@ -80,21 +88,20 @@ class SiloLimit(abc.ABC):
         :return: the conditional method object
         """
 
-        available_modes = frozenset(itertools.chain([SiloMode.MONOLITH], self.modes, extra_modes))
-
         def override(*args: Any, **kwargs: Any) -> Any:
             # It's important to do this check inside the override, so that tests
             # using `override_settings` or a similar context can change the value of
             # settings.SILO_MODE effectively. Otherwise, availability would be
             # immutably determined when the decorator is first evaluated.
-            current_mode = SiloMode.get_current_mode()
-            is_available = current_mode in available_modes
+            is_available = self.is_available(extra_modes)
 
             if is_available:
                 return original_method(*args, **kwargs)
             else:
                 handler = self.handle_when_unavailable(
-                    original_method, current_mode, available_modes
+                    original_method,
+                    SiloMode.get_current_mode(),
+                    itertools.chain([SiloMode.MONOLITH], self.modes, extra_modes),
                 )
                 return handler(*args, **kwargs)
 

+ 2 - 2
src/sentry/utils/silo/audit_silo_decorators.py

@@ -34,7 +34,7 @@ def _create_model_table(app_label):
     for model_class in django.apps.apps.get_models():
         if model_class._meta.app_label != app_label:
             continue
-        limit = getattr(model_class._meta, "_ModelSiloLimit__silo_limit", None)
+        limit = getattr(model_class._meta, "silo_limit", None)
         key = (limit.modes, limit.read_only) if limit else None
         table[key].append(model_class)
     return table
@@ -57,7 +57,7 @@ def _create_endpoint_table(app_label):
     for endpoint_class in get_endpoint_classes():
         if not endpoint_class.__module__.startswith(app_label):
             continue
-        limit = getattr(endpoint_class, "__silo_limit", None)
+        limit = getattr(endpoint_class, "silo_limit", None)
         key = frozenset(limit.modes if limit else ())
         table[key].append(endpoint_class)
 

+ 1 - 5
src/sentry/utils/silo/decorate_models_by_relation.py

@@ -40,11 +40,7 @@ def decorate_models_by_relation(
     control_classes = model_classes.difference(region_classes)
 
     def filtered_names(classes):
-        return (
-            (c.__module__, c.__name__)
-            for c in classes
-            if not hasattr(c._meta, "_ModelSiloLimit__silo_limit")
-        )
+        return ((c.__module__, c.__name__) for c in classes if not hasattr(c._meta, "silo_limit"))
 
     apply_decorators(
         "control_silo_model",

+ 28 - 1
tests/sentry/models/test_base.py

@@ -1,7 +1,9 @@
+from unittest.mock import MagicMock
+
 from django.test import override_settings
 from pytest import raises
 
-from sentry.db.models.base import Model, ModelSiloLimit
+from sentry.db.models.base import Model, ModelSiloLimit, get_model_if_available
 from sentry.silo import SiloMode
 from sentry.testutils import TestCase
 
@@ -71,3 +73,28 @@ class AvailableOnTest(TestCase):
             self.ReadOnlyModel.objects.create()
         with raises(ModelSiloLimit.AvailabilityError):
             self.ReadOnlyModel.objects.filter(id=1).delete()
+
+    def test_get_model_if_available(self):
+        test_models = {
+            m.__name__: m
+            for m in (
+                self.ControlModel,
+                self.RegionModel,
+                self.ReadOnlyModel,
+                self.ModelOnMonolith,
+            )
+        }
+        app_config = MagicMock()
+        app_config.get_model.side_effect = test_models.get
+
+        with override_settings(SILO_MODE=SiloMode.REGION):
+            assert get_model_if_available(app_config, "ControlModel") is None
+            assert get_model_if_available(app_config, "RegionModel") is self.RegionModel
+            assert get_model_if_available(app_config, "ReadOnlyModel") is None
+            assert get_model_if_available(app_config, "ModelOnMonolith") is self.ModelOnMonolith
+
+    def test_get_model_with_nonexistent_name(self):
+        app_config = MagicMock()
+        app_config.get_model.side_effect = LookupError
+        assert get_model_if_available(app_config, "BogusModel") is None
+        app_config.get_model.assert_called_with("BogusModel")

Some files were not shown because too many files changed in this diff