PYTHON-5947 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.

Author: blink1073

pymongo/_otel.py

@@ -68,6 +68,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
 
@@ -123,6 +132,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.
 
@@ -369,6 +400,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:
@@ -378,8 +423,14 @@ def end_command_span_success(span: Optional[Span], reply: _DocumentOut) -> None:
         # A cursor id of 0 means the cursor is already exhausted, i.e. there is
         # no cursor left to track, so per the OTel spec the attribute is
         # omitted, never set to 0, when a cursor-creating command's reply
-        # returns 0.
-        span.set_attribute("db.mongodb.cursor_id", cursor["id"])
+        # returns 0. A getMore keeps the id it sent, set in start_command_span,
+        # which this deliberately does not overwrite with a 0 reply id.
+        cursor_id = cursor["id"]
+        span.set_attribute("db.mongodb.cursor_id", cursor_id)
+        # The enclosing operation span carries the same attribute. For a
+        # cursor-creating command that is this reply's id; for a getMore it is
+        # the id already set from the sent value, which this repeats unchanged.
+        _set_operation_cursor_id(cursor_id)
     span.end()
 
 
@@ -451,6 +502,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.
 
@@ -469,6 +521,11 @@ def start_operation_span(
     avoid a concurrently-running unrelated session's operations picking up
     this transaction by accident. Pass None outside of a transaction.
 
+    ``cursor_id`` sets ``db.mongodb.cursor_id`` up front, for an operation
+    reading a cursor that already exists: the id is known before the command is
+    even built, and the operation span needs it even if the operation fails
+    before any command span exists.
+
     With ``set_current=False`` the span is created but not made current, and
     the operation-name contextvar is left alone. That suits a span created
     outside the ``_retry_internal`` call it covers, where the caller makes it
@@ -489,6 +546,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()
 

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: