PYTHON-978 - Use getMore command for CommandCursor

Author: aherlihy

pymongo/command_cursor.py

@@ -106,13 +106,17 @@ def __send_message(self, operation):
 
  cmd_duration = response.duration
  rqst_id = response.request_id
+ from_command = response.from_command
+
  if publish:
  start = datetime.datetime.now()
  try:
  doc = helpers._unpack_response(response.data,
  self.__id,
  self.__collection.codec_options)
- # TODO: check_command_response when getMore works with agg cursor
+ if from_command:
+ helpers._check_command_response(doc['data'][0])
+
  except OperationFailure as exc:
  self.__killed = True
 
@@ -142,22 +146,30 @@ def __send_message(self, operation):
  self.__address)
  raise
 
+ if from_command:
+ cursor = doc['data'][0]['cursor']
+ documents = cursor['nextBatch']
+ self.__id = cursor['id']
+ self.__retrieved += len(documents)
+ else:
+ documents = doc["data"]
+ self.__id = doc["cursor_id"]
+ self.__retrieved += doc["number_returned"]
+
  if publish:
  duration = (datetime.datetime.now() - start) + cmd_duration
  # Must publish in getMore command response format.
- res = {"cursor": {"id": doc["cursor_id"],
+ res = {"cursor": {"id": self.__id,
  "ns": self.__collection.full_name,
- "nextBatch": doc["data"]},
+ "nextBatch": documents},
  "ok": 1}
  listeners.publish_command_success(
  duration, res, "getMore", rqst_id, self.__address)
 
- self.__id = doc["cursor_id"]
  if self.__id == 0:
  self.__killed = True
+ self.__data = deque(documents)
 
- self.__retrieved += doc["number_returned"]
- self.__data = deque(doc["data"])
 
  def _refresh(self):
  """Refreshes the cursor with more data from the server.
@@ -175,8 +187,7 @@ def _refresh(self):
  self.__collection.name,
  self.__batch_size,
  self.__id,
- self.__collection.codec_options,
- cmd_cursor=True))
+ self.__collection.codec_options))
 
  else: # Cursor id is zero nothing else to return
  self.__killed = True

pymongo/message.py

@@ -275,21 +275,18 @@ class _GetMore(object):
  """A getmore operation."""
 
  __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_time_ms',
- 'codec_options', 'cmd_cursor')
+ 'codec_options')
 
  name = 'getMore'
 
  def __init__(self, db, coll, ntoreturn, cursor_id, codec_options,
- max_time_ms=None, cmd_cursor=False):
+ max_time_ms=None):
  self.db = db
  self.coll = coll
  self.ntoreturn = ntoreturn
  self.cursor_id = cursor_id
  self.codec_options = codec_options
  self.max_time_ms = max_time_ms
- # XXX: Temporarily keep track of if this getMore is for a command cursor
- # so we can use OP_KILLCURSORS until find support for mongos is completed.
- self.cmd_cursor = cmd_cursor
 
  def as_command(self):
  """Return a getMore command document for this query."""
@@ -301,7 +298,7 @@ def get_message(self, dummy0, dummy1, use_cmd=False):
 
  ns = _UJOIN % (self.db, self.coll)
 
- if use_cmd and not self.cmd_cursor:
+ if use_cmd:
  ns = _UJOIN % (self.db, "$cmd")
  spec = self.as_command()[0]