Refs #29898 -- Made ProjectState encapsulate alterations in relations registry.

Thanks Simon Charette and Chris Jerdonek for reviews.

Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>

1 files changed, 78 insertions, 20 deletions

diff --git a/django/db/migrations/state.py b/django/db/migrations/state.py
index 392708134d..e04213b656 100644
--- a/django/db/migrations/state.py
+++ b/django/db/migrations/state.py
@@ -98,18 +98,34 @@ class ProjectState:
         self.real_apps = real_apps
         self.is_delayed = False
         # {remote_model_key: {model_key: [(field_name, field)]}}
-        self.relations = None
+        self._relations = None
+
+    @property
+    def relations(self):
+        if self._relations is None:
+            self.resolve_fields_and_relations()
+        return self._relations
 
     def add_model(self, model_state):
-        app_label, model_name = model_state.app_label, model_state.name_lower
-        self.models[(app_label, model_name)] = model_state
+        model_key = model_state.app_label, model_state.name_lower
+        self.models[model_key] = model_state
+        if self._relations is not None:
+            self.resolve_model_relations(model_key)
         if 'apps' in self.__dict__:  # hasattr would cache the property
-            self.reload_model(app_label, model_name)
+            self.reload_model(*model_key)
 
     def remove_model(self, app_label, model_name):
-        del self.models[app_label, model_name]
+        model_key = app_label, model_name
+        del self.models[model_key]
+        if self._relations is not None:
+            self._relations.pop(model_key, None)
+            # Call list() since _relations can change size during iteration.
+            for related_model_key, model_relations in list(self._relations.items()):
+                model_relations.pop(model_key, None)
+                if not model_relations:
+                    del self._relations[related_model_key]
         if 'apps' in self.__dict__:  # hasattr would cache the property
-            self.apps.unregister_model(app_label, model_name)
+            self.apps.unregister_model(*model_key)
             # Need to do this explicitly since unregister_model() doesn't clear
             # the cache automatically (#24513)
             self.apps.clear_cache()
@@ -137,6 +153,14 @@ class ProjectState:
             if changed_field:
                 model_state.fields[name] = changed_field
                 to_reload.add((model_state.app_label, model_state.name_lower))
+        if self._relations is not None:
+            old_name_key = app_label, old_name_lower
+            new_name_key = app_label, new_name_lower
+            if old_name_key in self._relations:
+                self._relations[new_name_key] = self._relations.pop(old_name_key)
+            for model_relations in self._relations.values():
+                if old_name_key in model_relations:
+                    model_relations[new_name_key] = model_relations.pop(old_name_key)
         # Reload models related to old model before removing the old model.
         self.reload_models(to_reload, delay=True)
         # Remove the old model.
@@ -187,17 +211,23 @@ class ProjectState:
             field.default = NOT_PROVIDED
         else:
             field = field
-        self.models[app_label, model_name].fields[name] = field
+        model_key = app_label, model_name
+        self.models[model_key].fields[name] = field
+        if self._relations is not None:
+            self.resolve_model_field_relations(model_key, name, field)
         # Delay rendering of relationships if it's not a relational field.
         delay = not field.is_relation
-        self.reload_model(app_label, model_name, delay=delay)
+        self.reload_model(*model_key, delay=delay)
 
     def remove_field(self, app_label, model_name, name):
-        model_state = self.models[app_label, model_name]
+        model_key = app_label, model_name
+        model_state = self.models[model_key]
         old_field = model_state.fields.pop(name)
+        if self._relations is not None:
+            self.resolve_model_field_relations(model_key, name, old_field)
         # Delay rendering of relationships if it's not a relational field.
         delay = not old_field.is_relation
-        self.reload_model(app_label, model_name, delay=delay)
+        self.reload_model(*model_key, delay=delay)
 
     def alter_field(self, app_label, model_name, name, field, preserve_default):
         if not preserve_default:
@@ -205,20 +235,30 @@ class ProjectState:
             field.default = NOT_PROVIDED
         else:
             field = field
