perf(migrations): Keep pending deletions unrendered until their DELETE runs

src/sentry/new_migrations/migrations.py

@@ -35,18 +35,6 @@ def apply(
 
         return super().apply(project_state, schema_editor, collect_sql)
 
-    def mutate_state(self, project_state: ProjectState, preserve: bool = True) -> ProjectState:
-        # `SafeDeleteModel`/`SafeRemoveField` with `MOVE_TO_PENDING` render `state.apps`
-        # to look up the model. If that rendered registry is left on the state, Django
-        # re-renders models on every later state change, so building state from the full
-        # migration history (`migrate`, unapplying migrations in tests) takes minutes.
-        # Drop it again at the end of the migration unless it was already rendered.
-        was_rendered = "apps" in project_state.__dict__
-        new_state = super().mutate_state(project_state, preserve)
-        if not was_rendered:
-            new_state.__dict__.pop("apps", None)
-        return new_state
-
 
 def validate_operation(op: Operation) -> None:
     if isinstance(op, RunSQL) and not isinstance(op, SafeRunSQL):

src/sentry/new_migrations/monkey/fields.py

@@ -37,9 +37,7 @@ def __init__(self, model_name: str, name: str, deletion_action: DeletionAction)
 
     def state_forwards(self, app_label: str, state: SentryProjectState) -> None:  # type: ignore[override]
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
-            field = state.apps.get_model(app_label, self.model_name_lower)._meta.get_field(
-                self.name_lower
-            )
+            field = state.get_field_state(app_label, self.model_name_lower, self.name_lower)
             if getattr(field, "db_constraint", False):
                 raise UnsafeOperationException(
                     f"Foreign key db constraint must be removed before dropping {app_label}.{self.model_name_lower}.{self.name}. "

src/sentry/new_migrations/monkey/models.py

@@ -1,5 +1,6 @@
 from django.db import router
 from django.db.migrations import DeleteModel
+from django.db.models import ManyToManyField
 from django_zero_downtime_migrations.backends.postgres.schema import UnsafeOperationException
 
 from sentry.db.postgres.schema import SafePostgresDatabaseSchemaEditor
@@ -14,9 +15,11 @@ def __init__(self, name: str, deletion_action: DeletionAction) -> None:
 
     def state_forwards(self, app_label: str, state: SentryProjectState) -> None:  # type: ignore[override]
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
-            model = state.apps.get_model(app_label, self.name)
+            model_state = state.get_model_state(app_label, self.name)
             fields_with_constraints = [
-                f.name for f in model._meta.fields if getattr(f, "db_constraint", False)
+                name
+                for name, f in model_state.fields.items()
+                if not isinstance(f, ManyToManyField) and getattr(f, "db_constraint", False)
             ]
             if fields_with_constraints:
                 raise UnsafeOperationException(

src/sentry/new_migrations/monkey/state.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
-from copy import deepcopy
+from collections.abc import Iterable
 from enum import Enum
 
+from django.core.exceptions import FieldDoesNotExist
 from django.db.migrations.state import ModelState, ProjectState
+
+# django-stubs still declares this in `django.db.migrations.operations.utils`.
+from django.db.migrations.utils import resolve_relation  # type: ignore[attr-defined]
 from django.db.models import Field, Model
 from django_zero_downtime_migrations.backends.postgres.schema import UnsafeOperationException
 
@@ -14,14 +18,39 @@ class DeletionAction(Enum):
 
 
 class SentryProjectState(ProjectState):
+    """
+    Keeps the models and fields that `DeletionAction.MOVE_TO_PENDING` removed, so the later
+    `DeletionAction.DELETE` can drop them from the database.
+
+    They're kept as unrendered `ModelState`s and fields: rendering `apps` while replaying the
+    migration history makes every later operation re-render models. They're only rendered
+    when a delete step asks for them.
+    """
+
     def __init__(
         self,
         models: dict[tuple[str, str], ModelState] | None = None,
         real_apps: set[str] | None = None,
     ) -> None:
         super().__init__(models=models, real_apps=real_apps)
-        self.pending_deletion_models: dict[tuple[str, str], type[Model]] = {}
+        self.pending_deletion_models: dict[tuple[str, str], ModelState] = {}
         self.pending_deletion_fields: dict[tuple[str, str, str], Field[object, object]] = {}
+        # Models whose DELETE already ran. Pending deletions can still point at them, like an
+        # M2M `through` that is deleted before the model declaring it, and need them to render.
+        self.deleted_models: dict[tuple[str, str], ModelState] = {}
+
+    def get_model_state(self, app_label: str, model_name: str) -> ModelState:
+        try:
+            return self.models[app_label, model_name.lower()]
+        except KeyError:
+            raise LookupError(f"App '{app_label}' doesn't have a '{model_name}' model.") from None
+
+    def get_field_state(self, app_label: str, model_name: str, name: str) -> Field[object, object]:
+        model_state = self.get_model_state(app_label, model_name)
+        try:
+            return model_state.fields[name]
+        except KeyError:
+            raise FieldDoesNotExist(f"{model_state.name} has no field named '{name}'") from None
 
     def get_pending_deletion_model(self, app_label: str, model_name: str) -> type[Model]:
         model_key = (app_label.lower(), model_name.lower())
@@ -30,7 +59,7 @@ def get_pending_deletion_model(self, app_label: str, model_name: str) -> type[Mo
                 "Model must be in the pending deletion state before full deletion. "
                 "More info: https://develop.sentry.dev/api-server/application-domains/database-migrations/#deleting-tables"
             )
-        return self.pending_deletion_models[model_key]
+        return self._render_pending_deletions([model_key], []).apps.get_model(*model_key)
 
     def get_pending_deletion_field(
         self, app_label: str, model_name: str, field_name: str
@@ -41,7 +70,72 @@ def get_pending_deletion_field(
                 "Field must be in the pending deletion state before full deletion. "
                 "More info: https://develop.sentry.dev/api-server/application-domains/database-migrations/#deleting-columns"
             )
-        return self.pending_deletion_fields[field_key]
+        model = self._render_pending_deletions([], [field_key]).apps.get_model(*field_key[:2])
+        return model._meta.get_field(field_key[2])
+
+    def render_pending_deletions(self) -> tuple[list[Field[object, object]], list[type[Model]]]:
+        """Render every pending deletion at once, fields first, each in the order it went pending."""
+        state = self._render_pending_deletions(
+            self.pending_deletion_models, self.pending_deletion_fields
+        )
+        fields = [
+            state.apps.get_model(app_label, model_name)._meta.get_field(name)
+            for app_label, model_name, name in self.pending_deletion_fields
+        ]
+        models = [state.apps.get_model(*key) for key in self.pending_deletion_models]
+        return fields, models
+
+    def _render_pending_deletions(
+        self,
+        model_keys: Iterable[tuple[str, str]],
+        field_keys: Iterable[tuple[str, str, str]],
+    ) -> SentryProjectState:
+        """
+        Return a throwaway copy of this state with the given pending deletions added back.
+
+        Pending or deleted models that those deletions point at (a field's model, a relation's
+        target or `through`) are added too, so their relations resolve.
+
+        A pending model replaces a live model re-created under the same name, so its DELETE
+        drops the pending model's table and never the live one.
+        """
+        field_keys = list(field_keys)
+        to_add: dict[tuple[str, str], ModelState] = {}
+
+        def visit_relations(field: Field[object, object], key: tuple[str, str]) -> None:
+            if (remote_field := field.remote_field) is not None:
+                visit(resolve_relation(remote_field.model, *key))
+                if (through := getattr(remote_field, "through", None)) is not None:
+                    visit(resolve_relation(through, *key))
+
+        def visit(key: tuple[str, str], replace_live: bool = False) -> None:
+            if key in to_add or (key in self.models and not replace_live):
+                return
+            model_state = self.pending_deletion_models.get(key) or self.deleted_models.get(key)
+            if model_state is None:
+                return
+            to_add[key] = model_state
+            for field in model_state.fields.values():
+                visit_relations(field, key)
+
+        for model_key in model_keys:
+            visit(model_key, replace_live=True)
+        for field_key in field_keys:
+            visit(field_key[:2])
+            visit_relations(self.pending_deletion_fields[field_key], field_key[:2])
+
+        # Render this state once so the copy clones its registry and only re-renders the
+        # models related to what's added back.
+        _ = self.apps
+        state = self.clone()
+        for key, model_state in to_add.items():
+            if key in state.models:
+                state.remove_model(*key)
+            state.add_model(model_state.clone())
+        for app_label, model_name, name in field_keys:
+            field = self.pending_deletion_fields[app_label, model_name, name]
+            state.add_field(app_label, model_name, name, field, preserve_default=True)
+        return state
 
     def remove_model(
         self, app_label: str, model_name: str, deletion_action: DeletionAction | None = None
@@ -53,15 +147,15 @@ def remove_model(
                     "Model must be in the pending deletion state before full deletion. "
                     "More info: https://develop.sentry.dev/api-server/application-domains/database-migrations/#deleting-tables"
                 )
-            del self.pending_deletion_models[model_key]
+            self.deleted_models[model_key] = self.pending_deletion_models.pop(model_key)
             return
         if deletion_action == DeletionAction.MOVE_TO_PENDING:
             if model_key in self.pending_deletion_models:
                 raise UnsafeOperationException(
                     f"{app_label}.{model_name} is already pending deletion. Use DeletionAction.DELETE to delete"
                     "More info: https://develop.sentry.dev/api-server/application-domains/database-migrations/#deleting-tables"
                 )
-            self.pending_deletion_models[model_key] = self.apps.get_model(app_label, model_name)
+            self.pending_deletion_models[model_key] = self.get_model_state(app_label, model_name)
         super().remove_model(app_label, model_name)
 
     def remove_field(
@@ -87,14 +181,16 @@ def remove_field(
                     f"{app_label}.{model_name}.{name} is already pending deletion. Use DeletionAction.DELETE to delete"
                     "More info: https://develop.sentry.dev/api-server/application-domains/database-migrations/#deleting-columns"
                 )
-            self.pending_deletion_fields[field_key] = self.apps.get_model(
-                app_label, model_name
-            )._meta.get_field(name)
+            self.pending_deletion_fields[field_key] = self.get_field_state(
+                app_label, model_name, name
+            )
 
         super().remove_field(app_label, model_name, name)
 
     def clone(self) -> SentryProjectState:
         new_state = super().clone()
-        new_state.pending_deletion_models = deepcopy(self.pending_deletion_models)  # type: ignore[attr-defined]
-        new_state.pending_deletion_fields = deepcopy(self.pending_deletion_fields)  # type: ignore[attr-defined]
+        # These model and field states are never mutated, so the copies can share them.
+        new_state.pending_deletion_models = dict(self.pending_deletion_models)  # type: ignore[attr-defined]
+        new_state.pending_deletion_fields = dict(self.pending_deletion_fields)  # type: ignore[attr-defined]
+        new_state.deleted_models = dict(self.deleted_models)  # type: ignore[attr-defined]
         return new_state  # type: ignore[return-value]

tests/sentry/db/postgres/schema/safe_migrations/integration/test_migrations.py

@@ -307,11 +307,10 @@ class DeletionModelBadDeleteDoublePendingTest(BaseSafeMigrationTest):
     migrate_to = "0003"
 
     def test(self) -> None:
-        # The app's only model is already pending deletion, so the freshly rendered
-        # state doesn't contain the app at all.
+        # The model is already pending deletion, so it's no longer in the state.
         with pytest.raises(
             LookupError,
-            match="No installed app with label 'bad_flow_delete_model_double_pending_app'",
+            match="App 'bad_flow_delete_model_double_pending_app' doesn't have a 'TestTable' model",
         ):
             self.run_migration()
 

tests/sentry/new_migrations/monkey/test_models.py

@@ -7,7 +7,8 @@
 from unittest.mock import Mock
 
 import pytest
-from django.db import connection
+from django.db import connection, models
+from django.db.migrations.state import ModelState
 
 from sentry.db.postgres.schema import SafePostgresDatabaseSchemaEditor
 from sentry.new_migrations.monkey.models import SafeDeleteModel
@@ -27,21 +28,18 @@ def test_delete_model_without_historical_assignment_fails(self) -> None:
         cannot be found, SafeDeleteModel should raise a ValueError in test
         environments.
         """
-        fake_meta = Mock()
-        fake_meta.db_table = "sentry_fake_deleted_table_not_in_router"
-        fake_meta.app_label = "sentry"
-        fake_meta.model_name = "fakedeletedmodel"
-
-        FakeDeletedModel = Mock()
-        FakeDeletedModel._meta = fake_meta
-
         from_state = SentryProjectState()
         to_state = SentryProjectState()
 
         # Manually add the model to pending deletion in the from_state
         # This simulates what happens when a model was marked for pending deletion
         # and is now being deleted
-        from_state.pending_deletion_models[("sentry", "fakedeletedmodel")] = FakeDeletedModel
+        from_state.pending_deletion_models[("sentry", "fakedeletedmodel")] = ModelState(
+            "sentry",
+            "FakeDeletedModel",
+            [("id", models.BigAutoField(primary_key=True))],
+            options={"db_table": "sentry_fake_deleted_table_not_in_router"},
+        )
 
         operation = SafeDeleteModel(name="FakeDeletedModel", deletion_action=DeletionAction.DELETE)