From 54d276f6a77060c7eb021583a7f4a79a75514a71 Mon Sep 17 00:00:00 2001
From: Harry Marr <harry.marr@gmail.com>
Date: Fri, 8 Jan 2010 04:49:14 +0000
Subject: [PATCH] Added index for _types

---
 mongoengine/queryset.py | 18 +++++++++++-------
 tests/document.py       |  3 ++-
 tests/queryset.py       | 21 +++++++++++++++++++++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 18a26f9d..760f413c 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -54,13 +54,6 @@ class QuerySet(object):
         """Filter the selected documents by calling the 
         :class:`~mongoengine.QuerySet` with a query.
         """
-        
-        # ensure document-defined indexes are created
-        if self._document._meta['indexes']:
-            for key_or_list in self._document._meta['indexes']:
-                # print "key", key_or_list
-                self.ensure_index(key_or_list)
-        
         query = QuerySet._transform_query(_doc_cls=self._document, **query)
         self._query.update(query)
         return self
@@ -68,7 +61,18 @@ class QuerySet(object):
     @property
     def _cursor(self):
         if not self._cursor_obj:
+            # Ensure document-defined indexes are created
+            if self._document._meta['indexes']:
+                for key_or_list in self._document._meta['indexes']:
+                    # print "key", key_or_list
+                    self.ensure_index(key_or_list)
+
+            # If _types is being used (for polymorphism), it needs an index
+            if '_types' in self._query:
+                self._collection.ensure_index('_types')
+
             self._cursor_obj = self._collection.find(self._query)
+
         return self._cursor_obj
 
     @classmethod
diff --git a/tests/document.py b/tests/document.py
index 5783ae05..412c89a9 100644
--- a/tests/document.py
+++ b/tests/document.py
@@ -239,7 +239,8 @@ class DocumentTest(unittest.TestCase):
         info = BlogPost.objects._collection.index_information()
         self.assertEqual(len(info), 0)
 
-        BlogPost.objects()
+        # Indexes are lazy so use list() to perform query
+        list(BlogPost.objects)
         info = BlogPost.objects._collection.index_information()
         self.assertTrue([('category', 1), ('addDate', -1)] in info.values())
         # Even though descending order was specified, single-key indexes use 1
diff --git a/tests/queryset.py b/tests/queryset.py
index e7e79ccc..7ea3da26 100644
--- a/tests/queryset.py
+++ b/tests/queryset.py
@@ -326,6 +326,27 @@ class QuerySetTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+
+    def test_types_index(self):
+        """Ensure that and index is used when '_types' is being used in a
+        query.
+        """
+        # Indexes are lazy so use list() to perform query
+        list(self.Person.objects)
+        info = self.Person.objects._collection.index_information()
+        self.assertTrue([('_types', 1)] in info.values())
+
+        class BlogPost(Document):
+            title = StringField()
+            meta = {'allow_inheritance': False}
+
+        # _types is not used on objects where allow_inheritance is False
+        list(BlogPost.objects)
+        info = BlogPost.objects._collection.index_information()
+        self.assertFalse([('_types', 1)] in info.values())
+
+        BlogPost.drop_collection()
+
     def tearDown(self):
         self.Person.drop_collection()