-        model_state = self.models[app_label, model_name]
-        model_state.fields[name] = field
+        model_key = app_label, model_name
+        fields = self.models[model_key].fields
+        if self._relations is not None:
+            old_field = fields.pop(name)
+            if old_field.is_relation:
+                self.resolve_model_field_relations(model_key, name, old_field)
+            fields[name] = field
+            if field.is_relation:
+                self.resolve_model_field_relations(model_key, name, field)
+        else:
+            fields[name] = field
         # TODO: investigate if old relational fields must be reloaded or if
         # it's sufficient if the new field is (#27737).
         # Delay rendering of relationships if it's not a relational field and
         # not referenced by a foreign key.
         delay = (
             not field.is_relation and
-            not field_is_referenced(self, (app_label, model_name), (name, field))
+            not field_is_referenced(self, model_key, (name, field))
         )
-        self.reload_model(app_label, model_name, delay=delay)
+        self.reload_model(*model_key, delay=delay)
 
     def rename_field(self, app_label, model_name, old_name, new_name):
-        model_state = self.models[app_label, model_name]
+        model_key = app_label, model_name
+        model_state = self.models[model_key]
         # Rename the field.
         fields = model_state.fields
         try:
@@ -246,7 +286,7 @@ class ProjectState:
                 ]
         # Fix to_fields to refer to the new field.
         delay = True
-        references = get_references(self, (app_label, model_name), (old_name, found))
+        references = get_references(self, model_key, (old_name, found))
         for *_, field, reference in references:
             delay = False
             if reference.to:
@@ -258,7 +298,19 @@ class ProjectState:
                         new_name if to_field_name == old_name else to_field_name
                         for to_field_name in to_fields
                     ])
-        self.reload_model(app_label, model_name, delay=delay)
+        if self._relations is not None:
+            old_name_lower = old_name.lower()
+            new_name_lower = new_name.lower()
+            for to_model in self._relations.values():
+                # It's safe to modify the same collection that is iterated
+                # because `break` is called right after.
+                for field_name, field in to_model[model_key]:
+                    if field_name == old_name_lower:
+                        field.name = new_name_lower
+                        to_model[model_key].remove((old_name_lower, field))
+                        to_model[model_key].append((new_name_lower, field))
+                        break
+        self.reload_model(*model_key, delay=delay)
 
     def _find_reload_model(self, app_label, model_name, delay=False):
         if delay:
@@ -352,7 +404,13 @@ class ProjectState:
         remote_model_key = resolve_relation(model, *model_key)
         if remote_model_key[0] not in self.real_apps and remote_model_key in concretes:
             remote_model_key = concretes[remote_model_key]
-        self.relations[remote_model_key][model_key].append((field_name, field))
+        relations_to_remote_model = self._relations[remote_model_key]
+        if field_name in self.models[model_key].fields:
+            relations_to_remote_model[model_key].append((field_name, field))
+        else:
+            relations_to_remote_model[model_key].remove((field_name, field))
+            if not relations_to_remote_model[model_key]:
+                del relations_to_remote_model[model_key]
 
     def resolve_model_field_relations(
         self, model_key, field_name, field, concretes=None,
@@ -387,14 +445,14 @@ class ProjectState:
                 field.name = field_name
         # Resolve relations.
         # {remote_model_key: {model_key: [(field_name, field)]}}
-        self.relations = defaultdict(partial(defaultdict, list))
+        self._relations = defaultdict(partial(defaultdict, list))
         concretes, proxies = self._get_concrete_models_mapping_and_proxy_models()
 
         for model_key in concretes:
             self.resolve_model_relations(model_key, concretes)
 
         for model_key in proxies:
-            self.relations[model_key] = self.relations[concretes[model_key]]
+            self._relations[model_key] = self._relations[concretes[model_key]]
 
     def get_concrete_model_key(self, model):
         concrete_models_mapping, _ = self._get_concrete_models_mapping_and_proxy_models()