INTPYTHON-698 Add EmbeddedFieldIndex and EmbeddedFieldUniqueConstraint

django_mongodb_backend/constraints.py

@@ -1,10 +1,13 @@
 from collections import defaultdict
 
+from django.core import checks
+from django.core.exceptions import FieldDoesNotExist
 from django.db.models import UniqueConstraint
 from pymongo import ASCENDING
 from pymongo.operations import IndexModel
 
-from .indexes import _get_condition_mql
+from .fields import EmbeddedModelArrayField
+from .indexes import EmbeddedFieldIndexMixin, _get_condition_mql, get_field
 
 
 def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefix=""):
@@ -25,23 +28,53 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefi
             filter_expression[column].update({"$type": field.db_type(schema_editor.connection)})
         else:
             for field_name in self.fields:
-                field_ = model._meta.get_field(field_name)
+                field_ = get_field(model, field_name)
                 filter_expression[field_.column].update(
-                    {"$type": field_.db_type(schema_editor.connection)}
+                    {"$type": field_.field.db_type(schema_editor.connection)}
                 )
     if filter_expression:
         kwargs["partialFilterExpression"] = filter_expression
     index_orders = (
         [(column_prefix + field.column, ASCENDING)]
         if field
         else [
-            (column_prefix + model._meta.get_field(field_name).column, ASCENDING)
+            (column_prefix + get_field(model, field_name).column, ASCENDING)
             for field_name in self.fields
         ]
     )
     return IndexModel(index_orders, name=self.name, unique=True, **kwargs)
 
 
+class EmbeddedFieldUniqueConstraint(EmbeddedFieldIndexMixin, UniqueConstraint):
+    meta_option_name = "constraints"
+
+    def check(self, model, connection):
+        errors = super().check(model, connection)
+        # Due to MongoDB limitation (https://jira.mongodb.org/browse/SERVER-17853),
+        # EmbeddedFieldUniqueConstraint cannot reference subfields of
+        # EmbeddedModelArrayField unless nulls_distinct=False (so that the
+        # index doesn't need to be created with a partialFilterExpression).
+        if self.nulls_distinct is not False:
+            for field_name in self.fields:
+                # Get the top-level field, removing paths to embedded fields.
+                local_field_name, *_ = field_name.split(".")
+                try:
+                    field = model._meta.get_field(local_field_name)
+                except FieldDoesNotExist:
+                    continue
+                if isinstance(field, EmbeddedModelArrayField):
+                    errors.append(
+                        checks.Error(
+                            f"EmbeddedFieldUniqueConstraint {self.name!r} must "
+                            "have nulls_distinct=False since it references "
+                            f"EmbeddedModelArrayField '{local_field_name}'.",
+                            obj=model,
+                            id="django_mongodb_backend.constraints.E001",
+                        )
+                    )
+        return errors
+
+
 def register_constraints():
     UniqueConstraint.get_pymongo_index_model = get_pymongo_index_model
     UniqueConstraint._get_condition_mql = _get_condition_mql

django_mongodb_backend/indexes.py

@@ -1,16 +1,22 @@
 import itertools
-from collections import defaultdict
+from collections import defaultdict, namedtuple
 
 from django.core.checks import Error, Warning
+from django.core.exceptions import FieldDoesNotExist
 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
 from django.db.models.sql.where import AND, XOR, WhereNode
 from pymongo import ASCENDING, DESCENDING
 from pymongo.operations import IndexModel, SearchIndexModel
 
-from django_mongodb_backend.fields import ArrayField
+from django_mongodb_backend.fields import (
+    ArrayField,
+    EmbeddedModelArrayField,
+    EmbeddedModelField,
+)
 
 from .query_utils import process_rhs
 
@@ -23,6 +29,8 @@
     "in": "$in",
 }
 
+FieldColumn = namedtuple("FieldColumn", ["field", "column"])
+
 
 def _get_condition_mql(self, model, schema_editor):
     """Analogous to Index._get_condition_sql()."""
@@ -44,6 +52,25 @@ def builtin_lookup_idx(self, compiler, connection):
     return {lhs_mql: {operator: value}}
 
 
+def get_field(model, field_name):
+    """
+    A version of Model_.meta.get_field() that can retrieve embedded model
+    fields.
+    """
+    path = []
+    base_model = model
+    *parents, leaf = field_name.split(".")
+    for name in parents:
+        field = model._meta.get_field(name)
+        path.append(field.column)
+        if not isinstance(field, EmbeddedModelField | EmbeddedModelArrayField):
+            raise FieldDoesNotExist(f"{base_model.__name__} has no field named '{field_name}'.")
+        model = field.embedded_model
+    field = model._meta.get_field(leaf)
+    path.append(field.column)
+    return FieldColumn(field, ".".join(path))
+
+
 def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefix=""):
     """Return a pymongo IndexModel for this Django Index."""
     if self.contains_expressions:
@@ -61,7 +88,7 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefi
             # 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_field(model, field_name).column,
                 ASCENDING if order == "" else DESCENDING,
             )
             for field_name, order in self.fields_orders
