Use as_mql_idx when creating indexes.

Author: WaVEV

django_mongodb/base.py

@@ -119,6 +119,15 @@ def _isnull_operator(a, b):
         "iregex": lambda a, b: regex_match(a, b, insensitive=True),
     }
 
+    mongo_operators_idx = {
+        "exact": lambda a, b: {a: {"$eq": b}},
+        "gt": lambda a, b: {a: {"$gt": b}},
+        "gte": lambda a, b: {a: {"$gte": b}},
+        "lt": lambda a, b: {a: {"$lt": b}},
+        "lte": lambda a, b: {a: {"$lte": b}},
+        "in": lambda a, b: {a: {"$in": b}},
+    }
+
     display_name = "MongoDB"
     vendor = "mongodb"
     Database = Database

django_mongodb/indexes.py

@@ -7,16 +7,7 @@ def _get_condition_mql(self, model, schema_editor):
     query = Query(model=model, alias_cols=False)
     where = query.build_where(self.condition)
     compiler = query.get_compiler(connection=schema_editor.connection)
-    mql_ = where.as_mql(compiler, schema_editor.connection)
-    # Transform aggregate() query syntax into find() syntax.
-    mql = {}
-    for key in mql_:
-        col, value = mql_[key]
-        # multiple conditions don't work yet
-        if isinstance(col, dict):
-            return {}
-        mql[col.lstrip("$")] = {key: value}
-    return mql
+    return where.as_mql_idx(compiler, schema_editor.connection)
 
 
 def register_indexes():

django_mongodb/lookups.py

@@ -1,4 +1,5 @@
 from django.db import NotSupportedError
+from django.db.models.expressions import Col
 from django.db.models.fields.related_lookups import In, MultiColSource, RelatedIn
 from django.db.models.lookups import (
     BuiltinLookup,
@@ -17,6 +18,14 @@ def builtin_lookup(self, compiler, connection):
     return connection.mongo_operators[self.lookup_name](lhs_mql, value)
 
 
+def builtin_lookup_idx(self, compiler, connection):
+    if not isinstance(self.lhs, Col):
+        raise NotSupportedError("Expressions as indexes are not supported.")
+    lhs_mql = self.lhs.target.column
+    value = process_rhs(self, compiler, connection)
+    return connection.mongo_operators_idx[self.lookup_name](lhs_mql, value)
+
+
 _field_resolve_expression_parameter = FieldGetDbPrepValueIterableMixin.resolve_expression_parameter
 
 
@@ -93,6 +102,7 @@ def uuid_text_mixin(self, compiler, connection):  # noqa: ARG001
 
 def register_lookups():
     BuiltinLookup.as_mql = builtin_lookup
+    BuiltinLookup.as_mql_idx = builtin_lookup_idx
     FieldGetDbPrepValueIterableMixin.resolve_expression_parameter = (
         field_resolve_expression_parameter
     )

django_mongodb/query.py

@@ -297,8 +297,52 @@ def where_node(self, compiler, connection):
     return mql
 
 
+def where_node_idx(self, compiler, connection):
+    if self.connector == AND:
+        full_needed, empty_needed = len(self.children), 1
+    else:
+        full_needed, empty_needed = 1, len(self.children)
+    if self.connector == AND:
+        operator = "$and"
+    elif self.connector == XOR:
+        raise NotSupportedError("Xor in indexes is not supported")
+    else:
+        operator = "$or"
+    if self.negated:
+        raise NotSupportedError("Negated field in indexes is not supported")
+    children_mql = []
+    for child in self.children:
+        try:
+            mql = child.as_mql_idx(compiler, connection)
+        except EmptyResultSet:
+            empty_needed -= 1
+        except FullResultSet:
+            full_needed -= 1
+        else:
+            if mql:
+                children_mql.append(mql)
+            else:
+                full_needed -= 1
+
+    if empty_needed == 0:
+        raise (FullResultSet if self.negated else EmptyResultSet)
+    if full_needed == 0:
+        raise (EmptyResultSet if self.negated else FullResultSet)
+
+    if len(children_mql) == 1:
+        mql = children_mql[0]
+    elif len(children_mql) > 1:
+        mql = {operator: children_mql}
+    else:
+        mql = {}
+    if not mql:
+        raise FullResultSet
+    return mql
+
+
 def register_nodes():
     ExtraWhere.as_mql = extra_where
     Join.as_mql = join
     NothingNode.as_mql = NothingNode.as_sql
     WhereNode.as_mql = where_node
+    WhereNode.as_mql_idx = where_node_idx