PYTHON-1065 - Support 1 for batch_size with MongoDB >= 3.2

Author: behackett

pymongo/cursor.py

@@ -426,7 +426,7 @@ def batch_size(self, batch_size):
             raise ValueError("batch_size must be >= 0")
         self.__check_okay_to_chain()
 
-        self.__batch_size = batch_size == 1 and 2 or batch_size
+        self.__batch_size = batch_size
         return self
 
     def skip(self, skip):
@@ -999,17 +999,10 @@ def _refresh(self):
             return len(self.__data)
 
         if self.__id is None:  # Query
-            ntoreturn = self.__batch_size
-            if self.__limit:
-                if self.__batch_size:
-                    ntoreturn = min(self.__limit, self.__batch_size)
-                else:
-                    ntoreturn = self.__limit
             self.__send_message(_Query(self.__query_flags,
                                        self.__collection.database.name,
                                        self.__collection.name,
                                        self.__skip,
-                                       ntoreturn,
                                        self.__query_spec(),
                                        self.__projection,
                                        self.__codec_options,

pymongo/helpers.py

@@ -240,8 +240,8 @@ def _first_batch(sock_info, db, coll, query, ntoreturn,
                  slave_ok, codec_options, read_preference, cmd, listeners):
     """Simple query helper for retrieving a first (and possibly only) batch."""
     query = _Query(
-        0, db, coll, 0, ntoreturn, query, None,
-        codec_options, read_preference, 0, 0, DEFAULT_READ_CONCERN)
+        0, db, coll, 0, query, None,
+        codec_options, read_preference, ntoreturn, 0, DEFAULT_READ_CONCERN)
 
     name = next(iter(cmd))
     duration = None

pymongo/message.py

@@ -214,18 +214,17 @@ def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms):
 class _Query(object):
     """A query operation."""
 
-    __slots__ = ('flags', 'db', 'coll', 'ntoskip', 'ntoreturn', 'spec',
+    __slots__ = ('flags', 'db', 'coll', 'ntoskip', 'spec',
                  'fields', 'codec_options', 'read_preference', 'limit',
                  'batch_size', 'name', 'read_concern')
 
-    def __init__(self, flags, db, coll, ntoskip, ntoreturn, spec, fields,
+    def __init__(self, flags, db, coll, ntoskip, spec, fields,
                  codec_options, read_preference, limit,
                  batch_size, read_concern):
         self.flags = flags
         self.db = db
         self.coll = coll
         self.ntoskip = ntoskip
-        self.ntoreturn = ntoreturn
         self.spec = spec
         self.fields = fields
         self.codec_options = codec_options
@@ -260,12 +259,21 @@ def get_message(self, set_slave_ok, is_mongos, use_cmd=False):
 
         ns = _UJOIN % (self.db, self.coll)
         spec = self.spec
-        ntoreturn = self.ntoreturn
 
         if use_cmd:
             ns = _UJOIN % (self.db, "$cmd")
             spec = self.as_command()[0]
             ntoreturn = -1  # All DB commands return 1 document
+        else:
+            # OP_QUERY treats ntoreturn of -1 and 1 the same, return
+            # one document and close the cursor. We have to use 2 for
+            # batch size if 1 is specified.
+            ntoreturn = self.batch_size == 1 and 2 or self.batch_size
+            if self.limit:
+                if ntoreturn:
+                    ntoreturn = min(self.limit, ntoreturn)
+                else:
+                    ntoreturn = self.limit
 
         if is_mongos:
             spec = _maybe_add_read_preference(spec,

test/test_cursor.py

@@ -540,6 +540,21 @@ def cursor_count(cursor, expected_count):
         cursor_count(db.test.find().batch_size(100).limit(10), 10)
         cursor_count(db.test.find().batch_size(500).limit(10), 10)
 
+        cur = db.test.find().batch_size(1)
+        next(cur)
+        if client_context.version.at_least(3, 1, 9):
+            # find command batchSize should be 1
+            self.assertEqual(0, len(cur._Cursor__data))
+        else:
+            # OP_QUERY ntoreturn should be 2
+            self.assertEqual(1, len(cur._Cursor__data))
+        next(cur)
+        self.assertEqual(0, len(cur._Cursor__data))
+        next(cur)
+        self.assertEqual(0, len(cur._Cursor__data))
+        next(cur)
+        self.assertEqual(0, len(cur._Cursor__data))
+
     def test_limit_and_batch_size(self):
         db = self.db
         db.test.drop()