perf(migrations): speed up state replay and snapshots

CHANGELOG.rst

@@ -19,6 +19,7 @@ Fixed
 ^^^^^
 - ``QuerySet.distinct().count()`` no longer counts rows duplicated by a join (e.g. filtering on a m2m relation); it now counts distinct primary keys and matches the number of rows the query returns. (#2255)
 - BlackSheep: ``register_tortoise`` now enables the global connection fallback, so database access works when BlackSheep runs handlers in tasks other than the one that initialized the ORM. (#2248)
+- Migration execution now snapshots rendered model registries and rebuilds querysets only for changed models instead of rebuilding every historical model before each operation, substantially improving large migration plans while preserving old/new operation states.
 
 1.1.8
 -----

tests/migrations/test_state_performance.py

@@ -9,8 +9,14 @@
 from tortoise import fields
 from tortoise.fields.relational import ForeignKeyFieldInstance
 from tortoise.migrations.migration import Migration
-from tortoise.migrations.operations import CreateModel, Operation
-from tortoise.migrations.schema_generator.state import State
+from tortoise.migrations.operations import (
+    AlterField,
+    CreateModel,
+    DeleteModel,
+    Operation,
+    RemoveField,
+)
+from tortoise.migrations.schema_generator.state import ModelState, State
 from tortoise.migrations.schema_generator.state_apps import StateApps
 
 
@@ -63,25 +69,25 @@ async def test_state_building_performance_200_models():
 
 
 @pytest.mark.asyncio
-async def test_apply_dry_run_does_not_clone_state():
-    """Verify that apply(dry_run=True) never calls State.clone()."""
+async def test_apply_dry_run_does_not_snapshot_state():
+    """Verify that apply(dry_run=True) never calls State.snapshot()."""
     migrations = _build_migrations(10)
     state = State(models={}, apps=StateApps())
 
-    clone_calls = 0
-    original_clone = State.clone
+    snapshot_calls = 0
+    original_snapshot = State.snapshot
 
-    def counting_clone(self):
-        nonlocal clone_calls
-        clone_calls += 1
-        return original_clone(self)
+    def counting_snapshot(self):
+        nonlocal snapshot_calls
+        snapshot_calls += 1
+        return original_snapshot(self)
 
-    with patch.object(State, "clone", counting_clone):
+    with patch.object(State, "snapshot", counting_snapshot):
         for migration in migrations:
             await migration.apply(state, dry_run=True, schema_editor=None)
 
     assert len(state.models) == 10
-    assert clone_calls == 0, f"State.clone() was called {clone_calls} times during dry_run"
+    assert snapshot_calls == 0, f"State.snapshot() was called {snapshot_calls} times during dry_run"
 
 
 def test_state_clone_produces_independent_copy():
@@ -146,3 +152,117 @@ def test_state_clone_preserves_relations():
     fk = child_model._meta.fields_map["parent"]
     assert isinstance(fk, ForeignKeyFieldInstance)
     assert fk.related_model is parent_model
+
+
+def test_state_snapshot_does_not_render_unchanged_models():
+    state = State(models={}, apps=StateApps())
+    for i in range(50):
+        _make_create_model_op(i).state_forward("app", state)
+
+    with patch.object(ModelState, "render", wraps=ModelState.render) as render:
+        snapshot = state.snapshot()
+
+    assert render.call_count == 0
+    for model_name, model in state.apps.apps["app"].items():
+        assert snapshot.apps.get_model("app", model_name) is model
+
+
+def test_state_snapshot_keeps_old_related_models_intact():
+    state = State(models={}, apps=StateApps())
+    CreateModel(
+        name="Parent",
+        fields=[
+            ("id", fields.IntField(primary_key=True)),
+            ("name", fields.CharField(max_length=50)),
+        ],
+    ).state_forward("app", state)
+    CreateModel(
+        name="Child",
+        fields=[
+            ("id", fields.IntField(primary_key=True)),
+            ("parent", fields.ForeignKeyField("app.Parent", related_name="children")),
+        ],
+    ).state_forward("app", state)
+    old_state = state.snapshot()
+    old_parent = old_state.apps.get_model("app", "Parent")
+    old_child = old_state.apps.get_model("app", "Child")
+
+    AlterField(
+        model_name="Parent",
+        name="name",
+        field=fields.TextField(),
+    ).state_forward("app", state)
+
+    new_parent = state.apps.get_model("app", "Parent")
+    new_child = state.apps.get_model("app", "Child")
+    assert old_parent is not new_parent
+    assert old_child is not new_child
+    assert isinstance(old_parent._meta.fields_map["name"], fields.CharField)
+    assert isinstance(new_parent._meta.fields_map["name"], fields.TextField)
+    assert old_child._meta.fields_map["parent"].related_model is old_parent
+    assert new_child._meta.fields_map["parent"].related_model is new_parent
+    assert old_parent._meta.app == "app"
+    assert old_child._meta.app == "app"
+
+
+def test_state_reload_builds_querysets_only_for_reloaded_models():
+    state = State(models={}, apps=StateApps())
+    CreateModel(
+        name="Changed",
+        fields=[
+            ("id", fields.IntField(primary_key=True)),
+            ("name", fields.CharField(max_length=50)),
+        ],
+    ).state_forward("app", state)
+    CreateModel(
+        name="Unchanged",
+        fields=[("id", fields.IntField(primary_key=True))],
+    ).state_forward("app", state)
+
+    with patch.object(state.apps, "_build_initial_querysets") as build_querysets:
+        AlterField(
+            model_name="Changed",
+            name="name",
+            field=fields.TextField(),
+        ).state_forward("app", state)
+
+    reloaded_models = build_querysets.call_args.args[0]
+    assert {model.__name__ for model in reloaded_models} == {"Changed"}
+
+
+def test_state_snapshot_keeps_removed_field_in_old_state():
+    state = State(models={}, apps=StateApps())
+    CreateModel(
+        name="Article",
+        fields=[
+            ("id", fields.IntField(primary_key=True)),
+            ("title", fields.CharField(max_length=100)),
+        ],
+    ).state_forward("app", state)
+    old_state = state.snapshot()
+    old_model = old_state.apps.get_model("app", "Article")
+
+    RemoveField(model_name="Article", name="title").state_forward("app", state)
+
+    new_model = state.apps.get_model("app", "Article")
+    assert "title" in old_state.models[("app", "Article")].fields
+    assert "title" not in state.models[("app", "Article")].fields
+    assert "title" in old_model._meta.fields_map
+    assert "title" not in new_model._meta.fields_map
+    assert old_model is not new_model
+
+
+def test_state_snapshot_survives_model_deletion():
+    state = State(models={}, apps=StateApps())
+    CreateModel(
+        name="Obsolete",
+        fields=[("id", fields.IntField(primary_key=True))],
+    ).state_forward("app", state)
+    old_state = state.snapshot()
+    old_model = old_state.apps.get_model("app", "Obsolete")
+
+    DeleteModel("Obsolete").state_forward("app", state)
+
+    assert old_state.apps.get_model("app", "Obsolete") is old_model
+    assert old_model._meta.app == "app"
+    assert "Obsolete" not in state.apps.apps["app"]

tortoise/migrations/executor.py

@@ -244,7 +244,7 @@ async def _project_state_cache(self, applied: set[MigrationKey]) -> dict[Migrati
         for key in self._full_plan():
             if key not in applied:
                 continue
-            cache[key] = state.clone()
+            cache[key] = state.snapshot()
             migration = self.loader.graph.nodes[key]
             if migration is None:
                 raise ValueError(f"Missing migration for {key}")

tortoise/migrations/migration.py

@@ -48,7 +48,7 @@ async def apply(
             not dry_run and schema_editor is not None
         )
         for operation in self.operations:
-            old_state = state.clone() if need_old_state else None
+            old_state = state.snapshot() if need_old_state else None
             operation.state_forward(self.app_label, state)
             if collect_sql and schema_editor:
                 schema_editor.collected_sql.append("--")
@@ -100,13 +100,13 @@ async def unapply(
         new_state = state
         if not need_old_state and self.operations:
             # Single working copy so state_forward doesn't mutate the original
-            new_state = state.clone()
+            new_state = state.snapshot()
         for operation in self.operations:
             if not getattr(operation, "reversible", True):
                 raise ValueError(f"Operation {operation} in {self} is not reversible")
             if need_old_state:
-                new_state = new_state.clone()
-                old_state = new_state.clone()
+                new_state = new_state.snapshot()
+                old_state = new_state.snapshot()
                 operation.state_forward(self.app_label, new_state)
                 to_run.insert(0, (operation, old_state, new_state))
             else:

tortoise/migrations/operations.py

@@ -131,7 +131,7 @@ async def run(
         dry_run: bool,
         state_editor: BaseSchemaEditor | None = None,
     ) -> None:
-        old_state = state.clone() if (not dry_run and state_editor) else None
+        old_state = state.snapshot() if (not dry_run and state_editor) else None
         self.state_forward(app_label, state)
         if dry_run or not state_editor:
             return
@@ -236,7 +236,7 @@ def state_forward(self, app_label: str, state: State) -> None:
         if not model_state_to_change:
             raise IncompatibleStateError()
 
-        state.apps.unregister_model(app_label, self.old_name)
+        state.apps.unregister_model(app_label, self.old_name, detach=False)
 
         old_table = model_state_to_change.table
         model_state_to_change.name = self.new_name
@@ -334,7 +334,7 @@ def state_forward(self, app_label: str, state: State) -> None:
 
             models_to_reload.add(state.apps.split_reference(field.model_name))
 
-        state.apps.unregister_model(app_label, self.name)
+        state.apps.unregister_model(app_label, self.name, detach=False)
         state.reload_models(models_to_reload)
 
     async def database_forward(

tortoise/migrations/schema_generator/state.py

@@ -166,14 +166,23 @@ def _find_related_models(self, app_label: str, model_name: str) -> set[tuple[str
         return related_models
 
     def _reload(self, models_to_reload: set[tuple[str, str]]) -> None:
+        reloaded_models: list[type[Model]] = []
         for app_label, model_name in models_to_reload:
-            self.apps.unregister_model(app_label, model_name)
-            model_state = self.models[(app_label, model_name)]
+            model_state = self.models.get((app_label, model_name))
+            if model_state is None:
+                # A deleted model can still be reachable from a rendered
+                # relation retained by an older snapshot. It must not be added
+                # back to the current state while related models are reloaded.
+                continue
+            # Old migration-state snapshots may still reference this rendered
+            # model, so remove it from the current registry without mutating it.
+            self.apps.unregister_model(app_label, model_name, detach=False)
             model = model_state.render(self.apps)
             self.apps.register_model(app_label, model)
+            reloaded_models.append(model)
 
         self.apps._init_relations()
-        self.apps._build_initial_querysets()
+        self.apps._build_initial_querysets(reloaded_models)
 
     def reload_model(self, app_label: str, model_name: str) -> None:
         model_state = self.models.get((app_label, model_name))
@@ -218,3 +227,13 @@ def validate_relations_initialized(self) -> None:
     def clone(self) -> State:
         models = {key: model.clone() for key, model in self.models.items()}
         return self.__class__(models=models, apps=self.apps.clone(model_states=models))
+
+    def snapshot(self) -> State:
+        """Return an isolated state snapshot without re-rendering every model.
+
+        Model states are copied because operations mutate their fields and
+        options. Rendered model classes can be shared: changed models are replaced
+        non-destructively by ``_reload``, leaving this snapshot's classes intact.
+        """
+        models = {key: model.clone() for key, model in self.models.items()}
+        return self.__class__(models=models, apps=self.apps.snapshot())

tortoise/migrations/schema_generator/state_apps.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, cast
 
 if TYPE_CHECKING:
@@ -90,33 +91,49 @@ def _init_relations(self) -> None:
         for model in models_with_missing_refs:
             model._meta._inited = False
 
-    def _build_initial_querysets(self) -> None:
+    def _build_initial_querysets(self, models: Iterable[type[Model]] | None = None) -> None:
         # Skip building querysets when no DB config is available (state-only mode)
         # This allows pure state operations to work without database connections
         if self._connections._db_config is None:
             return
 
-        for app in self.apps.values():
-            for model in app.values():
-                if model._meta.default_connection is None:
-                    continue
-                if not model._meta._inited:
-                    continue
-                model._meta.finalise_model()
-                model._meta.basetable = Table(name=model._meta.db_table, schema=model._meta.schema)
-                basequery = model._meta.db.query_class.from_(model._meta.basetable)
-                model._meta.basequery = cast(Query, basequery)
-                model._meta.basequery_all_fields = cast(
-                    Query, basequery.select(*model._meta.db_fields)
-                )
-
-    def unregister_model(self, app_label: str, model_name: str) -> None:
+        if models is None:
+            models = (model for app in self.apps.values() for model in app.values())
+
+        for model in models:
+            if model._meta.default_connection is None:
+                continue
+            if not model._meta._inited:
+                continue
+            model._meta.finalise_model()
+            model._meta.basetable = Table(name=model._meta.db_table, schema=model._meta.schema)
+            basequery = model._meta.db.query_class.from_(model._meta.basetable)
+            model._meta.basequery = cast(Query, basequery)
+            model._meta.basequery_all_fields = cast(Query, basequery.select(*model._meta.db_fields))
+
+    def unregister_model(self, app_label: str, model_name: str, *, detach: bool = True) -> None:
         try:
             model = self.apps[app_label].pop(model_name)
-            model._meta.app = None
+            if detach:
+                model._meta.app = None
         except KeyError:
             return
 
+    def snapshot(self) -> StateApps:
+        """Return a cheap snapshot of the rendered migration models.
+
+        Rendered model classes are immutable versions of a migration model. State
+        operations replace changed models (and their related models) instead of
+        mutating them, so an apps registry only needs its mappings copied to keep
+        the previous version available to database operations.
+        """
+        state_apps = self.__class__(
+            default_connections=dict(self._default_connections),
+            connections=self._connections,
+        )
+        state_apps.apps = {app_label: dict(models) for app_label, models in self.apps.items()}
+        return state_apps
+
     def split_reference(self, reference: str | type[Model]) -> tuple[str, str]:
         if not isinstance(reference, str):
             model_class = reference