refactor.

Author: WaVEV

django_mongodb_backend/expressions/builtins.py

@@ -14,7 +14,6 @@
     Exists,
     ExpressionList,
     ExpressionWrapper,
-    Func,
     NegatedExpression,
     OrderBy,
     RawSQL,
@@ -25,12 +24,9 @@
     Value,
     When,
 )
-from django.db.models.fields.json import KeyTransform
 from django.db.models.sql import Query
 
-from django_mongodb_backend.fields.array import Array
-
-from ..query_utils import is_direct_value, process_lhs
+from ..query_utils import process_lhs
 
 
 def case(self, compiler, connection, as_path=False):
@@ -235,34 +231,14 @@ def value(self, compiler, connection, as_path=False):  # noqa: ARG001
     return value
 
 
-@staticmethod
-def _is_constant_value(value):
-    if isinstance(value, list | Array):
-        iterable = value.get_source_expressions() if isinstance(value, Array) else value
-        return all(_is_constant_value(e) for e in iterable)
-    if is_direct_value(value):
-        return True
-    return isinstance(value, Func | Value) and not (
-        value.contains_aggregate
-        or value.contains_over_clause
-        or value.contains_column_references
-        or value.contains_subquery
-    )
-
-
-@staticmethod
-def _is_simple_column(lhs):
-    while isinstance(lhs, KeyTransform):
-        if "." in getattr(lhs, "key_name", ""):
-            return False
-        lhs = lhs.lhs
-    col = lhs.source if isinstance(lhs, Ref) else lhs
-    # Foreign columns from parent cannot be addressed as single match
-    return isinstance(col, Col) and col.alias is not None
-
+def base_expression(self, compiler, connection, as_path=False):
+    if as_path and getattr(self, "is_simple_expression", lambda: False)():
+        self.is_simple_expression()
+        if hasattr(self, "as_mql_path"):
+            return self.as_mql_path(compiler, connection)
 
-def _is_simple_expression(self):
-    return self.is_simple_column(self.lhs) and self.is_constant_value(self.rhs)
+    expr = self.as_mql_expr(compiler, connection)
+    return {"$expr": expr} if as_path else expr
 
 
 def register_expressions():
@@ -283,6 +259,4 @@ def register_expressions():
     Subquery.as_mql = subquery
     When.as_mql = when
     Value.as_mql = value
-    BaseExpression.is_simple_expression = _is_simple_expression
-    BaseExpression.is_simple_column = _is_simple_column
-    BaseExpression.is_constant_value = _is_constant_value
+    BaseExpression.as_mql = base_expression

django_mongodb_backend/fields/json.py

@@ -16,8 +16,8 @@
     KeyTransformNumericLookupMixin,
 )
 
-from ..lookups import builtin_lookup
-from ..query_utils import process_lhs, process_rhs
+from ..lookups import builtin_lookup_expr, builtin_lookup_path
+from ..query_utils import is_simple_column, is_simple_expression, process_lhs, process_rhs
 
 
 def build_json_mql_path(lhs, key_transforms, as_path=False):
@@ -71,28 +71,35 @@ def _has_key_predicate(path, root_column=None, negated=False, as_path=False):
     return result
 
 
+def has_key_check_simple_expression(self):
+    return is_simple_expression(self) and all("." not in v for v in self.rhs)
+
+
 def has_key_lookup(self, compiler, connection, as_path=False):
     """Return MQL to check for the existence of a key."""
     rhs = self.rhs
     if not isinstance(rhs, (list, tuple)):
         rhs = [rhs]
-    as_path = as_path and self.is_simple_expression() and all("." not in v for v in rhs)
     lhs = process_lhs(self, compiler, connection, as_path=as_path)
     paths = []
     # Transform any "raw" keys into KeyTransforms to allow consistent handling
     # in the code that follows.
-
     for key in rhs:
         rhs_json_path = key if isinstance(key, KeyTransform) else KeyTransform(key, self.lhs)
         paths.append(rhs_json_path.as_mql(compiler, connection, as_path=as_path))
     keys = []
     for path in paths:
         keys.append(_has_key_predicate(path, lhs, as_path=as_path))
 
-    result = keys[0] if self.mongo_operator is None else {self.mongo_operator: keys}
-    if not as_path:
-        result = {"$expr": result}
-    return result
+    return keys[0] if self.mongo_operator is None else {self.mongo_operator: keys}
+
+
+def has_key_lookup_path(self, compiler, connection):
+    return has_key_lookup(self, compiler, connection, as_path=True)
+
+
+def has_key_lookup_expr(self, compiler, connection):
+    return has_key_lookup(self, compiler, connection, as_path=False)
 
 
 _process_rhs = JSONExact.process_rhs
