Support simple expressions in indexes.

Author: WaVEV

django_mongodb_backend/expressions/builtins.py

@@ -23,6 +23,7 @@
     Value,
     When,
 )
+from django.db.models.indexes import IndexExpression
 from django.db.models.sql import Query
 
 from ..query_utils import process_lhs
@@ -95,12 +96,25 @@ def expression_wrapper(self, compiler, connection):
     return self.expression.as_mql(compiler, connection)
 
 
+def index_expression(self, compiler, connection):
+    result = []
+    for expr in self.get_source_expressions():
+        if expr is None:
+            continue
+        for sub_expr in expr.get_source_expressions():
+            try:
+                result.append(sub_expr.as_mql(compiler, connection, as_path=True))
+            except FullResultSet:
+                result.append(Value(True).as_mql(compiler, connection))
+    return result
+
+
 def negated_expression(self, compiler, connection):
     return {"$not": expression_wrapper(self, compiler, connection)}
 
 
-def order_by(self, compiler, connection):
-    return self.expression.as_mql(compiler, connection)
+def order_by(self, compiler, connection, **extra_args):
+    return self.expression.as_mql(compiler, connection, **extra_args)
 
 
 def query(self, compiler, connection, get_wrapping_pipeline=None):
@@ -217,6 +231,7 @@ def register_expressions():
     Exists.as_mql = exists
     ExpressionList.as_mql = process_lhs
     ExpressionWrapper.as_mql = expression_wrapper
+    IndexExpression.as_mql = index_expression
     NegatedExpression.as_mql = negated_expression
     OrderBy.as_mql = order_by
     Query.as_mql = query

django_mongodb_backend/features.py

@@ -29,7 +29,7 @@ class DatabaseFeatures(GISFeatures, BaseDatabaseFeatures):
     supports_deferrable_unique_constraints = False
     supports_explaining_query_execution = True
     supports_expression_defaults = False
-    supports_expression_indexes = False
+    supports_expression_indexes = True
     supports_foreign_keys = False
     supports_ignore_conflicts = False
     supports_json_field_contains = False

django_mongodb_backend/fields/embedded_model.py

@@ -155,16 +155,6 @@ def formfield(self, **kwargs):
             }
         )
 
-    def contribute_to_class(self, cls, name):
-        super().contribute_to_class(cls, name)
-        for field in self.embedded_model._meta.local_fields:
-            embedded_field_name = f"{name}.{field.name}"
-            field = field.clone()
-            field.unique = False
-            field.db_index = False
-            field.name = embedded_field_name
-            models.Field.contribute_to_class(field, cls, embedded_field_name, private_only=False)
-
 
 class KeyTransform(Transform):
     def __init__(self, key_name, ref_field, *args, **kwargs):
@@ -213,6 +203,9 @@ def as_mql(self, compiler, connection, as_path=False):
     def output_field(self):
         return self.ref_field
 
+    def db_type(self, connection):
+        return self.output_field.db_type(connection)
+
 
 class KeyTransformFactory:
     def __init__(self, key_name, ref_field):

django_mongodb_backend/indexes.py

@@ -4,6 +4,7 @@
 from django.core.checks import Error, Warning
 from django.db import NotSupportedError
 from django.db.models import FloatField, Index, IntegerField
+from django.db.models.expressions import F, OrderBy
 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
@@ -46,10 +47,31 @@ def builtin_lookup_idx(self, compiler, connection):
 
 def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False, column_prefix=""):
     """Return a pymongo IndexModel for this Django Index."""
+    filter_expression = defaultdict(dict)
+    expressions_fields = []
     if self.contains_expressions:
-        return None
+        for expression in self.expressions:
+            if not isinstance(expression, F | OrderBy):
+                continue
+            query = Query(model=model, alias_cols=False)
+            field_ = expression.resolve_expression(query)
+            compiler = query.get_compiler(connection=schema_editor.connection)
+            column = field_.as_mql(compiler, schema_editor.connection, as_path=True)
+            db_type = (
+                field_.expression.db_type(schema_editor.connection)
+                if isinstance(field_, OrderBy)
+                else field_.db_type(schema_editor.connection)
+            )
+            if unique:
+                filter_expression[column].update({"$type": db_type})
+            order = (
+                DESCENDING
+                if isinstance(expression, OrderBy) and expression.descending
+                else ASCENDING
+            )
+            expressions_fields.append((column, order))
+
     kwargs = {}
