_INSERT: b'\x04documents\x00\x00\x00\x00\x00',

_UPDATE: b'\x04updates\x00\x00\x00\x00\x00',

_DELETE: b'\x04deletes\x00\x00\x00\x00\x00',

'insert': 'documents',

'update': 'updates',

'delete': 'deletes'

"""Add $readPreference to spec when appropriate."""

mode = read_preference.mode

tag_sets = read_preference.tag_sets

max_staleness = read_preference.max_staleness

# Only add $readPreference if it's something other than primary to avoid

# problems with mongos versions that don't support read preferences. Also,

# for maximum backwards compatibility, don't add $readPreference for

# secondaryPreferred unless tags or maxStalenessSeconds are in use (setting

# the slaveOkay bit has the same effect).

if mode and (

mode != ReadPreference.SECONDARY_PREFERRED.mode

or tag_sets != [{}]

or max_staleness != -1):

if "$query" not in spec:

spec = SON([("$query", spec)])

spec["$readPreference"] = read_preference.document

"""Convert a legacy write result to write commmand format."""

# Based on _merge_legacy from bulk.py

affected = result.get("n", 0)

res = {"ok": 1, "n": affected}

errmsg = result.get("errmsg", result.get("err", ""))

if errmsg:

# The write was successful on at least the primary so don't return.

if result.get("wtimeout"):

res["writeConcernError"] = {"errmsg": errmsg,

"code": 64,

"errInfo": {"wtimeout": True}}

else:

# The write failed.

error = {"index": 0,

"code": result.get("code", 8),

"errmsg": errmsg}

if "errInfo" in result:

error["errInfo"] = result["errInfo"]

res["writeErrors"] = [error]

return res

if operation == "insert":

# GLE result for insert is always 0 in most MongoDB versions.

res["n"] = len(command['documents'])

elif operation == "update":

if "upserted" in result:

res["upserted"] = [{"index": 0, "_id": result["upserted"]}]

# Versions of MongoDB before 2.6 don't return the _id for an

# upsert if _id is not an ObjectId.

elif result.get("updatedExisting") is False and affected == 1:

# If _id is in both the update document *and* the query spec

# the update document _id takes precedence.

update = command['updates'][0]

_id = update["u"].get("_id", update["q"].get("_id"))

res["upserted"] = [{"index": 0, "_id": _id}]

('tailable', 2),

('oplogReplay', 8),

('noCursorTimeout', 16),

('awaitData', 32),

('$query', 'filter'),

('$orderby', 'sort'),

('$hint', 'hint'),

('$comment', 'comment'),

('$maxScan', 'maxScan'),

('$maxTimeMS', 'maxTimeMS'),

('$max', 'max'),

('$min', 'min'),

('$returnKey', 'returnKey'),

('$showRecordId', 'showRecordId'),

('$showDiskLoc', 'showRecordId'), # <= MongoDb 3.0

read_concern, collation=None, session=None):

"""Generate a find command document."""

cmd = SON([('find', coll)])

if '$query' in spec:

cmd.update([(_MODIFIERS[key], val) if key in _MODIFIERS else (key, val)

for key, val in spec.items()])

if '$explain' in cmd:

cmd.pop('$explain')

if '$readPreference' in cmd:

cmd.pop('$readPreference')

else:

cmd['filter'] = spec

if projection:

cmd['projection'] = projection

if skip:

cmd['skip'] = skip

if limit:

cmd['limit'] = abs(limit)

if limit < 0:

cmd['singleBatch'] = True

if batch_size:

cmd['batchSize'] = batch_size

if read_concern.level and not (session and session.in_transaction):

cmd['readConcern'] = read_concern.document

if collation:

cmd['collation'] = collation

if options:

cmd.update([(opt, True)

for opt, val in _OPTIONS.items()

if options & val])

"""Generate a getMore command document."""

cmd = SON([('getMore', cursor_id),

('collection', coll)])

if batch_size:

