Merge pull request #6 from iainb/batch-size

feat: noticket - add batch_size support to QuerySet

Author: iainb

mongoengine/__init__.py

@@ -12,7 +12,7 @@
 __all__ = (document.__all__ + fields.__all__ + connection.__all__ +
            queryset.__all__ + signals.__all__)
 
-VERSION = (0, 6, 21)
+VERSION = (0, 6, 22)
 
 
 def get_version():

mongoengine/queryset.py

@@ -351,6 +351,7 @@ def __init__(self, document, collection):
         self._limit = None
         self._skip = None
         self._hint = -1  # Using -1 as None is a valid value for hint
+        self._batch_size = None
 
     def clone(self):
         """Creates a copy of the current :class:`~mongoengine.queryset.QuerySet`
@@ -361,7 +362,7 @@ def clone(self):
 
         copy_props = ('_initial_query', '_query_obj', '_where_clause',
                     '_loaded_fields', '_ordering',
-                    '_limit', '_skip',  '_hint',
+                    '_limit', '_skip',  '_hint', '_batch_size',
                     '_read_preference',)
 
         for prop in copy_props:
@@ -522,6 +523,9 @@ def _cursor(self):
             if self._hint != -1:
                 self._cursor_obj.hint(self._hint)
 
+            if self._batch_size is not None:
+                self._cursor_obj.batch_size(self._batch_size)
+
         return self._cursor_obj
 
     @classmethod
@@ -1014,6 +1018,21 @@ def hint(self, index=None):
         self._hint = index
         return self
 
+    def batch_size(self, size):
+        """Limit the number of documents returned in a single batch (each batch
+        requires a round trip to the server).
+
+
+        See http://api.mongodb.com/python/current/api/pymongo/cursor.html#pymongo.cursor.Cursor.batch_size
+        for details.
+        :param size: desired size of each batch.
+        """
+        if self._cursor_obj is not None:
+            self._cursor_obj.batch_size(size)
+
+        self._batch_size = size
+        return self
+
     def __getitem__(self, key):
         """Support skip and limit using getitem and slicing syntax.
         """

tests/queryset.py

@@ -2892,6 +2892,64 @@ def test_scalar_cursor_behaviour(self):
         pks = self.Person.objects.order_by('age').scalar('pk')[1:3]
         self.assertEqual("[u'A1', u'A2']",  "%s" % sorted(self.Person.objects.scalar('name').in_bulk(list(pks)).values()))
 
+    def test_batch_size(self):
+        """Ensure that batch_size works."""
+        class A(Document):
+            s = StringField()
+
+        A.drop_collection()
+
+        for i in range(100):
+            A.objects.create(s=str(i))
+
+        # test iterating over the result set
+        cnt = 0
+        for a in A.objects.batch_size(10):
+            cnt += 1
+        self.assertEqual(cnt, 100)
+
+        # test chaining
+        qs = A.objects.all()
+        qs = qs.limit(10).batch_size(20).skip(91)
+        cnt = 0
+        for a in qs:
+            cnt += 1
+        self.assertEqual(cnt, 9)
+
+        # test chaining and iteration with small batch size
+        qs = A.objects.all()
+        qs = qs.limit(10).batch_size(4).skip(91)
+        cnt = 0
+        for a in qs:
+            cnt += 1
+        self.assertEqual(cnt, 9)
+
+        # test invalid batch size
+        qs = A.objects.batch_size(-1)
+        with self.assertRaises(ValueError):
+            list(qs)
+
+    def test_modify_evaluated_query_with_batch_size(self):
+        class Number(Document):
+            n = IntField()
+
+        Number.drop_collection()
+
+        for i in xrange(1, 101):
+            t = Number(n=i)
+            t.save()
+
+        test = Number.objects
+        self.assertEqual(test.count(), 100)
+
+        test = test.filter(n__gt=11)
+        self.assertEqual(test.count(), 89)
+
+        test = test.batch_size(10)
+        self.assertEqual(len(list(test)), 89)
+
+        Number.drop_collection()
+
 
 class QTest(unittest.TestCase):