PYTHON-5993 Add OpenTelemetry operation spans for cursor getMores

pymongo/_otel.py

@@ -65,6 +65,12 @@
     "_CURRENT_OPERATION_NAME", default=None
 )
 
+# True while the driver is draining a cursor of its own to build one public API
+# call's return value. See internal_cursor_iteration.
+_INTERNAL_CURSOR_ITERATION: ContextVar[bool] = ContextVar(
+    "_INTERNAL_CURSOR_ITERATION", default=False
+)
+
 if TYPE_CHECKING:
     from opentelemetry.trace import Span, Tracer
 
@@ -127,6 +133,26 @@ 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.
+
+    Everything the block sends, every getMore included, belongs to the enclosing
+    method's one operation span. Outside such a block a getMore is assumed to be
+    caller-driven and gets an operation span of its own.
+    """
+    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.
 
@@ -295,6 +321,11 @@ def start_command_span(
         return None
 
     collection = _extract_collection_name(command_name, dbname, cmd)
+    # The id sent, not the reply's, which is 0 once the cursor is exhausted
+    # while the attribute is still required.
+    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
     # Runs per attempt, not per operation: an attempt that fails before building
     # a command never reaches here, so a retry may be where the span learns its
     # namespace.
@@ -308,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
@@ -329,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)
@@ -345,15 +380,29 @@ 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, since the current span could
+    otherwise be an unrelated one 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"])
+        # Omitted rather than set to 0, so a getMore keeps the id it sent.
+        cursor_id = cursor["id"]
+        span.set_attribute("db.mongodb.cursor_id", cursor_id)
+        _set_operation_cursor_id(cursor_id)
     span.end()
 
 
@@ -421,6 +470,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.
 
@@ -433,6 +483,9 @@ 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, since a getMore knows
+    it before the command is built and needs it even if the operation fails.
+
     ``set_current=False`` leaves the span and the operation-name contextvar
     alone, for a caller that makes it current with ``use_operation_span``.
     """
@@ -451,6 +504,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,9 @@ async def _run_aggregation_cmd(
             result_processor=self._process_result,
             comment=self._comment,
         )
+        # No operation span is attached to the resulting cursor: a change stream
+        # can tail indefinitely, so a span covering its whole lifetime would
+        # never end. Each getMore gets its own sibling span instead.
         return await self._client._retryable_read(
             cmd.get_cursor,
             self._target._read_preference_for(session),

pymongo/asynchronous/client_bulk.py

@@ -336,6 +336,9 @@ async def _process_results_cursor(
                 session=session,
                 comment=self.comment,
             )
+            # These getMores run inside the enclosing bulkWrite operation span,
+            # so a getMore operation span of their own would be spurious.
+            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,34 @@ 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 on every failure path and close()
+                # ends it once exhausted, leaving only this case.
+                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,10 @@ async def _refresh(self) -> int:
                 collection=self._collection.name,
                 set_current=False,
             )
-            await self._send_message_in_operation_span(q)
+            # The query's span stays open only when the call that created this
+            # cursor drains it itself, 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 +1115,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: