Add support for AddEmbeddedIndex

django/db/migrations/autodetector.py

@@ -5,6 +5,9 @@
 from graphlib import TopologicalSorter
 from itertools import chain
 
+from django_mongodb_backend.fields import EmbeddedModelField
+from django_mongodb_backend.models import EMBEDDED
+
 from django.conf import settings
 from django.db import models
 from django.db.migrations import operations
@@ -145,18 +148,23 @@
         # resolve dependencies caused by M2Ms and FKs.
         self.generated_operations = {}
         self.altered_indexes = {}
+        self.altered_embedded_indexes = {}
         self.altered_constraints = {}
         self.renamed_fields = {}
 
         # Prepare some old/new state and model lists, separating
         # proxy models and ignoring unmigrated apps.
+        self.old_embedded_keys = set()
         self.old_model_keys = set()
         self.old_proxy_keys = set()
         self.old_unmanaged_keys = set()
+        self.new_embedded_keys = set()
         self.new_model_keys = set()
         self.new_proxy_keys = set()
         self.new_unmanaged_keys = set()
         for (app_label, model_name), model_state in self.from_state.models.items():
+            if model_state.options.get("db_table") is EMBEDDED:
+                self.old_embedded_keys.add((app_label, model_name))
             if not model_state.options.get("managed", True):
                 self.old_unmanaged_keys.add((app_label, model_name))
             elif app_label not in self.from_state.real_apps:
@@ -166,7 +174,9 @@
                     self.old_model_keys.add((app_label, model_name))
 
         for (app_label, model_name), model_state in self.to_state.models.items():
-            if not model_state.options.get("managed", True):
+            if model_state.options.get("db_table") is EMBEDDED:
+                self.new_embedded_keys.add((app_label, model_name))
+            elif not model_state.options.get("managed", True):
                 self.new_unmanaged_keys.add((app_label, model_name))
             elif app_label not in self.from_state.real_apps or (
                 convert_apps and app_label in convert_apps
@@ -204,6 +214,7 @@
         # This avoids the same computation in generate_removed_indexes()
         # and generate_added_indexes().
         self.create_altered_indexes()
+        self.create_altered_embedded_indexes()
         self.create_altered_constraints()
         # Generate index removal operations before field is removed
         self.generate_removed_constraints()
@@ -1462,6 +1473,81 @@
                 }
             )
 
+    def create_altered_embedded_indexes(self, column_prefix=None, parent_model=None):
+        option_name = operations.AddEmbeddedIndex.option_name
+        for app_label, model_name in sorted(self.kept_model_keys):
+            # old_model_name = self.renamed_models.get(
+            #     (app_label, model_name), model_name
+            # )
+            # old_parent_model_state = self.from_state.models[app_label, old_model_name]
+            new_parent_model_state = self.to_state.models[app_label, model_name]
+
+            for field_name in new_parent_model_state.fields:
+                field = new_parent_model_state.get_field(field_name)
+                if isinstance(field, EmbeddedModelField):
+                    parent_model = new_parent_model_state.name
+                    embedded_model = field.embedded_model
+                    column_prefix = f"{field_name}."
+                    embedded_model_name = embedded_model._meta.model_name
+
+                    # TODO: handle renamed embedded models
+                    old_model_state = self.from_state.models[
+                        embedded_model._meta.app_label, embedded_model_name
+                    ]
+                    new_model_state = self.to_state.models[
+                        embedded_model._meta.app_label, embedded_model_name
+                    ]
+
+                    old_indexes = old_model_state.options[option_name]
+                    new_indexes = new_model_state.options[option_name]
+                    added_indexes = [
+                        idx for idx in new_indexes if idx not in old_indexes
+                    ]
+                    # removed_indexes = [
+                    #     idx for idx in old_indexes if idx not in new_indexes
+                    # ]
+                    # renamed_indexes = []
+                    # Find renamed indexes.
+                    remove_from_added = []
+                    # remove_from_removed = []
+                    # for new_index in added_indexes:
+                    #     new_index_dec = new_index.deconstruct()
+                    #     new_index_name = new_index_dec[2].pop("name")
+                    #     for old_index in removed_indexes:
+                    #         old_index_dec = old_index.deconstruct()
+                    #         old_index_name = old_index_dec[2].pop("name")
+                    #         # Indexes are the same except for the names.
+                    #         if (
+                    #             new_index_dec == old_index_dec
+                    #             and new_index_name != old_index_name
+                    #         ):
+                    #             renamed_indexes.append((old_index_name, new_index_name, None))  # noqa [temp line length]
+                    #             remove_from_added.append(new_index)
+                    #             remove_from_removed.append(old_index)
+                    # Remove renamed indexes from the lists of added and removed
+                    # indexes.
+                    added_indexes = [
+                        idx for idx in added_indexes if idx not in remove_from_added
+                    ]
+                    # removed_indexes = [
+                    #     idx for idx in removed_indexes if idx not in remove_from_removed  # noqa [temp line length]
+                    # ]
+
+                    self.altered_embedded_indexes.update(
+                        {
+                            (
+                                app_label,
+                                embedded_model_name,
+                                column_prefix,
+                                parent_model,
+                            ): {
+                                "added_indexes": added_indexes,
+                                # "removed_indexes": removed_indexes,
+                                # "renamed_indexes": renamed_indexes,
+                            }
+                        }
+                    )
+
     def generate_added_indexes(self):
         for (app_label, model_name), alt_indexes in self.altered_indexes.items():
             dependencies = self._get_dependencies_for_model(app_label, model_name)
