PYTHON-5903 : Retry attempts should share a single stable operation_id

pymongo/_telemetry.py

@@ -86,7 +86,7 @@ def _emit_log(self, message: _CommandStatusMessage, **extra: Any) -> None:
             commandName=self._name,
             databaseName=self._dbname,
             requestId=self._request_id,
-            operationId=self._request_id,
+            operationId=self._op_id if self._op_id is not None else self._request_id,
             driverConnectionId=self._conn.id,
             serverConnectionId=self._conn.server_connection_id,
             serverHost=self._conn.address[0],

pymongo/asynchronous/command_runner.py

@@ -284,6 +284,7 @@ async def run_cursor_command(
     more_to_come: bool = False,
     unpack_res: Optional[Callable[..., Any]] = None,
     cursor_id: Optional[int] = None,
+    op_id: Optional[int] = None,
 ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], datetime.timedelta]:
     """Run a cursor ``find``/``getMore`` operation over ``conn``.
 
@@ -304,6 +305,7 @@ async def run_cursor_command(
     :param unpack_res: A callable decoding the wire response; when ``None`` the
         reply's own ``unpack_response`` is used.
     :param cursor_id: The cursor id passed to ``unpack_res``.
+    :param op_id: The APM operation id; defaults to ``request_id``.
     """
     topology_id = client._topology_id if client is not None else None
     return await _run_command(
@@ -325,6 +327,7 @@ async def run_cursor_command(
         more_to_come=more_to_come,
         unpack_res=unpack_res,
         cursor_id=cursor_id,
+        op_id=op_id,
     )
 
 
@@ -348,6 +351,7 @@ async def run_command(
     user_fields: Optional[Mapping[str, Any]] = None,
     exhaust_allowed: bool = False,
     write_concern: Optional[WriteConcern] = None,
+    op_id: Optional[int] = None,
 ) -> _DocumentType:
     """Encode and execute a command over ``conn``, or raise socket.error.
 
@@ -376,6 +380,7 @@ async def run_command(
         passed to ``bson._decode_all_selective``.
     :param exhaust_allowed: True if we should enable OP_MSG exhaustAllowed.
     :param write_concern: The write concern for this command. Applied via CSOT.
+    :param op_id: The APM operation id; defaults to ``request_id``.
     """
     name = next(iter(spec))
 
@@ -428,6 +433,7 @@ async def run_command(
         codec_options=codec_options,
         user_fields=user_fields,
         orig=orig,
+        op_id=op_id,
         check=check,
         allowable_errors=allowable_errors,
         parse_write_concern_error=parse_write_concern_error,

pymongo/asynchronous/mongo_client.py

@@ -95,7 +95,7 @@
     _log_client_error,
     _log_or_warn,
 )
-from pymongo.message import _CursorAddress, _GetMore, _Query
+from pymongo.message import _CursorAddress, _GetMore, _Query, _randint
 from pymongo.monitoring import ConnectionClosedReason, _EventListeners
 from pymongo.operations import (
     DeleteMany,
@@ -1837,6 +1837,8 @@ async def _select_server(
             be pinned to a mongos server address.
           - `address` (optional): Address when sending a message
             to a specific server, used for getMore.
+          - `operation_id` (optional): Stable operation id shared across retries,
+            used for command monitoring.
         """
         try:
             topology = await self._get_topology()
@@ -1932,6 +1934,7 @@ async def _run_operation(
             async with operation.conn_mgr._lock:
                 async with _MongoClientErrorHandler(self, server, operation.session) as err_handler:  # type: ignore[arg-type]
                     err_handler.contribute_socket(operation.conn_mgr.conn)
+                    operation.conn_mgr.conn.op_id = _randint()
                     return await server.run_operation(
                         operation.conn_mgr.conn,
                         operation,
@@ -2023,6 +2026,7 @@ async def _retry_internal(
         :param retryable: If the operation should be retried once, defaults to None
         :param is_run_command: If this is a runCommand operation, defaults to False
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
+        :param operation_id: Stable operation id shared across retries, defaults to None
 
         :return: Output of the calling func()
         """
@@ -2069,6 +2073,7 @@ async def _retryable_read(
             (may not always be supported even if supplied), defaults to False
         :param is_run_command: If this is a runCommand operation, defaults to False.
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
+        :param operation_id: Stable operation id shared across retries, defaults to None
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2112,6 +2117,7 @@ async def _retryable_write(
         :param session: Client session we will use to execute write operation
         :param operation: The name of the operation that the server is being selected for
         :param bulk: bulk abstraction to execute operations in bulk, defaults to None
+        :param operation_id: Stable operation id shared across retries, defaults to None
         """
         async with self._tmp_session(session) as s:
             return await self._retry_with_session(retryable, func, s, bulk, operation, operation_id)
@@ -2795,7 +2801,7 @@ def __init__(
         self._server: Server = None  # type: ignore
         self._deprioritized_servers: list[Server] = []
         self._operation = operation
-        self._operation_id = operation_id
+        self._operation_id = operation_id if operation_id is not None else _randint()
         self._attempt_number = 0
         self._is_run_command = is_run_command
         self._is_aggregate_write = is_aggregate_write
@@ -3001,6 +3007,7 @@ async def _write(self) -> T:
             is_mongos = False
             self._server = await self._get_server()
             async with self._client._checkout(self._server, self._session) as conn:
+                conn.op_id = self._operation_id
                 max_wire_version = conn.max_wire_version
                 sessions_supported = (
                     self._session
@@ -3040,6 +3047,7 @@ async def _read(self) -> T:
             conn,
             read_pref,
         ):
+            conn.op_id = self._operation_id
             if self._retrying and not self._retryable and not self._always_retryable:
                 self._check_last_error()
             if self._retrying:

pymongo/asynchronous/pool.py

@@ -177,6 +177,8 @@ def __init__(
         self.creation_time = time.monotonic()
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
+        # Stable operation id for the operation currently using this connection.
+        self.op_id: Optional[int] = None
 
     def set_conn_timeout(self, timeout: Optional[float]) -> None:
         """Cache last timeout to avoid duplicate calls to conn.settimeout."""
@@ -416,6 +418,7 @@ async def command(
                 user_fields=user_fields,
                 exhaust_allowed=exhaust_allowed,
                 write_concern=write_concern,
+                op_id=self.op_id,
             )
         except (OperationFailure, NotPrimaryError):
             raise
@@ -1319,6 +1322,7 @@ async def checkin(self, conn: AsyncConnection) -> None:
         txn = conn.pinned_txn
         cursor = conn.pinned_cursor
         conn.active = False
+        conn.op_id = None
         conn.pinned_txn = False
         conn.pinned_cursor = False
         self.__pinned_sockets.discard(conn)

pymongo/asynchronous/server.py

@@ -185,6 +185,7 @@ async def run_operation(
             more_to_come=more_to_come,
             unpack_res=unpack_res,
             cursor_id=operation.cursor_id,
+            op_id=conn.op_id,
         )
         assert reply is not None
 

pymongo/synchronous/command_runner.py

@@ -284,6 +284,7 @@ def run_cursor_command(
     more_to_come: bool = False,
     unpack_res: Optional[Callable[..., Any]] = None,
     cursor_id: Optional[int] = None,
+    op_id: Optional[int] = None,
 ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], datetime.timedelta]:
     """Run a cursor ``find``/``getMore`` operation over ``conn``.
 
@@ -304,6 +305,7 @@ def run_cursor_command(
     :param unpack_res: A callable decoding the wire response; when ``None`` the
         reply's own ``unpack_response`` is used.
     :param cursor_id: The cursor id passed to ``unpack_res``.
+    :param op_id: The APM operation id; defaults to ``request_id``.
     """
     topology_id = client._topology_id if client is not None else None
     return _run_command(
@@ -325,6 +327,7 @@ def run_cursor_command(
         more_to_come=more_to_come,
         unpack_res=unpack_res,
         cursor_id=cursor_id,
+        op_id=op_id,
     )
 
 
@@ -348,6 +351,7 @@ def run_command(
     user_fields: Optional[Mapping[str, Any]] = None,
     exhaust_allowed: bool = False,
     write_concern: Optional[WriteConcern] = None,
+    op_id: Optional[int] = None,
 ) -> _DocumentType:
     """Encode and execute a command over ``conn``, or raise socket.error.
 
@@ -376,6 +380,7 @@ def run_command(
         passed to ``bson._decode_all_selective``.
     :param exhaust_allowed: True if we should enable OP_MSG exhaustAllowed.
     :param write_concern: The write concern for this command. Applied via CSOT.
+    :param op_id: The APM operation id; defaults to ``request_id``.
     """
     name = next(iter(spec))
 
@@ -428,6 +433,7 @@ def run_command(
         codec_options=codec_options,
         user_fields=user_fields,
         orig=orig,
+        op_id=op_id,
         check=check,
         allowable_errors=allowable_errors,
         parse_write_concern_error=parse_write_concern_error,

pymongo/synchronous/mongo_client.py

@@ -85,7 +85,7 @@
     _log_client_error,
     _log_or_warn,
 )
-from pymongo.message import _CursorAddress, _GetMore, _Query
+from pymongo.message import _CursorAddress, _GetMore, _Query, _randint
 from pymongo.monitoring import ConnectionClosedReason, _EventListeners
 from pymongo.operations import (
     DeleteMany,
@@ -1834,6 +1834,8 @@ def _select_server(
             be pinned to a mongos server address.
           - `address` (optional): Address when sending a message
             to a specific server, used for getMore.
+          - `operation_id` (optional): Stable operation id shared across retries,
+            used for command monitoring.
         """
         try:
             topology = self._get_topology()
@@ -1929,6 +1931,7 @@ def _run_operation(
             with operation.conn_mgr._lock:
                 with _MongoClientErrorHandler(self, server, operation.session) as err_handler:  # type: ignore[arg-type]
                     err_handler.contribute_socket(operation.conn_mgr.conn)
+                    operation.conn_mgr.conn.op_id = _randint()
                     return server.run_operation(
                         operation.conn_mgr.conn,
                         operation,
@@ -2020,6 +2023,7 @@ def _retry_internal(
         :param retryable: If the operation should be retried once, defaults to None
         :param is_run_command: If this is a runCommand operation, defaults to False
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
+        :param operation_id: Stable operation id shared across retries, defaults to None
 
         :return: Output of the calling func()
         """
@@ -2066,6 +2070,7 @@ def _retryable_read(
             (may not always be supported even if supplied), defaults to False
         :param is_run_command: If this is a runCommand operation, defaults to False.
         :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
+        :param operation_id: Stable operation id shared across retries, defaults to None
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2109,6 +2114,7 @@ def _retryable_write(
         :param session: Client session we will use to execute write operation
         :param operation: The name of the operation that the server is being selected for
         :param bulk: bulk abstraction to execute operations in bulk, defaults to None
+        :param operation_id: Stable operation id shared across retries, defaults to None
         """
         with self._tmp_session(session) as s:
             return self._retry_with_session(retryable, func, s, bulk, operation, operation_id)
@@ -2786,7 +2792,7 @@ def __init__(
         self._server: Server = None  # type: ignore
         self._deprioritized_servers: list[Server] = []
         self._operation = operation
-        self._operation_id = operation_id
+        self._operation_id = operation_id if operation_id is not None else _randint()
         self._attempt_number = 0
         self._is_run_command = is_run_command
         self._is_aggregate_write = is_aggregate_write
@@ -2992,6 +2998,7 @@ def _write(self) -> T:
             is_mongos = False
             self._server = self._get_server()
             with self._client._checkout(self._server, self._session) as conn:
+                conn.op_id = self._operation_id
                 max_wire_version = conn.max_wire_version
                 sessions_supported = (
                     self._session
@@ -3031,6 +3038,7 @@ def _read(self) -> T:
             conn,
             read_pref,
         ):
+            conn.op_id = self._operation_id
             if self._retrying and not self._retryable and not self._always_retryable:
                 self._check_last_error()
             if self._retrying:

pymongo/synchronous/pool.py

@@ -177,6 +177,8 @@ def __init__(
         self.creation_time = time.monotonic()
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
+        # Stable operation id for the operation currently using this connection.
+        self.op_id: Optional[int] = None
 
     def set_conn_timeout(self, timeout: Optional[float]) -> None:
         """Cache last timeout to avoid duplicate calls to conn.settimeout."""
@@ -416,6 +418,7 @@ def command(
                 user_fields=user_fields,
                 exhaust_allowed=exhaust_allowed,
                 write_concern=write_concern,
+                op_id=self.op_id,
             )
         except (OperationFailure, NotPrimaryError):
             raise
@@ -1315,6 +1318,7 @@ def checkin(self, conn: Connection) -> None:
         txn = conn.pinned_txn
         cursor = conn.pinned_cursor
         conn.active = False
+        conn.op_id = None
         conn.pinned_txn = False
         conn.pinned_cursor = False
         self.__pinned_sockets.discard(conn)

pymongo/synchronous/server.py

@@ -185,6 +185,7 @@ def run_operation(
             more_to_come=more_to_come,
             unpack_res=unpack_res,
             cursor_id=operation.cursor_id,
+            op_id=conn.op_id,
         )
         assert reply is not None
 

test/__init__.py

@@ -699,7 +699,7 @@ def require_failCommand_fail_point(self, func: Any) -> Any:
             func=func,
         )
 
-    def require_failCommand_appName(self, func):
+    def require_failCommand_appName(self, func: Any) -> Any:
         """Run a test only if the server supports the failCommand appName."""
         # SERVER-47195 and SERVER-49336.
         return self._require(

test/asynchronous/__init__.py

@@ -699,7 +699,7 @@ def require_failCommand_fail_point(self, func: Any) -> Any:
             func=func,
         )
 
-    def require_failCommand_appName(self, func):
+    def require_failCommand_appName(self, func: Any) -> Any:
         """Run a test only if the server supports the failCommand appName."""
         # SERVER-47195 and SERVER-49336.
         return self._require(