cmd['batchSize'] = batch_size

if max_await_time_ms is not None:

cmd['maxTimeMS'] = max_await_time_ms

"""A query operation."""

__slots__ = ('flags', 'db', 'coll', 'ntoskip', 'spec',

'fields', 'codec_options', 'read_preference', 'limit',

'batch_size', 'name', 'read_concern', 'collation',

'session', 'client', '_as_command')

# For compatibility with the _GetMore class.

exhaust_mgr = None

cursor_id = None

def __init__(self, flags, db, coll, ntoskip, spec, fields,

codec_options, read_preference, limit,

batch_size, read_concern, collation, session, client):

self.flags = flags

self.db = db

self.coll = coll

self.ntoskip = ntoskip

self.spec = spec

self.fields = fields

self.codec_options = codec_options

self.read_preference = read_preference

self.read_concern = read_concern

self.limit = limit

self.batch_size = batch_size

self.collation = collation

self.session = session

self.client = client

self.name = 'find'

self._as_command = None

def namespace(self):

return _UJOIN % (self.db, self.coll)

def use_command(self, sock_info, exhaust):

use_find_cmd = False

if sock_info.max_wire_version >= 4:

if not exhaust:

use_find_cmd = True

elif not self.read_concern.ok_for_legacy:

raise ConfigurationError(

'read concern level of %s is not valid '

'with a max wire version of %d.'

% (self.read_concern.level,

sock_info.max_wire_version))

if sock_info.max_wire_version < 5 and self.collation is not None:

raise ConfigurationError(

'Specifying a collation is unsupported with a max wire '

'version of %d.' % (sock_info.max_wire_version,))

sock_info.validate_session(self.client, self.session)

return use_find_cmd

def as_command(self, sock_info):

"""Return a find command document for this query."""

# We use the command twice: on the wire and for command monitoring.

# Generate it once, for speed and to avoid repeating side-effects.

if self._as_command is not None:

return self._as_command

explain = '$explain' in self.spec

cmd = _gen_find_command(

self.coll, self.spec, self.fields, self.ntoskip,

self.limit, self.batch_size, self.flags, self.read_concern,

self.collation, self.session)

if explain:

self.name = 'explain'

cmd = SON([('explain', cmd)])

session = self.session

if session:

session._apply_to(cmd, False, self.read_preference)

# Explain does not support readConcern.

if (not explain and session.options.causal_consistency

and session.operation_time is not None

and not session.in_transaction):

cmd.setdefault(

'readConcern', {})[

'afterClusterTime'] = session.operation_time

sock_info.send_cluster_time(cmd, session, self.client)

# Support auto encryption

client = self.client

if (client._encrypter and

not client._encrypter._bypass_auto_encryption):

cmd = client._encrypter.encrypt(

self.db, cmd, False, self.codec_options)

self._as_command = cmd, self.db

return self._as_command

def get_message(self, set_slave_ok, sock_info, use_cmd=False):

"""Get a query message, possibly setting the slaveOk bit."""

if set_slave_ok:

# Set the slaveOk bit.

flags = self.flags | 4

else:

flags = self.flags

ns = self.namespace()

spec = self.spec

if use_cmd:

spec = self.as_command(sock_info)[0]

if sock_info.op_msg_enabled:

request_id, msg, size, _ = _op_msg(

0, spec, self.db, self.read_preference,

set_slave_ok, False, self.codec_options,

ctx=sock_info.compression_context)

return request_id, msg, size

ns = _UJOIN % (self.db, "$cmd")

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 sock_info.is_mongos:

spec = _maybe_add_read_preference(spec,

self.read_preference)

return query(flags, ns, self.ntoskip, ntoreturn,

spec, None if use_cmd else self.fields,

"""A getmore operation."""

__slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_await_time_ms',

'codec_options', 'read_preference', 'session', 'client',

'exhaust_mgr', '_as_command')

name = 'getMore'

