From b412b2a6f18caf30ce0fdda6fe82d3c5ff7bc6a3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Lafr=C3=A9choux?=
 <jlafrechoux@nobatek.inef4.com>
Date: Mon, 12 Oct 2020 11:00:43 +0200
Subject: [PATCH 1/2] Allow passing an EmbeddedDocument in a query

---
 tests/test_query_mapper.py | 9 +++++++++
 umongo/query_mapper.py     | 3 +++
 2 files changed, 12 insertions(+)

diff --git a/tests/test_query_mapper.py b/tests/test_query_mapper.py
index 235a0bd7..f6608c85 100644
--- a/tests/test_query_mapper.py
+++ b/tests/test_query_mapper.py
@@ -100,6 +100,15 @@ class Book(Document):
             ]}
         }
 
+        # Test embedded document in query
+        query = map_query({
+            'author': Author(name='JRR Tolkien', birthday=dt.datetime(1892, 1, 3))
+        }, book_fields)
+        assert isinstance(query['a'], dict)
+        assert query == {
+            'a': {'name': 'JRR Tolkien', 'b': dt.datetime(1892, 1, 3)}
+        }
+
     def test_mix(self):
 
         @self.instance.register
diff --git a/umongo/query_mapper.py b/umongo/query_mapper.py
index c20224b3..272443c6 100644
--- a/umongo/query_mapper.py
+++ b/umongo/query_mapper.py
@@ -1,4 +1,5 @@
 from umongo.fields import ListField, EmbeddedField
+from umongo.embedded_document import EmbeddedDocumentImplementation
 
 
 def map_entry(entry, fields):
@@ -44,4 +45,6 @@ def map_query(query, fields):
         return mapped_query
     if isinstance(query, (list, tuple)):
         return [map_query(x, fields) for x in query]
+    if isinstance(query, EmbeddedDocumentImplementation):
+        return query.to_mongo()
     return query

From 86a742c63246ee37afe2b8f135b4eefc121191d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Lafr=C3=A9choux?=
 <jlafrechoux@nobatek.inef4.com>
Date: Mon, 12 Oct 2020 11:55:15 +0200
Subject: [PATCH 2/2] Allow passing a Document as reference in a query

---
 tests/test_query_mapper.py | 15 ++++++++++++++-
 umongo/query_mapper.py     |  4 ++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/tests/test_query_mapper.py b/tests/test_query_mapper.py
index f6608c85..e5d48153 100644
--- a/tests/test_query_mapper.py
+++ b/tests/test_query_mapper.py
@@ -1,5 +1,7 @@
 import datetime as dt
 
+from bson import ObjectId
+
 from umongo import Document, EmbeddedDocument, fields
 from umongo.query_mapper import map_query
 
@@ -10,6 +12,10 @@ class TestQueryMapper(BaseTest):
 
     def test_query_mapper(self):
 
+        @self.instance.register
+        class Editor(Document):
+            name = fields.StrField()
+
         @self.instance.register
         class Author(EmbeddedDocument):
             name = fields.StrField()
@@ -19,7 +25,6 @@ class Author(EmbeddedDocument):
         class Chapter(EmbeddedDocument):
             title = fields.StrField()
             pagination = fields.IntField(attribute='p')
-            # sub_chapters = fields.EmbeddedField('Chapter')
 
         @self.instance.register
         class Book(Document):
@@ -28,6 +33,7 @@ class Book(Document):
             author = fields.EmbeddedField(Author, attribute='a')
             chapters = fields.ListField(fields.EmbeddedField(Chapter))
             tags = fields.ListField(fields.StrField(), attribute='t')
+            editor = fields.ReferenceField(Editor, attribute='e')
 
         book_fields = Book.schema.fields
         # No changes needed
@@ -109,6 +115,13 @@ class Book(Document):
             'a': {'name': 'JRR Tolkien', 'b': dt.datetime(1892, 1, 3)}
         }
 
+        # Test document in query
+        editor = Editor(name='Allen & Unwin')
+        editor.id = ObjectId()
+        query = map_query({'editor': editor}, book_fields)
+        assert isinstance(query['e'], ObjectId)
+        assert query['e'] == editor.id
+
     def test_mix(self):
 
         @self.instance.register
diff --git a/umongo/query_mapper.py b/umongo/query_mapper.py
index 272443c6..6f28c278 100644
--- a/umongo/query_mapper.py
+++ b/umongo/query_mapper.py
@@ -1,4 +1,5 @@
 from umongo.fields import ListField, EmbeddedField
+from umongo.document import DocumentImplementation
 from umongo.embedded_document import EmbeddedDocumentImplementation
 
 
@@ -45,6 +46,9 @@ def map_query(query, fields):
         return mapped_query
     if isinstance(query, (list, tuple)):
         return [map_query(x, fields) for x in query]
+    # Passing a Document only makes sense in a Reference, let's query on ObjectId
+    if isinstance(query, DocumentImplementation):
+        return query.pk
     if isinstance(query, EmbeddedDocumentImplementation):
         return query.to_mongo()
     return query