From 75f4e2179dcaabbe6ebb3b5fa36f5cd291dea5dd Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Tue, 30 Jun 2026 15:46:07 +0000 Subject: [PATCH 01/14] fixed merge conflict --- pymongo/_telemetry.py | 2 +- pymongo/asynchronous/command_runner.py | 6 + pymongo/asynchronous/mongo_client.py | 12 +- pymongo/asynchronous/pool.py | 4 + pymongo/asynchronous/server.py | 1 + pymongo/synchronous/command_runner.py | 6 + pymongo/synchronous/mongo_client.py | 12 +- pymongo/synchronous/pool.py | 4 + pymongo/synchronous/server.py | 1 + test/__init__.py | 2 +- test/asynchronous/__init__.py | 2 +- test/asynchronous/test_operation_id_retry.py | 160 +++++++++++++++++++ test/test_operation_id_retry.py | 158 ++++++++++++++++++ tools/synchro.py | 1 + 14 files changed, 364 insertions(+), 7 deletions(-) create mode 100644 test/asynchronous/test_operation_id_retry.py create mode 100644 test/test_operation_id_retry.py diff --git a/pymongo/_telemetry.py b/pymongo/_telemetry.py index 471a013eb9..168d2a07fb 100644 --- a/pymongo/_telemetry.py +++ b/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], diff --git a/pymongo/asynchronous/command_runner.py b/pymongo/asynchronous/command_runner.py index 1f32cc744e..ac47aed798 100644 --- a/pymongo/asynchronous/command_runner.py +++ b/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, diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index a558f96356..bcbbc87e3d 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/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: diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 4ed3b85dbf..a008246f40 100644 --- a/pymongo/asynchronous/pool.py +++ b/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) diff --git a/pymongo/asynchronous/server.py b/pymongo/asynchronous/server.py index 57158dfc44..bceb080e17 100644 --- a/pymongo/asynchronous/server.py +++ b/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 diff --git a/pymongo/synchronous/command_runner.py b/pymongo/synchronous/command_runner.py index 077e0f9409..88574426f6 100644 --- a/pymongo/synchronous/command_runner.py +++ b/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, diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 5f321afe5c..42d912475a 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/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: diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 1006735444..422b07726b 100644 --- a/pymongo/synchronous/pool.py +++ b/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) diff --git a/pymongo/synchronous/server.py b/pymongo/synchronous/server.py index 09d8fb75e1..5ce6683f8c 100644 --- a/pymongo/synchronous/server.py +++ b/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 diff --git a/test/__init__.py b/test/__init__.py index f4ae7fe948..5d40b8dbac 100644 --- a/test/__init__.py +++ b/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( diff --git a/test/asynchronous/__init__.py b/test/asynchronous/__init__.py index 0699451d7e..bee0a2b8ce 100644 --- a/test/asynchronous/__init__.py +++ b/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( diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py new file mode 100644 index 0000000000..5941a2e7ad --- /dev/null +++ b/test/asynchronous/test_operation_id_retry.py @@ -0,0 +1,160 @@ +# Copyright 2024-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. + +"""Test that retry attempts reuse a single stable CommandStartedEvent.operation_id.""" + +from __future__ import annotations + +import sys + +sys.path[0:0] = [""] + +import pymongo +from pymongo.errors import ConnectionFailure +from pymongo.operations import InsertOne +from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest +from test.utils_shared import AllowListEventListener + +_IS_SYNC = False + +_APP_NAME = "operationIdRetryTest" + +# Each operation, paired with the wire command it issues and an awaitable action. +# These are all retryable; a stable operation_id must span every retry attempt. +_RETRYABLE_WRITES = [ + ("insert", lambda c: c.insert_one({"_id": 100})), + ("update", lambda c: c.update_one({"_id": 1}, {"$set": {"y": 1}})), + ("update", lambda c: c.replace_one({"_id": 2}, {"x": 9})), + ("delete", lambda c: c.delete_one({"_id": 3})), + ("findAndModify", lambda c: c.find_one_and_update({"_id": 4}, {"$set": {"y": 2}})), + ("insert", lambda c: c.bulk_write([InsertOne({"_id": 200}), InsertOne({"_id": 201})])), +] + + +_RETRYABLE_READS = [ + ("find", lambda c: c.find({"x": 1}).to_list()), + ("find", lambda c: c.find_one({"_id": 1})), + ("aggregate", lambda c: _agg(c)), + ("aggregate", lambda c: c.count_documents({"x": 1})), + ("distinct", lambda c: c.distinct("x")), + ("listIndexes", lambda c: _list_indexes(c)), +] + + +async def _agg(coll): + cursor = await coll.aggregate([{"$match": {"x": 1}}]) + return await cursor.to_list() + + +async def _list_indexes(coll): + cursor = await coll.list_indexes() + return await cursor.to_list() + + +class TestOperationIdRetry(AsyncIntegrationTest): + RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. + + @async_client_context.require_failCommand_fail_point + @async_client_context.require_failCommand_appName + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + + async def _seed(self, coll): + await coll.drop() + await coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) + await coll.create_index("x") + + async def _check_stable_operation_id(self, command_name, action, retries): + """Force ``retries`` retries of ``action`` and assert every command + event for ``command_name`` shares one integer operation_id.""" + listener = AllowListEventListener(command_name) + client = await self.async_rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) + coll = client.pymongo_test.test_operation_id_retry + await self._seed(coll) + listener.reset() + + fail_point = { + "mode": {"times": retries}, + "data": { + "failCommands": [command_name], + "closeConnection": True, + "appName": _APP_NAME, + }, + } + async with self.fail_point(fail_point): + # A CSOT timeout lets a single operation retry more than once. + with pymongo.timeout(60): + await action(coll) + + started = listener.started_events + failed = listener.failed_events + succeeded = listener.succeeded_events + op_ids = [e.operation_id for e in started + failed + succeeded] + + self.assertEqual(len(started), retries + 1, "expected one started event per attempt") + self.assertEqual(len(failed), retries) + self.assertEqual(len(succeeded), 1) + self.assertTrue(all(isinstance(op, int) for op in op_ids)) + self.assertEqual( + len(set(op_ids)), + 1, + f"operation_id not stable across retries for {command_name}: {op_ids}", + ) + + @async_client_context.require_no_standalone + async def test_retryable_writes_reuse_operation_id(self): + for command_name, action in _RETRYABLE_WRITES: + with self.subTest(command=command_name): + await self._check_stable_operation_id(command_name, action, self.RETRIES) + + async def test_retryable_reads_reuse_operation_id(self): + for command_name, action in _RETRYABLE_READS: + with self.subTest(command=command_name): + await self._check_stable_operation_id(command_name, action, self.RETRIES) + + @async_client_context.require_no_standalone + async def test_non_retryable_write_is_not_retried(self): + # Multi-document writes are not retryable: a single network error must + # surface immediately, with exactly one attempt. + for command_name, action in [ + ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), + ("delete", lambda c: c.delete_many({"x": 2})), + ]: + with self.subTest(command=command_name): + listener = AllowListEventListener(command_name) + client = await self.async_rs_or_single_client( + event_listeners=[listener], appname=_APP_NAME + ) + coll = client.pymongo_test.test_operation_id_retry + await self._seed(coll) + listener.reset() + + fail_point = { + "mode": {"times": 1}, + "data": { + "failCommands": [command_name], + "closeConnection": True, + "appName": _APP_NAME, + }, + } + async with self.fail_point(fail_point): + with self.assertRaises(ConnectionFailure): + await action(coll) + + self.assertEqual(len(listener.started_events), 1, "must not retry") + self.assertIsInstance(listener.started_events[0].operation_id, int) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py new file mode 100644 index 0000000000..313416af1a --- /dev/null +++ b/test/test_operation_id_retry.py @@ -0,0 +1,158 @@ +# Copyright 2024-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. + +"""Test that retry attempts reuse a single stable CommandStartedEvent.operation_id.""" + +from __future__ import annotations + +import sys + +sys.path[0:0] = [""] + +import pymongo +from pymongo.errors import ConnectionFailure +from pymongo.operations import InsertOne +from test import IntegrationTest, client_context, unittest +from test.utils_shared import AllowListEventListener + +_IS_SYNC = True + +_APP_NAME = "operationIdRetryTest" + +# Each operation, paired with the wire command it issues and an awaitable action. +# These are all retryable; a stable operation_id must span every retry attempt. +_RETRYABLE_WRITES = [ + ("insert", lambda c: c.insert_one({"_id": 100})), + ("update", lambda c: c.update_one({"_id": 1}, {"$set": {"y": 1}})), + ("update", lambda c: c.replace_one({"_id": 2}, {"x": 9})), + ("delete", lambda c: c.delete_one({"_id": 3})), + ("findAndModify", lambda c: c.find_one_and_update({"_id": 4}, {"$set": {"y": 2}})), + ("insert", lambda c: c.bulk_write([InsertOne({"_id": 200}), InsertOne({"_id": 201})])), +] + + +_RETRYABLE_READS = [ + ("find", lambda c: c.find({"x": 1}).to_list()), + ("find", lambda c: c.find_one({"_id": 1})), + ("aggregate", lambda c: _agg(c)), + ("aggregate", lambda c: c.count_documents({"x": 1})), + ("distinct", lambda c: c.distinct("x")), + ("listIndexes", lambda c: _list_indexes(c)), +] + + +def _agg(coll): + cursor = coll.aggregate([{"$match": {"x": 1}}]) + return cursor.to_list() + + +def _list_indexes(coll): + cursor = coll.list_indexes() + return cursor.to_list() + + +class TestOperationIdRetry(IntegrationTest): + RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. + + @client_context.require_failCommand_fail_point + @client_context.require_failCommand_appName + def setUp(self) -> None: + super().setUp() + + def _seed(self, coll): + coll.drop() + coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) + coll.create_index("x") + + def _check_stable_operation_id(self, command_name, action, retries): + """Force ``retries`` retries of ``action`` and assert every command + event for ``command_name`` shares one integer operation_id.""" + listener = AllowListEventListener(command_name) + client = self.rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) + coll = client.pymongo_test.test_operation_id_retry + self._seed(coll) + listener.reset() + + fail_point = { + "mode": {"times": retries}, + "data": { + "failCommands": [command_name], + "closeConnection": True, + "appName": _APP_NAME, + }, + } + with self.fail_point(fail_point): + # A CSOT timeout lets a single operation retry more than once. + with pymongo.timeout(60): + action(coll) + + started = listener.started_events + failed = listener.failed_events + succeeded = listener.succeeded_events + op_ids = [e.operation_id for e in started + failed + succeeded] + + self.assertEqual(len(started), retries + 1, "expected one started event per attempt") + self.assertEqual(len(failed), retries) + self.assertEqual(len(succeeded), 1) + self.assertTrue(all(isinstance(op, int) for op in op_ids)) + self.assertEqual( + len(set(op_ids)), + 1, + f"operation_id not stable across retries for {command_name}: {op_ids}", + ) + + @client_context.require_no_standalone + def test_retryable_writes_reuse_operation_id(self): + for command_name, action in _RETRYABLE_WRITES: + with self.subTest(command=command_name): + self._check_stable_operation_id(command_name, action, self.RETRIES) + + def test_retryable_reads_reuse_operation_id(self): + for command_name, action in _RETRYABLE_READS: + with self.subTest(command=command_name): + self._check_stable_operation_id(command_name, action, self.RETRIES) + + @client_context.require_no_standalone + def test_non_retryable_write_is_not_retried(self): + # Multi-document writes are not retryable: a single network error must + # surface immediately, with exactly one attempt. + for command_name, action in [ + ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), + ("delete", lambda c: c.delete_many({"x": 2})), + ]: + with self.subTest(command=command_name): + listener = AllowListEventListener(command_name) + client = self.rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) + coll = client.pymongo_test.test_operation_id_retry + self._seed(coll) + listener.reset() + + fail_point = { + "mode": {"times": 1}, + "data": { + "failCommands": [command_name], + "closeConnection": True, + "appName": _APP_NAME, + }, + } + with self.fail_point(fail_point): + with self.assertRaises(ConnectionFailure): + action(coll) + + self.assertEqual(len(listener.started_events), 1, "must not retry") + self.assertIsInstance(listener.started_events[0].operation_id, int) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/synchro.py b/tools/synchro.py index 8132167e71..9fa60f1120 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -253,6 +253,7 @@ def async_only_test(f: Path) -> bool: "test_monitoring.py", "test_mongos_load_balancing.py", "test_on_demand_csfle.py", + "test_operation_id_retry.py", "test_periodic_executor.py", "test_pooling.py", "test_raw_bson.py", From 7237bf149c90bc42cae54910ddadfb39db1a5e2d Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Tue, 7 Jul 2026 19:08:27 +0000 Subject: [PATCH 02/14] fixed stale op id for reauth and killCursors ops and added justification --- pymongo/asynchronous/helpers.py | 4 + pymongo/asynchronous/mongo_client.py | 2 + pymongo/asynchronous/pool.py | 6 +- pymongo/synchronous/helpers.py | 4 + pymongo/synchronous/mongo_client.py | 2 + pymongo/synchronous/pool.py | 6 +- test/__init__.py | 2 +- test/asynchronous/__init__.py | 2 +- test/asynchronous/test_operation_id_retry.py | 1 - test/test_operation_id_retry.py | 1 - tools/synchro.py | 90 +------------------- 11 files changed, 26 insertions(+), 94 deletions(-) diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py index 0ee99a441a..be5ceef9bb 100644 --- a/pymongo/asynchronous/helpers.py +++ b/pymongo/asynchronous/helpers.py @@ -68,7 +68,11 @@ async def inner(*args: Any, **kwargs: Any) -> Any: conn = arg.conn # type: ignore[assignment] break if conn: + # Don't let reauth's auth commands inherit the in-flight op's id. + prev_op_id = conn.op_id + conn.op_id = None await conn.authenticate(reauthenticate=True) + conn.op_id = prev_op_id else: raise return await func(*args, **kwargs) diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index bcbbc87e3d..b720457f35 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -2234,6 +2234,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: diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index a008246f40..2780aae427 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -177,7 +177,11 @@ 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. + # APM operation id of the op currently using this connection. Kept on the + # connection for simplicity: retries reuse it without threading an op_id + # kwarg through every command. Cleared on checkin; an op that runs a + # command directly on a pinned connection must clear it (see killCursors, + # reauth) so it isn't attributed to an unrelated command. self.op_id: Optional[int] = None def set_conn_timeout(self, timeout: Optional[float]) -> None: diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py index f24d72aea9..efcb29fb52 100644 --- a/pymongo/synchronous/helpers.py +++ b/pymongo/synchronous/helpers.py @@ -68,7 +68,11 @@ def inner(*args: Any, **kwargs: Any) -> Any: conn = arg.conn # type: ignore[assignment] break if conn: + # Don't let reauth's auth commands inherit the in-flight op's id. + prev_op_id = conn.op_id + conn.op_id = None conn.authenticate(reauthenticate=True) + conn.op_id = prev_op_id else: raise return func(*args, **kwargs) diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 42d912475a..59ebfc6de7 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -2231,6 +2231,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: diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 422b07726b..6ec1efb9b0 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -177,7 +177,11 @@ 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. + # APM operation id of the op currently using this connection. Kept on the + # connection for simplicity: retries reuse it without threading an op_id + # kwarg through every command. Cleared on checkin; an op that runs a + # command directly on a pinned connection must clear it (see killCursors, + # reauth) so it isn't attributed to an unrelated command. self.op_id: Optional[int] = None def set_conn_timeout(self, timeout: Optional[float]) -> None: diff --git a/test/__init__.py b/test/__init__.py index 5d40b8dbac..f4ae7fe948 100644 --- a/test/__init__.py +++ b/test/__init__.py @@ -699,7 +699,7 @@ def require_failCommand_fail_point(self, func: Any) -> Any: func=func, ) - def require_failCommand_appName(self, func: Any) -> Any: + def require_failCommand_appName(self, func): """Run a test only if the server supports the failCommand appName.""" # SERVER-47195 and SERVER-49336. return self._require( diff --git a/test/asynchronous/__init__.py b/test/asynchronous/__init__.py index bee0a2b8ce..0699451d7e 100644 --- a/test/asynchronous/__init__.py +++ b/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: Any) -> Any: + def require_failCommand_appName(self, func): """Run a test only if the server supports the failCommand appName.""" # SERVER-47195 and SERVER-49336. return self._require( diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 5941a2e7ad..27ca9f6ec4 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -66,7 +66,6 @@ class TestOperationIdRetry(AsyncIntegrationTest): RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. @async_client_context.require_failCommand_fail_point - @async_client_context.require_failCommand_appName async def asyncSetUp(self) -> None: await super().asyncSetUp() diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 313416af1a..58e3b533aa 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -66,7 +66,6 @@ class TestOperationIdRetry(IntegrationTest): RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. @client_context.require_failCommand_fail_point - @client_context.require_failCommand_appName def setUp(self) -> None: super().setUp() diff --git a/tools/synchro.py b/tools/synchro.py index 9fa60f1120..f9c3f23c8c 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -189,6 +189,7 @@ def async_only_test(f: Path) -> bool: "test_async_loop_safety.py", "test_async_contextvars_reset.py", "test_async_loop_unblocked.py", + "test_async_network_layer.py", ] @@ -198,92 +199,6 @@ def async_only_test(f: Path) -> bool: if f.is_file() and not async_only_test(f) ] -# Add each asynchronized test here as part of the converting PR -converted_tests = [ - "__init__.py", - "conftest.py", - "helpers.py", - "pymongo_mocks.py", - "utils_spec_runner.py", - "qcheck.py", - "test_auth.py", - "test_auth_oidc.py", - "test_auth_spec.py", - "test_bulk.py", - "test_causal_consistency.py", - "test_change_stream.py", - "test_client.py", - "test_client_backpressure.py", - "test_client_bulk_write.py", - "test_client_context.py", - "test_client_metadata.py", - "test_collation.py", - "test_collection.py", - "test_collection_management.py", - "test_command_logging.py", - "test_command_logging.py", - "test_command_monitoring.py", - "test_comment.py", - "test_common.py", - "test_connection_logging.py", - "test_connection_monitoring.py", - "test_connections_survive_primary_stepdown_spec.py", - "test_create_entities.py", - "test_crud_unified.py", - "test_csot.py", - "test_cursor.py", - "test_custom_types.py", - "test_database.py", - "test_discovery_and_monitoring.py", - "test_dns.py", - "test_encryption.py", - "test_examples.py", - "test_grid_file.py", - "test_gridfs.py", - "test_gridfs_bucket.py", - "test_gridfs_spec.py", - "test_handshake_unified.py", - "test_heartbeat_monitoring.py", - "test_index_management.py", - "test_json_util_integration.py", - "test_load_balancer.py", - "test_logger.py", - "test_max_staleness.py", - "test_monitor.py", - "test_monitoring.py", - "test_mongos_load_balancing.py", - "test_on_demand_csfle.py", - "test_operation_id_retry.py", - "test_periodic_executor.py", - "test_pooling.py", - "test_raw_bson.py", - "test_read_concern.py", - "test_read_preferences.py", - "test_read_write_concern_spec.py", - "test_retryable_reads.py", - "test_retryable_reads_unified.py", - "test_retryable_writes.py", - "test_retryable_writes_unified.py", - "test_run_command.py", - "test_sdam_monitoring_spec.py", - "test_server_selection.py", - "test_server_selection_in_window.py", - "test_server_selection_logging.py", - "test_server_selection_rtt.py", - "test_session.py", - "test_sessions_unified.py", - "test_srv_polling.py", - "test_ssl.py", - "test_streaming_protocol.py", - "test_transactions.py", - "test_transactions_unified.py", - "test_unified_format.py", - "test_versioned_api_integration.py", - "unified_format.py", - "utils_selection_tests.py", - "utils.py", -] - def process_files( files: list[str], docstring_translate_files: list[str], sync_test_files: list[str] @@ -467,13 +382,12 @@ def main() -> None: unasync_directory(filtered_tests, _test_base, _test_dest_base, replacements) # Derive generated output paths directly from filtered source paths. - converted_tests_set = set(converted_tests) generated_pymongo = [_pymongo_dest_base + Path(f).name for f in filtered_async] generated_gridfs = [_gridfs_dest_base + Path(f).name for f in filtered_gridfs] generated_tests = [ _test_dest_base + Path(f).name for f in filtered_tests - if Path(f).name in converted_tests_set and (Path(_test_dest_base) / Path(f).name).is_file() + if (Path(_test_dest_base) / Path(f).name).is_file() ] docstring_translate_files = generated_pymongo + generated_gridfs + generated_tests From 643ccef6d8d884bd2d0287cfc8d48e4648aece8d Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 8 Jul 2026 18:11:57 +0000 Subject: [PATCH 03/14] made requested changes --- pymongo/asynchronous/command_runner.py | 8 +- pymongo/asynchronous/cursor_base.py | 1 - pymongo/asynchronous/mongo_client.py | 5 +- pymongo/asynchronous/pool.py | 13 ++- pymongo/synchronous/command_runner.py | 8 +- pymongo/synchronous/cursor_base.py | 1 - pymongo/synchronous/mongo_client.py | 5 +- pymongo/synchronous/pool.py | 13 ++- test/asynchronous/test_operation_id_retry.py | 86 ++++++++++---------- test/test_operation_id_retry.py | 84 +++++++++---------- 10 files changed, 107 insertions(+), 117 deletions(-) diff --git a/pymongo/asynchronous/command_runner.py b/pymongo/asynchronous/command_runner.py index fd850d8c61..7900a6c958 100644 --- a/pymongo/asynchronous/command_runner.py +++ b/pymongo/asynchronous/command_runner.py @@ -284,7 +284,6 @@ 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``. @@ -305,7 +304,6 @@ 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( @@ -327,7 +325,7 @@ async def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, - op_id=op_id, + op_id=conn.op_id, ) @@ -351,7 +349,6 @@ 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. @@ -380,7 +377,6 @@ 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)) @@ -433,7 +429,7 @@ async def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=op_id, + op_id=conn.op_id, check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/asynchronous/cursor_base.py b/pymongo/asynchronous/cursor_base.py index 54d473a739..e8ac4c3139 100644 --- a/pymongo/asynchronous/cursor_base.py +++ b/pymongo/asynchronous/cursor_base.py @@ -136,7 +136,6 @@ async def _run_with_conn( more_to_come=more_to_come, unpack_res=self._unpack_response, cursor_id=operation.cursor_id, - op_id=conn.op_id, ) assert reply is not None if client._should_pin_cursor(operation.session) or operation.exhaust: # type: ignore[arg-type] diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index 51b9c99646..9c3cb06126 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -1797,6 +1797,7 @@ async def _checkout( yield session._pinned_connection 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 @@ -1935,8 +1936,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) - # Each getMore on a pinned connection is its own operation. - operation.conn_mgr.conn.op_id = _randint() + # 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 ) diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 2780aae427..426ffb624e 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -177,11 +177,12 @@ 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. Kept on the - # connection for simplicity: retries reuse it without threading an op_id - # kwarg through every command. Cleared on checkin; an op that runs a - # command directly on a pinned connection must clear it (see killCursors, - # reauth) so it isn't attributed to an unrelated command. + # 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: @@ -422,7 +423,6 @@ async def command( user_fields=user_fields, exhaust_allowed=exhaust_allowed, write_concern=write_concern, - op_id=self.op_id, ) except (OperationFailure, NotPrimaryError): raise @@ -1326,7 +1326,6 @@ 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) diff --git a/pymongo/synchronous/command_runner.py b/pymongo/synchronous/command_runner.py index 69686a0743..31fdd72ad8 100644 --- a/pymongo/synchronous/command_runner.py +++ b/pymongo/synchronous/command_runner.py @@ -284,7 +284,6 @@ 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``. @@ -305,7 +304,6 @@ 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( @@ -327,7 +325,7 @@ def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, - op_id=op_id, + op_id=conn.op_id, ) @@ -351,7 +349,6 @@ 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. @@ -380,7 +377,6 @@ 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)) @@ -433,7 +429,7 @@ def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=op_id, + op_id=conn.op_id, check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/synchronous/cursor_base.py b/pymongo/synchronous/cursor_base.py index acec222621..4cad4e0c09 100644 --- a/pymongo/synchronous/cursor_base.py +++ b/pymongo/synchronous/cursor_base.py @@ -136,7 +136,6 @@ def _run_with_conn( more_to_come=more_to_come, unpack_res=self._unpack_response, cursor_id=operation.cursor_id, - op_id=conn.op_id, ) assert reply is not None if client._should_pin_cursor(operation.session) or operation.exhaust: # type: ignore[arg-type] diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 764fdef1d3..22adf45eed 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -1794,6 +1794,7 @@ def _checkout( yield session._pinned_connection 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 @@ -1932,8 +1933,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) - # Each getMore on a pinned connection is its own operation. - operation.conn_mgr.conn.op_id = _randint() + # 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 ) diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 6ec1efb9b0..408a330c2f 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -177,11 +177,12 @@ 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. Kept on the - # connection for simplicity: retries reuse it without threading an op_id - # kwarg through every command. Cleared on checkin; an op that runs a - # command directly on a pinned connection must clear it (see killCursors, - # reauth) so it isn't attributed to an unrelated command. + # 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: @@ -422,7 +423,6 @@ def command( user_fields=user_fields, exhaust_allowed=exhaust_allowed, write_concern=write_concern, - op_id=self.op_id, ) except (OperationFailure, NotPrimaryError): raise @@ -1322,7 +1322,6 @@ 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) diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 27ca9f6ec4..9a57a483be 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -1,4 +1,4 @@ -# Copyright 2024-present MongoDB, Inc. +# 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. @@ -51,6 +51,9 @@ ("listIndexes", lambda c: _list_indexes(c)), ] +# Every command name the above operations issue, for the shared listener. +_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} + async def _agg(coll): cursor = await coll.aggregate([{"$match": {"x": 1}}]) @@ -63,28 +66,29 @@ async def _list_indexes(coll): class TestOperationIdRetry(AsyncIntegrationTest): - RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. + RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds. @async_client_context.require_failCommand_fail_point async def asyncSetUp(self) -> None: await super().asyncSetUp() - - async def _seed(self, coll): - await coll.drop() - await coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) - await coll.create_index("x") - - async def _check_stable_operation_id(self, command_name, action, retries): - """Force ``retries`` retries of ``action`` and assert every command - event for ``command_name`` shares one integer operation_id.""" - listener = AllowListEventListener(command_name) - client = await self.async_rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) - coll = client.pymongo_test.test_operation_id_retry - await self._seed(coll) - listener.reset() - + self.listener = AllowListEventListener(*_COMMANDS) + self.client = await self.async_rs_or_single_client( + event_listeners=[self.listener], appname=_APP_NAME + ) + self.coll = self.client.pymongo_test.test_operation_id_retry + + async def _seed(self): + await self.coll.drop() + await self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) + await self.coll.create_index("x") + + async def _run_under_failpoint(self, command_name, action, times, expected_error=None): + """Seed, force ``times`` closeConnection failures of ``command_name``, + run ``action``, and return its ``(started, failed, succeeded)`` events.""" + await self._seed() + self.listener.reset() fail_point = { - "mode": {"times": retries}, + "mode": {"times": times}, "data": { "failCommands": [command_name], "closeConnection": True, @@ -94,11 +98,24 @@ async def _check_stable_operation_id(self, command_name, action, retries): async with self.fail_point(fail_point): # A CSOT timeout lets a single operation retry more than once. with pymongo.timeout(60): - await action(coll) + if expected_error is not None: + with self.assertRaises(expected_error): + await action(self.coll) + else: + await action(self.coll) + + def of(events): + return [e for e in events if e.command_name == command_name] + + return ( + of(self.listener.started_events), + of(self.listener.failed_events), + of(self.listener.succeeded_events), + ) - started = listener.started_events - failed = listener.failed_events - succeeded = listener.succeeded_events + async def _check_stable_operation_id(self, command_name, action, retries): + """Assert every command event for ``command_name`` shares one integer operation_id.""" + started, failed, succeeded = await self._run_under_failpoint(command_name, action, retries) op_ids = [e.operation_id for e in started + failed + succeeded] self.assertEqual(len(started), retries + 1, "expected one started event per attempt") @@ -131,28 +148,11 @@ async def test_non_retryable_write_is_not_retried(self): ("delete", lambda c: c.delete_many({"x": 2})), ]: with self.subTest(command=command_name): - listener = AllowListEventListener(command_name) - client = await self.async_rs_or_single_client( - event_listeners=[listener], appname=_APP_NAME + started, _, _ = await self._run_under_failpoint( + command_name, action, times=1, expected_error=ConnectionFailure ) - coll = client.pymongo_test.test_operation_id_retry - await self._seed(coll) - listener.reset() - - fail_point = { - "mode": {"times": 1}, - "data": { - "failCommands": [command_name], - "closeConnection": True, - "appName": _APP_NAME, - }, - } - async with self.fail_point(fail_point): - with self.assertRaises(ConnectionFailure): - await action(coll) - - self.assertEqual(len(listener.started_events), 1, "must not retry") - self.assertIsInstance(listener.started_events[0].operation_id, int) + self.assertEqual(len(started), 1, "must not retry") + self.assertIsInstance(started[0].operation_id, int) if __name__ == "__main__": diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 58e3b533aa..534386aca6 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -1,4 +1,4 @@ -# Copyright 2024-present MongoDB, Inc. +# 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. @@ -51,6 +51,9 @@ ("listIndexes", lambda c: _list_indexes(c)), ] +# Every command name the above operations issue, for the shared listener. +_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} + def _agg(coll): cursor = coll.aggregate([{"$match": {"x": 1}}]) @@ -63,28 +66,27 @@ def _list_indexes(coll): class TestOperationIdRetry(IntegrationTest): - RETRIES = 5 # fail this many attempts; the (RETRIES + 1)th succeeds. + RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds. @client_context.require_failCommand_fail_point def setUp(self) -> None: super().setUp() - - def _seed(self, coll): - coll.drop() - coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) - coll.create_index("x") - - def _check_stable_operation_id(self, command_name, action, retries): - """Force ``retries`` retries of ``action`` and assert every command - event for ``command_name`` shares one integer operation_id.""" - listener = AllowListEventListener(command_name) - client = self.rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) - coll = client.pymongo_test.test_operation_id_retry - self._seed(coll) - listener.reset() - + self.listener = AllowListEventListener(*_COMMANDS) + self.client = self.rs_or_single_client(event_listeners=[self.listener], appname=_APP_NAME) + self.coll = self.client.pymongo_test.test_operation_id_retry + + def _seed(self): + self.coll.drop() + self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) + self.coll.create_index("x") + + def _run_under_failpoint(self, command_name, action, times, expected_error=None): + """Seed, force ``times`` closeConnection failures of ``command_name``, + run ``action``, and return its ``(started, failed, succeeded)`` events.""" + self._seed() + self.listener.reset() fail_point = { - "mode": {"times": retries}, + "mode": {"times": times}, "data": { "failCommands": [command_name], "closeConnection": True, @@ -94,11 +96,24 @@ def _check_stable_operation_id(self, command_name, action, retries): with self.fail_point(fail_point): # A CSOT timeout lets a single operation retry more than once. with pymongo.timeout(60): - action(coll) + if expected_error is not None: + with self.assertRaises(expected_error): + action(self.coll) + else: + action(self.coll) + + def of(events): + return [e for e in events if e.command_name == command_name] + + return ( + of(self.listener.started_events), + of(self.listener.failed_events), + of(self.listener.succeeded_events), + ) - started = listener.started_events - failed = listener.failed_events - succeeded = listener.succeeded_events + def _check_stable_operation_id(self, command_name, action, retries): + """Assert every command event for ``command_name`` shares one integer operation_id.""" + started, failed, succeeded = self._run_under_failpoint(command_name, action, retries) op_ids = [e.operation_id for e in started + failed + succeeded] self.assertEqual(len(started), retries + 1, "expected one started event per attempt") @@ -131,26 +146,11 @@ def test_non_retryable_write_is_not_retried(self): ("delete", lambda c: c.delete_many({"x": 2})), ]: with self.subTest(command=command_name): - listener = AllowListEventListener(command_name) - client = self.rs_or_single_client(event_listeners=[listener], appname=_APP_NAME) - coll = client.pymongo_test.test_operation_id_retry - self._seed(coll) - listener.reset() - - fail_point = { - "mode": {"times": 1}, - "data": { - "failCommands": [command_name], - "closeConnection": True, - "appName": _APP_NAME, - }, - } - with self.fail_point(fail_point): - with self.assertRaises(ConnectionFailure): - action(coll) - - self.assertEqual(len(listener.started_events), 1, "must not retry") - self.assertIsInstance(listener.started_events[0].operation_id, int) + started, _, _ = self._run_under_failpoint( + command_name, action, times=1, expected_error=ConnectionFailure + ) + self.assertEqual(len(started), 1, "must not retry") + self.assertIsInstance(started[0].operation_id, int) if __name__ == "__main__": From b36e20ba86cafe239c1277c0f03aa08d2f5fb41b Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 8 Jul 2026 20:31:33 +0000 Subject: [PATCH 04/14] made requested changes --- pymongo/asynchronous/helpers.py | 6 ++++-- pymongo/asynchronous/mongo_client.py | 6 ++++-- pymongo/synchronous/helpers.py | 6 ++++-- pymongo/synchronous/mongo_client.py | 6 ++++-- test/asynchronous/test_operation_id_retry.py | 3 --- test/test_operation_id_retry.py | 3 --- 6 files changed, 16 insertions(+), 14 deletions(-) diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py index be5ceef9bb..03215af785 100644 --- a/pymongo/asynchronous/helpers.py +++ b/pymongo/asynchronous/helpers.py @@ -71,8 +71,10 @@ async def inner(*args: Any, **kwargs: Any) -> Any: # Don't let reauth's auth commands inherit the in-flight op's id. prev_op_id = conn.op_id conn.op_id = None - await conn.authenticate(reauthenticate=True) - conn.op_id = prev_op_id + try: + await conn.authenticate(reauthenticate=True) + finally: + conn.op_id = prev_op_id else: raise return await func(*args, **kwargs) diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index 9c3cb06126..8cf6d4a05f 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -1793,8 +1793,10 @@ 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. diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py index efcb29fb52..877e2c2da5 100644 --- a/pymongo/synchronous/helpers.py +++ b/pymongo/synchronous/helpers.py @@ -71,8 +71,10 @@ def inner(*args: Any, **kwargs: Any) -> Any: # Don't let reauth's auth commands inherit the in-flight op's id. prev_op_id = conn.op_id conn.op_id = None - conn.authenticate(reauthenticate=True) - conn.op_id = prev_op_id + try: + conn.authenticate(reauthenticate=True) + finally: + conn.op_id = prev_op_id else: raise return func(*args, **kwargs) diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 22adf45eed..0ff845321d 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -1790,8 +1790,10 @@ 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. diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 9a57a483be..0fd91a2611 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -30,8 +30,6 @@ _APP_NAME = "operationIdRetryTest" -# Each operation, paired with the wire command it issues and an awaitable action. -# These are all retryable; a stable operation_id must span every retry attempt. _RETRYABLE_WRITES = [ ("insert", lambda c: c.insert_one({"_id": 100})), ("update", lambda c: c.update_one({"_id": 1}, {"$set": {"y": 1}})), @@ -51,7 +49,6 @@ ("listIndexes", lambda c: _list_indexes(c)), ] -# Every command name the above operations issue, for the shared listener. _COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 534386aca6..b73b4b24aa 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -30,8 +30,6 @@ _APP_NAME = "operationIdRetryTest" -# Each operation, paired with the wire command it issues and an awaitable action. -# These are all retryable; a stable operation_id must span every retry attempt. _RETRYABLE_WRITES = [ ("insert", lambda c: c.insert_one({"_id": 100})), ("update", lambda c: c.update_one({"_id": 1}, {"$set": {"y": 1}})), @@ -51,7 +49,6 @@ ("listIndexes", lambda c: _list_indexes(c)), ] -# Every command name the above operations issue, for the shared listener. _COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} From d8f29b9dc1d3372c196c7d78cecf7fdfcbb5dd81 Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Thu, 9 Jul 2026 15:56:13 +0000 Subject: [PATCH 05/14] implemented ContextVar approach for op_id --- pymongo/_op_id.py | 28 +++++++++++++++++++ pymongo/asynchronous/command_runner.py | 6 ++-- pymongo/asynchronous/helpers.py | 7 ++--- pymongo/asynchronous/mongo_client.py | 29 ++++++++++---------- pymongo/asynchronous/pool.py | 7 ----- pymongo/synchronous/command_runner.py | 6 ++-- pymongo/synchronous/helpers.py | 7 ++--- pymongo/synchronous/mongo_client.py | 29 ++++++++++---------- pymongo/synchronous/pool.py | 7 ----- test/asynchronous/test_operation_id_retry.py | 7 ++--- test/test_operation_id_retry.py | 7 ++--- 11 files changed, 74 insertions(+), 66 deletions(-) create mode 100644 pymongo/_op_id.py diff --git a/pymongo/_op_id.py b/pymongo/_op_id.py new file mode 100644 index 0000000000..dea80c92ea --- /dev/null +++ b/pymongo/_op_id.py @@ -0,0 +1,28 @@ +# 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 contextvars import ContextVar +from typing import Optional + +OP_ID: ContextVar[Optional[int]] = ContextVar("OP_ID", default=None) diff --git a/pymongo/asynchronous/command_runner.py b/pymongo/asynchronous/command_runner.py index 7900a6c958..f7858f1353 100644 --- a/pymongo/asynchronous/command_runner.py +++ b/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 @@ -325,7 +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, + op_id=_op_id.OP_ID.get(), ) @@ -429,7 +429,7 @@ async def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=conn.op_id, + op_id=_op_id.OP_ID.get(), check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py index 03215af785..984ac12301 100644 --- a/pymongo/asynchronous/helpers.py +++ b/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, @@ -69,12 +69,11 @@ async def inner(*args: Any, **kwargs: Any) -> Any: break if conn: # Don't let reauth's auth commands inherit the in-flight op's id. - prev_op_id = conn.op_id - conn.op_id = None + token = _op_id.OP_ID.set(None) try: await conn.authenticate(reauthenticate=True) finally: - conn.op_id = prev_op_id + _op_id.OP_ID.reset(token) else: raise return await func(*args, **kwargs) diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index 8cf6d4a05f..ec860e09b6 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/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.asynchronous import client_session, database, uri_parser from pymongo.asynchronous.change_stream import AsyncChangeStream, AsyncClusterChangeStream from pymongo.asynchronous.client_bulk import _AsyncClientBulk @@ -1793,13 +1793,10 @@ 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: - conn = session._pinned_connection - conn.op_id = None - err_handler.contribute_socket(conn) - yield conn + err_handler.contribute_socket(session._pinned_connection) + yield session._pinned_connection 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 @@ -1938,8 +1935,6 @@ 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 ) @@ -2227,8 +2222,6 @@ 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: @@ -3002,7 +2995,6 @@ 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 @@ -3023,7 +3015,12 @@ async def _write(self) -> T: commandName=self._operation, operationId=self._operation_id, ) - return await self._func(self._session, conn, self._retryable) # type: ignore + # Publish this op's id on every command of every attempt. + token = _op_id.OP_ID.set(self._operation_id) + try: + return await self._func(self._session, conn, self._retryable) # type: ignore + finally: + _op_id.OP_ID.reset(token) except PyMongoError as exc: if not self._retryable: raise @@ -3042,7 +3039,6 @@ 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: @@ -3053,7 +3049,12 @@ async def _read(self) -> T: commandName=self._operation, operationId=self._operation_id, ) - return await self._func(self._session, self._server, conn, read_pref) # type: ignore + # Publish this op's id on every command of every attempt. + token = _op_id.OP_ID.set(self._operation_id) + try: + return await self._func(self._session, self._server, conn, read_pref) # type: ignore + finally: + _op_id.OP_ID.reset(token) def _after_fork_child() -> None: diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index ab99eb150d..212f35f0d1 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -177,13 +177,6 @@ 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.""" diff --git a/pymongo/synchronous/command_runner.py b/pymongo/synchronous/command_runner.py index 31fdd72ad8..1e9dafbf08 100644 --- a/pymongo/synchronous/command_runner.py +++ b/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 @@ -325,7 +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, + op_id=_op_id.OP_ID.get(), ) @@ -429,7 +429,7 @@ def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=conn.op_id, + op_id=_op_id.OP_ID.get(), check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py index 877e2c2da5..8d0f402c3e 100644 --- a/pymongo/synchronous/helpers.py +++ b/pymongo/synchronous/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, @@ -69,12 +69,11 @@ def inner(*args: Any, **kwargs: Any) -> Any: break if conn: # Don't let reauth's auth commands inherit the in-flight op's id. - prev_op_id = conn.op_id - conn.op_id = None + token = _op_id.OP_ID.set(None) try: conn.authenticate(reauthenticate=True) finally: - conn.op_id = prev_op_id + _op_id.OP_ID.reset(token) else: raise return func(*args, **kwargs) diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 0ff845321d..56f128be44 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/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.client_options import ClientOptions from pymongo.driver_info import DriverInfo from pymongo.errors import ( @@ -1790,13 +1790,10 @@ 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: - conn = session._pinned_connection - conn.op_id = None - err_handler.contribute_socket(conn) - yield conn + err_handler.contribute_socket(session._pinned_connection) + yield session._pinned_connection 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 @@ -1935,8 +1932,6 @@ 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 ) @@ -2224,8 +2219,6 @@ 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: @@ -2993,7 +2986,6 @@ 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 @@ -3014,7 +3006,12 @@ def _write(self) -> T: commandName=self._operation, operationId=self._operation_id, ) - return self._func(self._session, conn, self._retryable) # type: ignore + # Publish this op's id on every command of every attempt. + token = _op_id.OP_ID.set(self._operation_id) + try: + return self._func(self._session, conn, self._retryable) # type: ignore + finally: + _op_id.OP_ID.reset(token) except PyMongoError as exc: if not self._retryable: raise @@ -3033,7 +3030,6 @@ 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: @@ -3044,7 +3040,12 @@ def _read(self) -> T: commandName=self._operation, operationId=self._operation_id, ) - return self._func(self._session, self._server, conn, read_pref) # type: ignore + # Publish this op's id on every command of every attempt. + token = _op_id.OP_ID.set(self._operation_id) + try: + return self._func(self._session, self._server, conn, read_pref) # type: ignore + finally: + _op_id.OP_ID.reset(token) def _after_fork_child() -> None: diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 84c7432079..f34a3341f3 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -177,13 +177,6 @@ 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.""" diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 0fd91a2611..3a6520c161 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -73,16 +73,13 @@ async def asyncSetUp(self) -> None: event_listeners=[self.listener], appname=_APP_NAME ) self.coll = self.client.pymongo_test.test_operation_id_retry - - async def _seed(self): await self.coll.drop() await self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) await self.coll.create_index("x") async def _run_under_failpoint(self, command_name, action, times, expected_error=None): - """Seed, force ``times`` closeConnection failures of ``command_name``, - run ``action``, and return its ``(started, failed, succeeded)`` events.""" - await self._seed() + """Force ``times`` closeConnection failures of ``command_name``, run + ``action``, and return its ``(started, failed, succeeded)`` events.""" self.listener.reset() fail_point = { "mode": {"times": times}, diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index b73b4b24aa..9bd336b8a1 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -71,16 +71,13 @@ def setUp(self) -> None: self.listener = AllowListEventListener(*_COMMANDS) self.client = self.rs_or_single_client(event_listeners=[self.listener], appname=_APP_NAME) self.coll = self.client.pymongo_test.test_operation_id_retry - - def _seed(self): self.coll.drop() self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) self.coll.create_index("x") def _run_under_failpoint(self, command_name, action, times, expected_error=None): - """Seed, force ``times`` closeConnection failures of ``command_name``, - run ``action``, and return its ``(started, failed, succeeded)`` events.""" - self._seed() + """Force ``times`` closeConnection failures of ``command_name``, run + ``action``, and return its ``(started, failed, succeeded)`` events.""" self.listener.reset() fail_point = { "mode": {"times": times}, From dc74bf1ccc7989be237fdd662117751ff0de6e7c Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Thu, 9 Jul 2026 19:01:53 +0000 Subject: [PATCH 06/14] made requested changes --- pymongo/_op_id.py | 24 +++++++++- pymongo/asynchronous/command_runner.py | 7 +-- pymongo/asynchronous/helpers.py | 5 +-- pymongo/asynchronous/mongo_client.py | 14 ++---- pymongo/periodic_executor.py | 5 ++- pymongo/synchronous/command_runner.py | 7 +-- pymongo/synchronous/helpers.py | 5 +-- pymongo/synchronous/mongo_client.py | 14 ++---- test/asynchronous/test_operation_id_retry.py | 47 ++++++++++---------- test/test_operation_id_retry.py | 47 ++++++++++---------- 10 files changed, 89 insertions(+), 86 deletions(-) diff --git a/pymongo/_op_id.py b/pymongo/_op_id.py index dea80c92ea..92d2d9a2de 100644 --- a/pymongo/_op_id.py +++ b/pymongo/_op_id.py @@ -22,7 +22,27 @@ from __future__ import annotations -from contextvars import ContextVar -from typing import Optional +from contextlib import AbstractContextManager +from contextvars import ContextVar, Token +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 + self._token: Optional[Token[Optional[int]]] = None + + 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: + if self._token: + OP_ID.reset(self._token) diff --git a/pymongo/asynchronous/command_runner.py b/pymongo/asynchronous/command_runner.py index f7858f1353..683dc0d8e8 100644 --- a/pymongo/asynchronous/command_runner.py +++ b/pymongo/asynchronous/command_runner.py @@ -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) @@ -325,7 +328,6 @@ async def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, - op_id=_op_id.OP_ID.get(), ) @@ -429,7 +431,6 @@ async def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=_op_id.OP_ID.get(), check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/asynchronous/helpers.py b/pymongo/asynchronous/helpers.py index 984ac12301..18b976cbe4 100644 --- a/pymongo/asynchronous/helpers.py +++ b/pymongo/asynchronous/helpers.py @@ -69,11 +69,8 @@ async def inner(*args: Any, **kwargs: Any) -> Any: break if conn: # Don't let reauth's auth commands inherit the in-flight op's id. - token = _op_id.OP_ID.set(None) - try: + with _op_id._OpIdContext(None): await conn.authenticate(reauthenticate=True) - finally: - _op_id.OP_ID.reset(token) else: raise return await func(*args, **kwargs) diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index ba0136a91c..80b1ab3499 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -3017,12 +3017,9 @@ async def _write(self) -> T: self._retryable = False if self._retrying: self._log_retry(is_write=True) - # Publish this op's id on every command of every attempt. - token = _op_id.OP_ID.set(self._operation_id) - try: + # 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 - finally: - _op_id.OP_ID.reset(token) except PyMongoError as exc: if not self._retryable: raise @@ -3045,12 +3042,9 @@ async def _read(self) -> T: self._check_last_error() if self._retrying: self._log_retry(is_write=False) - # Publish this op's id on every command of every attempt. - token = _op_id.OP_ID.set(self._operation_id) - try: + # 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 - finally: - _op_id.OP_ID.reset(token) def _after_fork_child() -> None: diff --git a/pymongo/periodic_executor.py b/pymongo/periodic_executor.py index 1e9022e721..4b979ca9f9 100644 --- a/pymongo/periodic_executor.py +++ b/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 diff --git a/pymongo/synchronous/command_runner.py b/pymongo/synchronous/command_runner.py index 1e9dafbf08..e81f514003 100644 --- a/pymongo/synchronous/command_runner.py +++ b/pymongo/synchronous/command_runner.py @@ -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) @@ -325,7 +328,6 @@ def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, - op_id=_op_id.OP_ID.get(), ) @@ -429,7 +431,6 @@ def run_command( codec_options=codec_options, user_fields=user_fields, orig=orig, - op_id=_op_id.OP_ID.get(), check=check, allowable_errors=allowable_errors, parse_write_concern_error=parse_write_concern_error, diff --git a/pymongo/synchronous/helpers.py b/pymongo/synchronous/helpers.py index 8d0f402c3e..89362645c0 100644 --- a/pymongo/synchronous/helpers.py +++ b/pymongo/synchronous/helpers.py @@ -69,11 +69,8 @@ def inner(*args: Any, **kwargs: Any) -> Any: break if conn: # Don't let reauth's auth commands inherit the in-flight op's id. - token = _op_id.OP_ID.set(None) - try: + with _op_id._OpIdContext(None): conn.authenticate(reauthenticate=True) - finally: - _op_id.OP_ID.reset(token) else: raise return func(*args, **kwargs) diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 0dcd947c22..81cb9b82c6 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -3008,12 +3008,9 @@ def _write(self) -> T: self._retryable = False if self._retrying: self._log_retry(is_write=True) - # Publish this op's id on every command of every attempt. - token = _op_id.OP_ID.set(self._operation_id) - try: + # One operation id across all attempts of this operation. + with _op_id._OpIdContext(self._operation_id): return self._func(self._session, conn, self._retryable) # type: ignore - finally: - _op_id.OP_ID.reset(token) except PyMongoError as exc: if not self._retryable: raise @@ -3036,12 +3033,9 @@ def _read(self) -> T: self._check_last_error() if self._retrying: self._log_retry(is_write=False) - # Publish this op's id on every command of every attempt. - token = _op_id.OP_ID.set(self._operation_id) - try: + # One operation id across all attempts of this operation. + with _op_id._OpIdContext(self._operation_id): return self._func(self._session, self._server, conn, read_pref) # type: ignore - finally: - _op_id.OP_ID.reset(token) def _after_fork_child() -> None: diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 3a6520c161..1ce07adfb3 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -49,8 +49,6 @@ ("listIndexes", lambda c: _list_indexes(c)), ] -_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} - async def _agg(coll): cursor = await coll.aggregate([{"$match": {"x": 1}}]) @@ -63,12 +61,14 @@ async def _list_indexes(coll): class TestOperationIdRetry(AsyncIntegrationTest): - RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds. + RETRIES = 2 @async_client_context.require_failCommand_fail_point async def asyncSetUp(self) -> None: await super().asyncSetUp() - self.listener = AllowListEventListener(*_COMMANDS) + self.listener = AllowListEventListener( + *{name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} + ) self.client = await self.async_rs_or_single_client( event_listeners=[self.listener], appname=_APP_NAME ) @@ -77,14 +77,13 @@ async def asyncSetUp(self) -> None: await self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) await self.coll.create_index("x") - async def _run_under_failpoint(self, command_name, action, times, expected_error=None): - """Force ``times`` closeConnection failures of ``command_name``, run - ``action``, and return its ``(started, failed, succeeded)`` events.""" + async def _run_under_failpoint(self, name, f, times, expected_error=None): + """Set a failpoint for the given command and return the corresponding events published during its execution.""" self.listener.reset() fail_point = { "mode": {"times": times}, "data": { - "failCommands": [command_name], + "failCommands": [name], "closeConnection": True, "appName": _APP_NAME, }, @@ -94,12 +93,12 @@ async def _run_under_failpoint(self, command_name, action, times, expected_error with pymongo.timeout(60): if expected_error is not None: with self.assertRaises(expected_error): - await action(self.coll) + await f(self.coll) else: - await action(self.coll) + await f(self.coll) def of(events): - return [e for e in events if e.command_name == command_name] + return [e for e in events if e.command_name == name] return ( of(self.listener.started_events), @@ -107,9 +106,9 @@ def of(events): of(self.listener.succeeded_events), ) - async def _check_stable_operation_id(self, command_name, action, retries): - """Assert every command event for ``command_name`` shares one integer operation_id.""" - started, failed, succeeded = await self._run_under_failpoint(command_name, action, retries) + async def _check_stable_operation_id(self, name, f, retries): + """Assert every command event for ``name`` shares one integer operation_id.""" + started, failed, succeeded = await self._run_under_failpoint(name, f, retries) op_ids = [e.operation_id for e in started + failed + succeeded] self.assertEqual(len(started), retries + 1, "expected one started event per attempt") @@ -119,31 +118,31 @@ async def _check_stable_operation_id(self, command_name, action, retries): self.assertEqual( len(set(op_ids)), 1, - f"operation_id not stable across retries for {command_name}: {op_ids}", + f"operation_id not stable across retries for {name}: {op_ids}", ) @async_client_context.require_no_standalone async def test_retryable_writes_reuse_operation_id(self): - for command_name, action in _RETRYABLE_WRITES: - with self.subTest(command=command_name): - await self._check_stable_operation_id(command_name, action, self.RETRIES) + for name, f in _RETRYABLE_WRITES: + with self.subTest(command=name): + await self._check_stable_operation_id(name, f, self.RETRIES) async def test_retryable_reads_reuse_operation_id(self): - for command_name, action in _RETRYABLE_READS: - with self.subTest(command=command_name): - await self._check_stable_operation_id(command_name, action, self.RETRIES) + for name, f in _RETRYABLE_READS: + with self.subTest(command=name): + await self._check_stable_operation_id(name, f, self.RETRIES) @async_client_context.require_no_standalone async def test_non_retryable_write_is_not_retried(self): # Multi-document writes are not retryable: a single network error must # surface immediately, with exactly one attempt. - for command_name, action in [ + for name, f in [ ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), ("delete", lambda c: c.delete_many({"x": 2})), ]: - with self.subTest(command=command_name): + with self.subTest(command=name): started, _, _ = await self._run_under_failpoint( - command_name, action, times=1, expected_error=ConnectionFailure + name, f, times=1, expected_error=ConnectionFailure ) self.assertEqual(len(started), 1, "must not retry") self.assertIsInstance(started[0].operation_id, int) diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 9bd336b8a1..0713e73514 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -49,8 +49,6 @@ ("listIndexes", lambda c: _list_indexes(c)), ] -_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} - def _agg(coll): cursor = coll.aggregate([{"$match": {"x": 1}}]) @@ -63,26 +61,27 @@ def _list_indexes(coll): class TestOperationIdRetry(IntegrationTest): - RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds. + RETRIES = 2 @client_context.require_failCommand_fail_point def setUp(self) -> None: super().setUp() - self.listener = AllowListEventListener(*_COMMANDS) + self.listener = AllowListEventListener( + *{name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS} + ) self.client = self.rs_or_single_client(event_listeners=[self.listener], appname=_APP_NAME) self.coll = self.client.pymongo_test.test_operation_id_retry self.coll.drop() self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) self.coll.create_index("x") - def _run_under_failpoint(self, command_name, action, times, expected_error=None): - """Force ``times`` closeConnection failures of ``command_name``, run - ``action``, and return its ``(started, failed, succeeded)`` events.""" + def _run_under_failpoint(self, name, f, times, expected_error=None): + """Set a failpoint for the given command and return the corresponding events published during its execution.""" self.listener.reset() fail_point = { "mode": {"times": times}, "data": { - "failCommands": [command_name], + "failCommands": [name], "closeConnection": True, "appName": _APP_NAME, }, @@ -92,12 +91,12 @@ def _run_under_failpoint(self, command_name, action, times, expected_error=None) with pymongo.timeout(60): if expected_error is not None: with self.assertRaises(expected_error): - action(self.coll) + f(self.coll) else: - action(self.coll) + f(self.coll) def of(events): - return [e for e in events if e.command_name == command_name] + return [e for e in events if e.command_name == name] return ( of(self.listener.started_events), @@ -105,9 +104,9 @@ def of(events): of(self.listener.succeeded_events), ) - def _check_stable_operation_id(self, command_name, action, retries): - """Assert every command event for ``command_name`` shares one integer operation_id.""" - started, failed, succeeded = self._run_under_failpoint(command_name, action, retries) + def _check_stable_operation_id(self, name, f, retries): + """Assert every command event for ``name`` shares one integer operation_id.""" + started, failed, succeeded = self._run_under_failpoint(name, f, retries) op_ids = [e.operation_id for e in started + failed + succeeded] self.assertEqual(len(started), retries + 1, "expected one started event per attempt") @@ -117,31 +116,31 @@ def _check_stable_operation_id(self, command_name, action, retries): self.assertEqual( len(set(op_ids)), 1, - f"operation_id not stable across retries for {command_name}: {op_ids}", + f"operation_id not stable across retries for {name}: {op_ids}", ) @client_context.require_no_standalone def test_retryable_writes_reuse_operation_id(self): - for command_name, action in _RETRYABLE_WRITES: - with self.subTest(command=command_name): - self._check_stable_operation_id(command_name, action, self.RETRIES) + for name, f in _RETRYABLE_WRITES: + with self.subTest(command=name): + self._check_stable_operation_id(name, f, self.RETRIES) def test_retryable_reads_reuse_operation_id(self): - for command_name, action in _RETRYABLE_READS: - with self.subTest(command=command_name): - self._check_stable_operation_id(command_name, action, self.RETRIES) + for name, f in _RETRYABLE_READS: + with self.subTest(command=name): + self._check_stable_operation_id(name, f, self.RETRIES) @client_context.require_no_standalone def test_non_retryable_write_is_not_retried(self): # Multi-document writes are not retryable: a single network error must # surface immediately, with exactly one attempt. - for command_name, action in [ + for name, f in [ ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), ("delete", lambda c: c.delete_many({"x": 2})), ]: - with self.subTest(command=command_name): + with self.subTest(command=name): started, _, _ = self._run_under_failpoint( - command_name, action, times=1, expected_error=ConnectionFailure + name, f, times=1, expected_error=ConnectionFailure ) self.assertEqual(len(started), 1, "must not retry") self.assertIsInstance(started[0].operation_id, int) From 64fff02f51187b75fafb652b01e2b62c8ce56233 Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Fri, 10 Jul 2026 18:41:37 +0000 Subject: [PATCH 07/14] fixed test coverage --- .../test_async_contextvars_reset.py | 2 +- test/asynchronous/test_operation_id_retry.py | 53 +++++++++++-------- test/test_operation_id_retry.py | 53 +++++++++++-------- 3 files changed, 65 insertions(+), 43 deletions(-) diff --git a/test/asynchronous/test_async_contextvars_reset.py b/test/asynchronous/test_async_contextvars_reset.py index c2a3cd1959..5d9cf69bb8 100644 --- a/test/asynchronous/test_async_contextvars_reset.py +++ b/test/asynchronous/test_async_contextvars_reset.py @@ -37,7 +37,7 @@ async def test_context_vars_are_reset_in_executor(self): for context in [ c for c in server._monitor._executor._task.get_context() - if c.name in ["TIMEOUT", "RTT", "DEADLINE"] + if c.name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"] ]: self.assertIn(context.get(), [None, float("inf"), 0.0]) await self.client.db.test.delete_many({}) diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 1ce07adfb3..7372f71a36 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -21,7 +21,11 @@ sys.path[0:0] = [""] import pymongo -from pymongo.errors import ConnectionFailure +from pymongo import _op_id +from pymongo.asynchronous.helpers import _handle_reauth +from pymongo.asynchronous.pool import AsyncConnection +from pymongo.errors import OperationFailure +from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE from pymongo.operations import InsertOne from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest from test.utils_shared import AllowListEventListener @@ -77,7 +81,7 @@ async def asyncSetUp(self) -> None: await self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) await self.coll.create_index("x") - async def _run_under_failpoint(self, name, f, times, expected_error=None): + async def _run_under_failpoint(self, name, f, times): """Set a failpoint for the given command and return the corresponding events published during its execution.""" self.listener.reset() fail_point = { @@ -91,11 +95,7 @@ async def _run_under_failpoint(self, name, f, times, expected_error=None): async with self.fail_point(fail_point): # A CSOT timeout lets a single operation retry more than once. with pymongo.timeout(60): - if expected_error is not None: - with self.assertRaises(expected_error): - await f(self.coll) - else: - await f(self.coll) + await f(self.coll) def of(events): return [e for e in events if e.command_name == name] @@ -132,20 +132,31 @@ async def test_retryable_reads_reuse_operation_id(self): with self.subTest(command=name): await self._check_stable_operation_id(name, f, self.RETRIES) - @async_client_context.require_no_standalone - async def test_non_retryable_write_is_not_retried(self): - # Multi-document writes are not retryable: a single network error must - # surface immediately, with exactly one attempt. - for name, f in [ - ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), - ("delete", lambda c: c.delete_many({"x": 2})), - ]: - with self.subTest(command=name): - started, _, _ = await self._run_under_failpoint( - name, f, times=1, expected_error=ConnectionFailure - ) - self.assertEqual(len(started), 1, "must not retry") - self.assertIsInstance(started[0].operation_id, int) + async def test_reauth_does_not_reuse_operation_id(self): + class FakeConnection(AsyncConnection): + def __init__(self): + self.auth_op_ids = [] + + async def authenticate(self, reauthenticate=False): + self.auth_op_ids.append(_op_id.OP_ID.get()) + + conn = FakeConnection() + attempt_op_ids = [] + + @_handle_reauth + async def func(conn): + attempt_op_ids.append(_op_id.OP_ID.get()) + if len(attempt_op_ids) == 1: + raise OperationFailure("reauth required", _REAUTHENTICATION_REQUIRED_CODE) + + op_id = 42 + with _op_id._OpIdContext(op_id): + await func(conn) + # Reauth's auth commands must not inherit the in-flight op's id. + self.assertEqual(conn.auth_op_ids, [None]) + # The op's id is restored for the retried command after reauth. + self.assertEqual(attempt_op_ids, [op_id, op_id]) + self.assertIsNone(_op_id.OP_ID.get()) if __name__ == "__main__": diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 0713e73514..cbf952cdfa 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -21,8 +21,12 @@ sys.path[0:0] = [""] import pymongo -from pymongo.errors import ConnectionFailure +from pymongo import _op_id +from pymongo.errors import OperationFailure +from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE from pymongo.operations import InsertOne +from pymongo.synchronous.helpers import _handle_reauth +from pymongo.synchronous.pool import Connection from test import IntegrationTest, client_context, unittest from test.utils_shared import AllowListEventListener @@ -75,7 +79,7 @@ def setUp(self) -> None: self.coll.insert_many([{"_id": i, "x": i % 3} for i in range(5)]) self.coll.create_index("x") - def _run_under_failpoint(self, name, f, times, expected_error=None): + def _run_under_failpoint(self, name, f, times): """Set a failpoint for the given command and return the corresponding events published during its execution.""" self.listener.reset() fail_point = { @@ -89,11 +93,7 @@ def _run_under_failpoint(self, name, f, times, expected_error=None): with self.fail_point(fail_point): # A CSOT timeout lets a single operation retry more than once. with pymongo.timeout(60): - if expected_error is not None: - with self.assertRaises(expected_error): - f(self.coll) - else: - f(self.coll) + f(self.coll) def of(events): return [e for e in events if e.command_name == name] @@ -130,20 +130,31 @@ def test_retryable_reads_reuse_operation_id(self): with self.subTest(command=name): self._check_stable_operation_id(name, f, self.RETRIES) - @client_context.require_no_standalone - def test_non_retryable_write_is_not_retried(self): - # Multi-document writes are not retryable: a single network error must - # surface immediately, with exactly one attempt. - for name, f in [ - ("update", lambda c: c.update_many({"x": 1}, {"$set": {"z": 1}})), - ("delete", lambda c: c.delete_many({"x": 2})), - ]: - with self.subTest(command=name): - started, _, _ = self._run_under_failpoint( - name, f, times=1, expected_error=ConnectionFailure - ) - self.assertEqual(len(started), 1, "must not retry") - self.assertIsInstance(started[0].operation_id, int) + def test_reauth_does_not_reuse_operation_id(self): + class FakeConnection(Connection): + def __init__(self): + self.auth_op_ids = [] + + def authenticate(self, reauthenticate=False): + self.auth_op_ids.append(_op_id.OP_ID.get()) + + conn = FakeConnection() + attempt_op_ids = [] + + @_handle_reauth + def func(conn): + attempt_op_ids.append(_op_id.OP_ID.get()) + if len(attempt_op_ids) == 1: + raise OperationFailure("reauth required", _REAUTHENTICATION_REQUIRED_CODE) + + op_id = 42 + with _op_id._OpIdContext(op_id): + func(conn) + # Reauth's auth commands must not inherit the in-flight op's id. + self.assertEqual(conn.auth_op_ids, [None]) + # The op's id is restored for the retried command after reauth. + self.assertEqual(attempt_op_ids, [op_id, op_id]) + self.assertIsNone(_op_id.OP_ID.get()) if __name__ == "__main__": From 816ecf932ce0c86028a19ca2fa99170f298672ce Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Tue, 14 Jul 2026 20:49:46 +0000 Subject: [PATCH 08/14] simplified _OpIdContext token handling --- pymongo/_op_id.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pymongo/_op_id.py b/pymongo/_op_id.py index 92d2d9a2de..d0de439789 100644 --- a/pymongo/_op_id.py +++ b/pymongo/_op_id.py @@ -23,7 +23,7 @@ from __future__ import annotations from contextlib import AbstractContextManager -from contextvars import ContextVar, Token +from contextvars import ContextVar from typing import Any, Optional OP_ID: ContextVar[Optional[int]] = ContextVar("OP_ID", default=None) @@ -38,11 +38,9 @@ class _OpIdContext(AbstractContextManager[Any]): def __init__(self, op_id: Optional[int]): self._op_id = op_id - self._token: Optional[Token[Optional[int]]] = None 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: - if self._token: - OP_ID.reset(self._token) + OP_ID.reset(self._token) From e430689dcec4bcb7aff7a4995df9b29dc113c30f Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 15 Jul 2026 20:12:03 +0000 Subject: [PATCH 09/14] updated changelog and test --- doc/changelog.rst | 5 +++++ test/asynchronous/test_operation_id_retry.py | 8 ++++---- test/test_operation_id_retry.py | 8 ++++---- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index b3e1c75319..8957ab2b9e 100644 --- a/doc/changelog.rst +++ b/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) -------------------------------------- diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 7372f71a36..410bb2a0ab 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -97,13 +97,13 @@ async def _run_under_failpoint(self, name, f, times): with pymongo.timeout(60): await f(self.coll) - def of(events): + def matching(events): return [e for e in events if e.command_name == name] return ( - of(self.listener.started_events), - of(self.listener.failed_events), - of(self.listener.succeeded_events), + matching(self.listener.started_events), + matching(self.listener.failed_events), + matching(self.listener.succeeded_events), ) async def _check_stable_operation_id(self, name, f, retries): diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index cbf952cdfa..e4d9ec875a 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -95,13 +95,13 @@ def _run_under_failpoint(self, name, f, times): with pymongo.timeout(60): f(self.coll) - def of(events): + def matching(events): return [e for e in events if e.command_name == name] return ( - of(self.listener.started_events), - of(self.listener.failed_events), - of(self.listener.succeeded_events), + matching(self.listener.started_events), + matching(self.listener.failed_events), + matching(self.listener.succeeded_events), ) def _check_stable_operation_id(self, name, f, retries): From c3fcca5674a24238e00554ec2078501937c9071d Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 15 Jul 2026 21:41:45 +0000 Subject: [PATCH 10/14] fixed contextvars reset test to read executor task context --- test/asynchronous/test_async_contextvars_reset.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/test/asynchronous/test_async_contextvars_reset.py b/test/asynchronous/test_async_contextvars_reset.py index 5d9cf69bb8..d9a6e877ca 100644 --- a/test/asynchronous/test_async_contextvars_reset.py +++ b/test/asynchronous/test_async_contextvars_reset.py @@ -34,10 +34,8 @@ async def test_context_vars_are_reset_in_executor(self): await self.client.db.test.insert_one({"x": 1}) for server in self.client._topology._servers.values(): - for context in [ - c - for c in server._monitor._executor._task.get_context() - if c.name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"] - ]: - self.assertIn(context.get(), [None, float("inf"), 0.0]) + ctx = server._monitor._executor._task.get_context() + for var in ctx: + if var.name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"]: + self.assertIn(ctx[var], [None, float("inf"), 0.0]) await self.client.db.test.delete_many({}) From 2a5581e8250baa630277f041525922ce25f5813f Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 15 Jul 2026 21:55:47 +0000 Subject: [PATCH 11/14] addressed subTest ambiguity and OP_ID reset assertion --- test/asynchronous/test_async_contextvars_reset.py | 5 ++++- test/asynchronous/test_operation_id_retry.py | 8 ++++---- test/test_operation_id_retry.py | 8 ++++---- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/test/asynchronous/test_async_contextvars_reset.py b/test/asynchronous/test_async_contextvars_reset.py index d9a6e877ca..d831c37630 100644 --- a/test/asynchronous/test_async_contextvars_reset.py +++ b/test/asynchronous/test_async_contextvars_reset.py @@ -36,6 +36,9 @@ async def test_context_vars_are_reset_in_executor(self): for server in self.client._topology._servers.values(): ctx = server._monitor._executor._task.get_context() for var in ctx: - if var.name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"]: + if var.name == "OP_ID": + # OP_ID resets to None; 0 is a valid op_id so don't accept it. + self.assertIsNone(ctx[var]) + elif var.name in ["TIMEOUT", "RTT", "DEADLINE"]: self.assertIn(ctx[var], [None, float("inf"), 0.0]) await self.client.db.test.delete_many({}) diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 410bb2a0ab..9e4dd46346 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -123,13 +123,13 @@ async def _check_stable_operation_id(self, name, f, retries): @async_client_context.require_no_standalone async def test_retryable_writes_reuse_operation_id(self): - for name, f in _RETRYABLE_WRITES: - with self.subTest(command=name): + for i, (name, f) in enumerate(_RETRYABLE_WRITES): + with self.subTest(command=name, index=i): await self._check_stable_operation_id(name, f, self.RETRIES) async def test_retryable_reads_reuse_operation_id(self): - for name, f in _RETRYABLE_READS: - with self.subTest(command=name): + for i, (name, f) in enumerate(_RETRYABLE_READS): + with self.subTest(command=name, index=i): await self._check_stable_operation_id(name, f, self.RETRIES) async def test_reauth_does_not_reuse_operation_id(self): diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index e4d9ec875a..d79a29df10 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -121,13 +121,13 @@ def _check_stable_operation_id(self, name, f, retries): @client_context.require_no_standalone def test_retryable_writes_reuse_operation_id(self): - for name, f in _RETRYABLE_WRITES: - with self.subTest(command=name): + for i, (name, f) in enumerate(_RETRYABLE_WRITES): + with self.subTest(command=name, index=i): self._check_stable_operation_id(name, f, self.RETRIES) def test_retryable_reads_reuse_operation_id(self): - for name, f in _RETRYABLE_READS: - with self.subTest(command=name): + for i, (name, f) in enumerate(_RETRYABLE_READS): + with self.subTest(command=name, index=i): self._check_stable_operation_id(name, f, self.RETRIES) def test_reauth_does_not_reuse_operation_id(self): From aa5fe73ac824b33a9e87790debf7fff0c210ca96 Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 15 Jul 2026 22:17:52 +0000 Subject: [PATCH 12/14] assert contextvars are present and reset in executor --- test/asynchronous/test_async_contextvars_reset.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/test/asynchronous/test_async_contextvars_reset.py b/test/asynchronous/test_async_contextvars_reset.py index d831c37630..3fa03650a9 100644 --- a/test/asynchronous/test_async_contextvars_reset.py +++ b/test/asynchronous/test_async_contextvars_reset.py @@ -35,10 +35,9 @@ async def test_context_vars_are_reset_in_executor(self): await self.client.db.test.insert_one({"x": 1}) for server in self.client._topology._servers.values(): ctx = server._monitor._executor._task.get_context() - for var in ctx: - if var.name == "OP_ID": - # OP_ID resets to None; 0 is a valid op_id so don't accept it. - self.assertIsNone(ctx[var]) - elif var.name in ["TIMEOUT", "RTT", "DEADLINE"]: - self.assertIn(ctx[var], [None, float("inf"), 0.0]) + values = {var.name: ctx[var] for var in ctx} + for name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"]: + # The executor resets these on startup, so each must be present. + self.assertIn(name, values) + self.assertIn(values[name], [None, float("inf"), 0.0]) await self.client.db.test.delete_many({}) From d9161e93f10fa5bb4f487b50691e7d345c08113f Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Wed, 15 Jul 2026 22:26:41 +0000 Subject: [PATCH 13/14] check contextvars against exact reset values --- test/asynchronous/test_async_contextvars_reset.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/asynchronous/test_async_contextvars_reset.py b/test/asynchronous/test_async_contextvars_reset.py index 3fa03650a9..95ef3fed84 100644 --- a/test/asynchronous/test_async_contextvars_reset.py +++ b/test/asynchronous/test_async_contextvars_reset.py @@ -33,11 +33,13 @@ async def test_context_vars_are_reset_in_executor(self): self.skipTest("Test requires asyncio.Task.get_context (added in Python 3.12)") await self.client.db.test.insert_one({"x": 1}) + # Value each contextvar is reset to at the start of the executor task. + expected = {"TIMEOUT": None, "RTT": 0.0, "DEADLINE": float("inf"), "OP_ID": None} for server in self.client._topology._servers.values(): ctx = server._monitor._executor._task.get_context() values = {var.name: ctx[var] for var in ctx} - for name in ["TIMEOUT", "RTT", "DEADLINE", "OP_ID"]: + for name, value in expected.items(): # The executor resets these on startup, so each must be present. self.assertIn(name, values) - self.assertIn(values[name], [None, float("inf"), 0.0]) + self.assertEqual(values[name], value) await self.client.db.test.delete_many({}) From 22905089bbbc527cc6ab38211ffb8b88fe718489 Mon Sep 17 00:00:00 2001 From: Shrey Varma Date: Thu, 16 Jul 2026 00:58:11 +0000 Subject: [PATCH 14/14] added op id guards for auto encryption --- pymongo/asynchronous/encryption.py | 10 ++++--- pymongo/synchronous/encryption.py | 10 ++++--- test/asynchronous/test_operation_id_retry.py | 28 ++++++++++++++++++++ test/test_operation_id_retry.py | 28 ++++++++++++++++++++ 4 files changed, 70 insertions(+), 6 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index c7540265f9..524ae45c11 100644 --- a/pymongo/asynchronous/encryption.py +++ b/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: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 8b9a7e3c27..014d162e2b 100644 --- a/pymongo/synchronous/encryption.py +++ b/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: diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index 9e4dd46346..91e3cd9011 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -21,7 +21,9 @@ sys.path[0:0] = [""] import pymongo +from bson.codec_options import DEFAULT_CODEC_OPTIONS from pymongo import _op_id +from pymongo.asynchronous.encryption import _Encrypter from pymongo.asynchronous.helpers import _handle_reauth from pymongo.asynchronous.pool import AsyncConnection from pymongo.errors import OperationFailure @@ -158,6 +160,32 @@ async def func(conn): self.assertEqual(attempt_op_ids, [op_id, op_id]) self.assertIsNone(_op_id.OP_ID.get()) + async def test_auto_encryption_does_not_reuse_operation_id(self): + class FakeAutoEncrypter: + def __init__(self): + self.op_ids = [] + + async def encrypt(self, database, cmd): + self.op_ids.append(_op_id.OP_ID.get()) + return cmd + + async def decrypt(self, response): + self.op_ids.append(_op_id.OP_ID.get()) + return response + + encrypter = _Encrypter.__new__(_Encrypter) + encrypter._closed = False + encrypter._auto_encrypter = FakeAutoEncrypter() + + op_id = 42 + with _op_id._OpIdContext(op_id): + await encrypter.encrypt("db", {"find": "test"}, DEFAULT_CODEC_OPTIONS) + await encrypter.decrypt(b"") + # The op's id is restored for the operation's own command. + self.assertEqual(_op_id.OP_ID.get(), op_id) + # Encryption's sub-commands must not inherit the in-flight op's id. + self.assertEqual(encrypter._auto_encrypter.op_ids, [None, None]) + if __name__ == "__main__": unittest.main() diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index d79a29df10..7224d0626a 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -21,10 +21,12 @@ sys.path[0:0] = [""] import pymongo +from bson.codec_options import DEFAULT_CODEC_OPTIONS from pymongo import _op_id from pymongo.errors import OperationFailure from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE from pymongo.operations import InsertOne +from pymongo.synchronous.encryption import _Encrypter from pymongo.synchronous.helpers import _handle_reauth from pymongo.synchronous.pool import Connection from test import IntegrationTest, client_context, unittest @@ -156,6 +158,32 @@ def func(conn): self.assertEqual(attempt_op_ids, [op_id, op_id]) self.assertIsNone(_op_id.OP_ID.get()) + def test_auto_encryption_does_not_reuse_operation_id(self): + class FakeAutoEncrypter: + def __init__(self): + self.op_ids = [] + + def encrypt(self, database, cmd): + self.op_ids.append(_op_id.OP_ID.get()) + return cmd + + def decrypt(self, response): + self.op_ids.append(_op_id.OP_ID.get()) + return response + + encrypter = _Encrypter.__new__(_Encrypter) + encrypter._closed = False + encrypter._auto_encrypter = FakeAutoEncrypter() + + op_id = 42 + with _op_id._OpIdContext(op_id): + encrypter.encrypt("db", {"find": "test"}, DEFAULT_CODEC_OPTIONS) + encrypter.decrypt(b"") + # The op's id is restored for the operation's own command. + self.assertEqual(_op_id.OP_ID.get(), op_id) + # Encryption's sub-commands must not inherit the in-flight op's id. + self.assertEqual(encrypter._auto_encrypter.op_ids, [None, None]) + if __name__ == "__main__": unittest.main()