WIP

Author: WaVEV

django_mongodb_backend/expressions/builtins.py

@@ -77,6 +77,11 @@ def col(self, compiler, connection, as_path=False):  # noqa: ARG001
     return f"{prefix}{self.target.column}"
 
 
+@property
+def col_is_simple_column(self):
+    return self.alias is not None
+
+
 def col_pairs(self, compiler, connection, as_path=False):
     cols = self.get_cols()
     if len(cols) > 1:
@@ -170,6 +175,11 @@ def ref(self, compiler, connection, as_path=False):  # noqa: ARG001
     return f"{prefix}{refs}"
 
 
+@property
+def ref_is_simple_column(self):
+    return isinstance(self.source, Col) and self.source.alias is not None
+
+
 def star(self, compiler, connection, as_path=False):  # noqa: ARG001
     return {"$literal": True}
 
@@ -229,8 +239,11 @@ def base_expression(self, compiler, connection, as_path=False, **extra):
 
 
 def register_expressions():
+    BaseExpression.as_mql = base_expression
+    BaseExpression.is_simple_column = False
     Case.as_mql_expr = case
     Col.as_mql = col
+    Col.is_simple_column = col_is_simple_column
     ColPairs.as_mql = col_pairs
     CombinedExpression.as_mql_expr = combined_expression
     Exists.as_mql_expr = exists
@@ -241,9 +254,9 @@ def register_expressions():
     Query.as_mql = query
     RawSQL.as_mql = raw_sql
     Ref.as_mql = ref
+    Ref.is_simple_column = ref_is_simple_column
     ResolvedOuterRef.as_mql = ResolvedOuterRef.as_sql
     Star.as_mql = star
     Subquery.as_mql_expr = subquery
     When.as_mql = when
     Value.as_mql = value
-    BaseExpression.as_mql = base_expression

django_mongodb_backend/expressions/search.py

@@ -933,12 +933,15 @@ def __str__(self):
     def __repr__(self):
         return f"SearchText({self.lhs}, {self.rhs})"
 
