EMF support json field.

Author: timgraham

django_mongodb_backend/fields/embedded_model.py

@@ -7,6 +7,7 @@
 from django.db.models.lookups import Transform
 
 from .. import forms
+from ..query_utils import key_transform_build_path
 
 
 class EmbeddedModelField(models.Field):
@@ -181,18 +182,24 @@ def get_transform(self, name):
         return result
 
     def preprocess_lhs(self, compiler, connection):
-        key_transforms = [self.key_name]
-        previous = self.lhs
+        previous = self
+        embedded_key_transforms = []
+        json_key_transforms = []
         while isinstance(previous, KeyTransform):
-            key_transforms.insert(0, previous.key_name)
+            if isinstance(previous.ref_field, EmbeddedModelField):
+                embedded_key_transforms.insert(0, previous.key_name)
+            else:
+                json_key_transforms.insert(0, previous.key_name)
             previous = previous.lhs
         mql = previous.as_mql(compiler, connection)
-        return mql, key_transforms
+        embedded_key_transforms.append(json_key_transforms.pop(0))
+        return mql, embedded_key_transforms, json_key_transforms
 
     def as_mql(self, compiler, connection):
-        mql, key_transforms = self.preprocess_lhs(compiler, connection)
+        mql, key_transforms, json_key_transforms = self.preprocess_lhs(compiler, connection)
         transforms = ".".join(key_transforms)
-        return f"{mql}.{transforms}"
+        result = f"{mql}.{transforms}"
+        return key_transform_build_path(json_key_transforms, result)
 
 
 class KeyTransformFactory:

django_mongodb_backend/fields/json.py

@@ -14,7 +14,7 @@
 )
 
 from ..lookups import builtin_lookup
-from ..query_utils import process_lhs, process_rhs
+from ..query_utils import key_transform_build_path, process_lhs, process_rhs
 
 
 def contained_by(self, compiler, connection):  # noqa: ARG001
@@ -89,23 +89,7 @@ def key_transform(self, compiler, connection):
         key_transforms.insert(0, previous.key_name)
         previous = previous.lhs
     lhs_mql = previous.as_mql(compiler, connection)
-    result = lhs_mql
-    # Build the MQL path using the collected key transforms.
-    for key in key_transforms:
-        get_field = {"$getField": {"input": result, "field": key}}
-        # Handle array indexing if the key is a digit. If key is something
-        # like '001', it's not an array index despite isdigit() returning True.
-        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
+    return key_transform_build_path(key_transforms, lhs_mql)
 
 
 def key_transform_in(self, compiler, connection):

django_mongodb_backend/query_utils.py

@@ -7,6 +7,26 @@ def is_direct_value(node):
     return not hasattr(node, "as_sql")
 
 
+def key_transform_build_path(key_trasnforms, lhs):
+    # Build the MQL path using the collected key transforms.
+    result = lhs
+    for key in key_trasnforms:
+        get_field = {"$getField": {"input": result, "field": key}}
+        # Handle array indexing if the key is a digit. If key is something
+        # like '001', it's not an array index despite isdigit() returning True.
+        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 process_lhs(node, compiler, connection):
     if not hasattr(node, "lhs"):
         # node is a Func or Expression, possibly with multiple source expressions.

tests/model_fields_/models.py

@@ -103,6 +103,7 @@ class Data(EmbeddedModel):
     integer = models.IntegerField(db_column="custom_column")
     auto_now = models.DateTimeField(auto_now=True)
     auto_now_add = models.DateTimeField(auto_now_add=True)
+    json_value = models.JSONField(default=dict)
 
 
 class Address(EmbeddedModel):

tests/model_fields_/test_embedded_model.py

@@ -139,6 +139,35 @@ def test_order_and_group_by_embedded_field(self):
             ],
         )
 
+    def test_embedded_with_json_field(self):
+        models = []
+        for i in range(4):
+            m = Holder.objects.create(
+                data=Data(json_value={"field1": i * 5, "field2": {"0": {"value": list(range(i))}}})
+            )
+            models.append(m)
+
+        all_models = Holder.objects.all()
+
+        self.assertCountEqual(
+            Holder.objects.filter(data__json_value__field2__0__value__0=0),
+            models[1:],
+        )
+        self.assertCountEqual(
+            Holder.objects.filter(data__json_value__field2__0__value__1=1),
+            models[2:],
+        )
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field2__0__value__1=5), [])
+
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field1__lt=100), all_models)
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field1__gt=100), [])
+        self.assertCountEqual(
+            Holder.objects.filter(
+                data__json_value__field1__gte=5, data__json_value__field1__lte=10
+            ),
+            models[1:3],
+        )
+
     def test_order_and_group_by_embedded_field_annotation(self):
         # Create repeated `data__integer` values.
         [Holder.objects.create(data=Data(integer=x)) for x in range(6)]