PYTHON-5903 Retry attempts should share a single stable operation_id by shrey-mongo · Pull Request #2901 · mongodb/mongo-python-driver

shrey-mongo · 2026-06-29T15:26:07Z

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

Changes in this PR

Command-monitoring events carry an operation_id to correlate the events of one logical operation. For single-document CRUD and cursor operations, each retry attempt got a different operation_id (it defaulted to the per-attempt request_id), so consumers couldn't group a retried operation's events. Bulk writes already shared one across batches.

_ClientConnectionRetryable mints the operation_id once per logical operation.
New internal pymongo/_op_id.py module holds an OP_ID ContextVar and an _OpIdContext context manager. _write/_read enter it around each attempt's command execution, so every attempt publishes the same id; on exit the previous value is restored.
command_runner._run_command falls back to OP_ID.get() when no explicit op_id is passed (bulk write pass their own), and threads it to the publish_command_* events and command logs.
Commands outside an operation's scope (handshake/auth during connection establishment, killCursors, pinned-cursor getMores, background tasks) read the default None and keep falling back to request_id.
_handle_reauth suppresses the in-flight id (_OpIdContext(None)) while re-authenticating, then restores it for the retried command.
AsyncPeriodicExecutor._run() resets OP_ID alongside _csot.reset_all() so executor tasks don't inherit an operation's id from their spawning context.

Test Plan

New test/asynchronous/test_operation_id_retry.py forces multiple retries via a failCommand failpoint across a read/write matrix and asserts all command events share one operation_id; verifies non-retryable multi-writes don't retry.
also unit-tests the reauth guard (id suppressed during auth, restored for the retry).
Extended test_async_contextvars_reset.py to assert OP_ID is reset inside executor tasks, alongside the CSOT vars.
Ran the new tests plus the existing command-monitoring and retryable reads/writes suites against a single-node replica set which passed.

Checklist

Checklist for Author

Did you update the changelog (if necessary)?
Is there test coverage?
Is any followup work tracked in a JIRA ticket? If so, add link(s).

Checklist for Reviewer

Does the title of the PR reference a JIRA Ticket?
Do you fully understand the implementation? (Would you be comfortable explaining how this code works to someone else?)
Is all relevant documentation (README or docstring) updated?

NoahStapp · 2026-06-29T16:46:08Z

Please fill out the pull request template completely.

Copilot

Pull request overview

This PR ensures command-monitoring operation_id remains stable across retry attempts for retryable reads/writes and relevant cursor operations, improving APM correlation for a single logical operation.

Changes:

Mint a single per-operation operation_id in the retryable-operation wrapper and stamp it onto the checked-out connection for each attempt.
Thread the connection’s op_id through pool/server command execution into command_runner so published command-monitoring events use it.
Add sync + async integration tests that force multiple retries via failCommand and assert all related command events share one operation_id.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
`tools/synchro.py`	Registers the new test file in the synchro conversion list.
`test/test_operation_id_retry.py`	Sync integration coverage asserting `operation_id` stability across retries (and no retry for non-retryable multi-writes).
`test/asynchronous/test_operation_id_retry.py`	Async equivalent of the retry `operation_id` stability tests.
`pymongo/synchronous/server.py`	Passes `conn.op_id` into cursor command execution for APM publishing.
`pymongo/synchronous/pool.py`	Adds `Connection.op_id`, forwards it into `run_command`, and clears it on check-in.
`pymongo/synchronous/mongo_client.py`	Mints/stabilizes per-operation ids across retries and stamps them on checked-out connections.
`pymongo/synchronous/command_runner.py`	Accepts/threads `op_id` through `run_command` / `run_cursor_command` into `_run_command`.
`pymongo/asynchronous/server.py`	Async counterpart: passes `conn.op_id` for cursor command execution.
`pymongo/asynchronous/pool.py`	Async counterpart: adds/forwards/clears `op_id` on connections.
`pymongo/asynchronous/mongo_client.py`	Async counterpart: stable per-operation id and per-attempt stamping onto connections.
`pymongo/asynchronous/command_runner.py`	Async counterpart: threads `op_id` through command execution to APM event publishing.

