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))
Exemplo n.º 2
0
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