querying support

Author: timgraham

django_mongodb_backend/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):
@@ -31,13 +32,8 @@ def _set_model(self, model):
         Resolve embedded model class once the field knows the model it belongs
         to.
 
-        If the model argument passed to __init__() was a string, resolve that
-        string to the corresponding model class, similar to relation fields.
-        However, we need to know our own model to generate a valid key
-        for the embedded model class lookup and EmbeddedModelFields are
-        not contributed_to_class if used in iterable fields. Thus the
-        collection field sets this field's "model" attribute in its
-        contribute_to_class().
+        If __init__()'s embedded_model argument is a string, resolve it to the
+        corresponding model class, similar to relation fields.
         """
         self._model = model
         if model is not None and isinstance(self.embedded_model, str):
@@ -54,8 +50,8 @@ def from_db_value(self, value, expression, connection):
 
     def to_python(self, value):
         """
-        Passes embedded model fields' values through embedded fields
-        to_python() and reinstiatates the embedded instance.
+        Pass embedded model fields' values through each field's to_python() and
+        reinstiatate the embedded instance.
         """
         if value is None:
             return None
@@ -76,14 +72,8 @@ def to_python(self, value):
 
     def get_db_prep_save(self, embedded_instance, connection):
         """
-        Apply pre_save() and get_db_prep_save() of embedded instance
-        fields and passes a field => value mapping down to database
-        type conversions.
-
-        The embedded instance will be saved as a column => value dict, but
-        because we need to apply database type conversions on embedded instance
-        fields' values and for these we need to know fields those values come
-        from, we need to entrust the database layer with creating the dict.
+        Apply pre_save() and get_db_prep_save() of embedded instance fields and
+        create the {field: value} dict to be saved.
         """
         if embedded_instance is None:
             return None
@@ -106,14 +96,48 @@ def get_db_prep_save(self, embedded_instance, connection):
                 continue
             field_values[field.attname] = value
         # This instance will exist in the database soon.
-        # TODO.XXX: Ensure that this doesn't cause race conditions.
+        # TODO: Ensure that this doesn't cause race conditions.
         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 as_mql(self, compiler, connection):
+        mql, key_transforms = self.preprocess_lhs(compiler, connection)
+        transforms = ".".join(key_transforms)
+        return f"{mql}.{transforms}"
+
+
+class KeyTransformFactory:
+    def __init__(self, key_name):
+        self.key_name = key_name
+
+    def __call__(self, *args, **kwargs):
+        return KeyTransform(self.key_name, *args, **kwargs)

tests/model_fields_/models.py

@@ -120,3 +120,19 @@ class EmbeddedModel(models.Model):
     someint = models.IntegerField(db_column="custom_column")
     auto_now = models.DateTimeField(auto_now=True)
     auto_now_add = models.DateTimeField(auto_now_add=True)
+
+
+class Address(models.Model):
+    city = models.CharField(max_length=20)
+    state = models.CharField(max_length=2)
+
+
+class Author(models.Model):
+    name = models.CharField(max_length=10)
+    age = models.IntegerField()
+    address = EmbeddedModelField(Address)
+
+
+class Book(models.Model):
+    name = models.CharField(max_length=100)
+    author = EmbeddedModelField(Author)

tests/model_fields_/test_embedded_model.py

@@ -6,6 +6,9 @@
 from django_mongodb_backend.fields import EmbeddedModelField
 
 from .models import (
+    Address,
+    Author,
+    Book,
     DecimalKey,
     DecimalParent,
     EmbeddedModel,
@@ -91,3 +94,43 @@ 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:]
+        )
+
+    def test_nested(self):
+        obj = Book.objects.create(
+            author=Author(name="Shakespeare", age=55, address=Address(city="NYC", state="NY"))
+        )
+        self.assertCountEqual(Book.objects.filter(author__address__city="NYC"), [obj])