diff --git a/src/sentry/new_migrations/migrations.py b/src/sentry/new_migrations/migrations.py index ab9c446e3f5f..a3f863f738c3 100644 --- a/src/sentry/new_migrations/migrations.py +++ b/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): diff --git a/src/sentry/new_migrations/monkey/fields.py b/src/sentry/new_migrations/monkey/fields.py index 71b7660a25a1..c24785a5c159 100644 --- a/src/sentry/new_migrations/monkey/fields.py +++ b/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}. " diff --git a/src/sentry/new_migrations/monkey/models.py b/src/sentry/new_migrations/monkey/models.py index ed7c6d202bd7..9b6ceb651e8b 100644 --- a/src/sentry/new_migrations/monkey/models.py +++ b/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( diff --git a/src/sentry/new_migrations/monkey/state.py b/src/sentry/new_migrations/monkey/state.py index c4a9f42c12a0..12dd20ea8b61 100644 --- a/src/sentry/new_migrations/monkey/state.py +++ b/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,7 +147,7 @@ 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: @@ -61,7 +155,7 @@ def remove_model( 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] diff --git a/tests/sentry/db/postgres/schema/safe_migrations/integration/test_migrations.py b/tests/sentry/db/postgres/schema/safe_migrations/integration/test_migrations.py index aea3f704831f..d9fabf7e2f61 100644 --- a/tests/sentry/db/postgres/schema/safe_migrations/integration/test_migrations.py +++ b/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() diff --git a/tests/sentry/new_migrations/monkey/test_models.py b/tests/sentry/new_migrations/monkey/test_models.py index 4260d2c48d67..ee4cf405f5b5 100644 --- a/tests/sentry/new_migrations/monkey/test_models.py +++ b/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) diff --git a/tests/sentry/new_migrations/test_migrations.py b/tests/sentry/new_migrations/test_migrations.py index 135151ddd413..50e23ff8159e 100644 --- a/tests/sentry/new_migrations/test_migrations.py +++ b/tests/sentry/new_migrations/test_migrations.py @@ -1,7 +1,7 @@ import pytest from django.core.exceptions import FieldDoesNotExist from django.db import connection, models -from django.db.migrations import Migration +from django.db.migrations import CreateModel, Migration from django.db.migrations.executor import MigrationExecutor from django.db.migrations.operations.base import Operation from django.db.migrations.state import ModelState @@ -134,13 +134,12 @@ def _state_with_m2m_through() -> SentryProjectState: def test_mutate_state_m2m_through_model_dangles_between_operations() -> None: # Mirrors `investigations/0003`: the through model goes to pending before the model that # declares the M2M, so between the two operations `Parent.children` points at a model that - # is no longer in the state. Rendered state may only be dropped at the migration boundary. + # is no longer in the state, and the state can't be rendered. first = SafeDeleteModel(name="Link", deletion_action=DeletionAction.MOVE_TO_PENDING) second = SafeDeleteModel(name="Parent", deletion_action=DeletionAction.MOVE_TO_PENDING) intermediate = _state_with_m2m_through() first.state_forwards("fakeapp", intermediate) - intermediate.__dict__.pop("apps") with pytest.raises(ValueError, match="lazy reference to fakeapp.link"): intermediate.apps @@ -153,6 +152,52 @@ def test_mutate_state_m2m_through_model_dangles_between_operations() -> None: } assert new_state.apps.get_model("fakeapp", "Child") + # Like `investigations/0008`, `Link` is deleted first, but the DELETE step for `Parent` + # still gets its M2M pointing at `Link`. + assert isinstance(new_state, SentryProjectState) + SafeDeleteModel(name="Link", deletion_action=DeletionAction.DELETE).state_forwards( + "fakeapp", new_state + ) + parent = new_state.get_pending_deletion_model("fakeapp", "Parent") + children = parent._meta.get_field("children") + assert isinstance(children, models.ManyToManyField) + through = children.remote_field.through + assert through is not None + assert through._meta.db_table == "fakeapp_link" + with pytest.raises(LookupError): + new_state.apps.get_model("fakeapp", "Parent") + + +@pytest.mark.parametrize( + "operation", + [ + SafeDeleteModel(name="Thing", deletion_action=DeletionAction.MOVE_TO_PENDING), + SafeRemoveField( + model_name="thing", name="flag", deletion_action=DeletionAction.MOVE_TO_PENDING + ), + ], + ids=["model", "field"], +) +def test_move_to_pending_does_not_render(operation: Operation) -> None: + # Rendering here would make every later operation in the history re-render models. + state = _state_with_model() + operation.state_forwards("fakeapp", state) + assert "apps" not in state.__dict__ + + +def test_pending_field_of_model_that_went_pending_later() -> None: + state = _state_with_model() + _migration( + SafeRemoveField( + model_name="thing", name="flag", deletion_action=DeletionAction.MOVE_TO_PENDING + ), + SafeDeleteModel(name="Thing", deletion_action=DeletionAction.MOVE_TO_PENDING), + ).mutate_state(state, preserve=False) + + field = state.get_pending_deletion_field("fakeapp", "thing", "flag") + assert field.column == "flag" + assert field.model._meta.db_table == "fakeapp_thing" + @override_settings( INSTALLED_APPS=( @@ -193,3 +238,42 @@ def test_executor_state_includes_migrations_applied_in_other_apps(self) -> None: assert (self.upstream, "upstreamlater") in executor_state.models assert (self.upstream, "upstreamlater") not in graph_state.models assert executor_state.apps.get_model(self.upstream, "UpstreamLater") + + +def test_pending_m2m_field_whose_through_model_was_deleted_first() -> None: + state = _state_with_m2m_through() + _migration( + SafeRemoveField( + model_name="parent", name="children", deletion_action=DeletionAction.MOVE_TO_PENDING + ), + SafeDeleteModel(name="Link", deletion_action=DeletionAction.MOVE_TO_PENDING), + SafeDeleteModel(name="Link", deletion_action=DeletionAction.DELETE), + ).mutate_state(state, preserve=False) + + field = state.get_pending_deletion_field("fakeapp", "parent", "children") + assert isinstance(field, models.ManyToManyField) + through = field.remote_field.through + assert through is not None + assert through._meta.db_table == "fakeapp_link" + + +def test_pending_model_wins_over_a_live_model_recreated_with_the_same_name() -> None: + state = _state_with_model() + _migration( + SafeDeleteModel(name="Thing", deletion_action=DeletionAction.MOVE_TO_PENDING), + CreateModel( + name="Thing", + fields=[("id", models.AutoField(primary_key=True))], + options={"db_table": "fakeapp_thing_v2"}, + ), + ).mutate_state(state, preserve=False) + + # The DELETE must drop the pending table, never the re-created model's. + pending = state.get_pending_deletion_model("fakeapp", "Thing") + assert pending._meta.db_table == "fakeapp_thing" + assert [f.column for f in pending._meta.fields] == ["id", "flag"] + _, [rendered] = state.render_pending_deletions() + assert rendered._meta.db_table == "fakeapp_thing" + + # The state itself still has the live model. + assert state.apps.get_model("fakeapp", "Thing")._meta.db_table == "fakeapp_thing_v2" diff --git a/tools/migrations/drop_pending_deletions.py b/tools/migrations/drop_pending_deletions.py index 55b717b230bd..4d4735a1b340 100644 --- a/tools/migrations/drop_pending_deletions.py +++ b/tools/migrations/drop_pending_deletions.py @@ -8,11 +8,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from django.db.backends.base.base import BaseDatabaseWrapper - from django.db.models import Field, ManyToManyField, Model + from django.db.models import ManyToManyField, Model from sentry.new_migrations.monkey.state import SentryProjectState @@ -45,8 +45,7 @@ def allow_migrate(model: type[Model]) -> bool: through_tables: list[str | None] = [] tables: list[str] = [] - # The state annotates these as `type[Field]`, but they're field instances. - fields = cast("list[Field[Any, Any]]", list(state.pending_deletion_fields.values())) + fields, models = state.render_pending_deletions() for field in fields: if not allow_migrate(field.model): continue @@ -55,7 +54,7 @@ def allow_migrate(model: type[Model]) -> bool: elif field.db_parameters(connection=connection)["type"] is not None: columns.append((field.model._meta.db_table, field.column)) - for model in state.pending_deletion_models.values(): + for model in models: if not allow_migrate(model): continue through_tables += [_auto_through_table(f) for f in model._meta.local_many_to_many]