PYTHON-5903 Retry attempts should share a single stable operation_id (#2901)

Co-authored-by: Jeffrey 'Alex' Clark <aclark@aclark.net>

Author: shrey-mongo

doc/changelog.rst

@@ -8,6 +8,11 @@ Changes in Version 4.18.0
   to the same server, avoiding a full handshake on each new connection.
   Session resumption is supported on all Python versions for synchronous clients
   and on Python 3.11+ for async clients.
+- Command monitoring events and command log messages for a single logical
+  operation now share one stable ``operation_id`` across all of its retry
+  attempts, so consumers can correlate a retried operation's events. As a
+  result, ``operation_id`` is no longer equal to the per-attempt ``request_id``
+  for these operations.
 
 Changes in Version 4.17.0 (2026/04/20)
 --------------------------------------

pymongo/_op_id.py

@@ -0,0 +1,46 @@
+# Copyright 2026-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you
+# may not use this file except in compliance with the License.  You
+# may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.  See the License for the specific language governing
+# permissions and limitations under the License.
+
+"""Internal helpers for the APM operation id.
+
+The retryable read/write logic sets OP_ID for the duration of each attempt so
+that every attempt of one logical operation publishes the same operation_id.
+Commands run outside that scope (handshake, auth, killCursors, pinned-cursor
+getMores) read the default None and fall back to their request_id.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from contextvars import ContextVar
+from typing import Any, Optional
+
+OP_ID: ContextVar[Optional[int]] = ContextVar("OP_ID", default=None)
+
+
+def reset() -> None:
+    OP_ID.set(None)
+
+
+class _OpIdContext(AbstractContextManager[Any]):
+    """Set OP_ID for the duration of a with block."""
+
+    def __init__(self, op_id: Optional[int]):
+        self._op_id = op_id
+
+    def __enter__(self) -> None:
+        self._token = OP_ID.set(self._op_id)
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        OP_ID.reset(self._token)

pymongo/_telemetry.py

@@ -111,7 +111,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

@@ -47,7 +47,7 @@
 )
 
 from bson import _decode_all_selective
-from pymongo import _csot, helpers_shared, message
+from pymongo import _csot, _op_id, helpers_shared, message
 from pymongo._telemetry import _CommandTelemetry
 from pymongo.compression_support import _NO_COMPRESSION
 from pymongo.errors import NotPrimaryError, OperationFailure
@@ -128,7 +128,8 @@ async def _run_command(
     :param orig: The command document published in the ``STARTED`` APM event;
         defaults to ``cmd`` (differs only when the wire command was mutated,
         e.g. with a read preference or after encryption).
-    :param op_id: The APM operation id; defaults to ``request_id``.
+    :param op_id: The APM operation id; defaults to the ``OP_ID`` contextvar,
+        then ``request_id``.
     :param command_name: The command name for the ``SUCCEEDED``/``FAILED`` APM
         events; defaults to the first key of ``cmd``.
     :param check: Raise OperationFailure on a command error.
@@ -158,6 +159,8 @@ async def _run_command(
         command_name = name
     if orig is None:
         orig = cmd
+    if op_id is None:
+        op_id = _op_id.OP_ID.get()
 
     telemetry = _CommandTelemetry(topology_id, conn, listeners, cmd, dbname, request_id, op_id)
     telemetry.started(orig, ensure_db)

pymongo/asynchronous/encryption.py

@@ -55,7 +55,7 @@
 from bson.codec_options import CodecOptions
 from bson.errors import BSONError
 from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson
-from pymongo import _csot
+from pymongo import _csot, _op_id
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.cursor import AsyncCursor
 from pymongo.asynchronous.database import AsyncDatabase
@@ -468,7 +468,9 @@ async def encrypt(
         self._check_closed()
         encoded_cmd = _dict_to_bson(cmd, False, codec_options)
         with _wrap_encryption_errors():
-            encrypted_cmd = await self._auto_encrypter.encrypt(database, encoded_cmd)
+            # Don't let encryption's sub-commands inherit the in-flight op's id.
+            with _op_id._OpIdContext(None):
+                encrypted_cmd = await self._auto_encrypter.encrypt(database, encoded_cmd)
             # TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
             return _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
 
@@ -481,7 +483,9 @@ async def decrypt(self, response: bytes | memoryview) -> Optional[bytes]:
         """
         self._check_closed()
         with _wrap_encryption_errors():
-            return cast(bytes, await self._auto_encrypter.decrypt(response))
+            # Don't let decryption's sub-commands inherit the in-flight op's id.
+            with _op_id._OpIdContext(None):
+                return cast(bytes, await self._auto_encrypter.decrypt(response))
 
     def _check_closed(self) -> None:
         if self._closed:

pymongo/asynchronous/helpers.py

@@ -30,7 +30,7 @@
     cast,
 )
 