shrey-mongo · 2026-06-29T17:12:33Z

Updated description to include full template

codecov-commenter · 2026-06-29T18:21:09Z

Codecov Report

❌ Patch coverage is 97.72727% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pymongo/_op_id.py	93.75%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

NoahStapp · 2026-06-30T15:05:07Z

There are merge conflicts, please resolve.

shrey-mongo · 2026-06-30T16:38:12Z

I fixed the merge conflict and the typing error that was causing the workflow to fail

NoahStapp · 2026-07-01T14:31:41Z

            func=func,
        )

-    def require_failCommand_appName(self, func):


What's the purpose of this change?

It was needed to fix an untyped-decorator mypy error that was occurring in my test. This matches the signature of require_failCommand_fail_point.

You don't need both the require_failCommand_fail_point and require_failCommand_appName annotations, the latter is a strict subset of the former.

NoahStapp · 2026-07-08T17:26:49Z

            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()


This will assign a new op_id to every call here even though nothing retryable goes through this path.

Good point, I changed it to conn.op_id = None so each get More falls back to its own request_id.

… into shrey-mongo/PYTHON-5903

NoahStapp · 2026-07-08T19:41:38Z

                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.


This check should be moved to before the if in_txn block to prevent a pinned connection from re-using an old op_id.

Added reset on this branch

NoahStapp

What do you think about using a ContextVar for the operationId instead of storing it on the connection? That would remove the current friction between a per-operation object being stored on a cross-operation object and simplify the cleanup significantly.

We already use this pattern in _csot.py and elsewhere for exactly the reason it might be useful here: to avoid threading through the entire call stack while still centralizing cleanup.

Something like this in _ClientConnectionRetryable.run:

token = _op_id.OP_ID.set(self._operation_id)
try:
   return await self._run()
finally:
   _op_id.OP_ID.reset(token)

And then replacing the current _handle_reauth block with this:

token = _op_id.OP_ID.set(None)
try:
   await conn.authenticate(reauthenticate=True)
finally:
   _op_id.OP_ID.reset(token)

The other cleanups could then be removed since they don't interact with _ClientConnectionRetryable.run() at all.

shrey-mongo · 2026-07-09T15:55:23Z

Implemented the ContextVar approach which helped with the cleanup, pool.py, cursor_base.py, and server.py match master and the explicit clears were removed for checkout, killCursors, and pinned getMore.

I set the token in _write/_read around self._func(...) instead of in run() since setting it in run would have the token live during connection checkout so the handshake/auth commands would inherit the operation's id.

… into shrey-mongo/PYTHON-5903 # Conflicts: # pymongo/asynchronous/mongo_client.py # pymongo/synchronous/mongo_client.py

… into shrey-mongo/PYTHON-5903

NoahStapp

OP_ID needs to be cleared inside AsyncPeriodicExecutor._run(), just as _csot.reset_all() is to prevent bleeding of op ids between async tasks.

Also look at _csot.py for an example of how to consolidate the OP_ID set/unset logic into a potential context manager.

NoahStapp · 2026-07-09T17:14:16Z

        codec_options=codec_options,
        user_fields=user_fields,
        orig=orig,
+        op_id=_op_id.OP_ID.get(),


We can consolidate these explicit passes of op_id to a new block within _run_command:

if op_id is None: op_id = _op_id.OP_ID.get()

NoahStapp · 2026-07-09T17:15:17Z

                if self._retrying:
                    self._log_retry(is_write=True)
-                return await self._func(self._session, conn, self._retryable)  # type: ignore
+                # Publish this op's id on every command of every attempt.


Suggested change

# Publish this op's id on every command of every attempt.

# One operation id across all attempts of this operation.