@@ -92,6 +119,82 @@ def where_node_idx(self, compiler, connection):
     return mql
 
 
+class EmbeddedFieldIndexMixin:
+    def check(self, model, connection):
+        # The parent check reports E012 (nonexistent fields) based on top-level
+        # fields. These errors are discarded and the checks redone, accounting
+        # for embedded model fields.
+        errors = [e for e in super().check(model, connection) if e.id != "models.E012"]
+        return errors + self._check_local_fields(model)
+
+    def _check_local_fields(self, model):
+        errors = []
+        forward_fields = self._get_forward_fields(model)
+        for field_name in self.fields:
+            if field_name not in forward_fields:
+                errors.append(
+                    Error(
+                        f"'{self.meta_option_name}' refers to the nonexistent "
+                        f"field '{field_name}'.",
+                        obj=model,
+                        id="models.E012",
+                    )
+                )
+        return errors
+
+    def _get_forward_fields(self, model):
+        """
+        Return the set of forward field paths for the given model, including
+        embedded fields of EmbeddedModelField/EmbeddedModelArrayField.
+        """
+        forward_fields = set()
+        for field in model._meta._get_fields(reverse=False):
+            # Recurse into embedded models and flatten their forward fields
+            # using dotted paths (e.g. "field.subfield").
+            if isinstance(field, EmbeddedModelField | EmbeddedModelArrayField):
+                sub_forward_fields = self._get_forward_fields(field.embedded_model)
+                for column in sub_forward_fields:
+                    forward_fields.add(f"{field.name}.{column}")
+                    if hasattr(field, "attname"):
+                        forward_fields.add(f"{field.attname}.{column}")
+            # For each field, both the public field name and its attribute name
+            # (attname), when present, are included in order to mirror Django's
+            # field resolution.
+            forward_fields.add(field.name)
+            if hasattr(field, "attname"):
+                forward_fields.add(field.attname)
+        return forward_fields
+
+
+class EmbeddedFieldIndex(EmbeddedFieldIndexMixin, Index):
+    meta_option_name = "indexes"
+
+    def set_name_with_model(self, model):
+        """
+        Generate a unique name for the index.
+
+        The name is divided into three parts: the table name, the field name,
+        and a hash plus suffix. Each part is truncated to fit within Django's
+        index name length constraints.
+
+        This method overrides Django's base implementation which uses
+        model._meta.get_field(field_name).column, which doesn't work for
+        embedded fields.
+        """
+        _, table_name = split_identifier(model._meta.db_table)
+        field_names_with_order = [
+            (f"-{column_name}" if order else column_name)
+            for column_name, order in self.fields_orders
+        ]
+        hash_data = [table_name, *field_names_with_order, self.suffix]
+        self.name = (
+            f"{table_name[:11]}_{self.fields_orders[0][0][:7]}_"
+            f"{names_digest(*hash_data, length=6)}_{self.suffix}"
+        )
+        if self.name[0] == "_" or self.name[0].isdigit():
+            self.name = f"D{self.name[1:]}"
+
+
 class SearchIndex(Index):
     suffix = "six"
     _error_id_prefix = "django_mongodb_backend.indexes.SearchIndex"

docs/ref/models/constraints.rst

@@ -22,3 +22,32 @@ If you wish to only allow one document with a ``NULL`` value, use a
 
     Support for :attr:`UniqueConstraint.nulls_distinct
     <django.db.models.UniqueConstraint.nulls_distinct>` was added.
+
+MongoDB-specific constraints
+============================
+
+.. module:: django_mongodb_backend.constraints
+   :synopsis: Database constraints for MongoDB.
+
+Some MongoDB-specific :doc:`constraints <django:ref/models/constraints>`, for
+use on a model's :attr:`Meta.constraints<django.db.models.Options.constraints>`
+option, are available in ``django_mongodb_backend.constraints``.
+
+Embedded field constraints
+==========================
+
+``EmbeddedFieldConstraint``
+---------------------------
+
+.. class:: EmbeddedFieldUniqueConstraint(**kwargs)
+
+    .. versionadded:: 6.0.2
+
+    Subclass of :class:`~django.db.models.UniqueConstraint` for use on a
+    top-level model in order to add a unique constraint on subfields of
+    :class:`~.fields.EmbeddedModelField` and
+    :class:`~.fields.EmbeddedModelArrayField`.
+
+    The ``fields`` argument uses dotted paths to reference embedded fields. For
+    examples, see :ref:`embedded-model-field-unique-constraints` and
+    :ref:`embedded-model-array-field-unique-constraints`.

docs/ref/models/fields.rst

@@ -319,8 +319,8 @@ These indexes use 0-based indexing.
     After you create a model with an ``EmbeddedModelField`` or add an
     ``EmbeddedModelField`` to an existing model, no further updates to the
     embedded model will be made. Using the models above as an example, if you
-    created these models and then added an indexed field to ``Address``,
-    the index created in the nested ``Book`` embed is not created.
+    created these models and then add or remove a field from ``Address``,
+    existing documents won't be updated.
 
 ``EmbeddedModelArrayField``
 ---------------------------

docs/ref/models/indexes.rst

