Standalone commands only

Author: NoahStapp

pymongo/asynchronous/network.py

@@ -18,6 +18,7 @@
 import datetime
 import logging
 import time
+from asyncio import streams, StreamReader
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -45,14 +46,14 @@
     _UNPACK_COMPRESSION_HEADER,
     _UNPACK_HEADER,
     async_receive_data,
-    async_sendall,
+    async_sendall, async_sendall_stream, async_receive_data_stream,
 )
 
 if TYPE_CHECKING:
     from bson import CodecOptions
     from pymongo.asynchronous.client_session import AsyncClientSession
     from pymongo.asynchronous.mongo_client import AsyncMongoClient
-    from pymongo.asynchronous.pool import AsyncConnection
+    from pymongo.asynchronous.pool import AsyncConnection, AsyncStreamConnection, AsyncConnectionStream
     from pymongo.compression_support import SnappyContext, ZlibContext, ZstdContext
     from pymongo.monitoring import _EventListeners
     from pymongo.read_concern import ReadConcern
@@ -114,6 +115,7 @@ async def command(
         bson._decode_all_selective.
     :param exhaust_allowed: True if we should enable OP_MSG exhaustAllowed.
     """
+    print("Running socket command!")
     name = next(iter(spec))
     ns = dbname + ".$cmd"
     speculative_hello = False
@@ -298,6 +300,243 @@ async def command(
 
     return response_doc  # type: ignore[return-value]
 
+async def command_stream(
+    conn: AsyncConnectionStream,
+    dbname: str,
+    spec: MutableMapping[str, Any],
+    is_mongos: bool,
+    read_preference: Optional[_ServerMode],
+    codec_options: CodecOptions[_DocumentType],
+    session: Optional[AsyncClientSession],
+    client: Optional[AsyncMongoClient],
+    check: bool = True,
+    allowable_errors: Optional[Sequence[Union[str, int]]] = None,
+    address: Optional[_Address] = None,
+    listeners: Optional[_EventListeners] = None,
+    max_bson_size: Optional[int] = None,
+    read_concern: Optional[ReadConcern] = None,
+    parse_write_concern_error: bool = False,
+    collation: Optional[_CollationIn] = None,
+    compression_ctx: Union[SnappyContext, ZlibContext, ZstdContext, None] = None,
+    use_op_msg: bool = False,
+    unacknowledged: bool = False,
+    user_fields: Optional[Mapping[str, Any]] = None,
+    exhaust_allowed: bool = False,
+    write_concern: Optional[WriteConcern] = None,
+) -> _DocumentType:
+    """Execute a command over the socket, or raise socket.error.
+
+    :param conn: a AsyncConnection instance
+    :param dbname: name of the database on which to run the command
+    :param spec: a command document as an ordered dict type, eg SON.
+    :param is_mongos: are we connected to a mongos?
+    :param read_preference: a read preference
+    :param codec_options: a CodecOptions instance
+    :param session: optional AsyncClientSession instance.
+    :param client: optional AsyncMongoClient instance for updating $clusterTime.
+    :param check: raise OperationFailure if there are errors
+    :param allowable_errors: errors to ignore if `check` is True
+    :param address: the (host, port) of `conn`
+    :param listeners: An instance of :class:`~pymongo.monitoring.EventListeners`
+    :param max_bson_size: The maximum encoded bson size for this server
+    :param read_concern: The read concern for this command.
+    :param parse_write_concern_error: Whether to parse the ``writeConcernError``
+        field in the command response.
+    :param collation: The collation for this command.
+    :param compression_ctx: optional compression Context.
+    :param use_op_msg: True if we should use OP_MSG.
+    :param unacknowledged: True if this is an unacknowledged command.
+    :param user_fields: Response fields that should be decoded
+        using the TypeDecoders from codec_options, passed to
+        bson._decode_all_selective.
+    :param exhaust_allowed: True if we should enable OP_MSG exhaustAllowed.
+    """
+    # print("Running stream command!")
+    name = next(iter(spec))
+    ns = dbname + ".$cmd"
+    speculative_hello = False
+
+    # Publish the original command document, perhaps with lsid and $clusterTime.
+    orig = spec
+    if is_mongos and not use_op_msg:
+        assert read_preference is not None
+        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, conn)
+    if collation is not None:
+        spec["collation"] = collation
+
+    publish = listeners is not None and listeners.enabled_for_commands
+    start = datetime.datetime.now()
+    if publish:
+        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 = await client._encrypter.encrypt(dbname, spec, codec_options)
+
+    # Support CSOT
+    if client:
+        conn.apply_timeout(client, spec)
+    _csot.apply_write_concern(spec, write_concern)
+
+    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, 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(
+            0, ns, 0, -1, spec, None, codec_options, 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 client is not None:
+        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
+            _debug_log(
+                _COMMAND_LOGGER,
+                clientId=client._topology_settings._topology_id,
+                message=_CommandStatusMessage.STARTED,
+                command=spec,
+                commandName=next(iter(spec)),
+                databaseName=dbname,
+                requestId=request_id,
+                operationId=request_id,
+                driverConnectionId=conn.id,
+                serverConnectionId=conn.server_connection_id,
+                serverHost=conn.address[0],
+                serverPort=conn.address[1],
+                serviceId=conn.service_id,
+            )
+    if publish:
+        assert listeners is not None
+        assert address is not None
+        listeners.publish_command_start(
+            orig,
+            dbname,
+            request_id,
+            address,
+            conn.server_connection_id,
+            service_id=conn.service_id,
+        )
+
+    try:
+        await async_sendall_stream(conn.conn[1], msg)
+        if use_op_msg and unacknowledged:
+            # Unacknowledged, fake a successful command response.
+            reply = None
+            response_doc: _DocumentOut = {"ok": 1}
+        else:
+            reply = await receive_message_stream(conn.conn[0], request_id)
+            conn.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:
+                await client._process_response(response_doc, session)
+            if check:
+                helpers_shared._check_command_response(
+                    response_doc,
+                    conn.max_wire_version,
+                    allowable_errors,
+                    parse_write_concern_error=parse_write_concern_error,
+                )
+    except Exception as exc:
+        duration = datetime.datetime.now() - start
+        if isinstance(exc, (NotPrimaryError, OperationFailure)):
+            failure: _DocumentOut = exc.details  # type: ignore[assignment]
+        else:
+            failure = message._convert_exception(exc)
+        if client is not None:
+            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
+                _debug_log(
+                    _COMMAND_LOGGER,
+                    clientId=client._topology_settings._topology_id,
+                    message=_CommandStatusMessage.FAILED,
+                    durationMS=duration,
+                    failure=failure,
+                    commandName=next(iter(spec)),
+                    databaseName=dbname,
+                    requestId=request_id,
+                    operationId=request_id,
+                    driverConnectionId=conn.id,
+                    serverConnectionId=conn.server_connection_id,
+                    serverHost=conn.address[0],
+                    serverPort=conn.address[1],
+                    serviceId=conn.service_id,
+                    isServerSideError=isinstance(exc, OperationFailure),
+                )
+        if publish:
+            assert listeners is not None
+            assert address is not None
+            listeners.publish_command_failure(
+                duration,
+                failure,
+                name,
+                request_id,
+                address,
+                conn.server_connection_id,
+                service_id=conn.service_id,
+                database_name=dbname,
+            )
+        raise
+    duration = datetime.datetime.now() - start
+    if client is not None:
+        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
+            _debug_log(
+                _COMMAND_LOGGER,
+                clientId=client._topology_settings._topology_id,
+                message=_CommandStatusMessage.SUCCEEDED,
+                durationMS=duration,
+                reply=response_doc,
+                commandName=next(iter(spec)),
+                databaseName=dbname,
+                requestId=request_id,
+                operationId=request_id,
+                driverConnectionId=conn.id,
+                serverConnectionId=conn.server_connection_id,
+                serverHost=conn.address[0],
+                serverPort=conn.address[1],
+                serviceId=conn.service_id,
+                speculative_authenticate="speculativeAuthenticate" in orig,
+            )
+    if publish:
+        assert listeners is not None
+        assert address is not None
+        listeners.publish_command_success(
+            duration,
+            response_doc,
+            name,
+            request_id,
+            address,
+            conn.server_connection_id,
+            service_id=conn.service_id,
+            speculative_hello=speculative_hello,
+            database_name=dbname,
+        )
+
+    if client and client._encrypter and reply:
+        decrypted = await client._encrypter.decrypt(reply.raw_command_response())
+        response_doc = cast(
+            "_DocumentOut", _decode_all_selective(decrypted, codec_options, user_fields)[0]
+        )
+
+    return response_doc  # type: ignore[return-value]
+
+
 
 async def receive_message(
     conn: AsyncConnection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
@@ -341,3 +580,48 @@ async def receive_message(
             f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
         ) from None
     return unpack_reply(data)
+
+async def receive_message_stream(
+    conn: StreamReader, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+) -> Union[_OpReply, _OpMsg]:
+    """Receive a raw BSON message or raise socket.error."""
+    # if _csot.get_timeout():
+    #     deadline = _csot.get_deadline()
+    # else:
+    #     timeout = conn.conn.gettimeout()
+    #     if timeout:
+    #         deadline = time.monotonic() + timeout
+    #     else:
+    #         deadline = None
+    deadline = None
+    # Ignore the response's request id.
+    length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data_stream(conn, 16, deadline))
+    # No request_id for exhaust cursor "getMore".
+    if request_id is not None:
+        if request_id != response_to:
+            raise ProtocolError(f"Got response id {response_to!r} but expected {request_id!r}")
+    if length <= 16:
+        raise ProtocolError(
+            f"Message length ({length!r}) not longer than standard message header size (16)"
+        )
+    if length > max_message_size:
+        raise ProtocolError(
+            f"Message length ({length!r}) is larger than server max "
+            f"message size ({max_message_size!r})"
+        )
+    if op_code == 2012:
+        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
+            await async_receive_data(conn, 9, deadline)
+        )
+        data = decompress(await async_receive_data_stream(conn, length - 25, deadline), compressor_id)
+    else:
+        data = await async_receive_data_stream(conn, length - 16, deadline)
+
+    try:
+        unpack_reply = _UNPACK_REPLY[op_code]
+    except KeyError:
+        raise ProtocolError(
+            f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
+        ) from None
+    return unpack_reply(data)
+