Refactor jsonfield as $expr.

Author: WaVEV

django_mongodb/__init__.py

@@ -7,13 +7,13 @@
 check_django_compatability()
 
 from .expressions import register_expressions  # noqa: E402
-from .fields import load_fields  # noqa: E402
+from .fields import register_fields  # noqa: E402
 from .functions import register_functions  # noqa: E402
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
 
 register_expressions()
+register_fields()
 register_functions()
 register_lookups()
 register_nodes()
-load_fields()

django_mongodb/base.py

@@ -44,6 +44,7 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         "IntegerField": "int",
         "BigIntegerField": "long",
         "GenericIPAddressField": "string",
+        "JSONField": "object",
         "OneToOneField": "int",
         "PositiveBigIntegerField": "int",
         "PositiveIntegerField": "long",

django_mongodb/expressions.py

@@ -1,14 +1,19 @@
-from django.db.models.expressions import Col, Value
+from django.db.models.expressions import Col, ExpressionWrapper, Value
 
 
 def col(self, compiler, connection):  # noqa: ARG001
     return f"${self.target.column}"
 
 
+def expression_wrapper(self, compiler, connection):
+    return self.expression.as_mql(compiler, connection)
+
+
 def value(self, compiler, connection):  # noqa: ARG001
     return {"$literal": self.value}
 
 
 def register_expressions():
     Col.as_mql = col
+    ExpressionWrapper.as_mql = expression_wrapper
     Value.as_mql = value

django_mongodb/features.py

@@ -386,11 +386,12 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "model_fields.test_jsonfield.JSONFieldTests.test_db_check_constraints",
         },
         "Mongodb's Null behaviour is different from sql's": {
-            "model_fields.test_jsonfield.TestQuerying.test_none_key_and_exact_lookup",
+            "model_fields.test_jsonfield.TestQuerying.test_expression_wrapper_key_transform",
             "model_fields.test_jsonfield.TestSaveLoad.test_json_null_different_from_sql_null",
-            "model_fields.test_jsonfield.TestQuerying.test_none_key",
             "model_fields.test_jsonfield.TestQuerying.test_lookup_exclude",
             "model_fields.test_jsonfield.TestQuerying.test_lookup_exclude_nonexistent_key",
+            "model_fields.test_jsonfield.TestQuerying.test_none_key",
+            "model_fields.test_jsonfield.TestQuerying.test_none_key_exclude",
         },
         "Pipeline filtering": {"model_fields.test_jsonfield.TestQuerying.test_icontains"},
     }

django_mongodb/fields/__init__.py

@@ -1,3 +1,4 @@
 from .auto import MongoAutoField
+from .json_field import register_fields
 
-__all__ = ["MongoAutoField"]
+__all__ = ["MongoAutoField", "register_fields"]

django_mongodb/fields/auto.py

@@ -1,25 +1,8 @@
 from bson import ObjectId, errors
 from django.core import exceptions
-from django.db import NotSupportedError
-from django.db.models import JSONField
 from django.db.models.fields import AutoField, Field
-from django.db.models.fields.json import (
-    ContainedBy,
-    DataContains,
-    HasAnyKeys,
-    HasKey,
-    HasKeyLookup,
-    HasKeys,
-    JSONExact,
-    KeyTransform,
-    KeyTransformIn,
-    KeyTransformIsNull,
-)
 from django.utils.translation import gettext_lazy as _
 
-from .base import DatabaseWrapper
-from .query_utils import process_lhs
-
 
 class MongoAutoField(AutoField):
     default_error_messages = {
@@ -49,130 +32,3 @@ def to_python(self, value):
                 code="invalid",
                 params={"value": value},
             ) from None
