5656
5757from bson .codec_options import DEFAULT_CODEC_OPTIONS , CodecOptions , TypeRegistry
5858from bson .timestamp import Timestamp
59- from pymongo import _csot , common , helpers_shared , periodic_executor
59+ from pymongo import _csot , _op_id , common , helpers_shared , periodic_executor
6060from pymongo ._telemetry import log_command_retry
6161from pymongo .asynchronous import client_session , database , uri_parser
6262from pymongo .asynchronous .change_stream import AsyncChangeStream , AsyncClusterChangeStream
9494 _log_client_error ,
9595 _log_or_warn ,
9696)
97- from pymongo .message import _CursorAddress , _GetMore , _Query
97+ from pymongo .message import _CursorAddress , _GetMore , _Query , _randint
9898from pymongo .monitoring import ConnectionClosedReason , _EventListeners
9999from pymongo .operations import (
100100 DeleteMany ,
@@ -1836,6 +1836,8 @@ async def _select_server(
18361836 be pinned to a mongos server address.
18371837 - `address` (optional): Address when sending a message
18381838 to a specific server, used for getMore.
1839+ - `operation_id` (optional): Stable operation id shared across retries,
1840+ used for command monitoring.
18391841 """
18401842 try :
18411843 topology = await self ._get_topology ()
@@ -2011,6 +2013,7 @@ async def _retry_internal(
20112013 :param retryable: If the operation should be retried once, defaults to None
20122014 :param is_run_command: If this is a runCommand operation, defaults to False
20132015 :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
2016+ :param operation_id: Stable operation id shared across retries, defaults to None
20142017
20152018 :return: Output of the calling func()
20162019 """
@@ -2057,6 +2060,7 @@ async def _retryable_read(
20572060 (may not always be supported even if supplied), defaults to False
20582061 :param is_run_command: If this is a runCommand operation, defaults to False.
20592062 :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
2063+ :param operation_id: Stable operation id shared across retries, defaults to None
20602064 """
20612065
20622066 # Ensure that the client supports retrying on reads and there is no session in
@@ -2100,6 +2104,7 @@ async def _retryable_write(
21002104 :param session: Client session we will use to execute write operation
21012105 :param operation: The name of the operation that the server is being selected for
21022106 :param bulk: bulk abstraction to execute operations in bulk, defaults to None
2107+ :param operation_id: Stable operation id shared across retries, defaults to None
21032108 """
21042109 async with self ._tmp_session (session ) as s :
21052110 return await self ._retry_with_session (retryable , func , s , bulk , operation , operation_id )
@@ -2783,7 +2788,7 @@ def __init__(
27832788 self ._server : Server = None # type: ignore
27842789 self ._deprioritized_servers : list [Server ] = []
27852790 self ._operation = operation
2786- self ._operation_id = operation_id
2791+ self ._operation_id = operation_id if operation_id is not None else _randint ()
27872792 self ._attempt_number = 0
27882793 self ._is_run_command = is_run_command
27892794 self ._is_aggregate_write = is_aggregate_write
@@ -3012,7 +3017,9 @@ async def _write(self) -> T:
30123017 self ._retryable = False
30133018 if self ._retrying :
30143019 self ._log_retry (is_write = True )
3015- return await self ._func (self ._session , conn , self ._retryable ) # type: ignore
3020+ # One operation id across all attempts of this operation.
3021+ with _op_id ._OpIdContext (self ._operation_id ):
3022+ return await self ._func (self ._session , conn , self ._retryable ) # type: ignore
30163023 except PyMongoError as exc :
30173024 if not self ._retryable :
30183025 raise
@@ -3035,7 +3042,9 @@ async def _read(self) -> T:
30353042 self ._check_last_error ()
30363043 if self ._retrying :
30373044 self ._log_retry (is_write = False )
3038- return await self ._func (self ._session , self ._server , conn , read_pref ) # type: ignore
3045+ # One operation id across all attempts of this operation.
3046+ with _op_id ._OpIdContext (self ._operation_id ):
3047+ return await self ._func (self ._session , self ._server , conn , read_pref ) # type: ignore
30393048
30403049
30413050def _after_fork_child () -> None :
0 commit comments