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