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

@@ -325,6 +325,7 @@ async def run_cursor_command(
         more_to_come=more_to_come,
         unpack_res=unpack_res,
         cursor_id=cursor_id,
+        op_id=conn.op_id,
     )
 
 
@@ -428,6 +429,7 @@ async def run_command(
         codec_options=codec_options,
         user_fields=user_fields,
         orig=orig,
+        op_id=conn.op_id,
         check=check,
         allowable_errors=allowable_errors,
         parse_write_concern_error=parse_write_concern_error,

pymongo/asynchronous/helpers.py

@@ -68,7 +68,13 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
                         conn = arg.conn  # type: ignore[assignment]
                         break
                 if conn:
-                    await conn.authenticate(reauthenticate=True)
+                    # Don't let reauth's auth commands inherit the in-flight op's id.
+                    prev_op_id = conn.op_id
+                    conn.op_id = None
+                    try:
+                        await conn.authenticate(reauthenticate=True)
+                    finally:
+                        conn.op_id = prev_op_id
                 else:
                     raise
                 return await func(*args, **kwargs)

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,
@@ -1793,10 +1793,13 @@ async def _checkout(
         async with _MongoClientErrorHandler(self, server, session) as err_handler:
             # Reuse the pinned connection, if it exists.
             if in_txn and session and session._pinned_connection:
-                err_handler.contribute_socket(session._pinned_connection)
-                yield session._pinned_connection
+                conn = session._pinned_connection
+                conn.op_id = None
+                err_handler.contribute_socket(conn)
+                yield conn
                 return
             async with await server.checkout(handler=err_handler) as conn:
+                conn.op_id = None  # Only retryable read/write logic sets an op_id.
                 # Pin this session to the selected server or connection.
                 if (
                     in_txn
@@ -1837,6 +1840,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()
@@ -1933,6 +1938,8 @@ 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)
+                    # Non-retryable getMore on a pinned conn; no shared op_id.
+                    operation.conn_mgr.conn.op_id = None
                     return await run_with_conn(
                         operation.conn_mgr.conn, operation, operation.read_preference
                     )
@@ -2012,6 +2019,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()
         """
@@ -2058,6 +2066,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
@@ -2101,6 +2110,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)
@@ -2217,6 +2227,8 @@ async def _kill_cursor_impl(
         namespace = address.namespace
         db, coll = namespace.split(".", 1)
         spec = {"killCursors": coll, "cursors": cursor_ids}
+        # killCursors is its own op; drop any op_id left on a pinned cursor conn.
+        conn.op_id = None
         await conn.command(db, spec, session=session, client=self)
 
     async def _process_kill_cursors(self) -> None:
@@ -2784,7 +2796,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
@@ -2990,6 +3002,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
@@ -3029,6 +3042,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,13 @@ def __init__(
         self.creation_time = time.monotonic()
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
+        # APM operation id of the op currently using this connection, read by
+        # command(). Kept here for simplicity: retries reuse it without threading
+        # an op_id kwarg through every command. Reset to None on checkout; only
+        # the retryable read/write logic sets a real value. An op running a
+        # command directly on a pinned connection (see killCursors, reauth) must
+        # clear it so it isn't attributed to an unrelated command.
+        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."""

pymongo/synchronous/command_runner.py

@@ -325,6 +325,7 @@ def run_cursor_command(
         more_to_come=more_to_come,
         unpack_res=unpack_res,
         cursor_id=cursor_id,
+        op_id=conn.op_id,
     )
 
 
@@ -428,6 +429,7 @@ def run_command(
         codec_options=codec_options,
         user_fields=user_fields,
         orig=orig,
+        op_id=conn.op_id,
         check=check,
         allowable_errors=allowable_errors,
         parse_write_concern_error=parse_write_concern_error,

pymongo/synchronous/helpers.py

@@ -68,7 +68,13 @@ def inner(*args: Any, **kwargs: Any) -> Any:
                         conn = arg.conn  # type: ignore[assignment]
                         break
                 if conn:
-                    conn.authenticate(reauthenticate=True)
+                    # Don't let reauth's auth commands inherit the in-flight op's id.
+                    prev_op_id = conn.op_id
+                    conn.op_id = None
+                    try:
+                        conn.authenticate(reauthenticate=True)
+                    finally:
+                        conn.op_id = prev_op_id
                 else:
                     raise
                 return func(*args, **kwargs)

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,
@@ -1790,10 +1790,13 @@ def _checkout(
         with _MongoClientErrorHandler(self, server, session) as err_handler:
             # Reuse the pinned connection, if it exists.
             if in_txn and session and session._pinned_connection:
-                err_handler.contribute_socket(session._pinned_connection)
-                yield session._pinned_connection
+                conn = session._pinned_connection
+                conn.op_id = None
+                err_handler.contribute_socket(conn)
+                yield conn
                 return
             with server.checkout(handler=err_handler) as conn:
+                conn.op_id = None  # Only retryable read/write logic sets an op_id.
                 # Pin this session to the selected server or connection.
                 if (
                     in_txn
@@ -1834,6 +1837,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()
@@ -1930,6 +1935,8 @@ 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)
+                    # Non-retryable getMore on a pinned conn; no shared op_id.
+                    operation.conn_mgr.conn.op_id = None
                     return run_with_conn(
                         operation.conn_mgr.conn, operation, operation.read_preference
                     )
@@ -2009,6 +2016,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()
         """
@@ -2055,6 +2063,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
@@ -2098,6 +2107,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)
@@ -2214,6 +2224,8 @@ def _kill_cursor_impl(
         namespace = address.namespace
         db, coll = namespace.split(".", 1)
         spec = {"killCursors": coll, "cursors": cursor_ids}
+        # killCursors is its own op; drop any op_id left on a pinned cursor conn.
+        conn.op_id = None
         conn.command(db, spec, session=session, client=self)
 
     def _process_kill_cursors(self) -> None:
@@ -2775,7 +2787,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
@@ -2981,6 +2993,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
@@ -3020,6 +3033,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,13 @@ def __init__(
         self.creation_time = time.monotonic()
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
+        # APM operation id of the op currently using this connection, read by
+        # command(). Kept here for simplicity: retries reuse it without threading
+        # an op_id kwarg through every command. Reset to None on checkout; only
+        # the retryable read/write logic sets a real value. An op running a
+        # command directly on a pinned connection (see killCursors, reauth) must
+        # clear it so it isn't attributed to an unrelated command.
+        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."""