-    def as_mql(self, compiler, connection, as_path=False):
-        lhs_mql = process_lhs(self, compiler, connection, as_path=as_path)
-        value = process_rhs(self, compiler, connection, as_path=as_path)
-        if as_path:
-            return {lhs_mql: {"$gte": value}}
-        return {"$expr": {"$gte": [lhs_mql, value]}}
+    def as_mql_expr(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection, as_path=False)
+        value = process_rhs(self, compiler, connection, as_path=False)
+        return {"$gte": [lhs_mql, value]}
+
+    def as_mql_path(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+        value = process_rhs(self, compiler, connection, as_path=True)
+        return {lhs_mql: {"$gte": value}}
 
 
 CharField.register_lookup(SearchTextLookup)

django_mongodb_backend/fields/array.py

@@ -230,9 +230,9 @@ def formfield(self, **kwargs):
 
 
 class Array(Func):
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_expr(self, compiler, connection):
         return [
-            expr.as_mql(compiler, connection, as_path=as_path)
+            expr.as_mql(compiler, connection, as_path=False)
             for expr in self.get_source_expressions()
         ]
 
@@ -254,24 +254,16 @@ def __init__(self, lhs, rhs):
 class ArrayContains(ArrayRHSMixin, FieldGetDbPrepValueMixin, Lookup):
     lookup_name = "contains"
 
-    def as_mql(self, compiler, connection, as_path=False):
-        if as_path and self.is_simple_expression():
-            lhs_mql = process_lhs(self, compiler, connection, as_path=as_path)
-            value = process_rhs(self, compiler, connection, as_path=as_path)
-            if value is None:
-                return False
-            return {lhs_mql: {"$all": value}}
+    def as_mql_path(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+        value = process_rhs(self, compiler, connection, as_path=True)
+        if value is None:
+            return False
+        return {lhs_mql: {"$all": value}}
+
+    def as_mql_expr(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection, as_path=False)
         value = process_rhs(self, compiler, connection, as_path=False)
-        expr = {
-            "$and": [
-                {"$ne": [lhs_mql, None]},
-                {"$ne": [value, None]},
-                {"$setIsSubset": [value, lhs_mql]},
-            ]
-        }
-        if as_path:
-            return {"$expr": expr}
         return {
             "$and": [
                 {"$ne": [lhs_mql, None]},
@@ -285,19 +277,16 @@ def as_mql(self, compiler, connection, as_path=False):
 class ArrayContainedBy(ArrayRHSMixin, FieldGetDbPrepValueMixin, Lookup):
     lookup_name = "contained_by"
 
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_expr(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection, as_path=False)
         value = process_rhs(self, compiler, connection, as_path=False)
-        expr = {
+        return {
             "$and": [
                 {"$ne": [lhs_mql, None]},
                 {"$ne": [value, None]},
                 {"$setIsSubset": [lhs_mql, value]},
             ]
         }
-        if as_path:
-            return {"$expr": expr}
-        return expr
 
 
 @ArrayField.register_lookup
@@ -344,36 +333,30 @@ def get_subquery_wrapping_pipeline(self, compiler, connection, field_name, expr)
             },
         ]
 
-    def as_mql(self, compiler, connection, as_path=False):
-        if as_path and self.is_simple_expression():
-            lhs_mql = process_lhs(self, compiler, connection, as_path=True)
-            value = process_rhs(self, compiler, connection, as_path=True)
-            return {lhs_mql: {"$in": value}}
+    def as_mql_path(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+        value = process_rhs(self, compiler, connection, as_path=True)
+        return {lhs_mql: {"$in": value}}
 
+    def as_mql_expr(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection, as_path=False)
         value = process_rhs(self, compiler, connection, as_path=False)
-        expr = {
+        return {
             "$and": [
                 {"$ne": [lhs_mql, None]},
                 {"$size": {"$setIntersection": [value, lhs_mql]}},
             ]
         }
-        if as_path:
-            return {"$expr": expr}
-        return expr
 
 
 @ArrayField.register_lookup
 class ArrayLenTransform(Transform):
     lookup_name = "len"
     output_field = IntegerField()
 
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_expr(self, compiler, connection, as_path=False):
         lhs_mql = process_lhs(self, compiler, connection, as_path=False)
-        expr = {"$cond": {"if": {"$isArray": lhs_mql}, "then": {"$size": lhs_mql}, "else": None}}
-        if as_path:
-            return {"$expr": expr}
-        return expr
+        return {"$cond": {"if": {"$isArray": lhs_mql}, "then": {"$size": lhs_mql}, "else": None}}
 
 
 @ArrayField.register_lookup
@@ -398,15 +381,20 @@ def __init__(self, index, base_field, *args, **kwargs):
         self.index = index
         self.base_field = base_field
 
-    def as_mql(self, compiler, connection, as_path=False):
-        if as_path and self.is_simple_column(self.lhs):
-            lhs_mql = process_lhs(self, compiler, connection, as_path=as_path)
-            return f"{lhs_mql}.{self.index}"
+    def is_simple_expression(self):
+        return self.is_simple_column
+
+    @property
+    def is_simple_column(self):
+        return self.lhs.is_simple_column
+
+    def as_mql_path(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+        return f"{lhs_mql}.{self.index}"
+
+    def as_mql_expr(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection, as_path=False)
-        expr = {"$arrayElemAt": [lhs_mql, self.index]}
-        if as_path:
-            return {"$expr": expr}
-        return expr
+        return {"$arrayElemAt": [lhs_mql, self.index]}
 
     @property
     def output_field(self):
@@ -428,7 +416,7 @@ def __init__(self, start, end, *args, **kwargs):
         self.start = start
         self.end = end
 
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_expr(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection)
         return {"$slice": [lhs_mql, self.start, self.end]}
 

django_mongodb_backend/fields/embedded_model.py

@@ -5,6 +5,7 @@
 from django.db import models
 from django.db.models.fields.related import lazy_related_operation
 from django.db.models.lookups import Transform
+from django.utils.functional import cached_property
 
 from .. import forms
 
@@ -165,6 +166,18 @@ def __init__(self, key_name, ref_field, *args, **kwargs):
     def get_lookup(self, name):
         return self.ref_field.get_lookup(name)
 
+    def is_simple_expression(self):
+        return self.is_simple_column
+
+    @cached_property
+    def is_simple_column(self):
+        previous = self
+        while isinstance(previous, KeyTransform):
+            if not previous.key_name.isalnum():
+                return False
+            previous = previous.lhs
+        return previous.is_simple_column
+
     def get_transform(self, name):
         """
         Validate that `name` is either a field of an embedded model or a
@@ -184,16 +197,22 @@ def get_transform(self, name):
             f"{suggestion}"
         )
 
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_path(self, compiler, connection):
+        previous = self
+        key_transforms = []
+        while isinstance(previous, KeyTransform):
+            key_transforms.insert(0, previous.key_name)
+            previous = previous.lhs
+        mql = previous.as_mql(compiler, connection, as_path=True)
+        mql_path = ".".join(key_transforms)
+        return f"{mql}.{mql_path}"
+
+    def as_mql_expr(self, compiler, connection):
         previous = self
         key_transforms = []
         while isinstance(previous, KeyTransform):
             key_transforms.insert(0, previous.key_name)
             previous = previous.lhs
-        if as_path:
-            mql = previous.as_mql(compiler, connection, as_path=True)
-            mql_path = ".".join(key_transforms)
-            return f"{mql}.{mql_path}"
         mql = previous.as_mql(compiler, connection)
         for key in key_transforms:
             mql = {"$getField": {"input": mql, "field": key}}

django_mongodb_backend/fields/embedded_model_array.py

@@ -5,6 +5,7 @@
 from django.db.models.expressions import Col
 from django.db.models.fields.related import lazy_related_operation
 from django.db.models.lookups import Lookup, Transform
+from django.utils.functional import cached_property
 
 from .. import forms
 from ..query_utils import process_lhs, process_rhs
@@ -128,7 +129,7 @@ def process_rhs(self, compiler, connection, as_path=False):
             for v in value
         ]
 
-    def as_mql(self, compiler, connection, as_path=False):
+    def as_mql_expr(self, compiler, connection):
         # Querying a subfield within the array elements (via nested
         # KeyTransform). Replicate MongoDB's implicit ANY-match by mapping over
         # the array and applying $in on the subfield.
@@ -138,7 +139,7 @@ def as_mql(self, compiler, connection, as_path=False):
         lhs_mql["$ifNull"][0]["$map"]["in"] = connection.mongo_operators_expr[self.lookup_name](
             inner_lhs_mql, values
         )
-        return {"$expr": {"$anyElementTrue": lhs_mql}}
+        return {"$anyElementTrue": lhs_mql}
 
 
 @_EmbeddedModelArrayOutputField.register_lookup
@@ -236,13 +237,25 @@ def __init__(self, key_name, array_field, *args, **kwargs):
         column_name = f"$item.{key_name}"
         column_target.db_column = column_name
         column_target.set_attributes_from_name(column_name)
-        self._lhs = Col(None, column_target)
+        self._lhs = Col(False, column_target)
         self._sub_transform = None
 
     def __call__(self, this, *args, **kwargs):
         self._lhs = self._sub_transform(self._lhs, *args, **kwargs)
         return self
 
+    def is_simple_expression(self):
+        return self.is_simple_column
+
+    @cached_property
+    def is_simple_column(self):
+        previous = self
+        while isinstance(previous, KeyTransform):
+            if not previous.key_name.isalnum():
+                return False
+            previous = previous.lhs
+        return previous.is_simple_column and self._lhs.is_simple_column
+
     def get_lookup(self, name):
         return self.output_field.get_lookup(name)
 
@@ -275,11 +288,12 @@ def get_transform(self, name):
             f"{suggestion}"
         )
 
-    def as_mql(self, compiler, connection, as_path=False):
-        if as_path:
-            inner_lhs_mql = self._lhs.as_mql(compiler, connection, as_path=True)
-            lhs_mql = process_lhs(self, compiler, connection, as_path=True)
-            return f"{inner_lhs_mql}.{lhs_mql}"
+    def as_mql_path(self, compiler, connection):
+        inner_lhs_mql = self._lhs.as_mql(compiler, connection, as_path=True).removeprefix("$item.")
+        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+        return f"{lhs_mql}.{inner_lhs_mql}"
+
+    def as_mql_expr(self, compiler, connection):
         inner_lhs_mql = self._lhs.as_mql(compiler, connection)
         lhs_mql = process_lhs(self, compiler, connection)
         return {

django_mongodb_backend/fields/json.py

@@ -18,7 +18,7 @@
 )
 
 from ..lookups import builtin_lookup_expr, builtin_lookup_path
-from ..query_utils import is_simple_column, is_simple_expression, process_lhs, process_rhs
+from ..query_utils import process_lhs, process_rhs
 
 
 def build_json_mql_path(lhs, key_transforms, as_path=False):
@@ -73,7 +73,8 @@ def _has_key_predicate(path, root_column=None, negated=False, as_path=False):
 
 
 def has_key_check_simple_expression(self):
-    return is_simple_expression(self) and all("." not in v for v in self.rhs)
+    rhs = [self.rhs] if not isinstance(self.rhs, (list, tuple)) else self.rhs
+    return self.is_simple_column and all(key.isalnum() for key in rhs)
 
 
 def has_key_lookup(self, compiler, connection, as_path=False):
@@ -225,17 +226,27 @@ def key_transform_exact_expr(self, compiler, connection):
     return builtin_lookup_expr(self, compiler, connection)
 
 
+def keytransform_is_simple_column(self):
+    previous = self
+    while isinstance(previous, KeyTransform):
+        if not previous.key_name.isalnum():
+            return False
+        previous = previous.lhs
+    return previous.is_simple_column
+
+
 def register_json_field():
     ContainedBy.as_mql = contained_by
     DataContains.as_mql = data_contains
     HasAnyKeys.mongo_operator = "$or"
     HasKey.mongo_operator = None
-    HasKeyLookup.is_simple_expression = has_key_check_simple_expression
     HasKeyLookup.as_mql_path = partialmethod(has_key_lookup, as_path=True)
     HasKeyLookup.as_mql_expr = partialmethod(has_key_lookup, as_path=False)
+    HasKeyLookup.is_simple_expression = has_key_check_simple_expression
     HasKeys.mongo_operator = "$and"
     JSONExact.process_rhs = json_exact_process_rhs
-    KeyTransform.is_simple_expression = is_simple_column
+    KeyTransform.is_simple_column = property(keytransform_is_simple_column)
+    KeyTransform.is_simple_expression = keytransform_is_simple_column
     KeyTransform.as_mql_path = partialmethod(key_transform, as_path=True)
     KeyTransform.as_mql_expr = partialmethod(key_transform, as_path=False)
     KeyTransformIn.as_mql_path = key_transform_in_path