Merge pull request #779 from DavidBord/fix-778

fix-#778: Add Support For MongoDB 2.6.X's maxTimeMS

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 
 Changes in 0.9.X - DEV
 ======================
+- Add Support For MongoDB 2.6.X's maxTimeMS #778
 - abstract shouldn't be inherited in EmbeddedDocument # 789
 - Allow specifying the '_cls' as a field for indexes #397
 - Stop ensure_indexes running on a secondaries unless connection is through mongos #746

mongoengine/queryset/base.py

@@ -82,6 +82,7 @@ class BaseQuerySet(object):
         self._skip = None
         self._hint = -1  # Using -1 as None is a valid value for hint
         self.only_fields = []
+        self._max_time_ms = None
 
     def __call__(self, q_obj=None, class_check=True, slave_okay=False,
                  read_preference=None, **query):
@@ -672,7 +673,7 @@ class BaseQuerySet(object):
                       '_timeout', '_class_check', '_slave_okay', '_read_preference',
                       '_iter', '_scalar', '_as_pymongo', '_as_pymongo_coerce',
                       '_limit', '_skip', '_hint', '_auto_dereference',
-                      '_search_text', '_include_text_scores', 'only_fields')
+                      '_search_text', '_include_text_scores', 'only_fields', '_max_time_ms')
 
         for prop in copy_props:
             val = getattr(self, prop)
@@ -969,6 +970,13 @@ class BaseQuerySet(object):
         queryset._as_pymongo_coerce = coerce_types
         return queryset
 
+    def max_time_ms(self, ms):
+        """Wait `ms` milliseconds before killing the query on the server
+
+        :param ms: the number of milliseconds before killing the query on the server
+        """
+        return self._chainable_method("max_time_ms", ms)
+
     # JSON Helpers
 
     def to_json(self, *args, **kwargs):
@@ -1700,6 +1708,13 @@ class BaseQuerySet(object):
                       code)
         return code
 
+    def _chainable_method(self, method_name, val):
+        queryset = self.clone()
+        method = getattr(queryset._cursor, method_name)
+        method(val)
+        setattr(queryset, "_" + method_name, val)
+        return queryset
+
     # Deprecated
     def ensure_index(self, **kwargs):
         """Deprecated use :func:`Document.ensure_index`"""

tests/queryset/queryset.py

@@ -4401,6 +4401,10 @@ class QuerySetTest(unittest.TestCase):
         self.Person.objects().delete()
         self.assertEqual(self.Person.objects().skip(1).delete(), 0)  # test Document delete without existing documents
 
+    def test_max_time_ms(self):
+        # 778: max_time_ms can get only int or None as input
+        self.assertRaises(TypeError, self.Person.objects(name="name").max_time_ms, "not a number")
+
 
 if __name__ == '__main__':
     unittest.main()