-from pymongo import _csot
+from pymongo import _csot, _op_id
 from pymongo.common import MAX_ADAPTIVE_RETRIES
 from pymongo.errors import (
     OperationFailure,
@@ -68,7 +68,9 @@ 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.
+                    with _op_id._OpIdContext(None):
+                        await conn.authenticate(reauthenticate=True)
                 else:
                     raise
                 return await func(*args, **kwargs)

pymongo/asynchronous/mongo_client.py

@@ -56,7 +56,7 @@
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions, TypeRegistry
 from bson.timestamp import Timestamp
-from pymongo import _csot, common, helpers_shared, periodic_executor
+from pymongo import _csot, _op_id, common, helpers_shared, periodic_executor
 from pymongo._telemetry import log_command_retry
 from pymongo.asynchronous import client_session, database, uri_parser
 from pymongo.asynchronous.change_stream import AsyncChangeStream, AsyncClusterChangeStream
@@ -94,7 +94,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,
@@ -1836,6 +1836,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()
@@ -2011,6 +2013,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()
         """
@@ -2057,6 +2060,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
@@ -2100,6 +2104,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)
@@ -2783,7 +2788,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
@@ -3012,7 +3017,9 @@ async def _write(self) -> T:
                     self._retryable = False
                 if self._retrying:
                     self._log_retry(is_write=True)
-                return await self._func(self._session, conn, self._retryable)  # type: ignore
+                # One operation id across all attempts of this operation.
+                with _op_id._OpIdContext(self._operation_id):
+                    return await self._func(self._session, conn, self._retryable)  # type: ignore
         except PyMongoError as exc:
             if not self._retryable:
                 raise
@@ -3035,7 +3042,9 @@ async def _read(self) -> T:
                 self._check_last_error()
             if self._retrying:
                 self._log_retry(is_write=False)
-            return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
+            # One operation id across all attempts of this operation.
+            with _op_id._OpIdContext(self._operation_id):
+                return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
 
 
 def _after_fork_child() -> None:

pymongo/periodic_executor.py

@@ -23,7 +23,7 @@
 import weakref
 from typing import Any, Optional
 
-from pymongo import _csot
+from pymongo import _csot, _op_id
 from pymongo._asyncio_task import create_task
 from pymongo.lock import _create_lock
 
@@ -94,8 +94,9 @@ def skip_sleep(self) -> None:
         self._skip_sleep = True
 
     async def _run(self) -> None:
-        # The CSOT contextvars must be cleared inside the executor task before execution begins
+        # The CSOT and op id contextvars must be cleared inside the executor task before execution begins
         _csot.reset_all()
+        _op_id.reset()
         while not self._stopped:
             if self._task and self._task.cancelling():  # type: ignore[unused-ignore, attr-defined]
                 raise asyncio.CancelledError

pymongo/synchronous/command_runner.py

@@ -47,7 +47,7 @@
 )
 
 from bson import _decode_all_selective
-from pymongo import _csot, helpers_shared, message
+from pymongo import _csot, _op_id, helpers_shared, message
 from pymongo._telemetry import _CommandTelemetry
 from pymongo.compression_support import _NO_COMPRESSION
 from pymongo.errors import NotPrimaryError, OperationFailure
@@ -128,7 +128,8 @@ def _run_command(
     :param orig: The command document published in the ``STARTED`` APM event;
         defaults to ``cmd`` (differs only when the wire command was mutated,
         e.g. with a read preference or after encryption).
-    :param op_id: The APM operation id; defaults to ``request_id``.
+    :param op_id: The APM operation id; defaults to the ``OP_ID`` contextvar,
+        then ``request_id``.
     :param command_name: The command name for the ``SUCCEEDED``/``FAILED`` APM
         events; defaults to the first key of ``cmd``.
     :param check: Raise OperationFailure on a command error.
@@ -158,6 +159,8 @@ def _run_command(
         command_name = name
     if orig is None:
         orig = cmd
+    if op_id is None:
+        op_id = _op_id.OP_ID.get()
 
     telemetry = _CommandTelemetry(topology_id, conn, listeners, cmd, dbname, request_id, op_id)
     telemetry.started(orig, ensure_db)

pymongo/synchronous/encryption.py

@@ -54,7 +54,7 @@
 from bson.codec_options import CodecOptions
 from bson.errors import BSONError
 from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson
-from pymongo import _csot
+from pymongo import _csot, _op_id
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import (
@@ -465,7 +465,9 @@ def encrypt(
         self._check_closed()
         encoded_cmd = _dict_to_bson(cmd, False, codec_options)
         with _wrap_encryption_errors():
-            encrypted_cmd = self._auto_encrypter.encrypt(database, encoded_cmd)
+            # Don't let encryption's sub-commands inherit the in-flight op's id.
+            with _op_id._OpIdContext(None):
+                encrypted_cmd = self._auto_encrypter.encrypt(database, encoded_cmd)
             # TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
             return _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
 
@@ -478,7 +480,9 @@ def decrypt(self, response: bytes | memoryview) -> Optional[bytes]:
         """
         self._check_closed()
         with _wrap_encryption_errors():
-            return cast(bytes, self._auto_encrypter.decrypt(response))
+            # Don't let decryption's sub-commands inherit the in-flight op's id.
+            with _op_id._OpIdContext(None):
+                return cast(bytes, self._auto_encrypter.decrypt(response))
 
     def _check_closed(self) -> None:
         if self._closed: