PYTHON-5676 Consolidate command execution logic (#2852)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Steven Silvester <steve.silvester@mongodb.com>
Co-authored-by: Noah Stapp <noah@noahstapp.com>

Author: 4 people

pymongo/asynchronous/bulk.py

@@ -20,8 +20,6 @@
 from __future__ import annotations
 
 import copy
-import datetime
-import logging
 from collections.abc import Iterator, Mapping, MutableMapping
 from itertools import islice
 from typing import (
@@ -34,9 +32,9 @@
 from bson.objectid import ObjectId
 from bson.raw_bson import RawBSONDocument
 from pymongo import _csot, common
-from pymongo.asynchronous.client_session import (
-    AsyncClientSession,
-    _validate_session_write_concern,
+from pymongo.asynchronous.client_session import AsyncClientSession, _validate_session_write_concern
+from pymongo.asynchronous.command_runner import (
+    run_bulk_write_command,
 )
 from pymongo.asynchronous.helpers import _handle_reauth
 from pymongo.bulk_shared import (
@@ -54,18 +52,14 @@
 from pymongo.errors import (
     ConfigurationError,
     InvalidOperation,
-    NotPrimaryError,
     OperationFailure,
 )
 from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES
-from pymongo.logger import _COMMAND_LOGGER, _CommandStatusMessage, _debug_log
 from pymongo.message import (
     _DELETE,
     _INSERT,
     _UPDATE,
     _BulkWriteContext,
-    _convert_exception,
-    _convert_write_result,
     _EncryptedBulkWriteContext,
     _randint,
 )
@@ -251,83 +245,16 @@ async def write_command(
         docs: list[Mapping[str, Any]],
         client: AsyncMongoClient[Any],
     ) -> dict[str, Any]:
-        """A proxy for SocketInfo.write_command that handles event publishing."""
+        """Run a batch write command, returning the response as a dict."""
         cmd[bwc.field] = docs
-        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-            _debug_log(
-                _COMMAND_LOGGER,
-                message=_CommandStatusMessage.STARTED,
-                clientId=client._topology_settings._topology_id,
-                command=cmd,
-                commandName=next(iter(cmd)),
-                databaseName=bwc.db_name,
-                requestId=request_id,
-                operationId=request_id,
-                driverConnectionId=bwc.conn.id,
-                serverConnectionId=bwc.conn.server_connection_id,
-                serverHost=bwc.conn.address[0],
-                serverPort=bwc.conn.address[1],
-                serviceId=bwc.conn.service_id,
-            )
-        if bwc.publish:
-            bwc._start(cmd, request_id, docs)
-        try:
-            if bwc.session is not None and bwc.session._starting_transaction:
-                bwc.session._transaction.set_in_progress()
-            reply = await bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc]
-            duration = datetime.datetime.now() - bwc.start_time
-            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-                _debug_log(
-                    _COMMAND_LOGGER,
-                    message=_CommandStatusMessage.SUCCEEDED,
-                    clientId=client._topology_settings._topology_id,
-                    durationMS=duration,
-                    reply=reply,
-                    commandName=next(iter(cmd)),
-                    databaseName=bwc.db_name,
-                    requestId=request_id,
-                    operationId=request_id,
-                    driverConnectionId=bwc.conn.id,
-                    serverConnectionId=bwc.conn.server_connection_id,
-                    serverHost=bwc.conn.address[0],
-                    serverPort=bwc.conn.address[1],
-                    serviceId=bwc.conn.service_id,
-                )
-            if bwc.publish:
-                bwc._succeed(request_id, reply, duration)  # type: ignore[arg-type]
-            await client._process_response(reply, bwc.session)  # type: ignore[arg-type]
-        except Exception as exc:
-            duration = datetime.datetime.now() - bwc.start_time
-            if isinstance(exc, (NotPrimaryError, OperationFailure)):
-                failure: _DocumentOut = exc.details  # type: ignore[assignment]
-            else:
-                failure = _convert_exception(exc)
-            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-                _debug_log(
-                    _COMMAND_LOGGER,
-                    message=_CommandStatusMessage.FAILED,
-                    clientId=client._topology_settings._topology_id,
-                    durationMS=duration,
-                    failure=failure,
-                    commandName=next(iter(cmd)),
-                    databaseName=bwc.db_name,
-                    requestId=request_id,
-                    operationId=request_id,
-                    driverConnectionId=bwc.conn.id,
-                    serverConnectionId=bwc.conn.server_connection_id,
-                    serverHost=bwc.conn.address[0],
-                    serverPort=bwc.conn.address[1],
-                    serviceId=bwc.conn.service_id,
-                    isServerSideError=isinstance(exc, OperationFailure),
-                )
-
-            if bwc.publish:
-                bwc._fail(request_id, failure, duration)
-            # Process the response from the server.
-            if isinstance(exc, (NotPrimaryError, OperationFailure)):
-                await client._process_response(exc.details, bwc.session)  # type: ignore[arg-type]
-            raise
-        return reply  # type: ignore[return-value]
+        result_docs, _, _ = await run_bulk_write_command(
+            bwc,
+            cmd,
+            request_id,
+            msg,
+            client=client,
+        )
+        return result_docs[0]
 
     async def unack_write(
         self,
@@ -339,83 +266,23 @@ async def unack_write(
         docs: list[Mapping[str, Any]],
         client: AsyncMongoClient[Any],
     ) -> Optional[Mapping[str, Any]]:
-        """A proxy for AsyncConnection.unack_write that handles event publishing."""
-        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-            _debug_log(
-                _COMMAND_LOGGER,
-                message=_CommandStatusMessage.STARTED,
-                clientId=client._topology_settings._topology_id,
-                command=cmd,
-                commandName=next(iter(cmd)),
-                databaseName=bwc.db_name,
-                requestId=request_id,
-                operationId=request_id,
-                driverConnectionId=bwc.conn.id,
-                serverConnectionId=bwc.conn.server_connection_id,
-                serverHost=bwc.conn.address[0],
-                serverPort=bwc.conn.address[1],
-                serviceId=bwc.conn.service_id,
-            )
-        if bwc.publish:
-            cmd = bwc._start(cmd, request_id, docs)
-        try:
-            result = await bwc.conn.unack_write(msg, max_doc_size)  # type: ignore[func-returns-value, misc, override]
-            duration = datetime.datetime.now() - bwc.start_time
-            if result is not None:
-                reply = _convert_write_result(bwc.name, cmd, result)  # type: ignore[arg-type]
-            else:
-                # Comply with APM spec.
-                reply = {"ok": 1}
-                if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-                    _debug_log(
-                        _COMMAND_LOGGER,
-                        message=_CommandStatusMessage.SUCCEEDED,
-                        clientId=client._topology_settings._topology_id,
-                        durationMS=duration,
-                        reply=reply,
-                        commandName=next(iter(cmd)),
-                        databaseName=bwc.db_name,
-                        requestId=request_id,
-                        operationId=request_id,
-                        driverConnectionId=bwc.conn.id,
-                        serverConnectionId=bwc.conn.server_connection_id,
-                        serverHost=bwc.conn.address[0],
-                        serverPort=bwc.conn.address[1],
-                        serviceId=bwc.conn.service_id,
-                    )
-            if bwc.publish:
-                bwc._succeed(request_id, reply, duration)
-        except Exception as exc:
-            duration = datetime.datetime.now() - bwc.start_time
-            if isinstance(exc, OperationFailure):
-                failure: _DocumentOut = _convert_write_result(bwc.name, cmd, exc.details)  # type: ignore[arg-type]
-            elif isinstance(exc, NotPrimaryError):
-                failure = exc.details  # type: ignore[assignment]
-            else:
-                failure = _convert_exception(exc)
-            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
-                _debug_log(
-                    _COMMAND_LOGGER,
-                    message=_CommandStatusMessage.FAILED,
-                    clientId=client._topology_settings._topology_id,
-                    durationMS=duration,
-                    failure=failure,
-                    commandName=next(iter(cmd)),
-                    databaseName=bwc.db_name,
-                    requestId=request_id,
-                    operationId=request_id,
-                    driverConnectionId=bwc.conn.id,
-                    serverConnectionId=bwc.conn.server_connection_id,
-                    serverHost=bwc.conn.address[0],
-                    serverPort=bwc.conn.address[1],
-                    serviceId=bwc.conn.service_id,
-                    isServerSideError=isinstance(exc, OperationFailure),
-                )
-            if bwc.publish:
-                assert bwc.start_time is not None
-                bwc._fail(request_id, failure, duration)
-            raise
-        return result  # type: ignore[return-value]
+        """Send an unacknowledged batch write command."""
+        # Historically the STARTED log omits the documents while the published
+        # CommandStartedEvent includes them, so log ``cmd`` but publish a copy
+        # carrying the ``docs`` field.
+        published = dict(cmd)
+        published[bwc.field] = docs
+        await run_bulk_write_command(
+            bwc,
+            cmd,
+            request_id,
+            msg,
+            client=client,
+            orig=published,
+            max_doc_size=max_doc_size,
+            unacknowledged=True,
+        )
+        return None
 
     async def _execute_batch_unack(
         self,
@@ -487,7 +354,7 @@ async def _execute_command(
         run = self.current_run
 
         # AsyncConnection.command validates the session, but we use
-        # AsyncConnection.write_command
+        # run_bulk_write_command.
         conn.validate_session(client, session)
         last_run = False