def __init__(self, db, coll, ntoreturn, cursor_id, codec_options,

read_preference, session, client, max_await_time_ms,

exhaust_mgr):

self.db = db

self.coll = coll

self.ntoreturn = ntoreturn

self.cursor_id = cursor_id

self.codec_options = codec_options

self.read_preference = read_preference

self.session = session

self.client = client

self.max_await_time_ms = max_await_time_ms

self.exhaust_mgr = exhaust_mgr

self._as_command = None

def namespace(self):

return _UJOIN % (self.db, self.coll)

def use_command(self, sock_info, exhaust):

sock_info.validate_session(self.client, self.session)

return sock_info.max_wire_version >= 4 and not exhaust

def as_command(self, sock_info):

"""Return a getMore command document for this query."""

# See _Query.as_command for an explanation of this caching.

if self._as_command is not None:

return self._as_command

cmd = _gen_get_more_command(self.cursor_id, self.coll,

self.ntoreturn,

self.max_await_time_ms)

if self.session:

self.session._apply_to(cmd, False, self.read_preference)

sock_info.send_cluster_time(cmd, self.session, self.client)

# Support auto encryption

client = self.client

if (client._encrypter and

not client._encrypter._bypass_auto_encryption):

cmd = client._encrypter.encrypt(

self.db, cmd, False, self.codec_options)

self._as_command = cmd, self.db

return self._as_command

def get_message(self, dummy0, sock_info, use_cmd=False):

"""Get a getmore message."""

ns = self.namespace()

ctx = sock_info.compression_context

if use_cmd:

spec = self.as_command(sock_info)[0]

if sock_info.op_msg_enabled:

request_id, msg, size, _ = _op_msg(

0, spec, self.db, ReadPreference.PRIMARY,

False, False, self.codec_options,

ctx=sock_info.compression_context)

return request_id, msg, size

ns = _UJOIN % (self.db, "$cmd")

return query(0, ns, 0, -1, spec, None, self.codec_options, ctx=ctx)

def use_command(self, socket_info, exhaust):

# Compatibility checks.

super(_RawBatchQuery, self).use_command(socket_info, exhaust)

return False

def get_message(self, set_slave_ok, sock_info, use_cmd=False):

# Always pass False for use_cmd.

