PYTHON-1674 Refactor Cursor/CommandCursor.__send_message

Move exhaust getMore out of Cursor. Move cursor command response decoding into Server so that all command listener events can be published in the same location.

Author: ShaneHarvey

pymongo/command_cursor.py

@@ -20,7 +20,7 @@
 
 from bson.py3compat import integer_types
 from pymongo import helpers
-from pymongo.errors import (AutoReconnect,
+from pymongo.errors import (ConnectionFailure,
  InvalidOperation,
  NotMasterError,
  OperationFailure)
@@ -127,87 +127,42 @@ def kill():
  self.__end_session(True)
 
  client = self.__collection.database.client
- listeners = client._event_listeners
- publish = listeners.enabled_for_commands
- start = datetime.datetime.now()
-
- def duration(): return datetime.datetime.now() - start
-
  try:
  response = client._send_message_with_response(
- operation, address=self.__address)
- except AutoReconnect:
+ operation, address=self.__address,
+ unpack_res=self._unpack_response)
+ except OperationFailure:
+ kill()
+ raise
+ except NotMasterError:
+ # Don't send kill cursors to another server after a "not master"
+ # error. It's completely pointless.
+ kill()
+ raise
+ except ConnectionFailure:
  # Don't try to send kill cursors on another socket
  # or to another server. It can cause a _pinValue
  # assertion on some server releases if we get here
  # due to a socket timeout.
  kill()
  raise
+ except Exception:
+ # Close the cursor
+ self.__die()
+ raise
 
- rqst_id = response.request_id
  from_command = response.from_command
  reply = response.data
-
- try:
- with client._reset_on_error(self.__address, self.__session):
- user_fields = None
- legacy_response = True
- if from_command:
- user_fields = {'cursor': {'nextBatch': 1}}
- legacy_response = False
- docs = self._unpack_response(
- reply, self.__id, self.__collection.codec_options,
- legacy_response=legacy_response, user_fields=user_fields)
- if from_command:
- first = docs[0]
- client._process_response(first, self.__session)
- helpers._check_command_response(first)
- except OperationFailure as exc:
- kill()
-
- if publish:
- listeners.publish_command_failure(
- duration(), exc.details, "getMore", rqst_id, self.__address)
-
- raise
- except NotMasterError as exc:
- # Don't send kill cursors to another server after a "not master"
- # error. It's completely pointless.
- kill()
-
- if publish:
- listeners.publish_command_failure(
- duration(), exc.details, "getMore", rqst_id, self.__address)
-
- raise
- except Exception as exc:
- if publish:
- listeners.publish_command_failure(
- duration(), _convert_exception(exc), "getMore", rqst_id,
- self.__address)
- raise
+ docs = response.docs
 
  if from_command:
  cursor = docs[0]['cursor']
  documents = cursor['nextBatch']
  self.__id = cursor['id']
- if publish:
- listeners.publish_command_success(
- duration(), docs[0], "getMore", rqst_id,
- self.__address)
  else:
  documents = docs
  self.__id = reply.cursor_id
 
- if publish:
- # Must publish in getMore command response format.
- res = {"cursor": {"id": self.__id,
- "ns": self.__collection.full_name,
- "nextBatch": documents},
- "ok": 1}
- listeners.publish_command_success(
- duration(), res, "getMore", rqst_id, self.__address)
-
  if self.__id == 0:
  kill()
  self.__data = deque(documents)
@@ -239,7 +194,8 @@ def _refresh(self):
  read_pref,
  self.__session,
  self.__collection.database.client,
- self.__max_await_time_ms))
+ self.__max_await_time_ms,
+ False))
  else: # Cursor id is zero nothing else to return
  self.__killed = True
  self.__end_session(True)

pymongo/cursor.py

@@ -15,7 +15,6 @@
 """Cursor class to iterate over Mongo query results."""
 
 import copy
-import datetime
 import warnings
 
 from collections import deque
@@ -29,13 +28,11 @@
 from pymongo import helpers
 from pymongo.common import validate_boolean, validate_is_mapping
 from pymongo.collation import validate_collation_or_none
