From 32fc4152a71f931f2e3a43ef215f8826488edb21 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Fri, 2 Mar 2012 13:42:24 +0000
Subject: [PATCH] Enable covered indexes for simple documents.

Refs #444
---
 docs/changelog.rst      |  1 +
 mongoengine/queryset.py | 13 ++++++++++---
 tests/document.py       | 29 ++++++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index e478befe..43eac27e 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -5,6 +5,7 @@ Changelog
 Changes in 0.6
 ==============
 
+- Added support for covered indexes when inheritance is off
 - No longer always upsert on save for items with a '_id'
 - Error raised if update doesn't have an operation
 - DeReferencing is now thread safe
diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 3eac1af1..85a5152a 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -273,16 +273,20 @@ class Q(QNode):
 
 class QueryFieldList(object):
     """Object that handles combinations of .only() and .exclude() calls"""
-    ONLY = True
-    EXCLUDE = False
+    ONLY = 1
+    EXCLUDE = 0
 
     def __init__(self, fields=[], value=ONLY, always_include=[]):
         self.value = value
         self.fields = set(fields)
         self.always_include = set(always_include)
+        self._id = None
 
     def as_dict(self):
-        return dict((field, self.value) for field in self.fields)
+        field_list = dict((field, self.value) for field in self.fields)
+        if self._id is not None:
+            field_list['_id'] = self._id
+        return field_list
 
     def __add__(self, f):
         if not self.fields:
@@ -298,6 +302,9 @@ class QueryFieldList(object):
             self.value = self.ONLY
             self.fields = f.fields - self.fields
 
+        if '_id' in f.fields:
+            self._id = f.value
+
         if self.always_include:
             if self.value is self.ONLY and self.fields:
                 self.fields = self.fields.union(self.always_include)
diff --git a/tests/document.py b/tests/document.py
index 252393cf..b6c72c5a 100644
--- a/tests/document.py
+++ b/tests/document.py
@@ -714,7 +714,6 @@ class DocumentTest(unittest.TestCase):
         self.assertEqual(info.keys(), ['_types_1_user_guid_1', '_id_', '_types_1_name_1'])
         Person.drop_collection()
 
-
     def test_embedded_document_index(self):
         """Tests settings an index on an embedded document
         """
@@ -788,6 +787,34 @@ class DocumentTest(unittest.TestCase):
 
         self.assertEquals(len(User._geo_indices()), 2)
 
+    def test_covered_index(self):
+        """Ensure that covered indexes can be used
+        """
+
+        class Test(Document):
+            a = IntField()
+
+            meta = {
+                'indexes': ['a'],
+                'allow_inheritance': False
+                }
+
+        Test.drop_collection()
+
+        obj = Test(a=1)
+        obj.save()
+
+        # Need to be explicit about covered indexes as mongoDB doesn't know if
+        # the documents returned might have more keys in that here.
+        query_plan = Test.objects(id=obj.id).exclude('a').explain()
+        self.assertFalse(query_plan['indexOnly'])
+
+        query_plan = Test.objects(id=obj.id).only('id').explain()
+        self.assertTrue(query_plan['indexOnly'])
+
+        query_plan = Test.objects(a=1).only('a').exclude('id').explain()
+        self.assertTrue(query_plan['indexOnly'])
+
     def test_hint(self):
 
         class BlogPost(Document):