return super(_RawBatchQuery, self).get_message(

def use_command(self, socket_info, exhaust):

return False

def get_message(self, set_slave_ok, sock_info, use_cmd=False):

# Always pass False for use_cmd.

return super(_RawBatchGetMore, self).get_message(

"""The server address (host, port) of a cursor, with namespace property."""

def __new__(cls, address, namespace):

self = tuple.__new__(cls, address)

self.__namespace = namespace

return self

@property

def namespace(self):

"""The namespace this cursor."""

return self.__namespace

def __hash__(self):

# Two _CursorAddress instances with different namespaces

# must not hash the same.

return (self + (self.__namespace,)).__hash__()

def __eq__(self, other):

if isinstance(other, _CursorAddress):

return (tuple(self) == tuple(other)

and self.namespace == other.namespace)

return NotImplemented

def __ne__(self, other):

"""Takes message data, compresses it, and adds an OP_COMPRESSED header."""

compressed = ctx.compress(data)

request_id = _randint()

header = _pack_compression_header(

_COMPRESSION_HEADER_SIZE + len(compressed), # Total message length

request_id, # Request id

0, # responseTo

2012, # operation id

operation, # original operation id

len(data), # uncompressed message length

ctx.compressor_id) # compressor id

"""Data to send to do a lastError.

cmd = SON([("getlasterror", 1)])

cmd.update(args)

splitns = namespace.split('.', 1)

return query(0, splitns[0] + '.$cmd', 0, -1, cmd,

"""Takes message data and adds a message header based on the operation.

rid = _randint()

message = _pack_header(16 + len(data), rid, 0, operation)

"""Get an OP_INSERT message"""

encode = _dict_to_bson # Make local. Uses extensions.

if len(docs) == 1:

encoded = encode(docs[0], check_keys, opts)

return b"".join([

b"\x00\x00\x00\x00", # Flags don't matter for one doc.

_make_c_string(collection_name),

encoded]), len(encoded)

encoded = [encode(doc, check_keys, opts) for doc in docs]

if not encoded:

raise InvalidOperation("cannot do an empty bulk insert")

return b"".join([

_pack_int(flags),

_make_c_string(collection_name),

collection_name, docs, check_keys, continue_on_error, opts, ctx):

"""Internal compressed unacknowledged insert message helper."""

op_insert, max_bson_size = _insert(

collection_name, docs, check_keys, continue_on_error, opts)

rid, msg = _compress(2002, op_insert, ctx)

safe, last_error_args, continue_on_error, opts):

"""Internal insert message helper."""

op_insert, max_bson_size = _insert(

collection_name, docs, check_keys, continue_on_error, opts)

rid, msg = __pack_message(2002, op_insert)

if safe:

rid, gle, _ = __last_error(collection_name, last_error_args)

return rid, msg + gle, max_bson_size

safe, last_error_args, continue_on_error, opts, ctx=None):

"""Get an **insert** message."""

if ctx:

return _insert_compressed(

collection_name, docs, check_keys, continue_on_error, opts, ctx)

return _insert_uncompressed(collection_name, docs, check_keys, safe,

"""Get an OP_UPDATE message."""

flags = 0

if upsert:

flags += 1

if multi:

flags += 2

encode = _dict_to_bson # Make local. Uses extensions.

encoded_update = encode(doc, check_keys, opts)

return b"".join([

_ZERO_32,

_make_c_string(collection_name),

_pack_int(flags),

encode(spec, False, opts),

collection_name, upsert, multi, spec, doc, check_keys, opts, ctx):

"""Internal compressed unacknowledged update message helper."""

op_update, max_bson_size = _update(

collection_name, upsert, multi, spec, doc, check_keys, opts)

rid, msg = _compress(2001, op_update, ctx)

doc, safe, last_error_args, check_keys, opts):

"""Internal update message helper."""

op_update, max_bson_size = _update(

collection_name, upsert, multi, spec, doc, check_keys, opts)

rid, msg = __pack_message(2001, op_update)

if safe:

rid, gle, _ = __last_error(collection_name, last_error_args)

return rid, msg + gle, max_bson_size

doc, safe, last_error_args, check_keys, opts, ctx=None):

"""Get an **update** message."""

if ctx:

return _update_compressed(

collection_name, upsert, multi, spec, doc, check_keys, opts, ctx)

return _update_uncompressed(collection_name, upsert, multi, spec,

"""Get a OP_MSG message.

# Encode the command document in payload 0 without checking keys.

encoded = _dict_to_bson(command, False, opts)

flags_type = _pack_op_msg_flags_type(flags, 0)

total_size = len(encoded)

max_doc_size = 0

if identifier:

type_one = _pack_byte(1)

cstring = _make_c_string(identifier)

encoded_docs = [_dict_to_bson(doc, check_keys, opts) for doc in docs]

size = len(cstring) + sum(len(doc) for doc in encoded_docs) + 4

encoded_size = _pack_int(size)

total_size += size

max_doc_size = max(len(doc) for doc in encoded_docs)

data = ([flags_type, encoded, type_one, encoded_size, cstring] +

encoded_docs)

else:

data = [flags_type, encoded]

ctx):

"""Internal OP_MSG message helper."""

msg, total_size, max_bson_size = _op_msg_no_header(

flags, command, identifier, docs, check_keys, opts)

rid, msg = _compress(2013, msg, ctx)

"""Internal compressed OP_MSG message helper."""

data, total_size, max_bson_size = _op_msg_no_header(

flags, command, identifier, docs, check_keys, opts)

request_id, op_message = __pack_message(2013, data)

opts, ctx=None):

"""Get a OP_MSG message."""

command['$db'] = dbname

if "$readPreference" not in command:

if slave_ok and not read_preference.mode:

command["$readPreference"] = (

ReadPreference.PRIMARY_PREFERRED.document)

else:

command["$readPreference"] = read_preference.document

name = next(iter(command))

try:

identifier = _FIELD_MAP.get(name)

docs = command.pop(identifier)

except KeyError:

identifier = ""

docs = None

try:

if ctx:

return _op_msg_compressed(

flags, command, identifier, docs, check_keys, opts, ctx)

return _op_msg_uncompressed(

flags, command, identifier, docs, check_keys, opts)

finally:

# Add the field back to the command.

if identifier:

num_to_return, query, field_selector, opts, check_keys):

"""Get an OP_QUERY message."""

encoded = _dict_to_bson(query, check_keys, opts)

if field_selector:

efs = _dict_to_bson(field_selector, False, opts)

else:

efs = b""

max_bson_size = max(len(encoded), len(efs))

return b"".join([

_pack_int(options),

_make_c_string(collection_name),

_pack_int(num_to_skip),

_pack_int(num_to_return),

encoded,

num_to_return, query, field_selector,

opts, check_keys=False, ctx=None):

"""Internal compressed query message helper."""

op_query, max_bson_size = _query(

options,

collection_name,

num_to_skip,

num_to_return,

query,

field_selector,

opts,

check_keys)

rid, msg = _compress(2004, op_query, ctx)

num_to_return, query, field_selector, opts, check_keys=False):

"""Internal query message helper."""

op_query, max_bson_size = _query(

options,

collection_name,

num_to_skip,

num_to_return,

query,

field_selector,

opts,

check_keys)

rid, msg = __pack_message(2004, op_query)

query, field_selector, opts, check_keys=False, ctx=None):

"""Get a **query** message."""

if ctx:

return _query_compressed(options, collection_name, num_to_skip,

num_to_return, query, field_selector,

opts, check_keys, ctx)

return _query_uncompressed(options, collection_name, num_to_skip,

num_to_return, query, field_selector, opts,

"""Get an OP_GET_MORE message."""

return b"".join([

_ZERO_32,

_make_c_string(collection_name),

_pack_int(num_to_return),

"""Get a **getMore** message."""

if ctx:

return _get_more_compressed(

collection_name, num_to_return, cursor_id, ctx)

"""Get an OP_DELETE message."""

encoded = _dict_to_bson(spec, False, opts) # Uses extensions.

return b"".join([

_ZERO_32,

_make_c_string(collection_name),

_pack_int(flags),

"""Internal compressed unacknowledged delete message helper."""

op_delete, max_bson_size = _delete(collection_name, spec, opts, flags)

rid, msg = _compress(2006, op_delete, ctx)

collection_name, spec, safe, last_error_args, opts, flags=0):

"""Internal delete message helper."""

op_delete, max_bson_size = _delete(collection_name, spec, opts, flags)

rid, msg = __pack_message(2006, op_delete)

if safe:

rid, gle, _ = __last_error(collection_name, last_error_args)

return rid, msg + gle, max_bson_size

collection_name, spec, safe, last_error_args, opts, flags=0, ctx=None):

"""Get a **delete** message.

if ctx:

return _delete_compressed(collection_name, spec, opts, flags, ctx)

return _delete_uncompressed(

"""Get a **killCursors** message.

num_cursors = len(cursor_ids)

pack = struct.Struct("<ii" + ("q" * num_cursors)).pack

op_kill_cursors = pack(0, num_cursors, *cursor_ids)

"""A wrapper around SocketInfo for use with write splitting functions."""

__slots__ = ('db_name', 'command', 'sock_info', 'op_id',

'name', 'field', 'publish', 'start_time', 'listeners',

'session', 'compress', 'op_type', 'codec')

def __init__(self, database_name, command, sock_info, operation_id,

listeners, session, op_type, codec):

self.db_name = database_name

self.command = command

self.sock_info = sock_info

self.op_id = operation_id

self.listeners = listeners

self.publish = listeners.enabled_for_commands

self.name = next(iter(command))

self.field = _FIELD_MAP[self.name]

self.start_time = datetime.datetime.now() if self.publish else None

self.session = session

self.compress = True if sock_info.compression_context else False

self.op_type = op_type

self.codec = codec

def _batch_command(self, docs):

namespace = self.db_name + '.$cmd'

request_id, msg, to_send = _do_bulk_write_command(

namespace, self.op_type, self.command, docs, self.check_keys,

self.codec, self)

if not to_send:

raise InvalidOperation("cannot do an empty bulk write")

return request_id, msg, to_send

def execute(self, docs, client):

request_id, msg, to_send = self._batch_command(docs)

result = self.write_command(request_id, msg, to_send)

client._process_response(result, self.session)

return result, to_send

def execute_unack(self, docs, client):

request_id, msg, to_send = self._batch_command(docs)

# Though this isn't strictly a "legacy" write, the helper

# handles publishing commands and sending our message

# without receiving a result. Send 0 for max_doc_size

# to disable size checking. Size checking is handled while

# the documents are encoded to BSON.

self.legacy_write(request_id, msg, 0, False, to_send)

return to_send

@property

def check_keys(self):

"""Should we check keys for this operation type?"""

return self.op_type == _INSERT

@property

def max_bson_size(self):

"""A proxy for SockInfo.max_bson_size."""

return self.sock_info.max_bson_size

@property

def max_message_size(self):

"""A proxy for SockInfo.max_message_size."""

if self.compress:

# Subtract 16 bytes for the message header.

return self.sock_info.max_message_size - 16

return self.sock_info.max_message_size

@property

def max_write_batch_size(self):

"""A proxy for SockInfo.max_write_batch_size."""

return self.sock_info.max_write_batch_size

@property

def max_split_size(self):

"""The maximum size of a BSON command before batch splitting."""

return self.max_bson_size

def legacy_bulk_insert(

self, request_id, msg, max_doc_size, acknowledged, docs, compress):

if compress:

request_id, msg = _compress(

2002, msg, self.sock_info.compression_context)

return self.legacy_write(

request_id, msg, max_doc_size, acknowledged, docs)

def legacy_write(self, request_id, msg, max_doc_size, acknowledged, docs):

"""A proxy for SocketInfo.legacy_write that handles event publishing.

if self.publish:

duration = datetime.datetime.now() - self.start_time

cmd = self._start(request_id, docs)

start = datetime.datetime.now()

try:

result = self.sock_info.legacy_write(

request_id, msg, max_doc_size, acknowledged)

if self.publish:

duration = (datetime.datetime.now() - start) + duration

if result is not None:

reply = _convert_write_result(self.name, cmd, result)

else:

# Comply with APM spec.

reply = {'ok': 1}

self._succeed(request_id, reply, duration)

except OperationFailure as exc:

if self.publish:

duration = (datetime.datetime.now() - start) + duration

self._fail(

request_id,

_convert_write_result(

self.name, cmd, exc.details),

duration)

raise

finally:

self.start_time = datetime.datetime.now()

return result

def write_command(self, request_id, msg, docs):

"""A proxy for SocketInfo.write_command that handles event publishing.

if self.publish:

duration = datetime.datetime.now() - self.start_time

self._start(request_id, docs)

start = datetime.datetime.now()

try:

reply = self.sock_info.write_command(request_id, msg)

if self.publish:

duration = (datetime.datetime.now() - start) + duration

self._succeed(request_id, reply, duration)

except OperationFailure as exc:

if self.publish:

duration = (datetime.datetime.now() - start) + duration

self._fail(request_id, exc.details, duration)

raise

finally:

self.start_time = datetime.datetime.now()

return reply

def _start(self, request_id, docs):