PYTHON-5993 Add OpenTelemetry operation spans for cursor getMores

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,15 +380,34 @@ 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:
         return
     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

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

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

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.

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 {}

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