@@ -107,7 +114,7 @@ def json_exact_process_rhs(self, compiler, connection):
     )
 
 
-def key_transform(self, compiler, connection, as_path=False):
+def key_transform_path(self, compiler, connection):
     """
     Return MQL for this KeyTransform (JSON path).
 
@@ -121,13 +128,19 @@ def key_transform(self, compiler, connection, as_path=False):
     while isinstance(previous, KeyTransform):
         key_transforms.insert(0, previous.key_name)
         previous = previous.lhs
-    if as_path and self.is_simple_column(self.lhs):
-        lhs_mql = previous.as_mql(compiler, connection, as_path=True)
-        return build_json_mql_path(lhs_mql, key_transforms, as_path=True)
+    # Collect all key transforms in order.
+    lhs_mql = previous.as_mql(compiler, connection, as_path=True)
+    return build_json_mql_path(lhs_mql, key_transforms, as_path=True)
+
+
+def key_transform_expr(self, compiler, connection):
+    key_transforms = [self.key_name]
+    previous = self.lhs
+    while isinstance(previous, KeyTransform):
+        key_transforms.insert(0, previous.key_name)
+        previous = previous.lhs
     # Collect all key transforms in order.
     lhs_mql = previous.as_mql(compiler, connection, as_path=False)
-    if as_path:
-        return {"$expr": build_json_mql_path(lhs_mql, key_transforms, as_path=False)}
     return build_json_mql_path(lhs_mql, key_transforms, as_path=False)
 
 
@@ -137,7 +150,7 @@ def key_transform_in(self, compiler, connection, as_path=False):
     set of specified values (rhs).
     """
     if as_path and self.is_simple_expression():
-        return builtin_lookup(self, compiler, connection, as_path=True)
+        return builtin_lookup_path(self, compiler, connection)
 
     lhs_mql = process_lhs(self, compiler, connection)
     # Traverse to the root column.
@@ -154,7 +167,24 @@ def key_transform_in(self, compiler, connection, as_path=False):
     return expr
 
 
