Allow querying an EmbeddedModelField by model instance

Author: timgraham

django_mongodb_backend/fields/embedded_model.py

@@ -3,10 +3,12 @@
 from django.core import checks
 from django.core.exceptions import FieldDoesNotExist
 from django.db import models
+from django.db.models import lookups
 from django.db.models.fields.related import lazy_related_operation
 from django.db.models.lookups import Transform
 
 from .. import forms
+from ..query_utils import process_lhs, process_rhs
 from .json import build_json_mql_path
 
 
@@ -149,6 +151,41 @@ def formfield(self, **kwargs):
         )
 
 
+@EmbeddedModelField.register_lookup
+class EMFExact(lookups.Exact):
+    def model_to_dict(self, instance):
+        """Return a dict containing the data in a model instance."""
+        data = {}
+        emf_data = {}
+        for f in instance._meta.concrete_fields:
+            value = f.value_from_object(instance)
+            if isinstance(f, EmbeddedModelField):
+                emf_data[f"{f.name}"] = self.model_to_dict(value)
+                continue
+            # Unless explicitly set, primary keys aren't included in embedded
+            # models.
+            if f.primary_key and value is None:
+                continue
+            data[f"{f.name}"] = value
+        return data, emf_data
+
+    def as_mql(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection)
+        value = process_rhs(self, compiler, connection)
+        if isinstance(value, models.Model):
+            value, emf_data = self.model_to_dict(value)
+            prefix = self.lhs.as_mql(compiler, connection)
+            conditions = [{"$eq": [f"{prefix}.{k}", v]} for k, v in value.items()]
+            # TODO: more tests for this logic. Might not work for another
+            # layer of embedding.
+            while emf_data:
+                for k, v in emf_data.items():
+                    v, emf_data = v
+                    conditions += [{"$eq": [f"{prefix}.{k}.{x}", y]} for x, y in v.items()]
+            return {"$and": conditions}
+        return connection.mongo_operators[self.lookup_name](lhs_mql, value)
+
+
 class KeyTransform(Transform):
     def __init__(self, key_name, ref_field, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -193,7 +230,13 @@ def preprocess_lhs(self, compiler, connection):
             previous = previous.lhs
         mql = previous.as_mql(compiler, connection)
         # The first json_key_transform is the field name.
-        embedded_key_transforms.append(json_key_transforms.pop(0))
+        try:
+            field_name = json_key_transforms.pop(0)
+        except IndexError:
+            # This is a lookup of the embedded model itself.
+            pass
+        else:
+            embedded_key_transforms.append(field_name)
         return mql, embedded_key_transforms, json_key_transforms
 
     def as_mql(self, compiler, connection):

docs/source/topics/embedded-models.rst

@@ -54,3 +54,12 @@ as relational fields. For example, to retrieve all customers who have an
 address with the city "New York"::
 
     >>> Customer.objects.filter(address__city="New York")
+
+You can also query using a model instance. Unlike a normal relational lookup
+which does the lookup by primary key, since embedded models typically don't
+have a primary key set, the query requires that every field match. For example,
+this query gives customers with addresses with the city "New York" and all
+other fields of address equal to their default (:attr:`Field.default
+<django.db.models.Field.default>`, ``None``, or an empty string).
+
+    >>> Customer.objects.filter(address=Address(city="New York"))

tests/model_fields_/test_embedded_model.py

@@ -2,7 +2,7 @@
 from datetime import timedelta
 
 from django.core.exceptions import FieldDoesNotExist, ValidationError
-from django.db import models
+from django.db import connection, models
 from django.db.models import (
     Exists,
     ExpressionWrapper,
@@ -117,6 +117,38 @@ 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_exact_with_model(self):
+        data = Holder.objects.first().data
+        self.assertEqual(
+            Holder.objects.filter(data=data).get().data.integer, self.objs[0].data.integer
+        )
+
+    def test_exact_with_model_ignores_key_order(self):
+        # Due to the possibility of schema changes or the reordering of a
+        # model's fields, a lookup must work if an embedded document has its
+        # keys in a different order than what's declared on the embedded model.
+        connection.get_collection("model_fields__holder").insert_one(
+            {
+                "data": {
+                    "auto_now": None,
+                    "auto_now_add": None,
+                    "json_value": None,
+                    "integer": 100,
+                }
+            }
+        )
+        self.assertEqual(Holder.objects.filter(data=Data(integer=100)).get().data.integer, 100)
+
+    def test_exact_with_nested_model(self):
+        address = Address(city="NYC", state="NY")
+        obj = Book.objects.create(author=Author(name="Shakespeare", age=55, address=address))
+        self.assertCountEqual(Book.objects.filter(author__address=address), [obj])
+
+    def test_exact_with_model_with_embedded_modelt(self):
+        author = Author(name="Shakespeare", age=55, address=Address(city="NYC", state="NY"))
+        obj = Book.objects.create(author=author)
+        self.assertCountEqual(Book.objects.filter(author=author), [obj])
+
     def test_embedded_json_field_lookups(self):
         objs = [
             Holder.objects.create(