-
-
-_from_db_value = JSONField.from_db_value
-
-
-def from_db_value(self, value, expression, connection):
-    return (
-        value
-        if isinstance(connection, DatabaseWrapper)
-        else _from_db_value(self, value, expression, connection)
-    )
-
-
-def json_process_rhs(node, compiler, connection):
-    rhs = node.rhs
-    if hasattr(rhs, "as_mql"):
-        return rhs.as_mql(compiler, connection)
-    _, value = node.process_rhs(compiler, connection)
-    lookup_name = node.lookup_name
-    if lookup_name not in ("in", "range"):
-        value = value[0] if len(value) > 0 else []
-
-    return value
-
-
-def key_transform(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
-    lhs_mql = previous.as_mql(compiler, connection)
-    return ".".join([lhs_mql, *key_transforms])
-
-
-def data_contains(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError("contains lookup is not supported on this database backend.")
-
-
-def contained_by(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError("contained_by lookup is not supported on this database backend.")
-
-
-def json_exact(self, compiler, connection):
-    lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
-    rhs_mql = json_process_rhs(self, compiler, connection)
-    if rhs_mql == "null":
-        return {"$or": [{lhs_mql: {"$eq": None}}, {lhs_mql: {"$exists": False}}]}
-    return {lhs_mql: {"$eq": rhs_mql, "$exists": True}}
-
-
-def key_transform_isnull(self, compiler, connection):
-    lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
-    rhs_mql = json_process_rhs(self, compiler, connection)
-
-    # https://code.djangoproject.com/ticket/32252
-    return {lhs_mql: {"$exists": not rhs_mql}}
-
-
-def key_transform_in(self, compiler, connection):
-    lhs_mql = key_transform_agg(self.lhs, compiler, connection)
-    bare_lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
-    value = json_process_rhs(self, compiler, connection)
-    expr = connection.mongo_aggregations[self.lookup_name](lhs_mql, value)
-    return {"$expr": expr, bare_lhs_mql: {"$exists": True}}
-
-
-def has_key_lookup(self, compiler, connection):
-    lhs = process_lhs(self, compiler, connection, bare_column_ref=True)
-    rhs = self.rhs
-    if not isinstance(rhs, (list | tuple)):
-        rhs = [rhs]
-    paths = []
-    for key in rhs:
-        if isinstance(key, KeyTransform):
-            *_, rhs_key_transforms = key.preprocess_lhs(compiler, connection)
-            rhs_key_transforms = ".".join(rhs_key_transforms)
-        else:
-            rhs_key_transforms = str(key)
-        rhs_json_path = f"{lhs}.{rhs_key_transforms}"
-        paths.append(rhs_json_path)
-
-    keys = []
-    for path in paths:
-        keys.append({path: {"$exists": True}})
-    if self.mongo_operator is None:
-        assert len(keys) == 1
-        return keys[0]
-    return {self.mongo_operator: keys}
-
-
-def key_transform_agg(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
-    lhs_mql = previous.as_mql(compiler, connection)
-    result = f"{lhs_mql}"
-    for key in key_transforms:
-        get_field = {"$getField": {"input": result, "field": key}}
-        if key.isdigit():
-            result = {
-                "$cond": {
-                    "if": {"$isArray": result},
-                    "then": {"$arrayElemAt": [result, int(key)]},
-                    "else": get_field,
-                }
-            }
-        else:
-            result = get_field
-    return result
-
-
-def load_fields():
-    JSONField.from_db_value = from_db_value
-    DataContains.as_mql = data_contains
-    KeyTransform.as_mql = key_transform
-    KeyTransform.as_mql_agg = key_transform_agg
-    JSONExact.as_mql = json_exact
-    ContainedBy.as_mql = contained_by
-    HasKeyLookup.as_mql = has_key_lookup
-    HasAnyKeys.mongo_operator = "$or"
-    HasKey.mongo_operator = None
-    HasKeys.mongo_operator = "$and"
-    KeyTransformIsNull.as_mql = key_transform_isnull
-    KeyTransformIn.as_mql = key_transform_in

django_mongodb/fields/json_field.py

@@ -0,0 +1,151 @@
+from django.db import NotSupportedError
+from django.db.models import JSONField
+from django.db.models.fields.json import (
+    ContainedBy,
+    DataContains,
+    HasAnyKeys,
+    HasKey,
+    HasKeyLookup,
+    HasKeys,
+    JSONExact,
+    KeyTransform,
+    KeyTransformIn,
+    KeyTransformIsNull,
+    KeyTransformNumericLookupMixin,
+)
+
+from django_mongodb.lookups import builtin_lookup
+from django_mongodb.query_utils import process_lhs, process_rhs
+
+
+def _key_transform_root(self, compiler, connection):
+    previous = self.lhs
+    while isinstance(previous, KeyTransform):
+        previous = previous.lhs
+    return previous.as_mql(compiler, connection)
+
+
+def contained_by(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError("contained_by lookup is not supported on this database backend.")
+
+
+def data_contains(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError("contains lookup is not supported on this database backend.")
+
+
+_from_db_value = JSONField.from_db_value
+
+
+def from_db_value(self, value, expression, connection):
+    """
+    Mongodb does not need to change the json value. It is store as it is.
+    """
+    return (
+        value
+        if connection.vendor == "mongodb"
+        else _from_db_value(self, value, expression, connection)
+    )
+
+
+def has_key_lookup(self, compiler, connection):
+    rhs = self.rhs
+    lhs = process_lhs(self, compiler, connection)
+    if not isinstance(rhs, (list | tuple)):
+        rhs = [rhs]
+    paths = []
+    for key in rhs:
+        if isinstance(key, KeyTransform):
+            rhs_json_path = key.as_mql(compiler, connection)
+        else:
+            rhs_json_path = KeyTransform(key, self.lhs).as_mql(compiler, connection)
+        paths.append(rhs_json_path)
+
+    keys = []
+    for path in paths:
+        keys.append({"$and": [{"$ne": [{"$type": path}, "missing"]}, {"$ne": [lhs, None]}]})
+    if self.mongo_operator is None:
+        assert len(keys) == 1
+        return keys[0]
+    return {self.mongo_operator: keys}
+
+
+__process_rhs = JSONExact.process_rhs
+
+
+def json_process_rhs(self, compiler, connection):
+    """
+    Django does some transformation over the parameters if it is null.
+    to avoid ambiguity, we keep using the parent process_rhs.
+    """
+    return (
+        super(JSONExact, self).process_rhs(compiler, connection)
+        if connection.vendor == "mongodb"
+        else __process_rhs(self, compiler, connection)
+    )
+
+
+def key_transform(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
+    lhs_mql = previous.as_mql(compiler, connection)
+    result = lhs_mql
+    for key in key_transforms:
+        get_field = {"$getField": {"input": result, "field": key}}
+        if key.isdigit() and str(int(key)) == key:
+            result = {
+                "$cond": {
+                    "if": {"$isArray": result},
+                    "then": {"$arrayElemAt": [result, int(key)]},
+                    "else": get_field,
+                }
+            }
+        else:
+            result = get_field
+    return result
+
+
+def key_transform_in(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    expr = connection.mongo_operators[self.lookup_name](lhs_mql, value)
+    return {"$and": [expr, {"$not": {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}}]}
+
+
+def key_transform_isnull(self, compiler, connection):
+    """
+    The KeyTransformIsNull lookup borrows the logic from HasKey for isnull=False.
+    If isnull=True, the query should only match objects that do not have the key.
+    # https://code.djangoproject.com/ticket/32252
+    """
+    lhs_mql = process_lhs(self, compiler, connection)
+    rhs_mql = process_rhs(self, compiler, connection)
+    root_column = _key_transform_root(self, compiler, connection)
+
+    result = {"$or": [{"$in": [{"$type": lhs_mql}, ["missing"]]}, {"$eq": [root_column, None]}]}
+    if not rhs_mql:
+        result = {"$not": result}
+    return result
+
+
+def key_transform_numeric_lookup_mixin(self, compiler, connection):
+    expr = builtin_lookup(self, compiler, connection)
+    lhs = process_lhs(self, compiler, connection)
+    return {"$and": [expr, {"$not": {"$in": [{"$type": lhs}, ["missing", "null"]]}}]}
+
+
+def register_fields():
+    ContainedBy.as_mql = contained_by
+    DataContains.as_mql = data_contains
+    HasAnyKeys.mongo_operator = "$or"
+    HasKey.mongo_operator = None
+    HasKeyLookup.as_mql = has_key_lookup
+    HasKeys.mongo_operator = "$and"
+    JSONExact.process_rhs = json_process_rhs
+    JSONField.from_db_value = from_db_value
+    KeyTransform.as_mql = key_transform
+    KeyTransformIn.as_mql = key_transform_in
+    KeyTransformIsNull.as_mql = key_transform_isnull
+    KeyTransformNumericLookupMixin.as_mql = key_transform_numeric_lookup_mixin

django_mongodb/functions.py

@@ -64,7 +64,10 @@ def cast(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)[0]
     if max_length := self.output_field.max_length:
         lhs_mql = {"$substrCP": [lhs_mql, 0, max_length]}
-    lhs_mql = {"$convert": {"input": lhs_mql, "to": output_type}}
+    if output_type == "object":
+        lhs_mql = {"$convert": {"input": lhs_mql, "to": output_type, "onError": lhs_mql}}
+    else:
+        lhs_mql = {"$convert": {"input": lhs_mql, "to": output_type}}
     if decimal_places := getattr(self.output_field, "decimal_places", None):
         lhs_mql = {"$trunc": [lhs_mql, decimal_places]}
     return lhs_mql

django_mongodb/lookups.py

@@ -1,6 +1,11 @@
 from django.db import NotSupportedError
 from django.db.models.fields.related_lookups import In, MultiColSource, RelatedIn
-from django.db.models.lookups import BuiltinLookup, IsNull, UUIDTextMixin
+from django.db.models.lookups import (
+    BuiltinLookup,
+    FieldGetDbPrepValueIterableMixin,
+    IsNull,
+    UUIDTextMixin,
+)
 
 from .query_utils import process_lhs, process_rhs