Merge branch 'main' into fix-redact-aws-session-token-repr

Author: aclark4life

.github/pull_request_template.md

@@ -1,6 +1,6 @@
 <!-- Thanks for contributing! -->
 <!-- Please ensure that the title of the PR is in the following form:
-[JIRA TICKET]: Issue Title
+PYTHON-XXXX Issue Title
 
 If you are an external contributor and there is no JIRA ticket associated with your change, then use your best judgement
 for the PR title. A MongoDB employee will create a JIRA ticket and edit the name and links as appropriate.
@@ -9,7 +9,9 @@ Note on AI Contributions:
 We only accept pull requests that are authored and submitted by human contributors who fully understand the changes they are proposing.
 All contributions must be written and understood by human contributors. Please read about our policy in our contributing guide.
 -->
-[JIRA TICKET]
+<!-- Replace XXXX with the ticket number. If you do not have a ticket (external contributors), you may leave PYTHON-XXXX;
+a MongoDB employee will create a JIRA ticket and update the PR title/links as appropriate. -->
+PYTHON-XXXX
 
 ## Changes in this PR
 <!-- What changes did you make to the code? What new APIs (public or private) were added, removed, or edited to generate

doc/changelog.rst

@@ -12,6 +12,11 @@ Changes in Version 4.18.0
   AWS session tokens, from the representations of
   :class:`~pymongo.synchronous.mongo_client.MongoClient` and
   :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`.
+- 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,
@@ -1838,6 +1838,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()
@@ -2013,6 +2015,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()
         """
@@ -2059,6 +2062,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
@@ -2102,6 +2106,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)
@@ -2785,7 +2790,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
@@ -3014,7 +3019,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
@@ -3037,7 +3044,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/asynchronous/pool.py

@@ -287,7 +287,9 @@ async def _hello(
             hello.logical_session_timeout_minutes is not None and hello.is_readable
         )
         self.logical_session_timeout_minutes: Optional[int] = hello.logical_session_timeout_minutes
-        self.hello_ok = hello.hello_ok
+        # hello_ok is set from helloOk, which is only returned from ismaster
+        # don't overwrite this when connection switches to hello
+        self.hello_ok = self.hello_ok or hello.hello_ok
         self.is_repl = hello.server_type in (
             SERVER_TYPE.RSPrimary,
             SERVER_TYPE.RSSecondary,

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