@@ -9,6 +9,25 @@ Some MongoDB-specific :doc:`indexes <django:ref/models/indexes>`, for use on a
 model's :attr:`Meta.indexes <django.db.models.Options.indexes>` option, are
 available in ``django_mongodb_backend.indexes``.
 
+Embedded field indexes
+======================
+
+``EmbeddedFieldIndex``
+----------------------
+
+.. class:: EmbeddedFieldIndex(**kwargs)
+
+    .. versionadded:: 6.0.2
+
+    Subclass of :class:`~django.db.models.Index` for use on a top-level model
+    in order to add an index on subfields of
+    :class:`~.fields.EmbeddedModelField` and
+    :class:`~.fields.EmbeddedModelArrayField`.
+
+    The ``fields`` argument uses dotted paths to reference embedded fields. For
+    examples, see :ref:`embedded-model-field-indexes` and
+    :ref:`embedded-model-array-field-indexes`.
+
 Search indexes
 ==============
 

docs/releases/6.0.x.rst

@@ -16,6 +16,16 @@ New features
 - Added support for :attr:`UniqueConstraint.nulls_distinct
   <django.db.models.UniqueConstraint.nulls_distinct>`. See
   :ref:`unique-constraints` for MongoDB-specific details.
+- Added :class:`~django_mongodb_backend.indexes.EmbeddedFieldIndex` and
+  :class:`~django_mongodb_backend.constraints.EmbeddedFieldUniqueConstraint`
+  to index and create unique constraints on subfields of
+  :class:`~.fields.EmbeddedModelField` and
+  :class:`~.fields.EmbeddedModelArrayField`.
+
+  See :ref:`embedded-model-field-indexes`,
+  :ref:`embedded-model-field-unique-constraints`,
+  :ref:`embedded-model-array-field-indexes`, and
+  :ref:`embedded-model-array-field-unique-constraints`.
 
 Bug fixes
 ---------

docs/topics/embedded-models.rst

@@ -64,6 +64,44 @@ address with the city "New York"::
 
     >>> Customer.objects.filter(address__city="New York")
 
+.. _embedded-model-field-indexes:
+
+Indexing ``EmbeddedModelField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 6.0.2
+
+You can create indexes on ``EmbeddedModelField`` subfields by using
+:class:`~django_mongodb_backend.indexes.EmbeddedFieldIndex` on the top-level
+model and referencing field names using dotted paths.
+
+For example, if model ``Customer`` has a ``EmbeddedModelField`` named
+``address`` and the embedded model has field ``zip_code``, use this on
+``Customer``\'s :attr:`Meta.indexes <django.db.models.Options.indexes>`
+option::
+
+    EmbeddedFieldIndex(fields=["address.zip_code"])]
+
+.. _embedded-model-field-unique-constraints:
+
+Unique constraints on ``EmbeddedModelField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 6.0.2
+
+You can enforce uniqueness of ``EmbeddedModelField`` subfields by using
+:class:`~django_mongodb_backend.constraints.EmbeddedFieldUniqueConstraint`
+on the top-level model and referencing field names using dotted paths.
+
+For example, if model ``Product`` has ``EmbeddedModelField`` ``product_info``
+and the embedded model has field ``sku``, use this on ``Product``\'s
+:attr:`Meta.constraints <django.db.models.Options.constraints>` option::
+
+    EmbeddedFieldUniqueConstraint(
+        fields=["product_info.sku"],
+        name="unique_product_sku",
+    )
+
 .. _embedded-model-array-field-example:
 
 ``EmbeddedModelArrayField``
@@ -184,6 +222,37 @@ For example, if the ``Tag`` model had an ``EmbeddedModelArrayField`` called
     ...
     ValueError: Cannot perform multiple levels of array traversal in a query.
 
+.. _embedded-model-array-field-indexes:
+
+Indexing ``EmbeddedModelArrayField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 6.0.2
+
+You can create indexes on ``EmbeddedModelArrayField`` subfields :ref:`in the
+same way <embedded-model-field-indexes>` as for ``EmbeddedModelField``.
+
+.. _embedded-model-array-field-unique-constraints:
+
+Unique constraints on ``EmbeddedModelArrayField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 6.0.2
+
+You can enforce uniqueness of ``EmbeddedModelArrayField`` subfields :ref:`in
+the same way <embedded-model-field-unique-constraints>` as for
+``EmbeddedModelField``. These constraints validate that each element across all
+of the arrays in all documents is unique (not just that the elements within
+each document's array are unique).
+
+The only difference from ``EmbeddedModelField`` unique constraints is that
+MongoDB doesn't support treating ``NULL`` values as distinct within
+``EmbeddedModelArrayField`` (the default behavior for ``EmbeddedModelField``
+constraints). You must use :attr:`UniqueConstraint(..., nulls_distinct=False)
+<django.db.models.UniqueConstraint.nulls_distinct>` for indexes that reference
+a subfield of ``EmbeddedModelArrayField``. Only one ``NULL`` value across all
+of the documents' arrays will be allowed.
+
 .. _polymorphic-embedded-model-field-example:
 
 ``PolymorphicEmbeddedModelField``