From b276906795c92f080daf802ba92aa509bf28523b Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 07:04:54 -0500 Subject: [PATCH] PYTHON-5993 Add OpenTelemetry operation spans for cursor getMores Give each caller-driven getMore an operation span of its own, as the specification requires: the application may do unrelated work between batches, so nesting them under the operation that created the cursor would misrepresent the timing. A public API call that creates a cursor and drains it itself, such as list_collection_names or index_information, is the exception. Those mark the block with internal_cursor_iteration(), and every getMore inside it belongs to that call's single operation span. The client bulk-write results cursor is a second exception, reusing the enclosing bulkWrite span rather than creating spurious siblings. Change streams deliberately get neither treatment: they can tail indefinitely, so a span covering the whole lifetime would never end. Also vendors the getMore spec fixture, the one fixture that needs this support. --- pymongo/_otel.py | 73 ++- pymongo/_telemetry.py | 7 + pymongo/asynchronous/change_stream.py | 8 + pymongo/asynchronous/client_bulk.py | 6 + pymongo/asynchronous/collection.py | 29 +- pymongo/asynchronous/command_cursor.py | 58 ++- pymongo/asynchronous/cursor.py | 19 +- pymongo/asynchronous/database.py | 10 +- pymongo/asynchronous/encryption.py | 10 +- pymongo/asynchronous/mongo_client.py | 49 +- pymongo/cursor_shared.py | 67 ++- pymongo/synchronous/change_stream.py | 8 + pymongo/synchronous/client_bulk.py | 6 + pymongo/synchronous/collection.py | 27 +- pymongo/synchronous/command_cursor.py | 60 ++- pymongo/synchronous/cursor.py | 19 +- pymongo/synchronous/database.py | 9 +- pymongo/synchronous/encryption.py | 8 +- pymongo/synchronous/mongo_client.py | 49 +- test/asynchronous/test_otel_getmore.py | 533 ++++++++++++++++++++ test/open_telemetry/operation/get_more.json | 318 ++++++++++++ test/test_otel_getmore.py | 533 ++++++++++++++++++++ 22 files changed, 1802 insertions(+), 104 deletions(-) create mode 100644 test/asynchronous/test_otel_getmore.py create mode 100644 test/open_telemetry/operation/get_more.json create mode 100644 test/test_otel_getmore.py diff --git a/pymongo/_otel.py b/pymongo/_otel.py index 64180a1d5b..caddcd2dd9 100644 --- a/pymongo/_otel.py +++ b/pymongo/_otel.py @@ -65,6 +65,15 @@ "_CURRENT_OPERATION_NAME", default=None ) +# True while the driver is iterating a cursor of its own to build the return +# value of one public API call (list_collection_names, index_information, ...). +# Such a call gets a single operation span covering every getMore it sends, +# whereas a cursor handed back to the caller gets a fresh operation span per +# caller-driven getMore. See internal_cursor_iteration. +_INTERNAL_CURSOR_ITERATION: ContextVar[bool] = ContextVar( + "_INTERNAL_CURSOR_ITERATION", default=False +) + if TYPE_CHECKING: from opentelemetry.trace import Span, Tracer @@ -116,6 +125,28 @@ def _env_truthy(name: str) -> bool: return os.getenv(name, "").strip().lower() in _TRUTHY +@contextlib.contextmanager +def internal_cursor_iteration() -> Iterator[None]: + """Mark the enclosing block as driver-internal cursor iteration. + + Wrap the block in which a public API method creates a cursor and drains it + itself to build its return value. Everything the block sends, including + every getMore, then belongs to that method's one operation span, as the + OTel spec requires. Outside such a block the cursor is assumed to reach the + caller, whose iteration is a separate operation per getMore. + """ + token = _INTERNAL_CURSOR_ITERATION.set(True) + try: + yield + finally: + _INTERNAL_CURSOR_ITERATION.reset(token) + + +def is_internal_cursor_iteration() -> bool: + """Return True inside an :func:`internal_cursor_iteration` block.""" + return _INTERNAL_CURSOR_ITERATION.get() + + def _is_tracing_enabled(tracing_options: Optional[TracingOptions]) -> bool: """Return True if spans should be created for this client. @@ -287,6 +318,14 @@ def start_command_span( return None collection = _extract_collection_name(command_name, dbname, cmd) + # A getMore's own command value is the id of the cursor being read, which is + # the value db.mongodb.cursor_id takes for a command operating on an + # existing cursor: the id sent, not whatever the reply comes back with. It + # has to be read here rather than from the reply because the reply is 0 once + # the cursor is exhausted, and the attribute is required even then. + sent_cursor_id = cmd.get(_GET_MORE) if command_name == _GET_MORE else None + if not isinstance(sent_cursor_id, int): + sent_cursor_id = None # Backfill the operation span's name/namespace/summary from the first command # built inside it. Before the sensitive-command return below, since the # operation span needs those attributes even when the command gets no span. @@ -300,6 +339,8 @@ def start_command_span( current_span.set_attribute("db.operation.summary", summary) if collection: current_span.set_attribute("db.collection.name", collection) + if sent_cursor_id: + current_span.set_attribute("db.mongodb.cursor_id", sent_cursor_id) if _is_sensitive_command(command_name, speculative_hello): return None @@ -321,6 +362,8 @@ def start_command_span( attributes["db.collection.name"] = collection if conn.server_connection_id is not None: attributes["db.mongodb.server_connection_id"] = conn.server_connection_id + if sent_cursor_id: + attributes["db.mongodb.cursor_id"] = sent_cursor_id lsid = cmd.get("lsid") if isinstance(lsid, Mapping): formatted_lsid = _format_lsid(lsid) @@ -337,6 +380,20 @@ def start_command_span( return _TRACER.start_span(command_name, kind=SpanKind.CLIENT, attributes=attributes) +def _set_operation_cursor_id(cursor_id: int) -> None: + """Set db.mongodb.cursor_id on the ambient operation span, if there is one. + + Guarded on the operation-name contextvar for the same reason + ``start_command_span``'s backfill is: without it the "current span" could be + an unrelated span belonging to the host application. + """ + if _CURRENT_OPERATION_NAME.get() is None: + return + current_span = trace.get_current_span() + if current_span.is_recording(): + current_span.set_attribute("db.mongodb.cursor_id", cursor_id) + + def end_command_span_success(span: Optional[Span], reply: _DocumentOut) -> None: """Set the cursor id (if any open cursor) and end the span.""" if span is None: @@ -344,8 +401,13 @@ def end_command_span_success(span: Optional[Span], reply: _DocumentOut) -> None: cursor = reply.get("cursor") if isinstance(cursor, Mapping) and cursor.get("id"): # Per the spec the attribute is omitted rather than set to 0, so a - # cursor-creating command that leaves no cursor open reports nothing. - span.set_attribute("db.mongodb.cursor_id", cursor["id"]) + # cursor-creating command that leaves no cursor open reports nothing. A + # getMore keeps the id it sent, which this does not overwrite with a 0. + cursor_id = cursor["id"] + span.set_attribute("db.mongodb.cursor_id", cursor_id) + # The operation span carries the same attribute: this reply's id for a + # cursor-creating command, or the already-set sent id for a getMore. + _set_operation_cursor_id(cursor_id) span.end() @@ -413,6 +475,7 @@ def start_operation_span( dbname: Optional[str] = None, collection: Optional[str] = None, set_current: bool = True, + cursor_id: Optional[int] = None, ) -> Optional[_OperationSpanHandle]: """Start a CLIENT-kind span for one logical operation, or None. @@ -425,6 +488,10 @@ def start_operation_span( ``parent_span`` becomes an *explicit* parent rather than being read from ambient context, so a concurrent unrelated session cannot be captured. + ``cursor_id`` sets ``db.mongodb.cursor_id`` up front, for an operation + reading an existing cursor: the id is known before the command is built and + is needed even if the operation fails before any command span exists. + ``set_current=False`` leaves the span and the operation-name contextvar alone, for a caller that makes it current with ``use_operation_span``. """ @@ -443,6 +510,8 @@ def start_operation_span( if collection: attributes["db.collection.name"] = collection attributes["db.operation.summary"] = name + if cursor_id: + attributes["db.mongodb.cursor_id"] = cursor_id if not set_current: span = _TRACER.start_span( name, kind=SpanKind.CLIENT, context=context, attributes=attributes diff --git a/pymongo/_telemetry.py b/pymongo/_telemetry.py index f4aea9d3b6..037ce18f9c 100644 --- a/pymongo/_telemetry.py +++ b/pymongo/_telemetry.py @@ -276,6 +276,9 @@ class _OperationTelemetry: That suits a span started outside the ``_retry_internal`` call it covers, such as a cursor-creating command's, whose span has to exist before the cursor does; that call makes it current with :meth:`use`. + + ``cursor_id`` presets ``db.mongodb.cursor_id`` for an operation reading a + cursor that already exists, whose id is known before the command is built. """ __slots__ = ("handle",) @@ -289,6 +292,7 @@ def __init__( dbname: Optional[str] = None, collection: Optional[str] = None, set_current: bool = True, + cursor_id: Optional[int] = None, ) -> None: parent_span = None if session is not None and session.in_transaction: @@ -300,6 +304,7 @@ def __init__( dbname=dbname, collection=collection, set_current=set_current, + cursor_id=cursor_id, ) def use(self) -> Any: @@ -330,6 +335,7 @@ def _operation_telemetry_or_none( dbname: Optional[str] = None, collection: Optional[str] = None, set_current: bool = True, + cursor_id: Optional[int] = None, ) -> Optional[_OperationTelemetry]: """Return an :class:`_OperationTelemetry`, or None if tracing is disabled. @@ -346,6 +352,7 @@ def _operation_telemetry_or_none( dbname=dbname, collection=collection, set_current=set_current, + cursor_id=cursor_id, ) diff --git a/pymongo/asynchronous/change_stream.py b/pymongo/asynchronous/change_stream.py index e9d588ac95..ef126938d3 100644 --- a/pymongo/asynchronous/change_stream.py +++ b/pymongo/asynchronous/change_stream.py @@ -250,6 +250,14 @@ async def _run_aggregation_cmd( result_processor=self._process_result, comment=self._comment, ) + # Deliberately no operation_telemetry is attached to the resulting + # cursor here: a change stream can tail indefinitely, so an operation + # span covering its whole lifetime (initial query + every getMore, + # like other command cursors) would never end while it's watching. + # Leaving it unattached means each getMore instead gets its own + # short-lived sibling "getMore" operation span, less ideal nesting, + # but not a leaked/never-exported span. Do not "fix" this without + # addressing that tradeoff. return await self._client._retryable_read( cmd.get_cursor, self._target._read_preference_for(session), diff --git a/pymongo/asynchronous/client_bulk.py b/pymongo/asynchronous/client_bulk.py index 8aa8892e5d..8f3c5645a2 100644 --- a/pymongo/asynchronous/client_bulk.py +++ b/pymongo/asynchronous/client_bulk.py @@ -335,6 +335,12 @@ async def _process_results_cursor( session=session, comment=self.comment, ) + # This cursor's getMores run inside the enclosing bulkWrite + # operation span, so their command spans belong under it directly; + # a getMore operation span of their own would be spurious. The + # cursor is also per-batch and never surfaces to the caller, so + # there is no cursor-lifetime span to own here. + cmd_cursor._reuse_current_span_for_getmore = True await cmd_cursor._maybe_pin_connection(conn) # Iterate the cursor to get individual write results. diff --git a/pymongo/asynchronous/collection.py b/pymongo/asynchronous/collection.py index f4c91c941d..14e41d0632 100644 --- a/pymongo/asynchronous/collection.py +++ b/pymongo/asynchronous/collection.py @@ -38,6 +38,7 @@ from bson.son import SON from bson.timestamp import Timestamp from pymongo import ASCENDING, _csot, common, helpers_shared, message +from pymongo._otel import internal_cursor_iteration from pymongo.asynchronous.aggregation import ( _CollectionAggregationCommand, _CollectionRawAggregationCommand, @@ -2639,12 +2640,13 @@ async def index_information( .. versionchanged:: 3.6 Added ``session`` parameter. """ - cursor = await self._list_indexes(session=session, comment=comment) - info = {} - async for index in cursor: - index["key"] = list(index["key"].items()) - index = dict(index) # noqa: PLW2901 - info[index.pop("name")] = index + with internal_cursor_iteration(): + cursor = await self._list_indexes(session=session, comment=comment) + info = {} + async for index in cursor: + index["key"] = list(index["key"].items()) + index = dict(index) # noqa: PLW2901 + info[index.pop("name")] = index return info async def list_search_indexes( @@ -2910,14 +2912,15 @@ async def options( self.write_concern, self.read_concern, ) - cursor = await dbo.list_collections( - session=session, filter={"name": self._name}, comment=comment - ) + with internal_cursor_iteration(): + cursor = await dbo.list_collections( + session=session, filter={"name": self._name}, comment=comment + ) - result = None - async for doc in cursor: - result = doc - break + result = None + async for doc in cursor: + result = doc + break if not result: return {} diff --git a/pymongo/asynchronous/command_cursor.py b/pymongo/asynchronous/command_cursor.py index 71404281a4..a15487bc9b 100644 --- a/pymongo/asynchronous/command_cursor.py +++ b/pymongo/asynchronous/command_cursor.py @@ -29,6 +29,7 @@ from pymongo.asynchronous.cursor_base import _AsyncCursorBase, _ConnectionManager from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure +from pymongo.helpers_shared import _split_namespace from pymongo.message import _GetMore, _OpMsg, _RawBatchGetMore from pymongo.response import PinnedResponse from pymongo.typings import _Address, _DocumentOut, _DocumentType @@ -173,9 +174,14 @@ async def _send_message(self, operation: _GetMore) -> None: client = self._collection.database.client try: response = await client._run_operation( - operation, self._run_with_conn, address=self._address + operation, + self._run_with_conn, + address=self._address, + operation_telemetry=self._operation_telemetry, + reuse_current_span=self._reuse_current_span_for_getmore, ) except OperationFailure as exc: + self._end_operation_telemetry(exc) if exc.code in _CURSOR_CLOSED_ERRORS: # Don't send killCursors because the cursor is already closed. self._killed = True @@ -185,13 +191,15 @@ async def _send_message(self, operation: _GetMore) -> None: # Return the session and pinned connection, if necessary. await self.close() raise - except ConnectionFailure: + except ConnectionFailure as exc: + self._end_operation_telemetry(exc) # Don't send killCursors because the cursor is already closed. self._killed = True # Return the session and pinned connection, if necessary. await self.close() raise - except Exception: + except Exception as exc: + self._end_operation_telemetry(exc) await self.close() raise @@ -218,24 +226,36 @@ async def _refresh(self) -> int: return len(self._data) if self._id: # Get More - dbname, collname = self._ns.split(".", 1) + dbname, collname = _split_namespace(self._ns) read_pref = self._collection._read_preference_for(self.session) - await self._send_message( - self._getmore_class( - dbname, - collname, - self._batch_size, - self._id, - self._collection.codec_options, - read_pref, - self._session, - self._collection.database.client, - self._max_await_time_ms, - self._sock_mgr, - False, - self._comment, - ) + getmore = self._getmore_class( + dbname, + collname, + self._batch_size, + self._id, + self._collection.codec_options, + read_pref, + self._session, + self._collection.database.client, + self._max_await_time_ms, + self._sock_mgr, + False, + self._comment, ) + own_span = self._start_getmore_operation_telemetry(dbname, collname) + if not own_span: + await self._send_message(getmore) + else: + # _send_message ends the span itself on every failure path, and + # an exhausted cursor's close() ends it on the way out; both are + # idempotent, so only a successful send leaving the cursor open + # is left to handle here. + try: + await self._send_message(getmore) + except BaseException as exc: + self._end_operation_telemetry(exc) + raise + self._end_operation_telemetry() else: # Cursor id is zero nothing else to return await self._die_lock() diff --git a/pymongo/asynchronous/cursor.py b/pymongo/asynchronous/cursor.py index b68e96a4bb..079588bc27 100644 --- a/pymongo/asynchronous/cursor.py +++ b/pymongo/asynchronous/cursor.py @@ -34,6 +34,7 @@ from bson.code import Code from bson.son import SON from pymongo import helpers_shared +from pymongo._otel import is_internal_cursor_iteration from pymongo._telemetry import _operation_telemetry_or_none from pymongo.asynchronous.cursor_base import _AsyncCursorBase, _ConnectionManager from pymongo.asynchronous.helpers import anext @@ -1088,7 +1089,11 @@ async def _refresh(self) -> int: collection=self._collection.name, set_current=False, ) - await self._send_message_in_operation_span(q) + # The query's span covers the query alone unless this cursor is + # being drained by the public API call that created it, in which + # case the span stays open to cover that call's getMores too. + own_span = not is_internal_cursor_iteration() + await self._send_message_in_operation_span(q, own_span) elif self._id: # Get More if self._limit: limit = self._limit - self._retrieved @@ -1111,18 +1116,24 @@ async def _refresh(self) -> int: self._exhaust, self._comment, ) - await self._send_message(g) + own_span = self._start_getmore_operation_telemetry(self._dbname, self._collname) + await self._send_message_in_operation_span(g, own_span) return len(self._data) - async def _send_message_in_operation_span(self, operation: Union[_Query, _GetMore]) -> None: - """Send ``operation``, ending the operation span once it completes. + async def _send_message_in_operation_span( + self, operation: Union[_Query, _GetMore], own_span: bool + ) -> None: + """Send ``operation``, ending the operation span after it when we own it. ``_send_message``'s own error handling already ends the span with the error on every failure path, and an exhausted cursor's close() ends it on the way out; both are idempotent, so this only has to cover the remaining case of a successful send that leaves the cursor open. """ + if not own_span: + await self._send_message(operation) + return try: await self._send_message(operation) except BaseException as exc: diff --git a/pymongo/asynchronous/database.py b/pymongo/asynchronous/database.py index 96971fdf2d..830a1ba9d1 100644 --- a/pymongo/asynchronous/database.py +++ b/pymongo/asynchronous/database.py @@ -33,6 +33,7 @@ from bson.dbref import DBRef from bson.timestamp import Timestamp from pymongo import _csot, common +from pymongo._otel import internal_cursor_iteration from pymongo.asynchronous.aggregation import _DatabaseAggregationCommand from pymongo.asynchronous.change_stream import AsyncDatabaseChangeStream from pymongo.asynchronous.collection import AsyncCollection @@ -1211,10 +1212,11 @@ async def _list_collection_names( if not filter or (len(filter) == 1 and "name" in filter): kwargs["nameOnly"] = True - return [ - result["name"] - async for result in await self._list_collections_helper(session=session, **kwargs) - ] + with internal_cursor_iteration(): + return [ + result["name"] + async for result in await self._list_collections_helper(session=session, **kwargs) + ] async def list_collection_names( self, diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 524ae45c11..dcabec39b8 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -56,6 +56,7 @@ from bson.errors import BSONError from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson from pymongo import _csot, _op_id +from pymongo._otel import internal_cursor_iteration from pymongo.asynchronous.collection import AsyncCollection from pymongo.asynchronous.cursor import AsyncCursor from pymongo.asynchronous.database import AsyncDatabase @@ -257,10 +258,11 @@ async def collection_info(self, database: str, filter: bytes) -> Optional[list[b :return: All documents from the listCollections command response as BSON. """ - async with await self.client_ref()[database].list_collections( - filter=RawBSONDocument(filter) - ) as cursor: - return [_dict_to_bson(doc, False, _DATA_KEY_OPTS) async for doc in cursor] + with internal_cursor_iteration(): + async with await self.client_ref()[database].list_collections( + filter=RawBSONDocument(filter) + ) as cursor: + return [_dict_to_bson(doc, False, _DATA_KEY_OPTS) async for doc in cursor] def spawn(self) -> None: """Spawn mongocryptd. diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index 51a6d02e91..a73cad9d5a 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -56,6 +56,7 @@ from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions, TypeRegistry from bson.timestamp import Timestamp from pymongo import _csot, _op_id, common, helpers_shared, periodic_executor +from pymongo._otel import is_internal_cursor_iteration from pymongo._telemetry import ( _generate_op_id_or_none, _operation_telemetry_or_none, @@ -1901,6 +1902,7 @@ async def _run_operation( run_with_conn: Callable, # type: ignore[type-arg] address: Optional[_Address] = None, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> Response: """Run a _Query/_GetMore operation and return a Response. @@ -1909,7 +1911,12 @@ async def _run_operation( that executes the operation on a given connection. :param address: Optional address when sending a message to a specific server, used for getMore. - :param operation_telemetry: The cursor's caller-owned operation span, or None. + :param operation_telemetry: The cursor's caller-owned operation span, shared + across its initial query and every getMore, or None. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. Mutually exclusive with ``operation_telemetry``. Defaults to + False. """ if operation.conn_mgr: server = await self._select_server( @@ -1927,8 +1934,8 @@ async def _run_operation( operation.conn_mgr.conn, ): # Exhaust/pinned cursors bypass _retry_internal, so make the - # caller's span current here to keep their command spans - # nested under it. + # caller's span current here to keep their getMore command + # spans nested under it. with ( operation_telemetry.use() if operation_telemetry @@ -1955,6 +1962,7 @@ async def _cmd( retryable=isinstance(operation, _Query), operation=operation.name, operation_telemetry=operation_telemetry, + reuse_current_span=reuse_current_span, ) async def _retry_with_session( @@ -2002,6 +2010,7 @@ async def _retry_internal( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> T: """Internal retryable helper for all client transactions. @@ -2020,6 +2029,13 @@ async def _retry_internal( (a cursor's, shared by its getMores). When given, this method neither creates nor ends a span; it only makes the caller's current for this call. Defaults to None, meaning this method owns a fresh span. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. For callers that know a suitable operation span is already + current, where a second one would be spurious (the client + bulk-write results cursor's getMores, which belong under the + enclosing bulkWrite span). Mutually exclusive with + ``operation_telemetry``. Defaults to False. :return: Output of the calling func() """ @@ -2037,6 +2053,7 @@ async def _retry_internal( is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, operation_telemetry=operation_telemetry, + reuse_current_span=reuse_current_span, ).run() async def _retryable_read( @@ -2051,6 +2068,7 @@ async def _retryable_read( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> T: """Execute an operation with consecutive retries if possible @@ -2071,6 +2089,10 @@ async def _retryable_read( :param operation_id: Stable operation id shared across retries, defaults to None :param operation_telemetry: A caller-owned operation span outliving this call, defaults to None, meaning this method owns a fresh span. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. Mutually exclusive with ``operation_telemetry``. Defaults to + False. """ # Ensure that the client supports retrying on reads and there is no session in @@ -2091,6 +2113,7 @@ async def _retryable_read( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + reuse_current_span=reuse_current_span, operation_telemetry=operation_telemetry, ) @@ -2116,7 +2139,11 @@ async def _retryable_read_cursor( call, before the cursor exists, so the span cannot be owned by the cursor the way a find cursor's is; create it here instead. - The span ends with the command that created the cursor. + The span ends with the command that created the cursor. Later getMores + belong to whoever drives iteration: each one the caller drives gets an + operation span of its own, so only a public API call that drains the + cursor itself (see ``_otel.internal_cursor_iteration``) keeps this one + open, by handing it to the cursor. """ operation_telemetry = _operation_telemetry_or_none( self.options.tracing, @@ -2143,7 +2170,11 @@ async def _retryable_read_cursor( if operation_telemetry is not None: operation_telemetry.failed(exc) raise - if operation_telemetry is not None: + if operation_telemetry is None: + pass + elif is_internal_cursor_iteration(): + cmd_cursor._attach_operation_telemetry(operation_telemetry) + else: operation_telemetry.succeeded() return cmd_cursor @@ -2985,6 +3016,7 @@ def __init__( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ): self._last_error: Optional[Exception] = None self._retrying = False @@ -3009,9 +3041,12 @@ def __init__( if operation_id is None: operation_id = _generate_op_id_or_none(self._client._event_listeners) self._operation_id = operation_id + if reuse_current_span and operation_telemetry is not None: + raise ValueError("reuse_current_span and operation_telemetry are mutually exclusive") # One span covering every attempt. A caller needing it to outlive this - # object (a cursor) passes its own and keeps ownership. - self._owns_telemetry = operation_telemetry is None + # object (a cursor) passes its own and keeps ownership; + # reuse_current_span means an enclosing span is already current. + self._owns_telemetry = operation_telemetry is None and not reuse_current_span if self._owns_telemetry: operation_telemetry = _operation_telemetry_or_none( mongo_client.options.tracing, operation, session, is_run_command=is_run_command diff --git a/pymongo/cursor_shared.py b/pymongo/cursor_shared.py index 5a26eef5cd..516673eb41 100644 --- a/pymongo/cursor_shared.py +++ b/pymongo/cursor_shared.py @@ -21,6 +21,8 @@ from collections.abc import Mapping, Sequence from typing import Any, Generic, Optional, Union +from pymongo import _otel +from pymongo._telemetry import _operation_telemetry_or_none from pymongo.message import _CursorAddress from pymongo.typings import _Address, _DocumentType @@ -56,6 +58,10 @@ class _AgnosticCursorBase(Generic[_DocumentType], ABC): _session: Optional[Any] _killed: bool _operation_telemetry: Optional[Any] = None + # Set by callers whose getMores belong under an operation span that is + # already current (the client bulk-write results cursor), rather than under + # a getMore operation span of their own. + _reuse_current_span_for_getmore: bool = False @abstractmethod def _get_namespace(self) -> str: @@ -119,11 +125,13 @@ def _prepare_to_die(self, already_killed: bool) -> tuple[int, Optional[_CursorAd def _end_operation_telemetry(self, exc: Optional[BaseException] = None) -> None: """End the operation span currently attached to this cursor, exactly once. - No span is scoped to the cursor's lifetime: the span of the operation - that created the cursor ends as soon as that creating command - completes, whether it succeeded, failed, or the cursor was abandoned - part-way and is being closed by close()/__del__. Idempotent, so every - one of those paths can call it unconditionally. + No span is ever scoped to the cursor's lifetime: a caller-driven getMore + attaches a span of its own and ends it as soon as that getMore + completes. What can outlive a single command is the span of a public API + call that drains the cursor itself (see + ``_otel.internal_cursor_iteration``), which ends when the cursor is + exhausted or, for a cursor abandoned part-way, at close()/__del__. + Idempotent, so every one of those paths can call it unconditionally. """ telemetry = self._operation_telemetry if telemetry is None: @@ -134,6 +142,55 @@ def _end_operation_telemetry(self, exc: Optional[BaseException] = None) -> None: else: telemetry.failed(exc) + def _start_getmore_operation_telemetry(self, dbname: str, collname: Optional[str]) -> bool: + """Give the getMore about to be sent an operation span of its own. + + The spec requires an operation span per caller-driven getMore, and + forbids nesting it under the operation that created the cursor, since + the application may do unrelated work between batches. + + Returns True when the caller now owns a span and must end it once the + getMore completes. Returns False when this getMore already belongs to + another operation, or when tracing is off. + """ + if self._operation_telemetry is not None or self._reuse_current_span_for_getmore: + return False + tracing_options = self._collection.database.client.options.tracing + if not _otel._is_tracing_enabled(tracing_options): + return False + # A cursor opened by a command (listCollections, listIndexes, a + # database-level aggregate) reports a namespace like + # "$cmd.listCollections", which names no user collection. + if _otel.is_command_namespace(collname): + collname = None + self._operation_telemetry = _operation_telemetry_or_none( + tracing_options, + "getMore", + self._session, + dbname=dbname, + collection=collname, + set_current=False, + cursor_id=self._id, + ) + return self._operation_telemetry is not None + + def _attach_operation_telemetry(self, telemetry: Any) -> None: + """Adopt the still-open operation span of the call that created this cursor. + + For command cursors only, and only when that call goes on to drain the + cursor itself, so its getMores belong to the same operation (see + ``_otel.internal_cursor_iteration``). A cursor returned to the caller + has its creating span ended right away and never gets here. + + A command cursor exhausted by its first batch is marked ``_killed`` in + ``__init__`` without calling ``close()``, so no getMore is sent and + neither ``_refresh()`` nor ``_die_lock()`` runs. Ending the span here + keeps that case prompt instead of leaving it to ``__del__``. + """ + self._operation_telemetry = telemetry + if self._killed: + self._end_operation_telemetry() + def _die_no_lock(self) -> None: """Closes this cursor without acquiring a lock.""" try: diff --git a/pymongo/synchronous/change_stream.py b/pymongo/synchronous/change_stream.py index f442cc6c67..8b5ac31a44 100644 --- a/pymongo/synchronous/change_stream.py +++ b/pymongo/synchronous/change_stream.py @@ -248,6 +248,14 @@ def _run_aggregation_cmd(self, session: Optional[ClientSession]) -> CommandCurso result_processor=self._process_result, comment=self._comment, ) + # Deliberately no operation_telemetry is attached to the resulting + # cursor here: a change stream can tail indefinitely, so an operation + # span covering its whole lifetime (initial query + every getMore, + # like other command cursors) would never end while it's watching. + # Leaving it unattached means each getMore instead gets its own + # short-lived sibling "getMore" operation span, less ideal nesting, + # but not a leaked/never-exported span. Do not "fix" this without + # addressing that tradeoff. return self._client._retryable_read( cmd.get_cursor, self._target._read_preference_for(session), diff --git a/pymongo/synchronous/client_bulk.py b/pymongo/synchronous/client_bulk.py index 7f7181d826..7f47dacc8c 100644 --- a/pymongo/synchronous/client_bulk.py +++ b/pymongo/synchronous/client_bulk.py @@ -333,6 +333,12 @@ def _process_results_cursor( session=session, comment=self.comment, ) + # This cursor's getMores run inside the enclosing bulkWrite + # operation span, so their command spans belong under it directly; + # a getMore operation span of their own would be spurious. The + # cursor is also per-batch and never surfaces to the caller, so + # there is no cursor-lifetime span to own here. + cmd_cursor._reuse_current_span_for_getmore = True cmd_cursor._maybe_pin_connection(conn) # Iterate the cursor to get individual write results. diff --git a/pymongo/synchronous/collection.py b/pymongo/synchronous/collection.py index 0821eb4573..a565061257 100644 --- a/pymongo/synchronous/collection.py +++ b/pymongo/synchronous/collection.py @@ -38,6 +38,7 @@ from bson.son import SON from bson.timestamp import Timestamp from pymongo import ASCENDING, _csot, common, helpers_shared, message +from pymongo._otel import internal_cursor_iteration from pymongo.collation import validate_collation_or_none from pymongo.common import _ecoc_coll_name, _esc_coll_name from pymongo.errors import ( @@ -2635,12 +2636,13 @@ def index_information( .. versionchanged:: 3.6 Added ``session`` parameter. """ - cursor = self._list_indexes(session=session, comment=comment) - info = {} - for index in cursor: - index["key"] = list(index["key"].items()) - index = dict(index) # noqa: PLW2901 - info[index.pop("name")] = index + with internal_cursor_iteration(): + cursor = self._list_indexes(session=session, comment=comment) + info = {} + for index in cursor: + index["key"] = list(index["key"].items()) + index = dict(index) # noqa: PLW2901 + info[index.pop("name")] = index return info def list_search_indexes( @@ -2906,12 +2908,15 @@ def options( self.write_concern, self.read_concern, ) - cursor = dbo.list_collections(session=session, filter={"name": self._name}, comment=comment) + with internal_cursor_iteration(): + cursor = dbo.list_collections( + session=session, filter={"name": self._name}, comment=comment + ) - result = None - for doc in cursor: - result = doc - break + result = None + for doc in cursor: + result = doc + break if not result: return {} diff --git a/pymongo/synchronous/command_cursor.py b/pymongo/synchronous/command_cursor.py index 8868d87939..bcc1492036 100644 --- a/pymongo/synchronous/command_cursor.py +++ b/pymongo/synchronous/command_cursor.py @@ -28,6 +28,7 @@ from bson import CodecOptions, _convert_raw_document_lists_to_streams from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure +from pymongo.helpers_shared import _split_namespace from pymongo.message import _GetMore, _OpMsg, _RawBatchGetMore from pymongo.response import PinnedResponse from pymongo.synchronous.cursor_base import _ConnectionManager, _CursorBase @@ -172,8 +173,15 @@ def _send_message(self, operation: _GetMore) -> None: """Send a getmore message and handle the response.""" client = self._collection.database.client try: - response = client._run_operation(operation, self._run_with_conn, address=self._address) + response = client._run_operation( + operation, + self._run_with_conn, + address=self._address, + operation_telemetry=self._operation_telemetry, + reuse_current_span=self._reuse_current_span_for_getmore, + ) except OperationFailure as exc: + self._end_operation_telemetry(exc) if exc.code in _CURSOR_CLOSED_ERRORS: # Don't send killCursors because the cursor is already closed. self._killed = True @@ -183,13 +191,15 @@ def _send_message(self, operation: _GetMore) -> None: # Return the session and pinned connection, if necessary. self.close() raise - except ConnectionFailure: + except ConnectionFailure as exc: + self._end_operation_telemetry(exc) # Don't send killCursors because the cursor is already closed. self._killed = True # Return the session and pinned connection, if necessary. self.close() raise - except Exception: + except Exception as exc: + self._end_operation_telemetry(exc) self.close() raise @@ -216,24 +226,36 @@ def _refresh(self) -> int: return len(self._data) if self._id: # Get More - dbname, collname = self._ns.split(".", 1) + dbname, collname = _split_namespace(self._ns) read_pref = self._collection._read_preference_for(self.session) - self._send_message( - self._getmore_class( - dbname, - collname, - self._batch_size, - self._id, - self._collection.codec_options, - read_pref, - self._session, - self._collection.database.client, - self._max_await_time_ms, - self._sock_mgr, - False, - self._comment, - ) + getmore = self._getmore_class( + dbname, + collname, + self._batch_size, + self._id, + self._collection.codec_options, + read_pref, + self._session, + self._collection.database.client, + self._max_await_time_ms, + self._sock_mgr, + False, + self._comment, ) + own_span = self._start_getmore_operation_telemetry(dbname, collname) + if not own_span: + self._send_message(getmore) + else: + # _send_message ends the span itself on every failure path, and + # an exhausted cursor's close() ends it on the way out; both are + # idempotent, so only a successful send leaving the cursor open + # is left to handle here. + try: + self._send_message(getmore) + except BaseException as exc: + self._end_operation_telemetry(exc) + raise + self._end_operation_telemetry() else: # Cursor id is zero nothing else to return self._die_lock() diff --git a/pymongo/synchronous/cursor.py b/pymongo/synchronous/cursor.py index 6554c62a09..9f4d3b59bb 100644 --- a/pymongo/synchronous/cursor.py +++ b/pymongo/synchronous/cursor.py @@ -34,6 +34,7 @@ from bson.code import Code from bson.son import SON from pymongo import helpers_shared +from pymongo._otel import is_internal_cursor_iteration from pymongo._telemetry import _operation_telemetry_or_none from pymongo.collation import validate_collation_or_none from pymongo.common import ( @@ -1086,7 +1087,11 @@ def _refresh(self) -> int: collection=self._collection.name, set_current=False, ) - self._send_message_in_operation_span(q) + # The query's span covers the query alone unless this cursor is + # being drained by the public API call that created it, in which + # case the span stays open to cover that call's getMores too. + own_span = not is_internal_cursor_iteration() + self._send_message_in_operation_span(q, own_span) elif self._id: # Get More if self._limit: limit = self._limit - self._retrieved @@ -1109,18 +1114,24 @@ def _refresh(self) -> int: self._exhaust, self._comment, ) - self._send_message(g) + own_span = self._start_getmore_operation_telemetry(self._dbname, self._collname) + self._send_message_in_operation_span(g, own_span) return len(self._data) - def _send_message_in_operation_span(self, operation: Union[_Query, _GetMore]) -> None: - """Send ``operation``, ending the operation span once it completes. + def _send_message_in_operation_span( + self, operation: Union[_Query, _GetMore], own_span: bool + ) -> None: + """Send ``operation``, ending the operation span after it when we own it. ``_send_message``'s own error handling already ends the span with the error on every failure path, and an exhausted cursor's close() ends it on the way out; both are idempotent, so this only has to cover the remaining case of a successful send that leaves the cursor open. """ + if not own_span: + self._send_message(operation) + return try: self._send_message(operation) except BaseException as exc: diff --git a/pymongo/synchronous/database.py b/pymongo/synchronous/database.py index 2f9b1a7eb8..216388c675 100644 --- a/pymongo/synchronous/database.py +++ b/pymongo/synchronous/database.py @@ -33,6 +33,7 @@ from bson.dbref import DBRef from bson.timestamp import Timestamp from pymongo import _csot, common +from pymongo._otel import internal_cursor_iteration from pymongo.common import _ecoc_coll_name, _esc_coll_name from pymongo.database_shared import _check_name, _CodecDocumentType from pymongo.errors import CollectionInvalid, InvalidOperation @@ -1209,9 +1210,11 @@ def _list_collection_names( if not filter or (len(filter) == 1 and "name" in filter): kwargs["nameOnly"] = True - return [ - result["name"] for result in self._list_collections_helper(session=session, **kwargs) - ] + with internal_cursor_iteration(): + return [ + result["name"] + for result in self._list_collections_helper(session=session, **kwargs) + ] def list_collection_names( self, diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 014d162e2b..2cb323c5b3 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -55,6 +55,7 @@ from bson.errors import BSONError from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson from pymongo import _csot, _op_id +from pymongo._otel import internal_cursor_iteration from pymongo.common import CONNECT_TIMEOUT from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( @@ -256,8 +257,11 @@ def collection_info(self, database: str, filter: bytes) -> Optional[list[bytes]] :return: All documents from the listCollections command response as BSON. """ - with self.client_ref()[database].list_collections(filter=RawBSONDocument(filter)) as cursor: - return [_dict_to_bson(doc, False, _DATA_KEY_OPTS) for doc in cursor] + with internal_cursor_iteration(): + with self.client_ref()[database].list_collections( + filter=RawBSONDocument(filter) + ) as cursor: + return [_dict_to_bson(doc, False, _DATA_KEY_OPTS) for doc in cursor] def spawn(self) -> None: """Spawn mongocryptd. diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 4e910dbb83..1e2bb3eb1d 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -56,6 +56,7 @@ from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions, TypeRegistry from bson.timestamp import Timestamp from pymongo import _csot, _op_id, common, helpers_shared, periodic_executor +from pymongo._otel import is_internal_cursor_iteration from pymongo._telemetry import ( _generate_op_id_or_none, _operation_telemetry_or_none, @@ -1896,6 +1897,7 @@ def _run_operation( run_with_conn: Callable, # type: ignore[type-arg] address: Optional[_Address] = None, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> Response: """Run a _Query/_GetMore operation and return a Response. @@ -1904,7 +1906,12 @@ def _run_operation( that executes the operation on a given connection. :param address: Optional address when sending a message to a specific server, used for getMore. - :param operation_telemetry: The cursor's caller-owned operation span, or None. + :param operation_telemetry: The cursor's caller-owned operation span, shared + across its initial query and every getMore, or None. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. Mutually exclusive with ``operation_telemetry``. Defaults to + False. """ if operation.conn_mgr: server = self._select_server( @@ -1922,8 +1929,8 @@ def _run_operation( operation.conn_mgr.conn, ): # Exhaust/pinned cursors bypass _retry_internal, so make the - # caller's span current here to keep their command spans - # nested under it. + # caller's span current here to keep their getMore command + # spans nested under it. with ( operation_telemetry.use() if operation_telemetry @@ -1950,6 +1957,7 @@ def _cmd( retryable=isinstance(operation, _Query), operation=operation.name, operation_telemetry=operation_telemetry, + reuse_current_span=reuse_current_span, ) def _retry_with_session( @@ -1997,6 +2005,7 @@ def _retry_internal( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> T: """Internal retryable helper for all client transactions. @@ -2015,6 +2024,13 @@ def _retry_internal( (a cursor's, shared by its getMores). When given, this method neither creates nor ends a span; it only makes the caller's current for this call. Defaults to None, meaning this method owns a fresh span. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. For callers that know a suitable operation span is already + current, where a second one would be spurious (the client + bulk-write results cursor's getMores, which belong under the + enclosing bulkWrite span). Mutually exclusive with + ``operation_telemetry``. Defaults to False. :return: Output of the calling func() """ @@ -2032,6 +2048,7 @@ def _retry_internal( is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, operation_telemetry=operation_telemetry, + reuse_current_span=reuse_current_span, ).run() def _retryable_read( @@ -2046,6 +2063,7 @@ def _retryable_read( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ) -> T: """Execute an operation with consecutive retries if possible @@ -2066,6 +2084,10 @@ def _retryable_read( :param operation_id: Stable operation id shared across retries, defaults to None :param operation_telemetry: A caller-owned operation span outliving this call, defaults to None, meaning this method owns a fresh span. + :param reuse_current_span: Create no operation span at all and leave the + ambient span in place as the parent for this operation's command + spans. Mutually exclusive with ``operation_telemetry``. Defaults to + False. """ # Ensure that the client supports retrying on reads and there is no session in @@ -2086,6 +2108,7 @@ def _retryable_read( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + reuse_current_span=reuse_current_span, operation_telemetry=operation_telemetry, ) @@ -2111,7 +2134,11 @@ def _retryable_read_cursor( call, before the cursor exists, so the span cannot be owned by the cursor the way a find cursor's is; create it here instead. - The span ends with the command that created the cursor. + The span ends with the command that created the cursor. Later getMores + belong to whoever drives iteration: each one the caller drives gets an + operation span of its own, so only a public API call that drains the + cursor itself (see ``_otel.internal_cursor_iteration``) keeps this one + open, by handing it to the cursor. """ operation_telemetry = _operation_telemetry_or_none( self.options.tracing, @@ -2138,7 +2165,11 @@ def _retryable_read_cursor( if operation_telemetry is not None: operation_telemetry.failed(exc) raise - if operation_telemetry is not None: + if operation_telemetry is None: + pass + elif is_internal_cursor_iteration(): + cmd_cursor._attach_operation_telemetry(operation_telemetry) + else: operation_telemetry.succeeded() return cmd_cursor @@ -2974,6 +3005,7 @@ def __init__( is_run_command: bool = False, is_aggregate_write: bool = False, operation_telemetry: Optional[_OperationTelemetry] = None, + reuse_current_span: bool = False, ): self._last_error: Optional[Exception] = None self._retrying = False @@ -2998,9 +3030,12 @@ def __init__( if operation_id is None: operation_id = _generate_op_id_or_none(self._client._event_listeners) self._operation_id = operation_id + if reuse_current_span and operation_telemetry is not None: + raise ValueError("reuse_current_span and operation_telemetry are mutually exclusive") # One span covering every attempt. A caller needing it to outlive this - # object (a cursor) passes its own and keeps ownership. - self._owns_telemetry = operation_telemetry is None + # object (a cursor) passes its own and keeps ownership; + # reuse_current_span means an enclosing span is already current. + self._owns_telemetry = operation_telemetry is None and not reuse_current_span if self._owns_telemetry: operation_telemetry = _operation_telemetry_or_none( mongo_client.options.tracing, operation, session, is_run_command=is_run_command diff --git a/test/asynchronous/test_otel_getmore.py b/test/asynchronous/test_otel_getmore.py new file mode 100644 index 0000000000..41ca9d09d5 --- /dev/null +++ b/test/asynchronous/test_otel_getmore.py @@ -0,0 +1,533 @@ +# 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. + +"""Test OpenTelemetry operation spans for cursor getMores.""" + +from __future__ import annotations + +import gc +import os +import sys +from typing import Optional +from unittest.mock import patch + +sys.path[0:0] = [""] + +import pytest + +import pymongo._otel as _otel +from pymongo import _telemetry, common +from pymongo._telemetry import _OperationTelemetry +from pymongo.errors import ( + ClientBulkWriteException, + ConfigurationError, + InvalidOperation, + OperationFailure, + ServerSelectionTimeoutError, +) +from pymongo.logger import _HELLO_COMMANDS +from pymongo.operations import InsertOne +from pymongo.read_preferences import ReadPreference +from pymongo.typings import _Address +from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest +from test.asynchronous.utils import async_wait_until +from test.unified_format_shared import _shared_test_provider + +_HAS_OTEL_TEST_DEPS = False +if _otel._HAS_OPENTELEMETRY: + try: + from opentelemetry import trace + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + _HAS_OTEL_TEST_DEPS = True + except ImportError: + pass + +_IS_SYNC = False + +pytestmark = pytest.mark.otel + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelGetMoreSpans(AsyncIntegrationTest): + """getMore spans, cursor lifetime, and change streams.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # See the matching comment in test/asynchronous/unified_format.py's + # UnifiedSpecTestMixinV1.tearDownClass: the span processor can never + # be removed from the shared process-wide TracerProvider, so without + # this shutdown() the exporter keeps accumulating every span from + # every client for the rest of the test run. + cls.exporter.shutdown() + super().tearDownClass() + + async def asyncSetUp(self): + await super().asyncSetUp() + self.exporter.clear() + + def spans(self, name: str | None = None): + finished = self.exporter.get_finished_spans() + if name is None: + return list(finished) + return [s for s in finished if s.name == name] + + @staticmethod + def operation_spans(finished, operation: str): + """Return the operation spans for ``operation``, excluding command spans. + + Only command spans carry db.command.name, so its absence is what tells + the two kinds apart when both name the same operation. + """ + return [ + s + for s in finished + if s.attributes.get("db.operation.name") == operation + and "db.command.name" not in s.attributes + ] + + @staticmethod + def command_spans(finished, command: str): + """Return the command spans for ``command``.""" + return [s for s in finished if s.attributes.get("db.command.name") == command] + + def ping_spans(self): + """Return the spans belonging to a ``ping`` run through ``db.command()``. + + For the tests that assert tracing produced *nothing*. Asserting the + exporter is empty would also catch spans no test asked for: a cursor + abandoned earlier in the class ends its operation span from a + finalizer, and on an interpreter that does not reference count, that + finalizer runs at an unpredictable point and lands in whichever test + happens to be running. Naming the ping's own spans keeps the assertion + about this client while staying immune to that. + """ + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "ping" + or s.attributes.get("db.operation.name") == "runCommand" + ] + + def _aggregate_operation_span(self): + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "aggregate" + ] + self.assertEqual(len(matching), 1) + return matching[0] + + async def test_span_created_for_get_more(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name].test_otel_getmore + await coll.drop() + await coll.insert_many([{"x": i} for i in range(5)]) + self.exporter.clear() + + docs = await coll.find({}, batch_size=2).to_list() + self.assertEqual(len(docs), 5) + + get_more_spans = self.spans("getMore") + self.assertGreater(len(get_more_spans), 0) + for span in get_more_spans: + self.assertEqual(span.attributes["db.collection.name"], "test_otel_getmore") + self.assertEqual(span.attributes["db.command.name"], "getMore") + + async def test_caller_driven_find_getmores_get_their_own_operation_spans(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.getmore_nesting + await coll.drop() + await coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + docs = await coll.find({}, batch_size=2).to_list() + self.assertEqual(len(docs), 10) + + finished = self.exporter.get_finished_spans() + # One operation span for the query that created the cursor. + find_op_spans = self.operation_spans(finished, "find") + self.assertEqual(len(find_op_spans), 1, [s.name for s in finished]) + find_op_span = find_op_spans[0] + self.assertEqual(find_op_span.name, "find pymongo_test.getmore_nesting") + self.assertTrue(find_op_span.attributes["db.mongodb.cursor_id"]) + + # One more, a sibling rather than a child, per getMore the caller drove. + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreater(len(getmore_op_spans), 1) + for op_span in getmore_op_spans: + self.assertEqual(op_span.name, "getMore pymongo_test.getmore_nesting") + self.assertNotEqual(op_span.parent, find_op_span.context) + self.assertEqual( + op_span.attributes["db.mongodb.cursor_id"], + find_op_span.attributes["db.mongodb.cursor_id"], + ) + + # Each getMore command span nests under its own operation span. + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), len(getmore_op_spans)) + parent_ids = {s.context.span_id for s in getmore_op_spans} + for cmd_span in getmore_cmd_spans: + self.assertIn(cmd_span.parent.span_id, parent_ids) + + async def test_caller_driven_aggregate_getmores_get_their_own_operation_spans(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.agg_nesting + await coll.drop() + await coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + docs = await (await coll.aggregate([{"$match": {}}], batchSize=2)).to_list() + self.assertEqual(len(docs), 10) + + finished = self.exporter.get_finished_spans() + agg_op_spans = self.operation_spans(finished, "aggregate") + self.assertEqual(len(agg_op_spans), 1, [s.name for s in finished]) + agg_op_span = agg_op_spans[0] + + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreater(len(getmore_op_spans), 1) + for op_span in getmore_op_spans: + self.assertEqual(op_span.name, "getMore pymongo_test.agg_nesting") + self.assertNotEqual(op_span.parent, agg_op_span.context) + + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), len(getmore_op_spans)) + parent_ids = {s.context.span_id for s in getmore_op_spans} + for cmd_span in getmore_cmd_spans: + self.assertIn(cmd_span.parent.span_id, parent_ids) + + async def test_internal_iteration_keeps_getmores_in_one_operation_span(self): + # list_collection_names drains its own listCollections cursor to build + # its return value, so the whole call is one operation and its getMores + # get no operation spans of their own. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test_internal_iteration + for i in range(6): + await db[f"coll{i}"].insert_one({}) + self.addAsyncCleanup(client.drop_database, db.name) + self.exporter.clear() + + names = await db.list_collection_names(cursor={"batchSize": 2}) + self.assertEqual(len(names), 6) + + finished = self.exporter.get_finished_spans() + op_spans = self.operation_spans(finished, "listCollections") + self.assertEqual(len(op_spans), 1, [s.name for s in finished]) + op_span = op_spans[0] + self.assertEqual(self.operation_spans(finished, "getMore"), []) + + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertGreater(len(getmore_cmd_spans), 0) + for cmd_span in getmore_cmd_spans: + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + async def test_single_batch_aggregate_ends_span_promptly_not_at_gc(self): + # A command cursor whose first batch exhausts it is marked _killed in + # __init__ without ever calling close(); no getMore is sent, so + # _refresh()/_die_lock() never run. Without explicit attachment its + # operation span would only be ended by __del__, i.e. whenever GC + # happens to run (or never, if the cursor is retained; Important #2). + # Assert the span is already ended while a reference to the + # cursor is still held, proving it ended at construction rather than + # waiting on GC. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.agg_single_batch + await coll.drop() + await coll.insert_many([{"i": i} for i in range(3)]) + self.exporter.clear() + + cursor = await coll.aggregate([{"$match": {}}]) + # Confirm this test actually exercises the single-batch path. + self.assertTrue(cursor._killed) + + finished = self.exporter.get_finished_spans() + agg_op_spans = [ + s + for s in finished + if s.attributes.get("db.operation.name") == "aggregate" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(agg_op_spans), 1, [s.name for s in finished]) + self.assertIsNotNone(agg_op_spans[0].end_time) + + # The cursor reference is kept alive through this assertion: if + # __del__ were the only thing ending the span, get_finished_spans() + # above would not have included it yet. + self.assertIsNotNone(cursor) + + async def test_abandoned_cursor_still_ends_operation_span(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.getmore_abandoned + await coll.drop() + await coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + cursor = coll.find({}, batch_size=2) + await cursor.next() # Leaves the cursor open with batches pending. + del cursor + gc.collect() + + find_op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "find" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(find_op_spans), 1) + + async def test_prose_3_get_more_records_sent_cursor_id_not_returned_cursor_id(self): + """Prose Test 3: getMore records the cursor id it sent, not the cursor id returned.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.prose3_getmore_cursor_id + await coll.drop() + await coll.insert_many([{"i": i} for i in range(3)]) + self.exporter.clear() + + cursor = coll.find({}, batch_size=2) + # Drain exactly the first batch (2 docs) without triggering a getMore + # yet, so the cursor id read here is the one `find` returned - the id + # about to be sent in the upcoming getMore - and not whatever that + # getMore's reply comes back with. + await cursor.next() + await cursor.next() + sent_cursor_id = cursor.cursor_id + self.assertIsNotNone(sent_cursor_id) + self.assertNotEqual(sent_cursor_id, 0) + + # Draining the last document sends exactly one getMore, which + # exhausts the cursor: the server's reply to that getMore returns a + # cursor id of 0. + remaining = await cursor.to_list() + self.assertEqual(len(remaining), 1) + self.assertEqual(cursor.cursor_id, 0) + + finished = self.exporter.get_finished_spans() + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), 1, [s.name for s in finished]) + + # Both the getMore operation span and the getMore command span must + # carry the id the driver sent, never the 0 the server's reply + # returned. + for span in (getmore_op_spans[0], getmore_cmd_spans[0]): + self.assertIn("db.mongodb.cursor_id", span.attributes) + self.assertNotEqual(span.attributes["db.mongodb.cursor_id"], 0) + self.assertEqual(span.attributes["db.mongodb.cursor_id"], sent_cursor_id) + + @async_client_context.require_transactions + async def test_prose_4_get_more_in_transaction_nests_under_transaction_span(self): + """Prose Test 4: getMore inside a transaction nests under the transaction span.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.prose4_getmore_in_txn + await coll.drop() + # Inserted outside the transaction, so the transaction below contains + # only the find and getMore. + await coll.insert_many([{"i": i} for i in range(3)]) + + async def callback(session): + docs = await coll.find({}, batch_size=2, session=session).to_list() + self.assertEqual(len(docs), 3) + + self.exporter.clear() + async with client.start_session() as session: + await session.with_transaction(callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + txn_span = txn_spans[0] + + find_op_spans = self.operation_spans(finished, "find") + self.assertEqual(len(find_op_spans), 1, [s.name for s in finished]) + find_op_span = find_op_spans[0] + + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_op_span = getmore_op_spans[0] + + # Both operation spans must nest directly under the transaction span... + self.assertEqual(find_op_span.parent.span_id, txn_span.context.span_id) + self.assertEqual(getmore_op_span.parent.span_id, txn_span.context.span_id) + # ...as siblings of each other, not one nested under the other. + self.assertNotEqual(getmore_op_span.parent.span_id, find_op_span.context.span_id) + + async def test_getmore_over_a_command_namespace_omits_the_collection(self): + """A cursor opened by a command targets no user collection. + + listCollections runs against the ".$cmd.listCollections" namespace, + and its getMore carries that in the command's "collection" field. That + is not a user collection, so per the spec db.collection.name is omitted + and the span is named "getMore " rather than + "getMore .$cmd.listCollections". + """ + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test_cmd_ns + await client.drop_database(db.name) + # Two collections with a batch size of one guarantees exactly one getMore. + await db.coll_one.insert_one({"x": 1}) + await db.coll_two.insert_one({"x": 1}) + self.exporter.clear() + + await (await db.list_collections(cursor={"batchSize": 1})).to_list() + + finished = self.exporter.get_finished_spans() + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreaterEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertGreaterEqual(len(getmore_cmd_spans), 1, [s.name for s in finished]) + + for span in getmore_op_spans + getmore_cmd_spans: + self.assertNotIn("db.collection.name", span.attributes) + # The operation span is named " " when no collection is + # targeted; the command span is named after the command alone, and + # carries the same " " form in db.query.summary. + for span in getmore_op_spans: + self.assertEqual(span.name, f"getMore {db.name}") + self.assertEqual(span.attributes["db.operation.summary"], f"getMore {db.name}") + for span in getmore_cmd_spans: + self.assertEqual(span.name, "getMore") + self.assertEqual(span.attributes["db.query.summary"], f"getMore {db.name}") + + await client.drop_database(db.name) + + @async_client_context.require_version_min(8, 0) + async def test_client_bulk_write_results_cursor_getmores_nest_under_bulk_write(self): + # A successful InsertOne's verbose result doc is tiny (~{"ok": 1, "idx": + # i, "n": 1}) regardless of the inserted document's size, and the driver + # never sends more than maxWriteBatchSize (100_000 by default) ops in one + # bulkWrite command, so plain successful inserts can never make the + # results cursor's first batch exceed the 16MB per-batch limit, no + # matter how many operations are given. Duplicate-key write errors, + # whose result docs embed the offending key (here padded to 3000 bytes), + # blow past that limit at a much smaller, fast-running operation count + # while still exercising the exact same code path (a real + # AsyncCommandCursor built and iterated by _process_results_cursor). + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.bulk_results_cursor + await coll.drop() + await coll.create_index("dup", unique=True) + dup_value = "d" * 3000 + models = [ + InsertOne(namespace=coll.full_name, document={"dup": dup_value}) for _ in range(10000) + ] + self.exporter.clear() + with self.assertRaises(ClientBulkWriteException): + await client.bulk_write(models, verbose_results=True, ordered=False) + + finished = self.exporter.get_finished_spans() + # Exactly one operation span, for the bulkWrite itself. + op_spans = [ + s + for s in finished + if "db.command.name" not in s.attributes + and s.attributes.get("db.operation.name") is not None + ] + self.assertEqual( + [s.attributes["db.operation.name"] for s in op_spans], + ["bulkWrite"], + [s.name for s in finished], + ) + (op_span,) = op_spans + + # Any getMore command spans parent directly to the bulkWrite span. + getmore_cmd_spans = [ + s for s in finished if s.attributes.get("db.command.name") == "getMore" + ] + self.assertGreater(len(getmore_cmd_spans), 0, "expected a multi-batch results cursor") + for cmd_span in getmore_cmd_spans: + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + async def test_caller_owned_operation_telemetry_is_not_ended_by_retry_internal(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + telemetry = _OperationTelemetry( + client.options.tracing, + "find", + None, + dbname="mydb", + collection="c", + set_current=False, + ) + self.exporter.clear() + + async def _noop_read(_session, _server, _conn, _read_pref): + return "ok" + + result = await client._retryable_read( + _noop_read, + ReadPreference.PRIMARY, + None, + operation="find", + operation_telemetry=telemetry, + ) + self.assertEqual(result, "ok") + # _retry_internal must not have ended the caller's span. + self.assertEqual( + [s for s in self.exporter.get_finished_spans() if s.name.startswith("find")], [] + ) + telemetry.succeeded() + self.assertEqual( + len([s for s in self.exporter.get_finished_spans() if s.name.startswith("find")]), + 1, + ) + + @async_client_context.require_version_min(4, 2, 0) + @async_client_context.require_change_streams + async def test_change_stream_collection_level_operation_span_has_full_namespace(self): + # A collection-level change stream's operation span carries both the + # database and the collection, derived from the aggregate command by + # _otel's lazy backfill. The database- and cluster-level cases below + # must omit db.collection.name, since neither targets one collection. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + coll = db.test_otel_change_stream_coll + await coll.drop() + self.exporter.clear() + async with await coll.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "pymongo_test") + self.assertEqual(span.attributes["db.collection.name"], "test_otel_change_stream_coll") + + @async_client_context.require_version_min(4, 2, 0) + @async_client_context.require_change_streams + async def test_change_stream_database_level_operation_span_omits_collection_name(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + self.exporter.clear() + async with await db.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "pymongo_test") + self.assertNotIn("db.collection.name", span.attributes) + + @async_client_context.require_version_min(4, 2, 0) + @async_client_context.require_change_streams + async def test_change_stream_cluster_level_operation_span_targets_admin(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + async with await client.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "admin") + self.assertNotIn("db.collection.name", span.attributes) diff --git a/test/open_telemetry/operation/get_more.json b/test/open_telemetry/operation/get_more.json new file mode 100644 index 0000000000..aa13a0afdd --- /dev/null +++ b/test/open_telemetry/operation/get_more.json @@ -0,0 +1,318 @@ +{ + "description": "operation getMore", + "schemaVersion": "1.27", + "createEntities": [ + { + "client": { + "id": "client0", + "useMultipleMongoses": false, + "observeTracingMessages": {} + } + }, + { + "database": { + "id": "database0", + "client": "client0", + "databaseName": "operation-get-more" + } + }, + { + "collection": { + "id": "collection0", + "database": "database0", + "collectionName": "test" + } + } + ], + "initialData": [ + { + "collectionName": "test", + "databaseName": "operation-get-more", + "documents": [ + { + "_id": 1 + }, + { + "_id": 2 + }, + { + "_id": 3 + } + ] + } + ], + "tests": [ + { + "description": "getMore is nested under its own operation span", + "operations": [ + { + "name": "createFindCursor", + "object": "collection0", + "arguments": { + "filter": {}, + "batchSize": 2 + }, + "saveResultAsEntity": "cursor0" + }, + { + "name": "iterateUntilDocumentOrError", + "object": "cursor0", + "expectResult": { + "_id": 1 + } + }, + { + "name": "iterateUntilDocumentOrError", + "object": "cursor0", + "expectResult": { + "_id": 2 + } + }, + { + "name": "iterateUntilDocumentOrError", + "object": "cursor0", + "expectResult": { + "_id": 3 + } + } + ], + "expectTracingMessages": [ + { + "client": "client0", + "ignoreExtraSpans": false, + "spans": [ + { + "name": "find operation-get-more.test", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.operation.name": "find", + "db.operation.summary": "find operation-get-more.test", + "db.mongodb.cursor_id": { + "$$type": [ + "int", + "long" + ] + } + }, + "nested": [ + { + "name": "find", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.command.name": "find", + "network.transport": "tcp", + "db.response.status_code": { + "$$exists": false + }, + "exception.message": { + "$$exists": false + }, + "exception.type": { + "$$exists": false + }, + "exception.stacktrace": { + "$$exists": false + }, + "server.address": { + "$$type": "string" + }, + "server.port": { + "$$type": [ + "int", + "long" + ] + }, + "db.query.summary": "find operation-get-more.test", + "db.mongodb.cursor_id": { + "$$type": [ + "int", + "long" + ] + }, + "db.mongodb.server_connection_id": { + "$$type": [ + "int", + "long" + ] + }, + "db.mongodb.driver_connection_id": { + "$$type": [ + "int", + "long" + ] + } + } + } + ] + }, + { + "name": "getMore operation-get-more.test", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.operation.name": "getMore", + "db.operation.summary": "getMore operation-get-more.test", + "db.mongodb.cursor_id": { + "$$type": [ + "int", + "long" + ] + } + }, + "nested": [ + { + "name": "getMore", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.command.name": "getMore", + "network.transport": "tcp", + "db.response.status_code": { + "$$exists": false + }, + "exception.message": { + "$$exists": false + }, + "exception.type": { + "$$exists": false + }, + "exception.stacktrace": { + "$$exists": false + }, + "server.address": { + "$$type": "string" + }, + "server.port": { + "$$type": [ + "int", + "long" + ] + }, + "db.query.summary": "getMore operation-get-more.test", + "db.mongodb.cursor_id": { + "$$type": [ + "int", + "long" + ] + }, + "db.mongodb.server_connection_id": { + "$$type": [ + "int", + "long" + ] + }, + "db.mongodb.driver_connection_id": { + "$$type": [ + "int", + "long" + ] + } + } + } + ] + } + ] + } + ] + }, + { + "description": "cursor_id is omitted when the server returns cursor id 0", + "operations": [ + { + "name": "find", + "object": "collection0", + "arguments": { + "filter": {} + }, + "expectResult": [ + { + "_id": 1 + }, + { + "_id": 2 + }, + { + "_id": 3 + } + ] + } + ], + "expectTracingMessages": [ + { + "client": "client0", + "ignoreExtraSpans": false, + "spans": [ + { + "name": "find operation-get-more.test", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.operation.name": "find", + "db.operation.summary": "find operation-get-more.test", + "db.mongodb.cursor_id": { + "$$exists": false + } + }, + "nested": [ + { + "name": "find", + "attributes": { + "db.system.name": "mongodb", + "db.namespace": "operation-get-more", + "db.collection.name": "test", + "db.command.name": "find", + "network.transport": "tcp", + "db.response.status_code": { + "$$exists": false + }, + "exception.message": { + "$$exists": false + }, + "exception.type": { + "$$exists": false + }, + "exception.stacktrace": { + "$$exists": false + }, + "server.address": { + "$$type": "string" + }, + "server.port": { + "$$type": [ + "int", + "long" + ] + }, + "db.query.summary": "find operation-get-more.test", + "db.mongodb.cursor_id": { + "$$exists": false + }, + "db.mongodb.server_connection_id": { + "$$type": [ + "int", + "long" + ] + }, + "db.mongodb.driver_connection_id": { + "$$type": [ + "int", + "long" + ] + } + } + } + ] + } + ] + } + ] + } + ] +} diff --git a/test/test_otel_getmore.py b/test/test_otel_getmore.py new file mode 100644 index 0000000000..1030537c01 --- /dev/null +++ b/test/test_otel_getmore.py @@ -0,0 +1,533 @@ +# 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. + +"""Test OpenTelemetry operation spans for cursor getMores.""" + +from __future__ import annotations + +import gc +import os +import sys +from typing import Optional +from unittest.mock import patch + +sys.path[0:0] = [""] + +import pytest + +import pymongo._otel as _otel +from pymongo import _telemetry, common +from pymongo._telemetry import _OperationTelemetry +from pymongo.errors import ( + ClientBulkWriteException, + ConfigurationError, + InvalidOperation, + OperationFailure, + ServerSelectionTimeoutError, +) +from pymongo.logger import _HELLO_COMMANDS +from pymongo.operations import InsertOne +from pymongo.read_preferences import ReadPreference +from pymongo.typings import _Address +from test import IntegrationTest, client_context, unittest +from test.unified_format_shared import _shared_test_provider +from test.utils import wait_until + +_HAS_OTEL_TEST_DEPS = False +if _otel._HAS_OPENTELEMETRY: + try: + from opentelemetry import trace + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + _HAS_OTEL_TEST_DEPS = True + except ImportError: + pass + +_IS_SYNC = True + +pytestmark = pytest.mark.otel + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelGetMoreSpans(IntegrationTest): + """getMore spans, cursor lifetime, and change streams.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # See the matching comment in test/synchronous/unified_format.py's + # UnifiedSpecTestMixinV1.tearDownClass: the span processor can never + # be removed from the shared process-wide TracerProvider, so without + # this shutdown() the exporter keeps accumulating every span from + # every client for the rest of the test run. + cls.exporter.shutdown() + super().tearDownClass() + + def setUp(self): + super().setUp() + self.exporter.clear() + + def spans(self, name: str | None = None): + finished = self.exporter.get_finished_spans() + if name is None: + return list(finished) + return [s for s in finished if s.name == name] + + @staticmethod + def operation_spans(finished, operation: str): + """Return the operation spans for ``operation``, excluding command spans. + + Only command spans carry db.command.name, so its absence is what tells + the two kinds apart when both name the same operation. + """ + return [ + s + for s in finished + if s.attributes.get("db.operation.name") == operation + and "db.command.name" not in s.attributes + ] + + @staticmethod + def command_spans(finished, command: str): + """Return the command spans for ``command``.""" + return [s for s in finished if s.attributes.get("db.command.name") == command] + + def ping_spans(self): + """Return the spans belonging to a ``ping`` run through ``db.command()``. + + For the tests that assert tracing produced *nothing*. Asserting the + exporter is empty would also catch spans no test asked for: a cursor + abandoned earlier in the class ends its operation span from a + finalizer, and on an interpreter that does not reference count, that + finalizer runs at an unpredictable point and lands in whichever test + happens to be running. Naming the ping's own spans keeps the assertion + about this client while staying immune to that. + """ + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "ping" + or s.attributes.get("db.operation.name") == "runCommand" + ] + + def _aggregate_operation_span(self): + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "aggregate" + ] + self.assertEqual(len(matching), 1) + return matching[0] + + def test_span_created_for_get_more(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name].test_otel_getmore + coll.drop() + coll.insert_many([{"x": i} for i in range(5)]) + self.exporter.clear() + + docs = coll.find({}, batch_size=2).to_list() + self.assertEqual(len(docs), 5) + + get_more_spans = self.spans("getMore") + self.assertGreater(len(get_more_spans), 0) + for span in get_more_spans: + self.assertEqual(span.attributes["db.collection.name"], "test_otel_getmore") + self.assertEqual(span.attributes["db.command.name"], "getMore") + + def test_caller_driven_find_getmores_get_their_own_operation_spans(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.getmore_nesting + coll.drop() + coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + docs = coll.find({}, batch_size=2).to_list() + self.assertEqual(len(docs), 10) + + finished = self.exporter.get_finished_spans() + # One operation span for the query that created the cursor. + find_op_spans = self.operation_spans(finished, "find") + self.assertEqual(len(find_op_spans), 1, [s.name for s in finished]) + find_op_span = find_op_spans[0] + self.assertEqual(find_op_span.name, "find pymongo_test.getmore_nesting") + self.assertTrue(find_op_span.attributes["db.mongodb.cursor_id"]) + + # One more, a sibling rather than a child, per getMore the caller drove. + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreater(len(getmore_op_spans), 1) + for op_span in getmore_op_spans: + self.assertEqual(op_span.name, "getMore pymongo_test.getmore_nesting") + self.assertNotEqual(op_span.parent, find_op_span.context) + self.assertEqual( + op_span.attributes["db.mongodb.cursor_id"], + find_op_span.attributes["db.mongodb.cursor_id"], + ) + + # Each getMore command span nests under its own operation span. + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), len(getmore_op_spans)) + parent_ids = {s.context.span_id for s in getmore_op_spans} + for cmd_span in getmore_cmd_spans: + self.assertIn(cmd_span.parent.span_id, parent_ids) + + def test_caller_driven_aggregate_getmores_get_their_own_operation_spans(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.agg_nesting + coll.drop() + coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + docs = (coll.aggregate([{"$match": {}}], batchSize=2)).to_list() + self.assertEqual(len(docs), 10) + + finished = self.exporter.get_finished_spans() + agg_op_spans = self.operation_spans(finished, "aggregate") + self.assertEqual(len(agg_op_spans), 1, [s.name for s in finished]) + agg_op_span = agg_op_spans[0] + + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreater(len(getmore_op_spans), 1) + for op_span in getmore_op_spans: + self.assertEqual(op_span.name, "getMore pymongo_test.agg_nesting") + self.assertNotEqual(op_span.parent, agg_op_span.context) + + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), len(getmore_op_spans)) + parent_ids = {s.context.span_id for s in getmore_op_spans} + for cmd_span in getmore_cmd_spans: + self.assertIn(cmd_span.parent.span_id, parent_ids) + + def test_internal_iteration_keeps_getmores_in_one_operation_span(self): + # list_collection_names drains its own listCollections cursor to build + # its return value, so the whole call is one operation and its getMores + # get no operation spans of their own. + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test_internal_iteration + for i in range(6): + db[f"coll{i}"].insert_one({}) + self.addCleanup(client.drop_database, db.name) + self.exporter.clear() + + names = db.list_collection_names(cursor={"batchSize": 2}) + self.assertEqual(len(names), 6) + + finished = self.exporter.get_finished_spans() + op_spans = self.operation_spans(finished, "listCollections") + self.assertEqual(len(op_spans), 1, [s.name for s in finished]) + op_span = op_spans[0] + self.assertEqual(self.operation_spans(finished, "getMore"), []) + + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertGreater(len(getmore_cmd_spans), 0) + for cmd_span in getmore_cmd_spans: + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + def test_single_batch_aggregate_ends_span_promptly_not_at_gc(self): + # A command cursor whose first batch exhausts it is marked _killed in + # __init__ without ever calling close(); no getMore is sent, so + # _refresh()/_die_lock() never run. Without explicit attachment its + # operation span would only be ended by __del__, i.e. whenever GC + # happens to run (or never, if the cursor is retained; Important #2). + # Assert the span is already ended while a reference to the + # cursor is still held, proving it ended at construction rather than + # waiting on GC. + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.agg_single_batch + coll.drop() + coll.insert_many([{"i": i} for i in range(3)]) + self.exporter.clear() + + cursor = coll.aggregate([{"$match": {}}]) + # Confirm this test actually exercises the single-batch path. + self.assertTrue(cursor._killed) + + finished = self.exporter.get_finished_spans() + agg_op_spans = [ + s + for s in finished + if s.attributes.get("db.operation.name") == "aggregate" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(agg_op_spans), 1, [s.name for s in finished]) + self.assertIsNotNone(agg_op_spans[0].end_time) + + # The cursor reference is kept alive through this assertion: if + # __del__ were the only thing ending the span, get_finished_spans() + # above would not have included it yet. + self.assertIsNotNone(cursor) + + def test_abandoned_cursor_still_ends_operation_span(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.getmore_abandoned + coll.drop() + coll.insert_many([{"i": i} for i in range(10)]) + self.exporter.clear() + + cursor = coll.find({}, batch_size=2) + cursor.next() # Leaves the cursor open with batches pending. + del cursor + gc.collect() + + find_op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "find" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(find_op_spans), 1) + + def test_prose_3_get_more_records_sent_cursor_id_not_returned_cursor_id(self): + """Prose Test 3: getMore records the cursor id it sent, not the cursor id returned.""" + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.prose3_getmore_cursor_id + coll.drop() + coll.insert_many([{"i": i} for i in range(3)]) + self.exporter.clear() + + cursor = coll.find({}, batch_size=2) + # Drain exactly the first batch (2 docs) without triggering a getMore + # yet, so the cursor id read here is the one `find` returned - the id + # about to be sent in the upcoming getMore - and not whatever that + # getMore's reply comes back with. + cursor.next() + cursor.next() + sent_cursor_id = cursor.cursor_id + self.assertIsNotNone(sent_cursor_id) + self.assertNotEqual(sent_cursor_id, 0) + + # Draining the last document sends exactly one getMore, which + # exhausts the cursor: the server's reply to that getMore returns a + # cursor id of 0. + remaining = cursor.to_list() + self.assertEqual(len(remaining), 1) + self.assertEqual(cursor.cursor_id, 0) + + finished = self.exporter.get_finished_spans() + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertEqual(len(getmore_cmd_spans), 1, [s.name for s in finished]) + + # Both the getMore operation span and the getMore command span must + # carry the id the driver sent, never the 0 the server's reply + # returned. + for span in (getmore_op_spans[0], getmore_cmd_spans[0]): + self.assertIn("db.mongodb.cursor_id", span.attributes) + self.assertNotEqual(span.attributes["db.mongodb.cursor_id"], 0) + self.assertEqual(span.attributes["db.mongodb.cursor_id"], sent_cursor_id) + + @client_context.require_transactions + def test_prose_4_get_more_in_transaction_nests_under_transaction_span(self): + """Prose Test 4: getMore inside a transaction nests under the transaction span.""" + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.prose4_getmore_in_txn + coll.drop() + # Inserted outside the transaction, so the transaction below contains + # only the find and getMore. + coll.insert_many([{"i": i} for i in range(3)]) + + def callback(session): + docs = coll.find({}, batch_size=2, session=session).to_list() + self.assertEqual(len(docs), 3) + + self.exporter.clear() + with client.start_session() as session: + session.with_transaction(callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + txn_span = txn_spans[0] + + find_op_spans = self.operation_spans(finished, "find") + self.assertEqual(len(find_op_spans), 1, [s.name for s in finished]) + find_op_span = find_op_spans[0] + + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_op_span = getmore_op_spans[0] + + # Both operation spans must nest directly under the transaction span... + self.assertEqual(find_op_span.parent.span_id, txn_span.context.span_id) + self.assertEqual(getmore_op_span.parent.span_id, txn_span.context.span_id) + # ...as siblings of each other, not one nested under the other. + self.assertNotEqual(getmore_op_span.parent.span_id, find_op_span.context.span_id) + + def test_getmore_over_a_command_namespace_omits_the_collection(self): + """A cursor opened by a command targets no user collection. + + listCollections runs against the ".$cmd.listCollections" namespace, + and its getMore carries that in the command's "collection" field. That + is not a user collection, so per the spec db.collection.name is omitted + and the span is named "getMore " rather than + "getMore .$cmd.listCollections". + """ + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test_cmd_ns + client.drop_database(db.name) + # Two collections with a batch size of one guarantees exactly one getMore. + db.coll_one.insert_one({"x": 1}) + db.coll_two.insert_one({"x": 1}) + self.exporter.clear() + + (db.list_collections(cursor={"batchSize": 1})).to_list() + + finished = self.exporter.get_finished_spans() + getmore_op_spans = self.operation_spans(finished, "getMore") + self.assertGreaterEqual(len(getmore_op_spans), 1, [s.name for s in finished]) + getmore_cmd_spans = self.command_spans(finished, "getMore") + self.assertGreaterEqual(len(getmore_cmd_spans), 1, [s.name for s in finished]) + + for span in getmore_op_spans + getmore_cmd_spans: + self.assertNotIn("db.collection.name", span.attributes) + # The operation span is named " " when no collection is + # targeted; the command span is named after the command alone, and + # carries the same " " form in db.query.summary. + for span in getmore_op_spans: + self.assertEqual(span.name, f"getMore {db.name}") + self.assertEqual(span.attributes["db.operation.summary"], f"getMore {db.name}") + for span in getmore_cmd_spans: + self.assertEqual(span.name, "getMore") + self.assertEqual(span.attributes["db.query.summary"], f"getMore {db.name}") + + client.drop_database(db.name) + + @client_context.require_version_min(8, 0) + def test_client_bulk_write_results_cursor_getmores_nest_under_bulk_write(self): + # A successful InsertOne's verbose result doc is tiny (~{"ok": 1, "idx": + # i, "n": 1}) regardless of the inserted document's size, and the driver + # never sends more than maxWriteBatchSize (100_000 by default) ops in one + # bulkWrite command, so plain successful inserts can never make the + # results cursor's first batch exceed the 16MB per-batch limit, no + # matter how many operations are given. Duplicate-key write errors, + # whose result docs embed the offending key (here padded to 3000 bytes), + # blow past that limit at a much smaller, fast-running operation count + # while still exercising the exact same code path (a real + # CommandCursor built and iterated by _process_results_cursor). + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.bulk_results_cursor + coll.drop() + coll.create_index("dup", unique=True) + dup_value = "d" * 3000 + models = [ + InsertOne(namespace=coll.full_name, document={"dup": dup_value}) for _ in range(10000) + ] + self.exporter.clear() + with self.assertRaises(ClientBulkWriteException): + client.bulk_write(models, verbose_results=True, ordered=False) + + finished = self.exporter.get_finished_spans() + # Exactly one operation span, for the bulkWrite itself. + op_spans = [ + s + for s in finished + if "db.command.name" not in s.attributes + and s.attributes.get("db.operation.name") is not None + ] + self.assertEqual( + [s.attributes["db.operation.name"] for s in op_spans], + ["bulkWrite"], + [s.name for s in finished], + ) + (op_span,) = op_spans + + # Any getMore command spans parent directly to the bulkWrite span. + getmore_cmd_spans = [ + s for s in finished if s.attributes.get("db.command.name") == "getMore" + ] + self.assertGreater(len(getmore_cmd_spans), 0, "expected a multi-batch results cursor") + for cmd_span in getmore_cmd_spans: + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + def test_caller_owned_operation_telemetry_is_not_ended_by_retry_internal(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + telemetry = _OperationTelemetry( + client.options.tracing, + "find", + None, + dbname="mydb", + collection="c", + set_current=False, + ) + self.exporter.clear() + + def _noop_read(_session, _server, _conn, _read_pref): + return "ok" + + result = client._retryable_read( + _noop_read, + ReadPreference.PRIMARY, + None, + operation="find", + operation_telemetry=telemetry, + ) + self.assertEqual(result, "ok") + # _retry_internal must not have ended the caller's span. + self.assertEqual( + [s for s in self.exporter.get_finished_spans() if s.name.startswith("find")], [] + ) + telemetry.succeeded() + self.assertEqual( + len([s for s in self.exporter.get_finished_spans() if s.name.startswith("find")]), + 1, + ) + + @client_context.require_version_min(4, 2, 0) + @client_context.require_change_streams + def test_change_stream_collection_level_operation_span_has_full_namespace(self): + # A collection-level change stream's operation span carries both the + # database and the collection, derived from the aggregate command by + # _otel's lazy backfill. The database- and cluster-level cases below + # must omit db.collection.name, since neither targets one collection. + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + coll = db.test_otel_change_stream_coll + coll.drop() + self.exporter.clear() + with coll.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "pymongo_test") + self.assertEqual(span.attributes["db.collection.name"], "test_otel_change_stream_coll") + + @client_context.require_version_min(4, 2, 0) + @client_context.require_change_streams + def test_change_stream_database_level_operation_span_omits_collection_name(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + self.exporter.clear() + with db.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "pymongo_test") + self.assertNotIn("db.collection.name", span.attributes) + + @client_context.require_version_min(4, 2, 0) + @client_context.require_change_streams + def test_change_stream_cluster_level_operation_span_targets_admin(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + with client.watch(): + pass + span = self._aggregate_operation_span() + self.assertEqual(span.attributes["db.namespace"], "admin") + self.assertNotIn("db.collection.name", span.attributes)