def test_socket_checker(self): s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.connect((client_context.host, client_context.port)) socket_checker = SocketChecker() # Socket has nothing to read. self.assertFalse(socket_checker.select(s, read=True)) self.assertFalse(socket_checker.select(s, read=True, timeout=0)) self.assertFalse(socket_checker.select(s, read=True, timeout=.05)) # Socket is writable. self.assertTrue(socket_checker.select(s, write=True, timeout=None)) self.assertTrue(socket_checker.select(s, write=True)) self.assertTrue(socket_checker.select(s, write=True, timeout=0)) self.assertTrue(socket_checker.select(s, write=True, timeout=.05)) # Make the socket readable _, msg, _ = message._query( 0, 'admin.$cmd', 0, -1, SON([('isMaster', 1)]), None, DEFAULT_CODEC_OPTIONS) s.sendall(msg) # Block until the socket is readable. self.assertTrue(socket_checker.select(s, read=True, timeout=None)) self.assertTrue(socket_checker.select(s, read=True)) self.assertTrue(socket_checker.select(s, read=True, timeout=0)) self.assertTrue(socket_checker.select(s, read=True, timeout=.05)) # Socket is still writable. self.assertTrue(socket_checker.select(s, write=True, timeout=None)) self.assertTrue(socket_checker.select(s, write=True)) self.assertTrue(socket_checker.select(s, write=True, timeout=0)) self.assertTrue(socket_checker.select(s, write=True, timeout=.05)) s.close() self.assertTrue(socket_checker.socket_closed(s))
def command(sock_info, dbname, spec, secondary_ok, is_mongos, read_preference, codec_options, session, client, check=True, allowable_errors=None, address=None, check_keys=False, listeners=None, max_bson_size=None, read_concern=None, parse_write_concern_error=False, collation=None, compression_ctx=None, use_op_msg=False, unacknowledged=False, user_fields=None, exhaust_allowed=False): """Execute a command over the socket, or raise socket.error. :Parameters: - `sock`: a raw socket instance - `dbname`: name of the database on which to run the command - `spec`: a command document as an ordered dict type, eg SON. - `secondary_ok`: whether to set the secondaryOkay wire protocol bit - `is_mongos`: are we connected to a mongos? - `read_preference`: a read preference - `codec_options`: a CodecOptions instance - `session`: optional ClientSession instance. - `client`: optional MongoClient instance for updating $clusterTime. - `check`: raise OperationFailure if there are errors - `allowable_errors`: errors to ignore if `check` is True - `address`: the (host, port) of `sock` - `check_keys`: if True, check `spec` for invalid keys - `listeners`: An instance of :class:`~pymongo.monitoring.EventListeners` - `max_bson_size`: The maximum encoded bson size for this server - `read_concern`: The read concern for this command. - `parse_write_concern_error`: Whether to parse the ``writeConcernError`` field in the command response. - `collation`: The collation for this command. - `compression_ctx`: optional compression Context. - `use_op_msg`: True if we should use OP_MSG. - `unacknowledged`: True if this is an unacknowledged command. - `user_fields` (optional): Response fields that should be decoded using the TypeDecoders from codec_options, passed to bson._decode_all_selective. - `exhaust_allowed`: True if we should enable OP_MSG exhaustAllowed. """ name = next(iter(spec)) ns = dbname + '.$cmd' flags = 4 if secondary_ok else 0 speculative_hello = False # Publish the original command document, perhaps with lsid and $clusterTime. orig = spec if is_mongos and not use_op_msg: spec = message._maybe_add_read_preference(spec, read_preference) if read_concern and not (session and session.in_transaction): if read_concern.level: spec['readConcern'] = read_concern.document if session: session._update_read_concern(spec, sock_info) if collation is not None: spec['collation'] = collation publish = listeners is not None and listeners.enabled_for_commands if publish: start = datetime.datetime.now() speculative_hello = _is_speculative_authenticate(name, spec) if compression_ctx and name.lower() in _NO_COMPRESSION: compression_ctx = None if (client and client._encrypter and not client._encrypter._bypass_auto_encryption): spec = orig = client._encrypter.encrypt( dbname, spec, check_keys, codec_options) # We already checked the keys, no need to do it again. check_keys = False if use_op_msg: flags = _OpMsg.MORE_TO_COME if unacknowledged else 0 flags |= _OpMsg.EXHAUST_ALLOWED if exhaust_allowed else 0 request_id, msg, size, max_doc_size = message._op_msg( flags, spec, dbname, read_preference, secondary_ok, check_keys, codec_options, ctx=compression_ctx) # If this is an unacknowledged write then make sure the encoded doc(s) # are small enough, otherwise rely on the server to return an error. if (unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size): message._raise_document_too_large(name, size, max_bson_size) else: request_id, msg, size = message._query( flags, ns, 0, -1, spec, None, codec_options, check_keys, compression_ctx) if (max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD): message._raise_document_too_large( name, size, max_bson_size + message._COMMAND_OVERHEAD) if publish: encoding_duration = datetime.datetime.now() - start listeners.publish_command_start(orig, dbname, request_id, address, service_id=sock_info.service_id) start = datetime.datetime.now() try: sock_info.sock.sendall(msg) if use_op_msg and unacknowledged: # Unacknowledged, fake a successful command response. reply = None response_doc = {"ok": 1} else: reply = receive_message(sock_info, request_id) sock_info.more_to_come = reply.more_to_come unpacked_docs = reply.unpack_response( codec_options=codec_options, user_fields=user_fields) response_doc = unpacked_docs[0] if client: client._process_response(response_doc, session) if check: helpers._check_command_response( response_doc, sock_info.max_wire_version, allowable_errors, parse_write_concern_error=parse_write_concern_error) except Exception as exc: if publish: duration = (datetime.datetime.now() - start) + encoding_duration if isinstance(exc, (NotPrimaryError, OperationFailure)): failure = exc.details else: failure = message._convert_exception(exc) listeners.publish_command_failure( duration, failure, name, request_id, address, service_id=sock_info.service_id) raise if publish: duration = (datetime.datetime.now() - start) + encoding_duration listeners.publish_command_success( duration, response_doc, name, request_id, address, service_id=sock_info.service_id, speculative_hello=speculative_hello) if client and client._encrypter and reply: decrypted = client._encrypter.decrypt(reply.raw_command_response()) response_doc = _decode_all_selective(decrypted, codec_options, user_fields)[0] return response_doc