PYTHON-978 - Use find/getMore commands

Author: aherlihy

pymongo/collection.py

@@ -1360,9 +1360,10 @@ def list_indexes(self):
  else:
  namespace = _UJOIN % (self.__database.name, "system.indexes")
  res = helpers._first_batch(
- sock_info, namespace, {"ns": self.__full_name},
- 0, slave_ok, codec_options, ReadPreference.PRIMARY,
- cmd, self.database.client._event_listeners)
+ sock_info, self.__database.name, "system.indexes",
+ {"ns": self.__full_name}, 0, slave_ok, codec_options,
+ ReadPreference.PRIMARY, cmd,
+ self.database.client._event_listeners)
  data = res["data"]
  cursor = {
  "id": res["cursor_id"],

pymongo/command_cursor.py

@@ -112,6 +112,7 @@ def __send_message(self, operation):
  doc = helpers._unpack_response(response.data,
  self.__id,
  self.__collection.codec_options)
+ # TODO: check_command_response when getMore works with agg cursor
  except OperationFailure as exc:
  self.__killed = True
 
@@ -170,7 +171,12 @@ def _refresh(self):
 
  if self.__id: # Get More
  self.__send_message(
- _GetMore(self.__ns, self.__batch_size, self.__id))
+ _GetMore(self.__collection.database.name,
+ self.__collection.name,
+ self.__batch_size,
+ self.__id,
+ self.__collection.codec_options,
+ cmd_cursor=True))
 
  else: # Cursor id is zero nothing else to return
  self.__killed = True

pymongo/cursor.py

@@ -805,6 +805,7 @@ def __send_message(self, operation):
  client = self.__collection.database.client
  listeners = client._event_listeners
  publish = listeners.enabled_for_commands
