def receive_message(sock,
operation,
request_id,
max_message_size=MAX_MESSAGE_SIZE):
"""Receive a raw BSON message or raise socket.error."""
header = _receive_data_on_socket(sock, 16)
length = _UNPACK_INT(header[:4])[0]
actual_op = _UNPACK_INT(header[12:])[0]
if operation != actual_op:
raise ProtocolError("Got opcode %r but expected "
"%r" % (actual_op, operation))
# No request_id for exhaust cursor "getMore".
if request_id is not None:
response_id = _UNPACK_INT(header[8:12])[0]
if request_id != response_id:
raise ProtocolError("Got response id %r but expected "
"%r" % (response_id, request_id))
if length <= 16:
raise ProtocolError("Message length (%r) not longer than standard "
"message header size (16)" % (length, ))
if length > max_message_size:
raise ProtocolError("Message length (%r) is larger than server max "
"message size (%r)" % (length, max_message_size))
return _receive_data_on_socket(sock, length - 16)
def receive_message(sock, request_id, max_message_size=MAX_MESSAGE_SIZE):
"""Receive a raw BSON message or raise socket.error."""
# Ignore the response's request id.
length, _, response_to, op_code = _UNPACK_HEADER(
_receive_data_on_socket(sock, 16))
# No request_id for exhaust cursor "getMore".
if request_id is not None:
if request_id != response_to:
raise ProtocolError("Got response id %r but expected "
"%r" % (response_to, request_id))
if length <= 16:
raise ProtocolError("Message length (%r) not longer than standard "
"message header size (16)" % (length, ))
if length > max_message_size:
raise ProtocolError("Message length (%r) is larger than server max "
"message size (%r)" % (length, max_message_size))
if op_code == 2012:
op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
_receive_data_on_socket(sock, 9))
data = decompress(_receive_data_on_socket(sock, length - 25),
compressor_id)
else:
data = _receive_data_on_socket(sock, length - 16)
try:
unpack_reply = _UNPACK_REPLY[op_code]
except KeyError:
raise ProtocolError("Got opcode %r but expected "
"%r" % (op_code, _UNPACK_REPLY.keys()))
return unpack_reply(data)
async def _read_loop_step(self) -> None:
header = await self.reader.readexactly(16)
length, = struct.unpack('<i', header[:4])
if length < 16:
raise ProtocolError('Message length ({}) not longer than standard '
'message header size (16)'.format(length))
response_id, = struct.unpack('<i', header[8:12])
if response_id not in self.__request_futures:
raise ProtocolError(
'Got response id {} but expected but request with such id was not sent.'
.format(response_id))
message_data = await self.reader.readexactly(length - 16)
ft = self.__request_futures.pop(response_id)
if not ft.cancelled():
ft.set_result(message_data)
def unpack(cls, msg):
"""Construct an _OpMsg from raw bytes."""
flags, first_payload_type, first_payload_size = cls.UNPACK_FROM(msg)
if flags != 0:
if flags & cls.CHECKSUM_PRESENT:
raise ProtocolError("Unsupported OP_MSG flag checksumPresent: "
"0x%x" % (flags, ))
if flags ^ cls.MORE_TO_COME:
raise ProtocolError("Unsupported OP_MSG flags: 0x%x" %
(flags, ))
if first_payload_type != 0:
raise ProtocolError("Unsupported OP_MSG payload type: "
"0x%x" % (first_payload_type, ))
if len(msg) != first_payload_size + 5:
raise ProtocolError("Unsupported OP_MSG reply: >1 section")
payload_document = msg[5:]
return cls(flags, payload_document)
def receive_message(sock, request_id, max_message_size=MAX_MESSAGE_SIZE):
"""Receive a raw BSON message or raise socket.error."""
# Ignore the response's request id.
length, _, response_to, op_code = _UNPACK_HEADER(
_receive_data_on_socket(sock, 16))
if op_code != _OpReply.OP_CODE:
raise ProtocolError("Got opcode %r but expected "
"%r" % (op_code, _OpReply.OP_CODE))
# No request_id for exhaust cursor "getMore".
if request_id is not None:
if request_id != response_to:
raise ProtocolError("Got response id %r but expected "
"%r" % (response_to, request_id))
if length <= 16:
raise ProtocolError("Message length (%r) not longer than standard "
"message header size (16)" % (length, ))
if length > max_message_size:
raise ProtocolError("Message length (%r) is larger than server max "
"message size (%r)" % (length, max_message_size))
return _OpReply.unpack(_receive_data_on_socket(sock, length - 16))
def _unpack_response(response,
cursor_id=None,
codec_options=_UNICODE_REPLACE_CODEC_OPTIONS):
"""Unpack a response from the database.
Check the response for errors and unpack, returning a dictionary
containing the response data.
Can raise CursorNotFound, NotMasterError, ExecutionTimeout, or
OperationFailure.
:Parameters:
- `response`: byte string as returned from the database
- `cursor_id` (optional): cursor_id we sent to get this response -
used for raising an informative exception when we get cursor id not
valid at server response
- `codec_options` (optional): an instance of
:class:`~bson.codec_options.CodecOptions`
"""
response_flag = struct.unpack("<i", response[:4])[0]
if response_flag & 1:
# Shouldn't get this response if we aren't doing a getMore
if cursor_id is None:
raise ProtocolError("No cursor id for getMore operation")
# Fake a getMore command response. OP_GET_MORE provides no document.
msg = "Cursor not found, cursor id: %d" % (cursor_id,)
errobj = {"ok": 0, "errmsg": msg, "code": 43}
raise CursorNotFound(msg, 43, errobj)
elif response_flag & 2:
error_object = bson.BSON(response[20:]).decode()
# Fake the ok field if it doesn't exist.
error_object.setdefault("ok", 0)
if error_object["$err"].startswith("not master"):
raise NotMasterError(error_object["$err"], error_object)
elif error_object.get("code") == 50:
raise ExecutionTimeout(error_object.get("$err"),
error_object.get("code"),
error_object)
raise OperationFailure("database error: %s" %
error_object.get("$err"),
error_object.get("code"),
error_object)
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
def raw_response(self, cursor_id=None, user_fields=None):
"""Check the response header from the database, without decoding BSON.
Check the response for errors and unpack.
Can raise CursorNotFound, NotPrimaryError, ExecutionTimeout, or
OperationFailure.
:Parameters:
- `cursor_id` (optional): cursor_id we sent to get this response -
used for raising an informative exception when we get cursor id not
valid at server response.
"""
if self.flags & 1:
# Shouldn't get this response if we aren't doing a getMore
if cursor_id is None:
raise ProtocolError("No cursor id for getMore operation")
# Fake a getMore command response. OP_GET_MORE provides no
# document.
msg = "Cursor not found, cursor id: %d" % (cursor_id, )
errobj = {"ok": 0, "errmsg": msg, "code": 43}
raise CursorNotFound(msg, 43, errobj)
elif self.flags & 2:
error_object: dict = bson.BSON(self.documents).decode()
# Fake the ok field if it doesn't exist.
error_object.setdefault("ok", 0)
if error_object["$err"].startswith(HelloCompat.LEGACY_ERROR):
raise NotPrimaryError(error_object["$err"], error_object)
elif error_object.get("code") == 50:
default_msg = "operation exceeded time limit"
raise ExecutionTimeout(error_object.get("$err", default_msg),
error_object.get("code"), error_object)
raise OperationFailure(
"database error: %s" % error_object.get("$err"),
error_object.get("code"),
error_object,
)
if self.documents:
return [self.documents]
return []