-def key_transform_is_null(self, compiler, connection, as_path=False):
+def key_transform_in_path(self, compiler, connection):
+    return builtin_lookup_path(self, compiler, connection)
+
+
+def key_transform_in_expr(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    # Traverse to the root column.
+    previous = self.lhs
+    while isinstance(previous, KeyTransform):
+        previous = previous.lhs
+    root_column = previous.as_mql(compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    # Construct the expression to check if lhs_mql values are in rhs values.
+    expr = connection.mongo_operators_expr[self.lookup_name](lhs_mql, value)
+    return {"$and": [_has_key_predicate(lhs_mql, root_column), expr]}
+
+
+def key_transform_is_null_path(self, compiler, connection):
     """
     Return MQL to check the nullability of a key.
 
@@ -164,63 +194,77 @@ def key_transform_is_null(self, compiler, connection, as_path=False):
 
     Reference: https://code.djangoproject.com/ticket/32252
     """
-    if as_path and self.is_simple_expression():
-        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
-        rhs_mql = process_rhs(self, compiler, connection)
-        return _has_key_predicate(lhs_mql, None, negated=rhs_mql, as_path=True)
-    # Get the root column.
+    lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+    rhs_mql = process_rhs(self, compiler, connection, as_path=True)
+    return _has_key_predicate(lhs_mql, None, negated=rhs_mql, as_path=True)
+
+
+def key_transform_is_null_expr(self, compiler, connection):
     previous = self.lhs
     while isinstance(previous, KeyTransform):
         previous = previous.lhs
     root_column = previous.as_mql(compiler, connection)
-    expr = _has_key_predicate(lhs_mql, root_column, negated=rhs_mql)
-    if as_path:
-        return {"$expr": expr}
-    return expr
+    lhs_mql = process_lhs(self, compiler, connection, as_path=False)
+    rhs_mql = process_rhs(self, compiler, connection)
+    return _has_key_predicate(lhs_mql, root_column, negated=rhs_mql)
 
 
-def key_transform_numeric_lookup_mixin(self, compiler, connection, as_path=False):
+def key_transform_numeric_lookup_mixin_path(self, compiler, connection):
     """
     Return MQL to check if the field exists (i.e., is not "missing" or "null")
     and that the field matches the given numeric lookup expression.
     """
-    if as_path and self.is_simple_expression():
-        return builtin_lookup(self, compiler, connection, as_path=True)
+    return builtin_lookup_path(self, compiler, connection)
+
 
+def key_transform_numeric_lookup_mixin_expr(self, compiler, connection):
+    """
+    Return MQL to check if the field exists (i.e., is not "missing" or "null")
+    and that the field matches the given numeric lookup expression.
+    """
     lhs = process_lhs(self, compiler, connection, as_path=False)
-    expr = builtin_lookup(self, compiler, connection, as_path=False)
+    expr = builtin_lookup_expr(self, compiler, connection)
     # Check if the type of lhs is not "missing" or "null".
     not_missing_or_null = {"$not": {"$in": [{"$type": lhs}, ["missing", "null"]]}}
-    expr = {"$and": [expr, not_missing_or_null]}
-    if as_path:
-        return {"$expr": expr}
-    return expr
+    return {"$and": [expr, not_missing_or_null]}
 
 
-def key_transform_exact(self, compiler, connection, as_path=False):
-    if as_path and self.is_simple_expression():
-        lhs_mql = process_lhs(self, compiler, connection, as_path=True)
-        return {
-            "$and": [
-                builtin_lookup(self, compiler, connection, as_path=True),
-                _has_key_predicate(lhs_mql, None, as_path=True),
-            ]
-        }
-    if as_path:
-        return {"$expr": builtin_lookup(self, compiler, connection, as_path=False)}
-    return builtin_lookup(self, compiler, connection, as_path=False)
+def key_transform_exact_path(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection, as_path=True)
+    return {
+        "$and": [
+            builtin_lookup_path(self, compiler, connection),
+            _has_key_predicate(lhs_mql, None, as_path=True),
+        ]
+    }
+
+
+def key_transform_exact_expr(self, compiler, connection):
+    return builtin_lookup_expr(self, compiler, connection)
 
 
 def register_json_field():
     ContainedBy.as_mql = contained_by
     DataContains.as_mql = data_contains
     HasAnyKeys.mongo_operator = "$or"
     HasKey.mongo_operator = None
-    HasKeyLookup.as_mql = has_key_lookup
+    # HasKeyLookup.as_mql = has_key_lookup
+    HasKeyLookup.is_simple_expression = has_key_check_simple_expression
+    HasKeyLookup.as_mql_path = has_key_lookup_path
+    HasKeyLookup.as_mql_expr = has_key_lookup_expr
     HasKeys.mongo_operator = "$and"
     JSONExact.process_rhs = json_exact_process_rhs
-    KeyTransform.as_mql = key_transform
-    KeyTransformIn.as_mql = key_transform_in
-    KeyTransformIsNull.as_mql = key_transform_is_null
-    KeyTransformNumericLookupMixin.as_mql = key_transform_numeric_lookup_mixin
-    KeyTransformExact.as_mql = key_transform_exact
+    KeyTransform.is_simple_expression = is_simple_column
+    # KeyTransform.as_mql = key_transform
+    KeyTransform.as_mql_path = key_transform_path
+    KeyTransform.as_mql_expr = key_transform_expr
+    # KeyTransformIn.as_mql = key_transform_in
+    KeyTransformIn.as_mql_path = key_transform_in_path
+    KeyTransformIn.as_mql_expr = key_transform_in_expr
+    # KeyTransformIsNull.as_mql = key_transform_is_null
+    KeyTransformIsNull.as_mql_path = key_transform_is_null_path
+    KeyTransformIsNull.as_mql_expr = key_transform_is_null_expr
+    KeyTransformNumericLookupMixin.as_mql_path = key_transform_numeric_lookup_mixin_path
+    KeyTransformNumericLookupMixin.as_mql_expr = key_transform_numeric_lookup_mixin_expr
+    KeyTransformExact.as_mql_expr = key_transform_exact_expr
+    KeyTransformExact.as_mql_path = key_transform_exact_path

django_mongodb_backend/functions.py

@@ -125,6 +125,18 @@ def func(self, compiler, connection, as_path=False):
     return {f"${operator}": lhs_mql}
 
 
+def func_path(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError(f"{self} May need an as_mql_path() method.")
+
+
+def func_expr(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection, as_path=False)
+    if self.function is None:
+        raise NotSupportedError(f"{self} may need an as_mql() method.")
+    operator = MONGO_OPERATORS.get(self.__class__, self.function.lower())
+    return {f"${operator}": lhs_mql}
+
+
 def left(self, compiler, connection, as_path=False):
     return self.get_substr().as_mql(compiler, connection, as_path=as_path)
 
@@ -312,7 +324,8 @@ def register_functions():
     ConcatPair.as_mql = concat_pair
     Cot.as_mql = cot
     Extract.as_mql = extract
-    Func.as_mql = func
+    Func.as_mql_path = func_path
+    Func.as_mql_expr = func_expr
     JSONArray.as_mql = process_lhs
     Left.as_mql = left
     Length.as_mql = length