Added index for _types

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

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

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()