From 5e7efcc8c2f4947d9d25bb3e40f0cfa1e759c419 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Fri, 17 Jun 2011 12:43:28 +0100
Subject: [PATCH] Added 'hint' support, telling Mongo the proper index to use
 for the query.

Judicious use of hints can greatly improve query performance. When doing a query
on multiple fields (at least one of which is indexed) pass the indexed field as
a hint to the query. Hinting will not do anything if the corresponding index
does not exist.  The last hint applied to this cursor takes precedence over all
others.

Closes #203
---
 docs/changelog.rst      |  1 +
 mongoengine/queryset.py | 21 ++++++++++++++++++++-
 tests/document.py       | 30 +++++++++++++++++++++++++++++-
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index ea926239..a9cfe328 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -5,6 +5,7 @@ Changelog
 Changes in dev
 ==============
 
+- Added hint() support, so cantell Mongo the proper index to use for the query
 - Fixed issue with inconsitent setting of _cls breaking inherited referencing
 - Added help_text and verbose_name to fields to help with some form libs
 - Updated item_frequencies to handle embedded document lookups
diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 92229a1c..bfa89b4f 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -347,6 +347,7 @@ class QuerySet(object):
         self._cursor_obj = None
         self._limit = None
         self._skip = None
+        self._hint = -1  # Using -1 as None is a valid value for hint
 
     def clone(self):
         """Creates a copy of the current :class:`~mongoengine.queryset.QuerySet`"""
@@ -354,7 +355,7 @@ class QuerySet(object):
 
         copy_props = ('_initial_query', '_query_obj', '_where_clause',
                     '_loaded_fields', '_ordering', '_snapshot',
-                    '_timeout', '_limit', '_skip', '_slave_okay')
+                    '_timeout', '_limit', '_skip', '_slave_okay', '_hint')
 
         for prop in copy_props:
             val = getattr(self, prop)
@@ -539,6 +540,9 @@ class QuerySet(object):
             if self._skip is not None:
                 self._cursor_obj.skip(self._skip)
 
+            if self._hint != -1:
+                self._cursor_obj.hint(self._hint)
+
         return self._cursor_obj
 
     @classmethod
@@ -965,6 +969,21 @@ class QuerySet(object):
         self._skip = n
         return self
 
+    def hint(self, index=None):
+        """Added 'hint' support, telling Mongo the proper index to use for the
+        query.
+
+        Judicious use of hints can greatly improve query performance. When doing
+        a query on multiple fields (at least one of which is indexed) pass the
+        indexed field as a hint to the query.
+
+        Hinting will not do anything if the corresponding index does not exist.
+        The last hint applied to this cursor takes precedence over all others.
+        """
+        self._cursor.hint(index)
+        self._hint = index
+        return self
+
     def __getitem__(self, key):
         """Support skip and limit using getitem and slicing syntax.
         """
diff --git a/tests/document.py b/tests/document.py
index b33f3fe7..5d44ca29 100644
--- a/tests/document.py
+++ b/tests/document.py
@@ -513,7 +513,6 @@ class DocumentTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
-
     def test_dictionary_indexes(self):
         """Ensure that indexes are used when meta[indexes] contains dictionaries
         instead of lists.
@@ -546,6 +545,35 @@ class DocumentTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+    def test_hint(self):
+
+        class BlogPost(Document):
+            tags = ListField(StringField())
+            meta = {
+                'indexes': [
+                    'tags',
+                ],
+            }
+
+        BlogPost.drop_collection()
+
+        for i in xrange(0, 10):
+            tags = [("tag %i" % n) for n in xrange(0, i % 2)]
+            BlogPost(tags=tags).save()
+
+        self.assertEquals(BlogPost.objects.count(), 10)
+        self.assertEquals(BlogPost.objects.hint().count(), 10)
+        self.assertEquals(BlogPost.objects.hint([('tags', 1)]).count(), 10)
+
+        self.assertEquals(BlogPost.objects.hint([('ZZ', 1)]).count(), 10)
+
+        def invalid_index():
+            BlogPost.objects.hint('tags')
+        self.assertRaises(TypeError, invalid_index)
+
+        def invalid_index_2():
+            return BlogPost.objects.hint(('tags', 1))
+        self.assertRaises(TypeError, invalid_index_2)
 
     def test_unique(self):
         """Ensure that uniqueness constraints are applied to fields.