NoahStapp · 2026-07-09T17:15:34Z

            if self._retrying:
                self._log_retry(is_write=False)
-            return await self._func(self._session, self._server, conn, read_pref)  # type: ignore
+            # Publish this op's id on every command of every attempt.


Suggested change

# Publish this op's id on every command of every attempt.

# One operation id across all attempts of this operation.

NoahStapp · 2026-07-09T17:18:07Z

+
+
+class TestOperationIdRetry(AsyncIntegrationTest):
+    RETRIES = 2  # fail this many attempts; the (RETRIES + 1)th succeeds.


Is this comment adding something non-obvious to a reader of the code?

NoahStapp · 2026-07-09T17:21:45Z

+        """Force ``times`` closeConnection failures of ``command_name``, run
+        ``action``, and return its ``(started, failed, succeeded)`` events."""


Suggested change

"""Force ``times`` closeConnection failures of ``command_name``, run

``action``, and return its ``(started, failed, succeeded)`` events."""

"""Set a failpoint for the given command and return the corresponding events published during its execution."""

NoahStapp · 2026-07-09T17:26:50Z

+        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):


Suggested change

async def _run_under_failpoint(self, command_name, action, times, expected_error=None):

async def _run_under_failpoint(self, name, f, times, expected_error=None):

NoahStapp · 2026-07-09T17:27:28Z

+    ("listIndexes", lambda c: _list_indexes(c)),
+]
+
+_COMMANDS = {name for name, _ in _RETRYABLE_WRITES + _RETRYABLE_READS}


If this isn't going to be reused, can it be defined inline at its use instead?

NoahStapp · 2026-07-09T17:28:36Z

+        for command_name, action in _RETRYABLE_WRITES:
+            with self.subTest(command=command_name):
+                await self._check_stable_operation_id(command_name, action, self.RETRIES)
+


Suggested change

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)

We use the name, f convention for similar tests elsewhere.

NoahStapp · 2026-07-09T17:29:30Z

+        for command_name, action in _RETRYABLE_READS:
+            with self.subTest(command=command_name):
+                await self._check_stable_operation_id(command_name, action, self.RETRIES)


Suggested change

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)

NoahStapp · 2026-07-09T18:22:18Z