-from pymongo.errors import (AutoReconnect,
- ConnectionFailure,
+from pymongo.errors import (ConnectionFailure,
  InvalidOperation,
  NotMasterError,
  OperationFailure)
-from pymongo.message import (_convert_exception,
- _CursorAddress,
+from pymongo.message import (_CursorAddress,
  _GetMore,
  _RawBatchGetMore,
  _Query,
@@ -50,7 +47,6 @@
  "await_data": 32,
  "exhaust": 64,
  "partial": 128}
-_CURSOR_DOC_FIELDS = {'cursor': {'firstBatch': 1, 'nextBatch': 1}}
 
 
 class CursorType(object):
@@ -941,129 +937,54 @@ def __send_message(self, operation):
  Can raise ConnectionFailure.
  """
  client = self.__collection.database.client
- listeners = client._event_listeners
- publish = listeners.enabled_for_commands
- from_command = False
- start = datetime.datetime.now()
-
- def duration(): return datetime.datetime.now() - start
-
- if operation:
- try:
- response = client._send_message_with_response(
- operation, exhaust=self.__exhaust, address=self.__address)
- self.__address = response.address
- if self.__exhaust:
- # 'response' is an ExhaustResponse.
- self.__exhaust_mgr = _SocketManager(response.socket_info,
- response.pool)
-
- cmd_name = operation.name
- reply = response.data
- 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
- # assertion on some server releases if we get here
- # due to a socket timeout.
- self.__killed = True
- self.__die()
- raise
- else:
- # Exhaust cursor - no getMore message.
- rqst_id = 0
- cmd_name = 'getMore'
- if publish:
- # Fake a getMore command.
- cmd = SON([('getMore', self.__id),
- ('collection', self.__collection.name)])
- if self.__batch_size:
- cmd['batchSize'] = self.__batch_size
- if self.__max_time_ms:
- cmd['maxTimeMS'] = self.__max_time_ms
- listeners.publish_command_start(
- cmd, self.__collection.database.name, 0, self.__address)
- try:
- reply = self.__exhaust_mgr.sock.receive_message(None)
- except Exception as exc:
- if publish:
- listeners.publish_command_failure(
- duration(), _convert_exception(exc), cmd_name, rqst_id,
- self.__address)
- if isinstance(exc, ConnectionFailure):
- self.__die()
- raise
-
  try:
- with client._reset_on_error(self.__address, self.__session):
- user_fields = None
- legacy_response = True
- if from_command:
- user_fields = _CURSOR_DOC_FIELDS
- legacy_response = False
- docs = self._unpack_response(
- reply, self.__id, self.__collection.codec_options,
- legacy_response=legacy_response, user_fields=user_fields)
- if from_command:
- first = docs[0]
- client._process_response(first, self.__session)
- helpers._check_command_response(first)
- except OperationFailure as exc:
+ response = client._send_message_with_response(
+ operation, exhaust=self.__exhaust, address=self.__address,
+ unpack_res=self._unpack_response)
+ except OperationFailure:
  self.__killed = True
 
  # Make sure exhaust socket is returned immediately, if necessary.
  self.__die()
 
- if publish:
- listeners.publish_command_failure(
- duration(), exc.details, cmd_name, rqst_id, self.__address)
-
  # If this is a tailable cursor the error is likely
  # due to capped collection roll over. Setting
  # self.__killed to True ensures Cursor.alive will be
  # False. No need to re-raise.
  if self.__query_flags & _QUERY_OPTIONS["tailable_cursor"]:
  return
  raise
- except NotMasterError as exc:
+ except NotMasterError:
  # Don't send kill cursors to another server after a "not master"
  # error. It's completely pointless.
  self.__killed = True
 
  # Make sure exhaust socket is returned immediately, if necessary.
  self.__die()
 
- if publish:
- listeners.publish_command_failure(
- duration(), exc.details, cmd_name, rqst_id, self.__address)
-
  raise
- except Exception as exc:
- if publish:
- listeners.publish_command_failure(
- duration(), _convert_exception(exc), cmd_name, rqst_id,
- self.__address)
+ except ConnectionFailure:
+ # Don't try to send kill cursors on another socket
+ # or to another server. It can cause a _pinValue
+ # assertion on some server releases if we get here
+ # due to a socket timeout.
+ self.__killed = True
+ self.__die()
+ raise
+ except Exception:
+ # Close the cursor
+ self.__die()
  raise
 
- if publish:
- # Must publish in find / getMore / explain command response format.
- if from_command:
- res = docs[0]
- elif cmd_name == "explain":
- res = docs[0] if docs else {}
- else:
- res = {"cursor": {"id": reply.cursor_id,
- "ns": self.__collection.full_name},
- "ok": 1}
- if cmd_name == "find":
- res["cursor"]["firstBatch"] = docs
- else:
- res["cursor"]["nextBatch"] = docs
- listeners.publish_command_success(
- duration(), res, cmd_name, rqst_id, self.__address)
+ self.__address = response.address
+ if self.__exhaust and not self.__exhaust_mgr:
+ # 'response' is an ExhaustResponse.
+ self.__exhaust_mgr = _SocketManager(response.socket_info,
+ response.pool)
 
- if from_command:
+ cmd_name = operation.name
+ docs = response.docs
+ if response.from_command:
  if cmd_name != "explain":
  cursor = docs[0]['cursor']
  self.__id = cursor['id']
@@ -1078,9 +999,9 @@ def duration(): return datetime.datetime.now() - start
  self.__data = deque(docs)
  self.__retrieved += len(docs)
  else:
- self.__id = reply.cursor_id
+ self.__id = response.data.cursor_id
  self.__data = deque(docs)
- self.__retrieved += reply.number_returned
+ self.__retrieved += response.data.number_returned
 
  if self.__id == 0:
  self.__killed = True
@@ -1147,19 +1068,17 @@ def _refresh(self):
  limit = self.__batch_size
 
  # Exhaust cursors don't send getMore messages.
- if self.__exhaust:
- self.__send_message(None)
- else:
- g = self._getmore_class(self.__collection.database.name,
- self.__collection.name,
- limit,
- self.__id,
- self.__codec_options,
- self._read_preference(),
- self.__session,
- self.__collection.database.client,
- self.__max_await_time_ms)
- self.__send_message(g)
+ g = self._getmore_class(self.__collection.database.name,
+ self.__collection.name,
+ limit,
+ self.__id,
+ self.__codec_options,
+ self._read_preference(),
+ self.__session,
+ self.__collection.database.client,
+ self.__max_await_time_ms,
+ self.__exhaust_mgr)
+ self.__send_message(g)
 
  return len(self.__data)