+ from_command = False
 
  if operation:
  kwargs = {
@@ -827,6 +828,7 @@ def __send_message(self, operation):
  data = response.data
  cmd_duration = response.duration
  rqst_id = response.request_id
+ from_command = response.from_command
  except AutoReconnect:
  # Don't try to send kill cursors on another socket
  # or to another server. It can cause a _pinValue
@@ -869,6 +871,8 @@ def __send_message(self, operation):
  doc = helpers._unpack_response(response=data,
  cursor_id=self.__id,
  codec_options=self.__codec_options)
+ if from_command:
+ helpers._check_command_response(doc['data'][0])
  except OperationFailure as exc:
  self.__killed = True
 
@@ -913,7 +917,9 @@ def __send_message(self, operation):
  if publish:
  duration = (datetime.datetime.now() - start) + cmd_duration
  # Must publish in find / getMore / explain command response format.
- if cmd_name == "explain":
+ if from_command:
+ res = doc['data'][0]
+ elif cmd_name == "explain":
  res = doc["data"][0] if doc["number_returned"] else {}
  else:
  res = {"cursor": {"id": doc["cursor_id"],
@@ -926,12 +932,23 @@ def __send_message(self, operation):
  listeners.publish_command_success(
  duration, res, cmd_name, rqst_id, self.__address)
 
- self.__id = doc["cursor_id"]
+ if from_command and cmd_name != "explain":
+ cursor = doc['data'][0]['cursor']
+ self.__id = cursor['id']
+ if cmd_name == 'find':
+ documents = cursor['firstBatch']
+ else:
+ documents = cursor['nextBatch']
+ self.__data = deque(documents)
+ self.__retrieved += len(documents)
+ else:
+ self.__id = doc["cursor_id"]
+ self.__data = deque(doc["data"])
+ self.__retrieved += doc["number_returned"]
+
  if self.__id == 0:
  self.__killed = True
 
- self.__retrieved += doc["number_returned"]
- self.__data = deque(doc["data"])
 
  if self.__limit and self.__id and self.__limit <= self.__retrieved:
  self.__die()
@@ -959,7 +976,8 @@ def _refresh(self):
  else:
  ntoreturn = self.__limit
  self.__send_message(_Query(self.__query_flags,
- self.__collection.full_name,
+ self.__collection.database.name,
+ self.__collection.name,
  self.__skip,
  ntoreturn,
  self.__query_spec(),
@@ -982,9 +1000,11 @@ def _refresh(self):
  if self.__exhaust:
  self.__send_message(None)
  else:
- self.__send_message(_GetMore(self.__collection.full_name,
+ self.__send_message(_GetMore(self.__collection.database.name,
+ self.__collection.name,
  limit,
  self.__id,
+ self.__codec_options,
  self.__max_time_ms))
 
  else: # Cursor id is zero nothing else to return

pymongo/database.py

@@ -466,7 +466,7 @@ def _list_collections(self, sock_info, slave_okay, criteria=None):
  return CommandCursor(coll, cursor, sock_info.address)
  else:
  coll = self["system.namespaces"]
- res = _first_batch(sock_info, coll.full_name,
+ res = _first_batch(sock_info, coll.database.name, coll.name,
  criteria, 0, slave_okay,
  CodecOptions(), ReadPreference.PRIMARY, cmd,
  self.client._event_listeners)
@@ -598,7 +598,7 @@ def current_op(self, include_all=False):
  return sock_info.command("admin", cmd)
  else:
  spec = {"$all": True} if include_all else {}
- x = helpers._first_batch(sock_info, "admin.$cmd.sys.inprog",
+ x = helpers._first_batch(sock_info, "admin", "$cmd.sys.inprog",
  spec, -1, True, self.codec_options,
  ReadPreference.PRIMARY, cmd, self.client._event_listeners)
  return x.get('data', [None])[0]

pymongo/helpers.py

@@ -126,11 +126,11 @@ def _unpack_response(response, cursor_id=None, codec_options=CodecOptions()):
  error_object.get("code"),
  error_object)
 
- result = {}
- result["cursor_id"] = struct.unpack("<q", response[4:12])[0]
- result["starting_from"] = struct.unpack("<i", response[12:16])[0]
- result["number_returned"] = struct.unpack("<i", response[16:20])[0]
- result["data"] = bson.decode_all(response[20:], codec_options)
+ result = {"cursor_id": struct.unpack("<q", response[4:12])[0],
+  "starting_from": struct.unpack("<i", response[12:16])[0],
+  "number_returned": struct.unpack("<i", response[16:20])[0],
+  "data": bson.decode_all(response[20:], codec_options)}
+
  assert len(result["data"]) == result["number_returned"]
  return result
 
@@ -188,6 +188,8 @@ def _check_command_response(response, msg=None, allowable_errors=None):
  raise DuplicateKeyError(errmsg, code, response)
  elif code == 50:
  raise ExecutionTimeout(errmsg, code, response)
+ elif code == 43:
+ raise CursorNotFound(errmsg, code, response)
 
  msg = msg or "%s"
  raise OperationFailure(msg % errmsg, code, response)
@@ -233,11 +235,11 @@ def _check_gle_response(response):
  raise OperationFailure(details["err"], code, result)
 
 
-def _first_batch(sock_info, namespace, query, ntoreturn,
+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, namespace, 0, ntoreturn, query, None,
+ 0, db, coll, 0, ntoreturn, query, None,
  codec_options, read_preference, 0, ntoreturn)
 
  name = next(iter(cmd))
@@ -252,7 +254,7 @@ def _first_batch(sock_info, namespace, query, ntoreturn,
  if publish:
  encoding_duration = datetime.datetime.now() - start
  listeners.publish_command_start(
- cmd, namespace.split('.', 1)[0], request_id, sock_info.address)
+ cmd, db, request_id, sock_info.address)
  start = datetime.datetime.now()
 
  sock_info.send_message(msg, max_doc_size)

pymongo/message.py

@@ -26,7 +26,7 @@
 
 import bson
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
-from bson.py3compat import b, StringIO
+from bson.py3compat import b, StringIO, u
 from bson.son import SON
 try:
  from pymongo import _cmessage
@@ -60,6 +60,8 @@
  _DELETE: b'\x04deletes\x00\x00\x00\x00\x00',
 }
 
+_UJOIN = u("%s.%s")
+
 
 def _randint():
  """Generate a pseudo random 32 bit integer."""
@@ -201,14 +203,15 @@ def _gen_get_more_command(cursor_id, coll, batch_size, max_time_ms):
 class _Query(object):
  """A query operation."""
 
- __slots__ = ('flags', 'ns', 'ntoskip', 'ntoreturn', 'spec', 'fields',
- 'codec_options', 'read_preference', 'limit', 'batch_size',
- 'name')
+ __slots__ = ('flags', 'db', 'coll', 'ntoskip', 'ntoreturn', 'spec',
+ 'fields', 'codec_options', 'read_preference', 'limit',
+ 'batch_size', 'name')
 
- def __init__(self, flags, ns, ntoskip, ntoreturn, spec, fields,
+ def __init__(self, flags, db, coll, ntoskip, ntoreturn, spec, fields,
  codec_options, read_preference, limit, batch_size):
  self.flags = flags
- self.ns = ns
+ self.db = db
+ self.coll = coll
  self.ntoskip = ntoskip
  self.ntoreturn = ntoreturn
  self.spec = spec
@@ -224,51 +227,77 @@ def as_command(self):
 
  Should be called *after* get_message.
  """
- dbn, coll = self.ns.split('.', 1)
  if '$explain' in self.spec:
  self.name = 'explain'
  return _gen_explain_command(
- coll, self.spec, self.fields, self.ntoskip,
- self.limit, self.batch_size, self.flags), dbn
- return _gen_find_command(coll, self.spec, self.fields, self.ntoskip,
- self.limit, self.batch_size, self.flags), dbn
+ self.coll, self.spec, self.fields, self.ntoskip,
+ self.limit, self.batch_size, self.flags), self.db
+ return _gen_find_command(
+ self.coll, self.spec, self.fields, self.ntoskip, self.limit,
+ self.batch_size, self.flags), self.db
 
- def get_message(self, set_slave_ok, is_mongos):
+ def get_message(self, set_slave_ok, is_mongos, use_cmd=False):
  """Get a query message, possibly setting the slaveOk bit."""
- if is_mongos:
- self.spec = _maybe_add_read_preference(self.spec,
- self.read_preference)
  if set_slave_ok:
  # Set the slaveOk bit.
  flags = self.flags | 4
  else:
  flags = self.flags
- return query(flags, self.ns, self.ntoskip, self.ntoreturn,
- self.spec, self.fields, self.codec_options)
+
+ 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
+
+ if is_mongos:
+ spec = _maybe_add_read_preference(spec,
+ self.read_preference)
+
+ return query(flags, ns, self.ntoskip, ntoreturn,
+ spec, self.fields, self.codec_options)
 
 
 class _GetMore(object):
  """A getmore operation."""
 
- __slots__ = ('ns', 'ntoreturn', 'cursor_id', 'max_time_ms')
+ __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_time_ms',
+ 'codec_options', 'cmd_cursor')
 
  name = 'getMore'
 
- def __init__(self, ns, ntoreturn, cursor_id, max_time_ms=None):
- self.ns = ns
+ def __init__(self, db, coll, ntoreturn, cursor_id, codec_options,
+ max_time_ms=None, cmd_cursor=False):
+ 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."""
- dbn, coll = self.ns.split('.', 1)
- return _gen_get_more_command(
- self.cursor_id, coll, self.ntoreturn, self.max_time_ms), dbn
+ return _gen_get_more_command(self.cursor_id, self.coll,
+ self.ntoreturn, self.max_time_ms), self.db
 
- def get_message(self, dummy0, dummy1):
+ def get_message(self, dummy0, dummy1, use_cmd=False):
  """Get a getmore message."""
- return get_more(self.ns, self.ntoreturn, self.cursor_id)
+
+ ns = _UJOIN % (self.db, self.coll)
+
+ if use_cmd and not self.cmd_cursor:
+ ns = _UJOIN % (self.db, "$cmd")
+ spec = self.as_command()[0]
+
+ return query(4, ns, 0, -1, spec, None, self.codec_options)
+
+ return get_more(ns, self.ntoreturn, self.cursor_id)
 
 
 class _CursorAddress(tuple):

pymongo/mongo_client.py

@@ -1102,7 +1102,7 @@ def unlock(self):
  if exc.code != 125:
  raise
  else:
- helpers._first_batch(sock_info, "admin.$cmd.sys.unlock",
+ helpers._first_batch(sock_info, "admin", "$cmd.sys.unlock",
  {}, -1, True, self.codec_options,
  ReadPreference.PRIMARY, cmd, self._event_listeners)
 

pymongo/response.py

@@ -16,21 +16,24 @@
 
 
 class Response(object):
- __slots__ = ('_data', '_address', '_request_id', '_duration')
+ __slots__ = ('_data', '_address', '_request_id', '_duration',
+ '_from_command')
 
- def __init__(self, data, address, request_id, duration):
+ def __init__(self, data, address, request_id, duration, from_command):
  """Represent a response from the server.
 
  :Parameters:
  - `data`: Raw BSON bytes.
  - `address`: (host, port) of the source server.
  - `request_id`: The request id of this operation.
  - `duration`: The duration of the operation.
+ - `from_command`: if the response is the result of a db command.
  """
  self._data = data
  self._address = address
  self._request_id = request_id
  self._duration = duration
+ self._from_command = from_command
 
  @property
  def data(self):
@@ -52,12 +55,17 @@ def duration(self):
  """The duration of the operation."""
  return self._duration
 
+ @property
+ def from_command(self):
+ """If the response is a result from a db command."""
+ return self._from_command
+
 
 class ExhaustResponse(Response):
  __slots__ = ('_socket_info', '_pool')
 
- def __init__(
- self, data, address, socket_info, pool, request_id, duration):
+ def __init__(self, data, address, socket_info, pool, request_id, duration,
+  from_command):
  """Represent a response to an exhaust cursor's initial query.
 
  :Parameters:
@@ -67,11 +75,13 @@ def __init__(
  - `pool`: The Pool from which the SocketInfo came.
  - `request_id`: The request id of this operation.
  - `duration`: The duration of the operation.
+ - `from_command`: If the response is the result of a db command.
  """
  super(ExhaustResponse, self).__init__(data,
  address,
  request_id,
- duration)
+ duration,
+ from_command)
  self._socket_info = socket_info
  self._pool = pool