add support for JSONField lookups in an embedded model

Author: WaVEV

django_mongodb_backend/fields/embedded_model.py

@@ -7,6 +7,7 @@
 from django.db.models.lookups import Transform
 
 from .. import forms
+from .json import build_json_mql_path
 
 
 class EmbeddedModelField(models.Field):
@@ -181,18 +182,27 @@ 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
+ # The first json_key_transform is the field name.
+ 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}"
+ if json_key_transforms:
+ result = build_json_mql_path(result, json_key_transforms)
+ return result
 
 
 class KeyTransformFactory:

django_mongodb_backend/fields/json.py

@@ -17,6 +17,26 @@
 from ..query_utils import process_lhs, process_rhs
 
 
+def build_json_mql_path(lhs, key_transforms):
+ # Build the MQL path using the collected key transforms.
+ result = lhs
+ 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
+
+
 def contained_by(self, compiler, connection): # noqa: ARG001
  raise NotSupportedError("contained_by lookup is not supported on this database backend.")
 
@@ -89,23 +109,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 build_json_mql_path(lhs_mql, key_transforms)
 
 
 def key_transform_in(self, compiler, connection):

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()
 
 
 class Address(EmbeddedModel):

tests/model_fields_/test_embedded_model.py

@@ -114,6 +114,31 @@ def test_order_by_embedded_field(self):
  qs = Holder.objects.filter(data__integer__gt=3).order_by("-data__integer")
  self.assertSequenceEqual(qs, list(reversed(self.objs[4:])))
 
+ def test_embedded_json_field_lookups(self):
+ objs = [
+ Holder.objects.create(
+ data=Data(json_value={"field1": i * 5, "field2": {"0": {"value": list(range(i))}}})
+ )
+ for i in range(4)
+ ]
+ self.assertCountEqual(
+ Holder.objects.filter(data__json_value__field2__0__value__0=0),
+ objs[1:],
+ )
+ self.assertCountEqual(
+ Holder.objects.filter(data__json_value__field2__0__value__1=1),
+ objs[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), objs)
+ 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
+ ),
+ objs[1:3],
+ )
+
  def test_order_and_group_by_embedded_field(self):
  # Create and sort test data by `data__integer`.
  expected_objs = sorted(