beginnings of querying

Author: timgraham

django_mongodb/fields/__init__.py

@@ -1,11 +1,12 @@
 from .auto import ObjectIdAutoField
 from .duration import register_duration_field
-from .embedded_model import EmbeddedModelField
+from .embedded_model import EmbeddedModelField, register_embedded_model_field
 from .json import register_json_field
 
 __all__ = ["register_fields", "EmbeddedModelField", "ObjectIdAutoField"]
 
 
 def register_fields():
     register_duration_field()
+    register_embedded_model_field()
     register_json_field()

django_mongodb/fields/embedded_model.py

@@ -1,5 +1,6 @@
 from django.db import models
 from django.db.models.fields.related import lazy_related_operation
+from django.db.models.lookups import Transform
 
 
 class EmbeddedModelField(models.Field):
@@ -108,10 +109,48 @@ def get_db_prep_save(self, embedded_instance, connection):
         embedded_instance._state.adding = False
         return field_values
 
+    def get_transform(self, name):
+        transform = super().get_transform(name)
+        if transform:
+            return transform
+        return KeyTransformFactory(name)
+
     def validate(self, value, model_instance):
         super().validate(value, model_instance)
         if self.embedded_model is None:
             return
         for field in self.embedded_model._meta.fields:
             attname = field.attname
             field.validate(getattr(value, attname), model_instance)
+
+
+class KeyTransform(Transform):
+    def __init__(self, key_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.key_name = str(key_name)
+
+    def preprocess_lhs(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
+        mql = previous.as_mql(compiler, connection)
+        return mql, key_transforms
+
+
+def key_transform(self, compiler, connection):
+    mql, key_transforms = self.preprocess_lhs(compiler, connection)
+    return f"{mql}.{key_transforms[0]}"
+
+
+class KeyTransformFactory:
+    def __init__(self, key_name):
+        self.key_name = key_name
+
+    def __call__(self, *args, **kwargs):
+        return KeyTransform(self.key_name, *args, **kwargs)
+
+
+def register_embedded_model_field():
+    KeyTransform.as_mql = key_transform

tests/model_fields_/test_embedded_model.py

@@ -91,3 +91,37 @@ def test_embedded_field_with_foreign_conversion(self):
         decimal = DecimalKey.objects.create(decimal=Decimal("1.5"))
         decimal_parent = DecimalParent.objects.create(child=decimal)
         EmbeddedModelFieldModel.objects.create(decimal_parent=decimal_parent)
+
+
+class QueryingTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.objs = [
+            EmbeddedModelFieldModel.objects.create(simple=EmbeddedModel(someint=x))
+            for x in range(6)
+        ]
+
+    def test_exact(self):
+        self.assertCountEqual(
+            EmbeddedModelFieldModel.objects.filter(simple__someint=3), [self.objs[3]]
+        )
+
+    def test_lt(self):
+        self.assertCountEqual(
+            EmbeddedModelFieldModel.objects.filter(simple__someint__lt=3), self.objs[:3]
+        )
+
+    def test_lte(self):
+        self.assertCountEqual(
+            EmbeddedModelFieldModel.objects.filter(simple__someint__lte=3), self.objs[:4]
+        )
+
+    def test_gt(self):
+        self.assertCountEqual(
+            EmbeddedModelFieldModel.objects.filter(simple__someint__gt=3), self.objs[4:]
+        )
+
+    def test_gte(self):
+        self.assertCountEqual(
+            EmbeddedModelFieldModel.objects.filter(simple__someint__gte=3), self.objs[3:]
+        )