@@ -1474,6 +1560,24 @@
                     ),
                     dependencies=dependencies,
                 )
+        for (
+            app_label,
+            model_name,
+            column_prefix,
+            parent_model_name,
+        ), alt_indexes in self.altered_embedded_indexes.items():
+            dependencies = self._get_dependencies_for_model(app_label, model_name)
+            for index in alt_indexes["added_indexes"]:
+                self.add_operation(
+                    app_label,
+                    operations.AddEmbeddedIndex(
+                        model_name=model_name,
+                        index=index,
+                        column_prefix=column_prefix,
+                        parent_model_name=parent_model_name,
+                    ),
+                    dependencies=dependencies,
+                )
 
     def generate_removed_indexes(self):
         for (app_label, model_name), alt_indexes in self.altered_indexes.items():

django/db/migrations/operations/__init__.py

@@ -1,6 +1,7 @@
 from .fields import AddField, AlterField, RemoveField, RenameField
 from .models import (
     AddConstraint,
+    AddEmbeddedIndex,
     AddIndex,
     AlterConstraint,
     AlterIndexTogether,
@@ -43,4 +44,5 @@
     "RunPython",
     "AlterOrderWithRespectTo",
     "AlterModelManagers",
+    "AddEmbeddedIndex",
 ]

django/db/migrations/operations/models.py

@@ -948,6 +948,38 @@ def reduce(self, operation, app_label):
         return super().reduce(operation, app_label)
 
 
+class AddEmbeddedIndex(AddIndex):
+
+    def __init__(self, model_name, index, column_prefix, parent_model_name):
+        super().__init__(model_name, index)
+        self.column_prefix = column_prefix
+        self.parent_model_name = parent_model_name
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        model = to_state.apps.get_model(app_label, self.model_name)
+        if self.parent_model_name:
+            parent_model = to_state.apps.get_model(app_label, self.parent_model_name)
+        else:
+            parent_model = None
+        if self.allow_migrate_model(schema_editor.connection.alias, parent_model):
+            schema_editor.add_index(
+                model,
+                self.index,
+                column_prefix=self.column_prefix,
+                parent_model=parent_model,
+            )
+
+    def deconstruct(self):
+        name, args, kwargs = super().deconstruct()
+        kwargs.update(
+            {
+                "column_prefix": self.column_prefix,
+                "parent_model_name": self.parent_model_name,
+            }
+        )
+        return name, args, kwargs
+
+
 class RemoveIndex(IndexOperation):
     """Remove an index from a model."""