Support index definition on Embedded Models in top level model.

Author: WaVEV

django_mongodb_backend/indexes.py

@@ -3,6 +3,7 @@
 
 from django.core.checks import Error, Warning
 from django.db import NotSupportedError
+from django.db.backends.utils import names_digest, split_identifier
 from django.db.models import FloatField, Index, IntegerField
 from django.db.models.lookups import BuiltinLookup
 from django.db.models.sql.query import Query
@@ -13,6 +14,7 @@
 from django_mongodb_backend.fields import ArrayField
 
 from .query_utils import process_rhs
+from .utils import get_column_path
 
 MONGO_INDEX_OPERATORS = {
     "exact": "$eq",
@@ -61,7 +63,7 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False
             filter_expression[column].update({"$type": field.db_type(schema_editor.connection)})
         else:
             for field_name, _ in self.fields_orders:
-                field_ = model._meta.get_field(field_name)
+                field_ = get_column_path(model, field_name)
                 filter_expression[field_.column].update(
                     {"$type": field_.db_type(schema_editor.connection)}
                 )
@@ -74,7 +76,7 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False
             # order is "" if ASCENDING or "DESC" if DESCENDING (see
             # django.db.models.indexes.Index.fields_orders).
             (
-                column_prefix + model._meta.get_field(field_name).column,
+                column_prefix + get_column_path(model, field_name).column,
                 ASCENDING if order == "" else DESCENDING,
             )
             for field_name, order in self.fields_orders
@@ -154,7 +156,7 @@ def get_pymongo_index_model(
         for field_name, _ in self.fields_orders:
             field = model._meta.get_field(field_name)
             type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
-            field_path = column_prefix + model._meta.get_field(field_name).column
+            field_path = column_prefix + get_column_path(model, field_name)
             fields[field_path] = {"type": type_}
         return SearchIndexModel(
             definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name
@@ -264,7 +266,7 @@ def get_pymongo_index_model(
         fields = []
         for field_name, _ in self.fields_orders:
             field_ = model._meta.get_field(field_name)
-            field_path = column_prefix + model._meta.get_field(field_name).column
+            field_path = column_prefix + get_column_path(model, field_name)
             mappings = {"path": field_path}
             if isinstance(field_, ArrayField):
                 mappings.update(
@@ -280,8 +282,40 @@ def get_pymongo_index_model(
         return SearchIndexModel(definition={"fields": fields}, name=self.name, type="vectorSearch")
 
 
+def set_name_with_model(self, model):
+    """
+    Generate a unique name for the index.
+
+    The name is divided into 3 parts - table name (12 chars), field name
+    (8 chars) and unique hash + suffix (10 chars). Each part is made to
+    fit its size by truncating the excess length.
+    """
+    _, table_name = split_identifier(model._meta.db_table)
+    column_names = [
+        get_column_path(model, field_name).column for field_name, order in self.fields_orders
+    ]
+    column_names_with_order = [
+        (f"-{column_name}" if order else column_name)
+        for column_name, (field_name, order) in zip(column_names, self.fields_orders, strict=False)
+    ]
+    # The length of the parts of the name is based on the default max
+    # length of 30 characters.
+    hash_data = [table_name, *column_names_with_order, self.suffix]
+    self.name = (
+        f"{table_name[:11]}_{column_names[0][:7]}_"
+        f"{names_digest(*hash_data, length=6)}_{self.suffix}"
+    )
+    if len(self.name) > self.max_name_length:
+        raise ValueError(
+            "Index too long for multiple database support. Is self.suffix longer than 3 characters?"
+        )
+    if self.name[0] == "_" or self.name[0].isdigit():
+        self.name = f"D{self.name[1:]}"
+
+
 def register_indexes():
     BuiltinLookup.as_mql_idx = builtin_lookup_idx
     Index._get_condition_mql = _get_condition_mql
     Index.get_pymongo_index_model = get_pymongo_index_model
+    Index.set_name_with_model = set_name_with_model
     WhereNode.as_mql_idx = where_node_idx

django_mongodb_backend/lookups.py

@@ -1,4 +1,6 @@
+import django.db.models.base as base
 from django.db import NotSupportedError
+from django.db.models.constants import LOOKUP_SEP
 from django.db.models.fields.related_lookups import In, RelatedIn
 from django.db.models.lookups import (
     BuiltinLookup,
@@ -8,6 +10,7 @@
     UUIDTextMixin,
 )
 
+from .fields import EmbeddedModelField
 from .query_utils import process_lhs, process_rhs
 
 
@@ -121,6 +124,26 @@ def uuid_text_mixin(self, compiler, connection):  # noqa: ARG001
     raise NotSupportedError("Pattern lookups on UUIDField are not supported.")
 
 
+class Options(base.Options):
+    def get_field(self, field_name):
+        if LOOKUP_SEP in field_name:
+            previous = self
+            keys = field_name.split(LOOKUP_SEP)
+            path = []
+            for field in keys:
+                field = base.Options.get_field(previous, field)
+                if isinstance(field, EmbeddedModelField):
+                    previous = field.embedded_model._meta
+                else:
+                    previous = field
+                path.append(field.column)
+            column = ".".join(path)
+            embedded_column = field.clone()
+            embedded_column.column = column
+            return embedded_column
+        return super().get_field(field_name)
+
+
 def register_lookups():
     BuiltinLookup.as_mql = builtin_lookup
     FieldGetDbPrepValueIterableMixin.resolve_expression_parameter = (
@@ -131,3 +154,4 @@ def register_lookups():
     IsNull.as_mql = is_null
     PatternLookup.prep_lookup_value_mongo = pattern_lookup_prep_lookup_value
     UUIDTextMixin.as_mql = uuid_text_mixin
+    base.Options = Options

django_mongodb_backend/schema.py

@@ -7,7 +7,7 @@
 from .fields import EmbeddedModelField
 from .gis.schema import GISSchemaEditor
 from .query import wrap_database_errors
-from .utils import OperationCollector
+from .utils import OperationCollector, get_column_path
 
 
 def ignore_embedded_models(func):
@@ -249,7 +249,7 @@ def alter_unique_together(
             )
         # Created uniques
         for field_names in news.difference(olds):
-            columns = [model._meta.get_field(field).column for field in field_names]
+            columns = [get_column_path(model, field).column for field in field_names]
             name = str(
                 self._unique_constraint_name(
                     model._meta.db_table, [column_prefix + col for col in columns]

django_mongodb_backend/utils.py

@@ -5,6 +5,7 @@
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured, ValidationError
 from django.db.backends.utils import logger
+from django.db.models.constants import LOOKUP_SEP
 from django.utils.functional import SimpleLazyObject
 from django.utils.text import format_lazy
 from django.utils.version import get_version_tuple
@@ -186,3 +187,21 @@ def wrapper(self, *args, **kwargs):
             self.log(method, args, kwargs)
 
         return wrapper
+
+
+def get_column_path(model, field_name):
+    from .fields import EmbeddedModelField  # noqa: PLC0415
+
+    if LOOKUP_SEP in field_name:
+        previous = model
+        keys = field_name.split(LOOKUP_SEP)
+        path = []
+        for field in keys:
+            field = previous._meta.get_field(field)
+            previous = field.embedded_model if isinstance(field, EmbeddedModelField) else field
+            path.append(field.column)
+        column = ".".join(path)
+        embedded_column = field.clone()
+        embedded_column.column = column
+        return embedded_column
+    return model._meta.get_field(field_name)

tests/schema_/test_embedded_model.py

@@ -519,6 +519,168 @@ class Meta:
         self.assertTableNotExists(Author)
 
 
+class EmbeddedModelsTopLevelIndexTest(TestMixin, TransactionTestCase):
+    @isolate_apps("schema_")
+    def test_unique_together(self):
+        """Meta.unique_together defined at the top-level for embedded fields."""
+
+        class Address(EmbeddedModel):
+            unique_together_one = models.CharField(max_length=10)
+            unique_together_two = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Author(EmbeddedModel):
+            address = EmbeddedModelField(Address)
+            unique_together_three = models.CharField(max_length=10)
+            unique_together_four = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Book(models.Model):
+            author = EmbeddedModelField(Author)
+
+            class Meta:
+                app_label = "schema_"
+                unique_together = [
+                    ("author__unique_together_three", "author__unique_together_four"),
+                    (
+                        "author__address__unique_together_one",
+                        "author__address__unique_together_two",
+                    ),
+                ]
+
+        with connection.schema_editor() as editor:
+            editor.create_model(Book)
+            self.assertTableExists(Book)
+            # Embedded uniques are created from top-level definition.
+            self.assertEqual(
+                self.get_constraints_for_columns(
+                    Book, ["author.unique_together_three", "author.unique_together_four"]
+                ),
+                [
+                    "schema__book_author.unique_together_three_author.unique_together_four_09a570b8_uniq"
+                ],
+            )
+            self.assertEqual(
+                self.get_constraints_for_columns(
+                    Book,
+                    ["author.address.unique_together_one", "author.address.unique_together_two"],
+                ),
+                [
+                    "schema__book_author.address.unique_together_one_author.address.unique_together_two_2c2d1477_uniq"
+                ],
+            )
+            editor.delete_model(Book)
+        self.assertTableNotExists(Book)
+
+    @isolate_apps("schema_")
+    def test_add_remove_field_indexes(self):
+        """AddField/RemoveField + EmbeddedModelField + Meta.indexes at top-level."""
+
+        class Address(EmbeddedModel):
+            indexed_one = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Author(EmbeddedModel):
+            address = EmbeddedModelField(Address)
+            indexed_two = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Book(models.Model):
+            author = EmbeddedModelField(Author)
+
+            class Meta:
+                app_label = "schema_"
+                indexes = [
+                    models.Index(fields=["author__indexed_two"]),
+                    models.Index(fields=["author__address__indexed_one"]),
+                ]
+
+        new_field = EmbeddedModelField(Author)
+        new_field.set_attributes_from_name("author")
+
+        with connection.schema_editor() as editor:
+            # Create the table and add the field.
+            editor.create_model(Book)
+            editor.add_field(Book, new_field)
+            # Embedded indexes are created.
+            self.assertEqual(
+                self.get_constraints_for_columns(Book, ["author.indexed_two"]),
+                ["schema__boo_author._333c90_idx"],
+            )
+            self.assertEqual(
+                self.get_constraints_for_columns(
+                    Book,
+                    ["author.address.indexed_one"],
+                ),
+                ["schema__boo_author._f54386_idx"],
+            )
+            editor.delete_model(Book)
+        self.assertTableNotExists(Book)
+
+    @isolate_apps("schema_")
+    def test_add_remove_field_constraints(self):
+        """AddField/RemoveField + EmbeddedModelField + Meta.constraints at top-level."""
+
+        class Address(EmbeddedModel):
+            unique_constraint_one = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Author(EmbeddedModel):
+            address = EmbeddedModelField(Address)
+            unique_constraint_two = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "schema_"
+
+        class Book(models.Model):
+            author = EmbeddedModelField(Author)
+
+            class Meta:
+                app_label = "schema_"
+                constraints = [
+                    models.UniqueConstraint(
+                        fields=["author__unique_constraint_two"],
+                        name="unique_two",
+                    ),
+                    models.UniqueConstraint(
+                        fields=["author__address__unique_constraint_one"],
+                        name="unique_one",
+                    ),
+                ]
+
+        new_field = EmbeddedModelField(Author)
+        new_field.set_attributes_from_name("author")
+
+        with connection.schema_editor() as editor:
+            # Create the table and add the field.
+            editor.create_model(Book)
+            editor.add_field(Book, new_field)
+            # Embedded constraints are created.
+            self.assertEqual(
+                self.get_constraints_for_columns(Book, ["author.unique_constraint_two"]),
+                ["unique_two"],
+            )
+            self.assertEqual(
+                self.get_constraints_for_columns(
+                    Book,
+                    ["author.address.unique_constraint_one"],
+                ),
+                ["unique_one"],
+            )
+            editor.delete_model(Book)
+        self.assertTableNotExists(Book)
+
+
 class EmbeddedModelsIgnoredTests(TestMixin, TransactionTestCase):
     def test_embedded_not_created(self):
         """create_model() and delete_model() ignore EmbeddedModel."""