def _gen_find_command(coll, spec, projection, skip, limit, batch_size, options, read_concern=DEFAULT_READ_CONCERN, collation=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: 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]) return cmd
def _command(self, sock_info, command, slave_ok=False, value=1, check=True, allowable_errors=None, read_preference=ReadPreference.PRIMARY, codec_options=DEFAULT_CODEC_OPTIONS, write_concern=None, parse_write_concern_error=False, **kwargs): """Internal command helper.""" if isinstance(command, string_type): command = SON([(command, value)]) if sock_info.max_wire_version >= 5 and write_concern: command['writeConcern'] = write_concern.document command.update(kwargs) return sock_info.command( self.__name, command, slave_ok, read_preference, codec_options, check, allowable_errors, parse_write_concern_error=parse_write_concern_error)
def _authenticate_scram_sha1(credentials, sock_info): """Authenticate using SCRAM-SHA-1.""" username = credentials.username password = credentials.password source = credentials.source # Make local _hmac = hmac.HMAC _sha1 = sha1 user = username.encode("utf-8").replace(b"=", b"=3D").replace(b",", b"=2C") nonce = standard_b64encode( (("%s" % (SystemRandom().random(), ))[2:]).encode("utf-8")) first_bare = b"n=" + user + b",r=" + nonce cmd = SON([('saslStart', 1), ('mechanism', 'SCRAM-SHA-1'), ('payload', Binary(b"n,," + first_bare)), ('autoAuthorize', 1)]) res = sock_info.command(source, cmd) server_first = res['payload'] parsed = _parse_scram_response(server_first) iterations = int(parsed[b'i']) salt = parsed[b's'] rnonce = parsed[b'r'] if not rnonce.startswith(nonce): raise OperationFailure("Server returned an invalid nonce.") without_proof = b"c=biws,r=" + rnonce salted_pass = _hi( _password_digest(username, password).encode("utf-8"), standard_b64decode(salt), iterations) client_key = _hmac(salted_pass, b"Client Key", _sha1).digest() stored_key = _sha1(client_key).digest() auth_msg = b",".join((first_bare, server_first, without_proof)) client_sig = _hmac(stored_key, auth_msg, _sha1).digest() client_proof = b"p=" + standard_b64encode(_xor(client_key, client_sig)) client_final = b",".join((without_proof, client_proof)) server_key = _hmac(salted_pass, b"Server Key", _sha1).digest() server_sig = standard_b64encode( _hmac(server_key, auth_msg, _sha1).digest()) cmd = SON([('saslContinue', 1), ('conversationId', res['conversationId']), ('payload', Binary(client_final))]) res = sock_info.command(source, cmd) parsed = _parse_scram_response(res['payload']) if not compare_digest(parsed[b'v'], server_sig): raise OperationFailure("Server returned an invalid signature.") # Depending on how it's configured, Cyrus SASL (which the server uses) # requires a third empty challenge. if not res['done']: cmd = SON([('saslContinue', 1), ('conversationId', res['conversationId']), ('payload', Binary(b''))]) res = sock_info.command(source, cmd) if not res['done']: raise OperationFailure('SASL conversation failed to complete.')
def transform_incoming(self, son, collection): """Move _id to the front if it's there. """ if not "_id" in son: return son transformed = SON({"_id": son["_id"]}) transformed.update(son) return transformed
def __last_error(namespace, args): """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, None, DEFAULT_CODEC_OPTIONS)
def _gen_explain_command( coll, spec, projection, skip, limit, batch_size, options, read_concern): """Generate an explain command document.""" cmd = _gen_find_command( coll, spec, projection, skip, limit, batch_size, options) if read_concern.level: return SON([('explain', cmd), ('readConcern', read_concern.document)]) return SON([('explain', cmd)])
def as_doc(self): """Get the SON document representation of this DBRef. Generally not needed by application developers """ doc = SON([("$ref", self.collection), ("$id", self.id)]) if self.database is not None: doc["$db"] = self.database doc.update(self.__kwargs) return doc
def _index_document(index_list): """Helper to generate an index specifying document. Takes a list of (key, direction) pairs. """ if isinstance(index_list, collections.Mapping): raise TypeError("passing a dict to sort/create_index/hint is not " "allowed - use a list of tuples instead. did you " "mean %r?" % list(iteritems(index_list))) elif not isinstance(index_list, (list, tuple)): raise TypeError("must use a list of (key, direction) pairs, " "not: " + repr(index_list)) if not len(index_list): raise ValueError("key_or_list must not be the empty list") index = SON() for (key, value) in index_list: if not isinstance(key, string_type): raise TypeError("first item in each key pair must be a string") if not isinstance(value, (string_type, int, collections.Mapping)): raise TypeError("second item in each key pair must be 1, -1, " "'2d', 'geoHaystack', or another valid MongoDB " "index specifier.") index[key] = value return index
def _list_collections(self, sock_info, slave_okay, criteria=None): """Internal listCollections helper.""" criteria = criteria or {} cmd = SON([("listCollections", 1), ("cursor", {})]) if criteria: cmd["filter"] = criteria if sock_info.max_wire_version > 2: coll = self["$cmd"] cursor = self._command(sock_info, cmd, slave_okay)["cursor"] return CommandCursor(coll, cursor, sock_info.address) else: coll = self["system.namespaces"] res = _first_batch(sock_info, coll.database.name, coll.name, criteria, 0, slave_okay, CodecOptions(), ReadPreference.PRIMARY, cmd, self.client._event_listeners) data = res["data"] cursor = { "id": res["cursor_id"], "firstBatch": data, "ns": coll.full_name, } # Need to tell the cursor how many docs were in the first batch. return CommandCursor(coll, cursor, sock_info.address, len(data))
def transform_value(value): if isinstance(value, DBRef): return self.database.dereference(value) elif isinstance(value, list): return [transform_value(v) for v in value] elif isinstance(value, collections.MutableMapping): return transform_dict(SON(value)) return value
def transform_value(value): if isinstance(value, collections.MutableMapping): if "_id" in value and "_ns" in value: return DBRef(value["_ns"], transform_value(value["_id"])) else: return transform_dict(SON(value)) elif isinstance(value, list): return [transform_value(v) for v in value] return value
def add_delete(self, selector, limit, collation=None): """Create a delete document and add it to the list of ops. """ cmd = SON([('q', selector), ('limit', limit)]) collation = validate_collation_or_none(collation) if collation is not None: self.uses_collation = True cmd['collation'] = collation self.ops.append((_DELETE, cmd))
def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms): """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 return cmd
def transform_outgoing(self, son, collection): """Manipulate an outgoing SON object. :Parameters: - `son`: the SON object being retrieved from the database - `collection`: the collection this object was stored in """ if self.will_copy(): return SON(son) return son
def _authenticate_plain(credentials, sock_info): """Authenticate using SASL PLAIN (RFC 4616) """ source = credentials.source username = credentials.username password = credentials.password payload = ('\x00%s\x00%s' % (username, password)).encode('utf-8') cmd = SON([('saslStart', 1), ('mechanism', 'PLAIN'), ('payload', Binary(payload)), ('autoAuthorize', 1)]) sock_info.command(source, cmd)
def transform_incoming(self, son, collection): """Manipulate an incoming SON object. :Parameters: - `son`: the SON object to be inserted into the database - `collection`: the collection the object is being inserted into """ if self.will_copy(): return SON(son) return son
def _authenticate_x509(credentials, sock_info): """Authenticate using MONGODB-X509. """ query = SON([('authenticate', 1), ('mechanism', 'MONGODB-X509')]) if credentials.username is not None: query['user'] = credentials.username elif sock_info.max_wire_version < 5: raise ConfigurationError( "A username is required for MONGODB-X509 authentication " "when connected to MongoDB versions older than 3.4.") sock_info.command('$external', query)
def add_replace(self, selector, replacement, upsert=False, collation=None): """Create a replace document and add it to the list of ops. """ validate_ok_for_replace(replacement) cmd = SON([('q', selector), ('u', replacement), ('multi', False), ('upsert', upsert)]) collation = validate_collation_or_none(collation) if collation is not None: self.uses_collation = True cmd['collation'] = collation self.ops.append((_UPDATE, cmd))
def _authenticate_cram_md5(credentials, sock_info): """Authenticate using CRAM-MD5 (RFC 2195) """ source = credentials.source username = credentials.username password = credentials.password # The password used as the mac key is the # same as what we use for MONGODB-CR passwd = _password_digest(username, password) cmd = SON([('saslStart', 1), ('mechanism', 'CRAM-MD5'), ('payload', Binary(b'')), ('autoAuthorize', 1)]) response = sock_info.command(source, cmd) # MD5 as implicit default digest for digestmod is deprecated # in python 3.4 mac = hmac.HMAC(key=passwd.encode('utf-8'), digestmod=md5) mac.update(response['payload']) challenge = username.encode('utf-8') + b' ' + b(mac.hexdigest()) cmd = SON([('saslContinue', 1), ('conversationId', response['conversationId']), ('payload', Binary(challenge))]) sock_info.command(source, cmd)
def __query_spec(self): """Get the spec to use for a query. """ operators = self.__modifiers.copy() if self.__ordering: operators["$orderby"] = self.__ordering if self.__explain: operators["$explain"] = True if self.__hint: operators["$hint"] = self.__hint if self.__comment: operators["$comment"] = self.__comment if self.__max_scan: operators["$maxScan"] = self.__max_scan if self.__max_time_ms is not None: operators["$maxTimeMS"] = self.__max_time_ms if self.__max: operators["$max"] = self.__max if self.__min: operators["$min"] = self.__min if self.__return_key: operators["$returnKey"] = self.__return_key if self.__show_record_id: # This is upgraded to showRecordId for MongoDB 3.2+ "find" command. operators["$showDiskLoc"] = self.__show_record_id if self.__snapshot: operators["$snapshot"] = self.__snapshot if operators: # Make a shallow copy so we can cleanly rewind or clone. spec = self.__spec.copy() # White-listed commands must be wrapped in $query. if "$query" not in spec: # $query has to come first spec = SON([("$query", spec)]) if not isinstance(spec, SON): # Ensure the spec is SON. As order is important this will # ensure its set before merging in any extra operators. spec = SON(spec) spec.update(operators) return spec # Have to wrap with $query if "query" is the first key. # We can't just use $query anytime "query" is a key as # that breaks commands like count and find_and_modify. # Checking spec.keys()[0] covers the case that the spec # was passed as an instance of SON or OrderedDict. elif ("query" in self.__spec and (len(self.__spec) == 1 or next(iter(self.__spec)) == "query")): return SON({"$query": self.__spec}) return self.__spec
def min(self, spec): """Adds `min` operator that specifies lower bound for specific index. :Parameters: - `spec`: a list of field, limit pairs specifying the inclusive lower bound for all keys of a specific index in order. .. versionadded:: 2.7 """ if not isinstance(spec, (list, tuple)): raise TypeError("spec must be an instance of list or tuple") self.__check_okay_to_chain() self.__min = SON(spec) return self
def _authenticate_mongo_cr(credentials, sock_info): """Authenticate using MONGODB-CR. """ source = credentials.source username = credentials.username password = credentials.password # Get a nonce response = sock_info.command(source, {'getnonce': 1}) nonce = response['nonce'] key = _auth_key(nonce, username, password) # Actually authenticate query = SON([('authenticate', 1), ('user', username), ('nonce', nonce), ('key', key)]) sock_info.command(source, query)
def count(self, with_limit_and_skip=False): """Get the size of the results set for this query. Returns the number of documents in the results set for this query. Does not take :meth:`limit` and :meth:`skip` into account by default - set `with_limit_and_skip` to ``True`` if that is the desired behavior. Raises :class:`~pymongo.errors.OperationFailure` on a database error. When used with MongoDB >= 2.6, :meth:`~count` uses any :meth:`~hint` applied to the query. In the following example the hint is passed to the count command: collection.find({'field': 'value'}).hint('field_1').count() The :meth:`count` method obeys the :attr:`~pymongo.collection.Collection.read_preference` of the :class:`~pymongo.collection.Collection` instance on which :meth:`~pymongo.collection.Collection.find` was called. :Parameters: - `with_limit_and_skip` (optional): take any :meth:`limit` or :meth:`skip` that has been applied to this cursor into account when getting the count .. note:: The `with_limit_and_skip` parameter requires server version **>= 1.1.4-** .. versionchanged:: 2.8 The :meth:`~count` method now supports :meth:`~hint`. """ validate_boolean("with_limit_and_skip", with_limit_and_skip) cmd = SON([("count", self.__collection.name), ("query", self.__spec)]) if self.__max_time_ms is not None: cmd["maxTimeMS"] = self.__max_time_ms if self.__comment: cmd["$comment"] = self.__comment if self.__hint is not None: cmd["hint"] = self.__hint if with_limit_and_skip: if self.__limit: cmd["limit"] = self.__limit if self.__skip: cmd["skip"] = self.__skip return self.__collection._count(cmd, self.__collation)
def current_op(self, include_all=False): """Get information on operations currently running. :Parameters: - `include_all` (optional): if ``True`` also list currently idle operations in the result """ cmd = SON([("currentOp", 1), ("$all", include_all)]) with self.__client._socket_for_writes() as sock_info: if sock_info.max_wire_version >= 4: return sock_info.command("admin", cmd) else: spec = {"$all": True} if include_all else {} 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]
def _check_with_socket(self, sock_info, metadata=None): """Return (IsMaster, round_trip_time). Can raise ConnectionFailure or OperationFailure. """ cmd = SON([('ismaster', 1)]) if metadata is not None: cmd['client'] = metadata start = _time() request_id, msg, max_doc_size = message.query( 0, 'admin.$cmd', 0, -1, cmd, None, DEFAULT_CODEC_OPTIONS) # TODO: use sock_info.command() sock_info.send_message(msg, max_doc_size) raw_response = sock_info.receive_message(1, request_id) result = helpers._unpack_response(raw_response) return IsMaster(result['data'][0]), _time() - start
def transform_outgoing(self, son, collection): """Replace DBRefs with embedded documents. """ def transform_value(value): if isinstance(value, DBRef): return self.database.dereference(value) elif isinstance(value, list): return [transform_value(v) for v in value] elif isinstance(value, collections.MutableMapping): return transform_dict(SON(value)) return value def transform_dict(object): for (key, value) in object.items(): object[key] = transform_value(value) return object return transform_dict(SON(son))
def execute_command(self, sock_info, generator, write_concern): """Execute using write commands. """ # nModified is only reported for write commands, not legacy ops. full_result = { "writeErrors": [], "writeConcernErrors": [], "nInserted": 0, "nUpserted": 0, "nMatched": 0, "nModified": 0, "nRemoved": 0, "upserted": [], } op_id = _randint() db_name = self.collection.database.name listeners = self.collection.database.client._event_listeners for run in generator: cmd = SON([(_COMMANDS[run.op_type], self.collection.name), ('ordered', self.ordered)]) if write_concern.document: cmd['writeConcern'] = write_concern.document if self.bypass_doc_val and sock_info.max_wire_version >= 4: cmd['bypassDocumentValidation'] = True bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id, listeners) results = _do_batched_write_command(self.namespace, run.op_type, cmd, run.ops, True, self.collection.codec_options, bwc) _merge_command(run, full_result, results) # We're supposed to continue if errors are # at the write concern level (e.g. wtimeout) if self.ordered and full_result['writeErrors']: break if full_result["writeErrors"] or full_result["writeConcernErrors"]: if full_result['writeErrors']: full_result['writeErrors'].sort( key=lambda error: error['index']) raise BulkWriteError(full_result) return full_result
def transform_incoming(self, son, collection): """Replace embedded documents with DBRefs. """ def transform_value(value): if isinstance(value, collections.MutableMapping): if "_id" in value and "_ns" in value: return DBRef(value["_ns"], transform_value(value["_id"])) else: return transform_dict(SON(value)) elif isinstance(value, list): return [transform_value(v) for v in value] return value def transform_dict(object): for (key, value) in object.items(): object[key] = transform_value(value) return object return transform_dict(SON(son))
def _maybe_add_read_preference(spec, read_preference): """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 return spec
def remove_user(self, name): """Remove user `name` from this :class:`Database`. User `name` will no longer have permissions to access this :class:`Database`. :Parameters: - `name`: the name of the user to remove """ try: cmd = SON([("dropUser", name)]) # Don't send {} as writeConcern. if self.write_concern.acknowledged and self.write_concern.document: cmd["writeConcern"] = self.write_concern.document self.command(cmd) except OperationFailure as exc: # See comment in add_user try / except above. if exc.code in common.COMMAND_NOT_FOUND_CODES: coll = self._collection_default_options('system.users') coll.delete_one({"user": name}) return raise