+    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 [


Suggested change

for command_name, action in [

for name, f in [

… into shrey-mongo/PYTHON-5903

NoahStapp · 2026-07-10T14:31:53Z

+                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):


This isn't testing any new behavior added in this PR, can we remove it?

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

shrey-mongo requested a review from a team as a code owner June 29, 2026 15:26

shrey-mongo requested a review from aclark4life June 29, 2026 15:26

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch 2 times, most recently from 247b26f to 682dbac Compare June 29, 2026 15:56

NoahStapp requested review from NoahStapp and Copilot and removed request for aclark4life June 29, 2026 16:44

Copilot started reviewing on behalf of NoahStapp June 29, 2026 16:44 View session

Copilot AI reviewed Jun 29, 2026

View reviewed changes

NoahStapp requested changes Jun 29, 2026

View reviewed changes

Comment thread test/asynchronous/test_operation_id_retry.py

Comment thread pymongo/asynchronous/command_runner.py Outdated

shrey-mongo requested a review from NoahStapp June 30, 2026 15:04

fixed merge conflict

75f4e21

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch from 5546477 to 75f4e21 Compare June 30, 2026 15:46

NoahStapp requested changes Jul 1, 2026

View reviewed changes

fixed stale op id for reauth and killCursors ops and added justification

7237bf1

shrey-mongo requested a review from NoahStapp July 7, 2026 19:10

PYTHON-5903 merged master and resolved op_id conflicts

63312b0

shrey-mongo force-pushed the shrey-mongo/PYTHON-5903 branch from 1707aa0 to 63312b0 Compare July 8, 2026 17:23

NoahStapp requested changes Jul 8, 2026

View reviewed changes

shrey-mongo added 2 commits July 8, 2026 18:11

made requested changes

643ccef

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

11223ce

… into shrey-mongo/PYTHON-5903

shrey-mongo requested a review from NoahStapp July 8, 2026 18:43

NoahStapp requested changes Jul 8, 2026

View reviewed changes

made requested changes

b36e20b

shrey-mongo requested a review from NoahStapp July 8, 2026 20:33

NoahStapp requested changes Jul 9, 2026

View reviewed changes

Comment thread test/asynchronous/test_operation_id_retry.py

shrey-mongo added 3 commits July 9, 2026 15:56

implemented ContextVar approach for op_id

d8f29b9

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

371b161

… into shrey-mongo/PYTHON-5903 # Conflicts: # pymongo/asynchronous/mongo_client.py # pymongo/synchronous/mongo_client.py

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

be5fa6e

… into shrey-mongo/PYTHON-5903

shrey-mongo requested a review from NoahStapp July 9, 2026 16:21

NoahStapp requested changes Jul 9, 2026

View reviewed changes

shrey-mongo added 2 commits July 9, 2026 19:01

made requested changes

dc74bf1

Merge branch 'master' of https://github.com/mongodb/mongo-python-driver…

f1e4285

… into shrey-mongo/PYTHON-5903

Jibola changed the title ~~PYTHON-5903 : Retry attempts should share a single stable operation_id~~ PYTHON-5903 Retry attempts should share a single stable operation_id Jul 9, 2026

shrey-mongo requested a review from NoahStapp July 9, 2026 19:18

NoahStapp requested changes Jul 10, 2026

View reviewed changes

fixed test coverage

64fff02

shrey-mongo requested a review from NoahStapp July 10, 2026 18:45

aclark4life requested review from aclark4life and Copilot and removed request for NoahStapp July 14, 2026 16:25

Copilot started reviewing on behalf of aclark4life July 14, 2026 16:25 View session

Merge branch 'main' into shrey-mongo/PYTHON-5903

005b6d4

Copilot AI reviewed Jul 14, 2026

View reviewed changes

Comment thread pymongo/_op_id.py

aclark4life reviewed Jul 14, 2026

View reviewed changes

Comment thread pymongo/_op_id.py Outdated

aclark4life requested a review from Copilot July 14, 2026 20:15

Copilot started reviewing on behalf of aclark4life July 14, 2026 20:16 View session

Copilot AI reviewed Jul 14, 2026

View reviewed changes

Comment thread pymongo/_op_id.py Outdated

simplified _OpIdContext token handling

816ecf9

	# Publish this op's id on every command of every attempt.
	# One operation id across all attempts of this operation.



		class TestOperationIdRetry(AsyncIntegrationTest):
		RETRIES = 2 # fail this many attempts; the (RETRIES + 1)th succeeds.

		"""Force ``times`` closeConnection failures of ``command_name``, run
		``action``, and return its ``(started, failed, succeeded)`` events."""

	"""Force ``times`` closeConnection failures of ``command_name``, run
	``action``, and return its ``(started, failed, succeeded)`` events."""
	"""Set a failpoint for the given command and return the corresponding events published during its execution."""

	async def _run_under_failpoint(self, command_name, action, times, expected_error=None):
	async def _run_under_failpoint(self, name, f, times, expected_error=None):

Uh oh!

Conversation

shrey-mongo commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes in this PR

Test Plan

Checklist

Checklist for Author

Checklist for Reviewer

Uh oh!

NoahStapp commented Jun 29, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

shrey-mongo commented Jun 29, 2026

Uh oh!

codecov-commenter commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

NoahStapp commented Jun 30, 2026

Uh oh!

shrey-mongo commented Jun 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

NoahStapp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

shrey-mongo commented Jul 9, 2026

Uh oh!

NoahStapp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shrey-mongo commented Jun 29, 2026 •

edited

Loading

codecov-commenter commented Jun 29, 2026 •

edited

Loading