-    filter_expression = defaultdict(dict)
     if self.condition:
         filter_expression.update(self._get_condition_mql(model, schema_editor))
     if unique:
@@ -80,7 +102,7 @@ def get_pymongo_index_model(self, model, schema_editor, field=None, unique=False
             for field_name, order in self.fields_orders
         ]
     )
-    return IndexModel(index_orders, name=self.name, **kwargs)
+    return IndexModel(expressions_fields + index_orders, name=self.name, **kwargs)
 
 
 def where_node_idx(self, compiler, connection):

django_mongodb_backend/schema.py

@@ -355,6 +355,7 @@ def add_constraint(self, model, constraint, field=None, column_prefix="", parent
             nulls_distinct=constraint.nulls_distinct,
         ):
             idx = Index(
+                *constraint.expressions,
                 fields=constraint.fields,
                 name=constraint.name,
                 condition=constraint.condition,

tests/schema_/test_embedded_model.py

@@ -1,6 +1,7 @@
 import itertools
 
 from django.db import connection, models
+from django.db.models.expressions import F
 from django.test import TransactionTestCase, skipUnlessDBFeature
 from django.test.utils import isolate_apps
 
@@ -544,11 +545,18 @@ class Book(models.Model):
 
             class Meta:
                 app_label = "schema_"
-                unique_together = [
-                    ("author.unique_together_three", "author.unique_together_four"),
+                constraints = [
+                    models.UniqueConstraint(
+                        F("author__unique_together_three").asc(),
+                        F("author__unique_together_four").desc(),
+                        name="unique_together_34",
+                    ),
                     (
-                        "author.address.unique_together_one",
-                        "author.address.unique_together_two",
+                        models.UniqueConstraint(
+                            F("author__address__unique_together_one"),
+                            F("author__address__unique_together_two").asc(),
+                            name="unique_together_12",
+                        )
                     ),
                 ]
 
@@ -560,18 +568,14 @@ class Meta:
                 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"
-                ],
+                ["unique_together_34"],
             )
             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"
-                ],
+                ["unique_together_12"],
             )
             editor.delete_model(Book)
         self.assertTableNotExists(Book)
@@ -599,8 +603,8 @@ class Book(models.Model):
             class Meta:
                 app_label = "schema_"
                 indexes = [
-                    models.Index(fields=["author.indexed_two"]),
-                    models.Index(fields=["author.address.indexed_one"]),
+                    models.Index(F("author__indexed_two").asc(), name="indexed_two"),
+                    models.Index(F("author__address__indexed_one").asc(), name="indexed_one"),
                 ]
 
         new_field = EmbeddedModelField(Author)
@@ -613,14 +617,14 @@ class Meta:
             # Embedded indexes are created.
             self.assertEqual(
                 self.get_constraints_for_columns(Book, ["author.indexed_two"]),
-                ["schema__boo_author._333c90_idx"],
+                ["indexed_two"],
             )
             self.assertEqual(
                 self.get_constraints_for_columns(
                     Book,
                     ["author.address.indexed_one"],
                 ),
-                ["schema__boo_author._f54386_idx"],
+                ["indexed_one"],
             )
             editor.delete_model(Book)
         self.assertTableNotExists(Book)
@@ -648,13 +652,9 @@ class Book(models.Model):
             class Meta:
                 app_label = "schema_"
                 constraints = [
+                    models.UniqueConstraint(F("author__unique_constraint_two"), name="unique_two"),
                     models.UniqueConstraint(
-                        fields=["author.unique_constraint_two"],
-                        name="unique_two",
-                    ),
-                    models.UniqueConstraint(
-                        fields=["author.address.unique_constraint_one"],
-                        name="unique_one",
+                        F("author__address__unique_constraint_one"), name="unique_one"
                     ),
                 ]