diff --git a/.evergreen/generated_configs/variants.yml b/.evergreen/generated_configs/variants.yml index de491b99be..bde14deada 100644 --- a/.evergreen/generated_configs/variants.yml +++ b/.evergreen/generated_configs/variants.yml @@ -445,7 +445,10 @@ buildvariants: # Otel tests - name: otel-rhel8 tasks: - - name: .test-non-standard .standalone-noauth-nossl + - name: .test-non-standard .replica_set-noauth-ssl + - name: .test-non-standard .sharded_cluster-auth-ssl .python-3.14 + - name: .test-non-standard .sharded_cluster-auth-ssl .python-pypy3.11 + - name: .test-non-standard .standalone-noauth-nossl .python-3.10 display_name: OTel RHEL8 run_on: - rhel87-small diff --git a/.evergreen/scripts/generate_config.py b/.evergreen/scripts/generate_config.py index 9dad14f286..a02bfc8043 100644 --- a/.evergreen/scripts/generate_config.py +++ b/.evergreen/scripts/generate_config.py @@ -458,7 +458,23 @@ def create_otel_variants(): expansions = dict(TEST_NAME="otel", COVERAGE="1") return [ create_variant( - [".test-non-standard .standalone-noauth-nossl"], + [ + # All three topologies, subset to keep the task count at 22. + # + # Replica set in full: the only topology where transaction spans + # run at all (they are skipped on standalone and sharded), and + # the only one covering free-threaded Python. + ".test-non-standard .replica_set-noauth-ssl", + # Sharded adds mongos, which rewrites commands and reports a + # different server.address, plus auth and ssl, which exercise + # sensitive-command redaction. Newest CPython across server + # versions, and PyPy for the alternate implementation. + ".test-non-standard .sharded_cluster-auth-ssl .python-3.14", + ".test-non-standard .sharded_cluster-auth-ssl .python-pypy3.11", + # Standalone only for its min-deps tasks, which resolve + # opentelemetry-api down to the floor in requirements/. + ".test-non-standard .standalone-noauth-nossl .python-3.10", + ], get_variant_name("OTel", host), host=host, tags=["pr"], diff --git a/.github/workflows/test-python.yml b/.github/workflows/test-python.yml index f665c5c694..22a814802b 100644 --- a/.github/workflows/test-python.yml +++ b/.github/workflows/test-python.yml @@ -305,3 +305,36 @@ jobs: source .venv/bin/activate uv pip install -e ".[test]" --resolution=lowest-direct --force-reinstall pytest -v test/test_srv_polling.py test/test_dns.py test/asynchronous/test_srv_polling.py test/asynchronous/test_dns.py + + # TEMPORARY, remove before merging the PYTHON-5947 stack. + # + # Evergreen only runs on PRs whose base is a configured branch, so the upper + # PRs in this stack (based on each other rather than on otel) get no Evergreen + # coverage at all, and the otel tests only run in Evergreen's otel variant. + # This job runs them on every branch in the stack instead. A replica set is + # requested because the transaction span tests require one. Actions are pinned + # to commit hashes here, unlike the rest of this file, to satisfy the + # mutable-action-tag scanners on newly added lines. + otel: + runs-on: ubuntu-latest + name: OTel Tests (temporary) + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + enable-cache: true + python-version: "3.11" + - id: setup-mongodb + uses: mongodb-labs/drivers-evergreen-tools@c70d29e2102cd048ad800932209776499553870c # v1.0.1 + with: + version: "8.0" + topology: "replica_set" + - name: Install just + run: uv tool install rust-just + - name: Setup tests + run: TEST_NAME=otel just setup-tests + - name: Run tests + run: just run-tests diff --git a/bson/json_util.py b/bson/json_util.py index e12e04ccb8..8b6948a76a 100644 --- a/bson/json_util.py +++ b/bson/json_util.py @@ -1115,20 +1115,18 @@ def _truncate_documents(obj: Any, max_length: int) -> tuple[Any, int]: if hasattr(obj, "items"): truncated: Any = {} for k, v in obj.items(): - truncated_v, remaining = _truncate_documents(v, remaining) - if truncated_v: - truncated[k] = truncated_v if remaining <= 0: break + truncated_v, remaining = _truncate_documents(v, remaining) + truncated[k] = truncated_v return truncated, remaining elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes)): truncated: Any = [] # type:ignore[no-redef] for v in obj: - truncated_v, remaining = _truncate_documents(v, remaining) - if truncated_v: - truncated.append(truncated_v) if remaining <= 0: break + truncated_v, remaining = _truncate_documents(v, remaining) + truncated.append(truncated_v) return truncated, remaining else: return _truncate(obj, remaining) diff --git a/doc/changelog.rst b/doc/changelog.rst index c5800602e7..728904fc89 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -29,11 +29,14 @@ PyMongo 4.18 brings a number of changes including: attempts, so consumers can correlate a retried operation's events. As a result, ``operation_id`` is no longer equal to the per-attempt ``request_id`` for these operations. -- Added optional OpenTelemetry command-span support, conforming to the +- Added optional OpenTelemetry tracing support, conforming to the `OpenTelemetry driver specification `_. - Enable it with the ``tracing`` :class:`~pymongo.mongo_client.MongoClient` - option or the ``OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED`` environment - variable. Install the ``opentelemetry-api`` package, or use the + Every public API call produces an operation span, which contains one span + per command sent to the server. Inside a transaction, those operation spans + nest under a ``transaction`` span. Enable it with the + ``tracing`` :class:`~pymongo.mongo_client.MongoClient` option or the + ``OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED`` environment variable. + Install the ``opentelemetry-api`` package, or use the ``pymongo[opentelemetry]`` extra, to enable this feature. - Fixed a potential out-of-bounds read in the C extension when decoding an array of BSON documents. An embedded document whose declared length exceeds diff --git a/pymongo/_otel.py b/pymongo/_otel.py index 713f963bbc..66d5a29153 100644 --- a/pymongo/_otel.py +++ b/pymongo/_otel.py @@ -12,17 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Optional OpenTelemetry command-span support. +"""Optional OpenTelemetry span support. Kept separate from :mod:`pymongo._telemetry` so that module stays free of ``opentelemetry`` import guards. Every function here is a no-op when ``opentelemetry`` isn't installed or tracing isn't enabled. + +This module also owns the specification's naming and attribute rules, such as +how span names, ``db.operation.name`` and ``db.query.summary`` are built. +:mod:`pymongo._telemetry` owns span lifecycles. A specification change to a +name or an attribute value stays here; one that changes when a span starts or +ends, or what it nests under, affects both. """ from __future__ import annotations +import contextlib +import enum import os -from collections.abc import Mapping, MutableMapping +import traceback +from collections.abc import Iterator, Mapping, MutableMapping +from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Optional, TypedDict from bson import json_util @@ -31,20 +41,30 @@ from pymongo.logger import _HELLO_COMMANDS, _JSON_OPTIONS, _SENSITIVE_COMMANDS try: - from opentelemetry import trace - from opentelemetry.trace import SpanKind, Status, StatusCode + from opentelemetry import context, trace # type:ignore[import-not-found,unused-ignore] + from opentelemetry.trace import ( # type:ignore[import-not-found,unused-ignore] + SpanKind, + Status, + StatusCode, + ) _HAS_OPENTELEMETRY = True - # Safe to cache at import time: opentelemetry.trace.get_tracer() returns a - # ProxyTracer when no real TracerProvider is registered yet, and that proxy - # transparently starts delegating to the real tracer once the application - # calls trace.set_tracer_provider() later, so this doesn't bind us to a - # permanently-inert no-op tracer. + # Safe to cache: get_tracer() returns a ProxyTracer when no provider is + # registered yet, and that proxy starts delegating once the application + # calls set_tracer_provider(), so this is not bound to a no-op tracer. _TRACER: Optional[Tracer] = trace.get_tracer("PyMongo", __version__) except ImportError: _HAS_OPENTELEMETRY = False _TRACER = None +# Name of the active operation span, so start_command_span can backfill that +# span's name and namespace attributes from the first command built inside it. +# Call sites that know the namespace up front pass it to start_operation_span; +# the generic retry path does not, and relies on this backfill. +_CURRENT_OPERATION_NAME: ContextVar[Optional[str]] = ContextVar( + "_CURRENT_OPERATION_NAME", default=None +) + if TYPE_CHECKING: from opentelemetry.trace import Span, Tracer @@ -55,9 +75,9 @@ class TracingOptions(TypedDict): """The shape of the ``MongoClient`` ``tracing`` option. - ``query_text_max_length`` is None when the client didn't configure it, so - the environment variable can be consulted; any explicit value (including - 0, to force ``db.query.text`` off) overrides the environment variable. + ``query_text_max_length`` is None as validated from user input; the options + a client holds have been through :func:`_resolve_tracing_options`, so both + fields are resolved. """ enabled: bool @@ -72,20 +92,24 @@ class TracingOptions(TypedDict): # from the equivalent CommandStartedEvent.command per the OpenTelemetry spec. _QUERY_TEXT_EXCLUDED_FIELDS = frozenset({"lsid", "$db", "$clusterTime", "signature"}) -# getMore's own command value is the cursor id, not the collection name; the -# collection lives under a separate "collection" key instead. -# See _gen_get_more_command in pymongo/message.py. +# getMore's command value is the cursor id, not a collection name; the collection +# lives under a separate "collection" key. See _gen_get_more_command in message.py. _GET_MORE = "getMore" -# explain wraps the real command (e.g. find/aggregate) rather than naming a -# collection directly: {"explain": {"find": "coll", ...}}. See _Query.as_command -# in pymongo/message.py. +# explain wraps the real command rather than naming a collection directly: +# {"explain": {"find": "coll", ...}}. See _Query.as_command in message.py. _EXPLAIN = "explain" # Commands against this database (e.g. user/role management, renameCollection) # never have a real collection name, even when their command value is a string. _ADMIN_DB = "admin" +# A cursor opened by a command rather than over a collection (listCollections, +# listIndexes, a database-level aggregate) has a namespace like +# "$cmd.listCollections", which names no user collection, so per the spec +# db.collection.name is omitted. +_CMD_NAMESPACE_PREFIX = "$cmd" + def _env_truthy(name: str) -> bool: """Return True if the environment variable ``name`` is set to "1", "true", or "yes".""" @@ -93,44 +117,47 @@ def _env_truthy(name: str) -> bool: def _is_tracing_enabled(tracing_options: Optional[TracingOptions]) -> bool: - """Return True if OTel command spans should be created for this client. + """Return True if spans should be created for this client. - The ``MongoClient`` ``tracing.enabled`` option and the - ``OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED`` environment variable both - gate enablement; either one being truthy is sufficient. + ClientOptions folds ``OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED`` into + ``tracing.enabled`` once at construction, so this is a lookup rather than + an os.environ read per command. None means there is no client to read the + option from, as for monitor and handshake connections, which are never + traced. """ - if not _HAS_OPENTELEMETRY: - return False - if tracing_options and tracing_options.get("enabled"): - return True - return _env_truthy(_OTEL_ENABLED_ENV) + return _HAS_OPENTELEMETRY and tracing_options is not None and tracing_options["enabled"] -def _get_query_text_max_length(tracing_options: Optional[TracingOptions]) -> int: - """Return the configured db.query.text truncation length, or 0 to omit the attribute. +def _resolve_tracing_options(tracing_options: TracingOptions) -> TracingOptions: + """Fold both environment variables into a client's validated tracing options. - An explicit client value (including 0) always wins; the environment - variable is only consulted when the client didn't configure it at all. + Called once when the client is built, so nothing re-reads the environment + per command. An explicit client value wins, including a + ``query_text_max_length`` of 0, which turns ``db.query.text`` off. """ - client_value = tracing_options.get("query_text_max_length") if tracing_options else None - if client_value is not None: - return max(0, client_value) - try: - return max(0, int(os.getenv(_OTEL_QUERY_TEXT_MAX_LENGTH_ENV, "0"))) - except ValueError: - return 0 + max_length = tracing_options["query_text_max_length"] + if max_length is None: + try: + max_length = int(os.getenv(_OTEL_QUERY_TEXT_MAX_LENGTH_ENV, "0")) + except ValueError: + max_length = 0 + return { + "enabled": tracing_options["enabled"] or _env_truthy(_OTEL_ENABLED_ENV), + "query_text_max_length": max(0, max_length), + } + + +def _get_query_text_max_length(tracing_options: Optional[TracingOptions]) -> int: + """Return the db.query.text truncation length, or 0 to omit the attribute.""" + return (tracing_options["query_text_max_length"] or 0) if tracing_options else 0 def _build_query_text(cmd: Mapping[str, Any], max_length: int) -> str: """Serialize ``cmd`` to extended JSON, redacted and truncated to ``max_length``. - Mirrors the truncation approach used for log messages: truncate field - values first, which usually keeps the result well-formed JSON (unlike a - blind cut of the fully-serialized string), then fall back to a hard - string cut as a safety net for whatever the field truncation's size - estimate still leaves over ``max_length``. The "..." marker is carved out - of the budget (not appended on top of it) so the result never exceeds - ``max_length``. + Mirrors log-message truncation: shorten field values first, which usually + keeps the result valid JSON, then hard-cut as a safety net. The "..." + marker comes out of the budget, so the result never exceeds ``max_length``. """ filtered = {k: v for k, v in cmd.items() if k not in _QUERY_TEXT_EXCLUDED_FIELDS} truncated_cmd = _truncate_documents(filtered, max_length)[0] @@ -162,7 +189,21 @@ def _extract_collection_name( return _extract_collection_name(inner_name, dbname, inner) key = "collection" if command_name == _GET_MORE else command_name value = cmd.get(key) - return value if isinstance(value, str) else None + if not isinstance(value, str) or is_command_namespace(value): + return None + return value + + +def is_command_namespace(collection: Optional[str]) -> bool: + """Return True if ``collection`` is a command pseudo-namespace, not a user collection. + + A cursor opened by a command (listCollections, a database-level aggregate) + reports something like "$cmd.listCollections", which targets no specific + collection, so per the spec ``db.collection.name`` is omitted. + """ + return collection is not None and ( + collection == _CMD_NAMESPACE_PREFIX or collection.startswith(_CMD_NAMESPACE_PREFIX + ".") + ) def _build_query_summary(command_name: str, dbname: str, collection: Optional[str]) -> str: @@ -172,6 +213,42 @@ def _build_query_summary(command_name: str, dbname: str, collection: Optional[st return f"{command_name} {dbname}" +# Some `_Op` values are the wire command name ("drop"/"create") rather than the +# spec's db.operation.name ("dropCollection"/"createCollection"). Translate here +# instead of renaming `_Op`: `_WRITES_WITH_CLUSTER_TIME` in operations.py matches +# these exact strings to pick which writes get afterClusterTime. +_OPERATION_NAME_OVERRIDES = { + "drop": "dropCollection", + "create": "createCollection", + "dropSearchIndexes": "dropSearchIndex", +} + +# The spec names anything sent through the generic `Database.command()` API "runCommand", +# not after the command it carries. +_RUN_COMMAND_OPERATION_NAME = "runCommand" + + +def _normalize_operation_name(operation: Any) -> str: + """Return the plain ``str`` form of an operation name. + + Call sites pass an `_Op` (a `str`-mixin enum), and Python 3.11 changed + ``Enum.__format__`` so ``str(_Op.INSERT)`` yields ``"_Op.INSERT"`` rather + than ``"insert"``. Normalizing once here keeps every span name and + attribute stable across versions. + """ + if isinstance(operation, enum.Enum): + return operation.value + return str(operation) + + +def _build_operation_name(operation: Any, is_run_command: bool = False) -> str: + """Return the ``db.operation.name`` the spec wants for this operation.""" + if is_run_command: + return _RUN_COMMAND_OPERATION_NAME + name = _normalize_operation_name(operation) + return _OPERATION_NAME_OVERRIDES.get(name, name) + + def _is_sensitive_command(command_name: str, speculative_hello: bool) -> bool: """Mirror the redaction rules in ``pymongo.logger.LogMessage._is_sensitive``.""" if command_name in _SENSITIVE_COMMANDS: @@ -200,15 +277,33 @@ def start_command_span( ) -> Optional[Span]: """Start and return a CLIENT-kind span for a server command, or None. - Returns None when tracing is disabled/unavailable or the command is - sensitive (mirroring the redaction applied to logs). + None when tracing is off or the command is sensitive, mirroring the + redaction applied to logs. One span per wire-protocol message, so a retried + operation produces one per attempt. Returned rather than made current: it is + a leaf under whichever operation span is current, and the caller holds it + for the command's duration. """ if not _is_tracing_enabled(tracing_options): return None + + collection = _extract_collection_name(command_name, dbname, cmd) + # 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. + current_operation = _CURRENT_OPERATION_NAME.get() + if current_operation is not None: + current_span = trace.get_current_span() + if current_span.is_recording(): + summary = _build_query_summary(current_operation, dbname, collection) + current_span.update_name(summary) + current_span.set_attribute("db.namespace", dbname) + current_span.set_attribute("db.operation.summary", summary) + if collection: + current_span.set_attribute("db.collection.name", collection) + if _is_sensitive_command(command_name, speculative_hello): return None - collection = _extract_collection_name(command_name, dbname, cmd) address = conn.address transport = "unix" if address[1] is None else "tcp" attributes: dict[str, Any] = { @@ -243,15 +338,35 @@ def start_command_span( def end_command_span_success(span: Optional[Span], reply: _DocumentOut) -> None: - """Set the cursor id (if any) and end the span.""" + """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 "id" in 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"]) span.end() +def _set_exception_attributes(span: Span, exc: BaseException) -> None: + """Set exception.type/exception.message/exception.stacktrace span attributes. + + ``record_exception`` attaches these to an "exception" *event* only, but the + spec requires them as span *attributes* too, for both command and operation + spans. Formatting mirrors ``record_exception``. + """ + module = type(exc).__module__ + qualname = type(exc).__qualname__ + exception_type = f"{module}.{qualname}" if module and module != "builtins" else qualname + span.set_attribute("exception.type", exception_type) + span.set_attribute("exception.message", str(exc)) + span.set_attribute( + "exception.stacktrace", + "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)), + ) + + def end_command_span_failure( span: Optional[Span], failure: _DocumentOut, @@ -261,8 +376,150 @@ def end_command_span_failure( if span is None: return span.record_exception(exc) + _set_exception_attributes(span, exc) code = failure.get("code") if code is not None: span.set_attribute("db.response.status_code", str(code)) span.set_status(Status(StatusCode.ERROR, description=failure.get("errmsg"))) span.end() + + +class _OperationSpanHandle: + """Bundles an operation span with what's needed to end it later. + + ``_cm`` is the ``start_as_current_span`` context manager, or None in + detached mode, where ``use_operation_span`` makes the span current per use. + """ + + __slots__ = ("_cm", "_name_token", "operation_name", "span") + + def __init__( + self, + span: Span, + cm: Any, + name_token: Any, + operation_name: str, + ) -> None: + self.span = span + self._cm = cm + self._name_token = name_token + self.operation_name = operation_name + + +def start_operation_span( + tracing_options: Optional[TracingOptions], + operation: str, + parent_span: Optional[Span], + dbname: Optional[str] = None, + collection: Optional[str] = None, + set_current: bool = True, +) -> Optional[_OperationSpanHandle]: + """Start a CLIENT-kind span for one logical operation, or None. + + Spans all retry attempts of one ``_retry_internal`` call. Namespace + attributes are set eagerly from ``dbname``/``collection`` so an operation + that fails before building any command, such as a server selection + timeout, still produces a conformant span; ``start_command_span`` + backfills the authoritative values once a command exists. + + ``parent_span`` becomes an *explicit* parent rather than being read from + ambient context, so a concurrent unrelated session cannot be captured. + + ``set_current=False`` leaves the span and the operation-name contextvar + alone, for a caller that makes it current with ``use_operation_span``. + """ + if not _is_tracing_enabled(tracing_options): + return None + assert _TRACER is not None # _is_tracing_enabled already checked _HAS_OPENTELEMETRY + context = trace.set_span_in_context(parent_span) if parent_span is not None else None + attributes: dict[str, Any] = { + "db.system.name": "mongodb", + "db.operation.name": operation, + } + name = operation + if dbname is not None: + name = _build_query_summary(operation, dbname, collection) + attributes["db.namespace"] = dbname + if collection: + attributes["db.collection.name"] = collection + attributes["db.operation.summary"] = name + if not set_current: + span = _TRACER.start_span( + name, kind=SpanKind.CLIENT, context=context, attributes=attributes + ) + return _OperationSpanHandle(span, None, None, operation) + cm = _TRACER.start_as_current_span( + name, + kind=SpanKind.CLIENT, + context=context, + attributes=attributes, + ) + span = cm.__enter__() + name_token = _CURRENT_OPERATION_NAME.set(operation) + return _OperationSpanHandle(span, cm, name_token, operation) + + +@contextlib.contextmanager +def use_operation_span(handle: Optional[_OperationSpanHandle]) -> Iterator[None]: + """Make a detached operation span current for the duration of the block. + + Does not end the span; its owner ends it explicitly. A no-op when + ``handle`` is None. + """ + if handle is None: + yield + return + token = _CURRENT_OPERATION_NAME.set(handle.operation_name) + try: + # Left on, these would auto-record any exception leaving the block and + # set ERROR status, duplicating the caller's end_operation_span_failure + # once the operation's final outcome is known. + with trace.use_span( + handle.span, + end_on_exit=False, + record_exception=False, + set_status_on_exception=False, + ): + yield + finally: + _CURRENT_OPERATION_NAME.reset(token) + + +def reset_context() -> None: + """Clear the OTel ambient span and operation-name contextvar. + + ``asyncio.create_task`` freezes the caller's context, so without this a + long-lived background task parents every span it emits under an unrelated, + long-ended operation. Attaching an empty context makes spans started + afterwards trace roots. Deliberately never detached: the task's context is + wrong for its whole life and dies with it. + """ + if not _HAS_OPENTELEMETRY: + return + _CURRENT_OPERATION_NAME.set(None) + context.attach(context.Context()) + + +def end_operation_span_success(handle: Optional[_OperationSpanHandle]) -> None: + """End the operation span with no error status.""" + if handle is None: + return + if handle._cm is None: + handle.span.end() + return + _CURRENT_OPERATION_NAME.reset(handle._name_token) + handle._cm.__exit__(None, None, None) + + +def end_operation_span_failure(handle: Optional[_OperationSpanHandle], exc: BaseException) -> None: + """Record the exception, set the error status, and end the operation span.""" + if handle is None: + return + handle.span.record_exception(exc) + _set_exception_attributes(handle.span, exc) + handle.span.set_status(Status(StatusCode.ERROR, description=str(exc))) + if handle._cm is None: + handle.span.end() + return + _CURRENT_OPERATION_NAME.reset(handle._name_token) + handle._cm.__exit__(None, None, None) diff --git a/pymongo/_telemetry.py b/pymongo/_telemetry.py index 84b11a7408..f9dde5520b 100644 --- a/pymongo/_telemetry.py +++ b/pymongo/_telemetry.py @@ -260,6 +260,92 @@ def failed( _otel.end_command_span_failure(self._span, failure, exc) +class _OperationTelemetry: + """One span-scoped context per logical operation (spanning all retry attempts). + + Construct once per call to ``_retry_internal``; call :meth:`succeeded` or + :meth:`failed` exactly once when the operation's outcome is known, or use + it as a context manager to do so automatically. A no-op throughout when + tracing is disabled. + + This span is shared by every attempt, while each attempt gets a command + span of its own underneath it. Retries are therefore visible as sibling + command spans rather than being collapsed into one. + + With ``set_current=False`` the span is not made current at construction. + 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`. + """ + + __slots__ = ("handle",) + + def __init__( + self, + tracing_options: Optional[_otel.TracingOptions], + operation: str, + session: Optional[Any], + is_run_command: bool = False, + dbname: Optional[str] = None, + collection: Optional[str] = None, + set_current: bool = True, + ) -> None: + self.handle = _otel.start_operation_span( + tracing_options, + _otel._build_operation_name(operation, is_run_command), + None, + dbname=dbname, + collection=collection, + set_current=set_current, + ) + + def use(self) -> Any: + """Make this operation's span current for the duration of a block.""" + return _otel.use_operation_span(self.handle) + + def succeeded(self) -> None: + _otel.end_operation_span_success(self.handle) + + def failed(self, exc: BaseException) -> None: + _otel.end_operation_span_failure(self.handle, exc) + + def __enter__(self) -> _OperationTelemetry: + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if exc_val is None: + self.succeeded() + else: + self.failed(exc_val) + + +def _operation_telemetry_or_none( + tracing_options: Optional[_otel.TracingOptions], + operation: str, + session: Optional[Any], + is_run_command: bool = False, + dbname: Optional[str] = None, + collection: Optional[str] = None, + set_current: bool = True, +) -> Optional[_OperationTelemetry]: + """Return an :class:`_OperationTelemetry`, or None if tracing is disabled. + + Every operation goes through here, so follow _CommandTelemetry's fast path + and skip the object rather than build one whose methods all do nothing. + """ + if not _otel._is_tracing_enabled(tracing_options): + return None + return _OperationTelemetry( + tracing_options, + operation, + session, + is_run_command=is_run_command, + dbname=dbname, + collection=collection, + set_current=set_current, + ) + + class _CmapTelemetry: """Combines CMAP structured logging and APM event publishing for pool and connection events.""" diff --git a/pymongo/asynchronous/aggregation.py b/pymongo/asynchronous/aggregation.py index f1f77acc73..724773a1a8 100644 --- a/pymongo/asynchronous/aggregation.py +++ b/pymongo/asynchronous/aggregation.py @@ -22,6 +22,7 @@ from pymongo import common from pymongo.collation import validate_collation_or_none from pymongo.errors import ConfigurationError +from pymongo.helpers_shared import _split_namespace from pymongo.read_preferences import ReadPreference, _AggWritePref if TYPE_CHECKING: @@ -251,5 +252,5 @@ def _cursor_collection(self, cursor: Mapping[str, Any]) -> AsyncCollection[Any]: # AsyncCollection level aggregate may not always return the "ns" field # according to our MockupDB tests. Let's handle that case for db level # aggregate too by defaulting to the .$cmd.aggregate namespace. - _, collname = cursor.get("ns", self._cursor_namespace).split(".", 1) + _, collname = _split_namespace(cursor.get("ns", self._cursor_namespace)) return self._database[collname] diff --git a/pymongo/asynchronous/client_bulk.py b/pymongo/asynchronous/client_bulk.py index 367fdd492f..8aa8892e5d 100644 --- a/pymongo/asynchronous/client_bulk.py +++ b/pymongo/asynchronous/client_bulk.py @@ -32,7 +32,7 @@ from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument from pymongo import _csot, common -from pymongo._telemetry import _generate_op_id_or_none +from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.asynchronous.client_session import ( AsyncClientSession, _validate_session_write_concern, @@ -630,13 +630,28 @@ async def execute( session = _validate_session_write_concern(session, self.write_concern) if not self.write_concern.acknowledged: - async with await self.client._conn_for_writes(session, operation) as connection: - if connection.max_wire_version < 25: - raise InvalidOperation( - "MongoClient.bulk_write requires MongoDB server version 8.0+." - ) - await self.execute_no_results(connection) - return ClientBulkWriteResult(None, False, False) # type: ignore[arg-type] + # This path never reaches the command-span code that would otherwise + # fill in the namespace, so pass it here. A client bulk write always + # runs against admin and spans multiple namespaces, so it reports no + # collection. + operation_telemetry = _operation_telemetry_or_none( + self.client.options.tracing, operation, session, dbname="admin" + ) + try: + async with await self.client._conn_for_writes(session, operation) as connection: + if connection.max_wire_version < 25: + raise InvalidOperation( + "MongoClient.bulk_write requires MongoDB server version 8.0+." + ) + await self.execute_no_results(connection) + except BaseException as exc: + if operation_telemetry is not None: + operation_telemetry.failed(exc) + raise + else: + if operation_telemetry is not None: + operation_telemetry.succeeded() + return ClientBulkWriteResult(None, False, False) # type: ignore[arg-type] result = await self.execute_command(session, operation) return ClientBulkWriteResult( diff --git a/pymongo/asynchronous/collection.py b/pymongo/asynchronous/collection.py index ff2df1be82..f4c91c941d 100644 --- a/pymongo/asynchronous/collection.py +++ b/pymongo/asynchronous/collection.py @@ -2596,8 +2596,13 @@ async def _cmd( return cmd_cursor async with self._database.client._tmp_session(session) as s: - return await self._database.client._retryable_read( - _cmd, read_pref, s, operation=_Op.LIST_INDEXES + return await self._database.client._retryable_read_cursor( + _cmd, + read_pref, + s, + operation=_Op.LIST_INDEXES, + dbname=self._database.name, + collection=self._name, ) async def index_information( @@ -2695,12 +2700,14 @@ async def list_search_indexes( user_fields={"cursor": {"firstBatch": 1}}, ) - return await self._database.client._retryable_read( + return await self._database.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(session), # type: ignore[arg-type] session, retryable=not cmd._performs_write, operation=_Op.LIST_SEARCH_INDEX, + dbname=self._database.name, + collection=self.name, ) async def create_search_index( @@ -2944,13 +2951,15 @@ async def _aggregate( user_fields={"cursor": {"firstBatch": 1}}, ) - return await self._database.client._retryable_read( + return await self._database.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(session), # type: ignore[arg-type] session, retryable=not cmd._performs_write, operation=_Op.AGGREGATE, is_aggregate_write=cmd._performs_write, + dbname=self._database.name, + collection=self._name, ) async def aggregate( diff --git a/pymongo/asynchronous/cursor.py b/pymongo/asynchronous/cursor.py index 1b2428a6d8..b68e96a4bb 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._telemetry import _operation_telemetry_or_none from pymongo.asynchronous.cursor_base import _AsyncCursorBase, _ConnectionManager from pymongo.asynchronous.helpers import anext from pymongo.collation import validate_collation_or_none @@ -974,9 +975,13 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None: 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, ) except OperationFailure as exc: + self._end_operation_telemetry(exc) if exc.code in _CURSOR_CLOSED_ERRORS or self._exhaust: # Don't send killCursors because the cursor is already closed. self._killed = True @@ -994,12 +999,14 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None: ): return raise - except ConnectionFailure: + except ConnectionFailure as exc: + self._end_operation_telemetry(exc) self._killed = True await self.close() raise # Catch KeyboardInterrupt, CancelledError, etc. and cleanup. - except BaseException: + except BaseException as exc: + self._end_operation_telemetry(exc) await self.close() raise self._address = response.address @@ -1017,7 +1024,7 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None: # Update the namespace used for future getMore commands. ns = cursor.get("ns") if ns: - self._dbname, self._collname = ns.split(".", 1) + self._dbname, self._collname = helpers_shared._split_namespace(ns) else: documents = cursor["nextBatch"] self._data = deque(documents) @@ -1072,7 +1079,16 @@ async def _refresh(self) -> int: self._allow_disk_use, self._exhaust, ) - await self._send_message(q) + client = self._collection.database.client + self._operation_telemetry = _operation_telemetry_or_none( + client.options.tracing, + q.name, + self._session, + dbname=self._collection.database.name, + collection=self._collection.name, + set_current=False, + ) + await self._send_message_in_operation_span(q) elif self._id: # Get More if self._limit: limit = self._limit - self._retrieved @@ -1099,6 +1115,21 @@ async def _refresh(self) -> int: 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. + + ``_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. + """ + try: + await self._send_message(operation) + except BaseException as exc: + self._end_operation_telemetry(exc) + raise + self._end_operation_telemetry() + async def rewind(self) -> AsyncCursor[_DocumentType]: """Rewind this cursor to its unevaluated state. diff --git a/pymongo/asynchronous/cursor_base.py b/pymongo/asynchronous/cursor_base.py index e8ac4c3139..870120f78a 100644 --- a/pymongo/asynchronous/cursor_base.py +++ b/pymongo/asynchronous/cursor_base.py @@ -175,6 +175,7 @@ async def _die_lock(self) -> None: # ___init__ did not run to completion (or at all). return + self._end_operation_telemetry() cursor_id, address = self._prepare_to_die(already_killed) await self._collection.database.client._cleanup_cursor_lock( cursor_id, diff --git a/pymongo/asynchronous/database.py b/pymongo/asynchronous/database.py index 88ca9cab5b..96971fdf2d 100644 --- a/pymongo/asynchronous/database.py +++ b/pymongo/asynchronous/database.py @@ -708,12 +708,13 @@ async def aggregate( kwargs, user_fields={"cursor": {"firstBatch": 1}}, ) - return await self.client._retryable_read( + return await self.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(s), # type: ignore[arg-type] s, retryable=not cmd._performs_write, operation=_Op.AGGREGATE, + dbname=self.name, ) @overload @@ -1051,8 +1052,8 @@ async def inner( else: raise InvalidOperation("Command does not return a cursor.") - return await self.client._retryable_read( - inner, read_preference, tmp_session, command_name, None, False + return await self.client._retryable_read_cursor( + inner, read_preference, tmp_session, command_name, None, False, dbname=self.name ) async def _retryable_read_command( @@ -1149,8 +1150,8 @@ async def _cmd( conn, session, read_preference=read_preference, **kwargs ) - return await self._client._retryable_read( - _cmd, read_pref, session, operation=_Op.LIST_COLLECTIONS + return await self._client._retryable_read_cursor( + _cmd, read_pref, session, operation=_Op.LIST_COLLECTIONS, dbname=self.name ) async def list_collections( diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py index aca51ced84..51a6d02e91 100644 --- a/pymongo/asynchronous/mongo_client.py +++ b/pymongo/asynchronous/mongo_client.py @@ -56,7 +56,12 @@ 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._telemetry import _generate_op_id_or_none, log_command_retry +from pymongo._telemetry import ( + _generate_op_id_or_none, + _operation_telemetry_or_none, + _OperationTelemetry, + log_command_retry, +) from pymongo.asynchronous import client_session, database, uri_parser from pymongo.asynchronous.change_stream import AsyncChangeStream, AsyncClusterChangeStream from pymongo.asynchronous.client_bulk import _AsyncClientBulk @@ -145,6 +150,7 @@ T = TypeVar("T") +_CommandCursor = TypeVar("_CommandCursor", bound=AsyncCommandCursor[Any]) _WriteCall = Callable[ [Optional["AsyncClientSession"], "AsyncConnection", bool], Coroutine[Any, Any, T] @@ -617,7 +623,8 @@ def __init__( | **OpenTelemetry options:** | (Requires the ``opentelemetry-api`` package; install with the ``pymongo[opentelemetry]`` extra.) - - `tracing`: (dict) Configuration for OpenTelemetry command spans, with keys: + - `tracing`: (dict) Configuration for OpenTelemetry command, operation, and + transaction spans, with keys: - ``enabled``: (boolean) Whether to create spans for server commands issued by this client. Defaults to ``False``. Also controlled by the @@ -633,7 +640,10 @@ def __init__( .. seealso:: The MongoDB documentation on `connections `_. .. versionchanged:: 4.18 - Added the ``tracing`` keyword argument. + Added the ``tracing`` keyword argument. Every public API call + produces an operation span, which contains one span per command + sent to the server. Inside a transaction, those operation spans + nest under a ``"transaction"`` span. .. versionchanged:: 4.17 Added the ``max_adaptive_retries`` and ``enable_overload_retargeting`` URI and keyword arguments. @@ -1735,7 +1745,13 @@ async def _end_sessions(self, session_ids: list[_ServerSession]) -> None: for i in range(0, len(session_ids), common._MAX_END_SESSIONS): spec = {"endSessions": session_ids[i : i + common._MAX_END_SESSIONS]} - await conn.command("admin", spec, read_preference=read_pref, client=self) + # endSessions bypasses _retry_internal (errors are ignored per + # spec, and it must not be retried), so start its span here. + telemetry = _operation_telemetry_or_none( + self.options.tracing, _Op.END_SESSIONS, None, dbname="admin" + ) + with telemetry or contextlib.nullcontext(): + await conn.command("admin", spec, read_preference=read_pref, client=self) except PyMongoError: # Drivers MUST ignore any errors returned by the endSessions # command. @@ -1884,6 +1900,7 @@ async def _run_operation( operation: Union[_Query, _GetMore], run_with_conn: Callable, # type: ignore[type-arg] address: Optional[_Address] = None, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> Response: """Run a _Query/_GetMore operation and return a Response. @@ -1892,6 +1909,7 @@ 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. """ if operation.conn_mgr: server = await self._select_server( @@ -1908,9 +1926,17 @@ async def _run_operation( operation.session, # type: ignore[arg-type] operation.conn_mgr.conn, ): - return await run_with_conn( - operation.conn_mgr.conn, operation, operation.read_preference - ) + # Exhaust/pinned cursors bypass _retry_internal, so make the + # caller's span current here to keep their command spans + # nested under it. + with ( + operation_telemetry.use() + if operation_telemetry + else contextlib.nullcontext() + ): + return await run_with_conn( + operation.conn_mgr.conn, operation, operation.read_preference + ) async def _cmd( _session: Optional[AsyncClientSession], @@ -1928,6 +1954,7 @@ async def _cmd( address=address, retryable=isinstance(operation, _Query), operation=operation.name, + operation_telemetry=operation_telemetry, ) async def _retry_with_session( @@ -1974,6 +2001,7 @@ async def _retry_internal( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> T: """Internal retryable helper for all client transactions. @@ -1988,6 +2016,10 @@ async def _retry_internal( :param is_run_command: If this is a runCommand operation, defaults to False :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False. :param operation_id: Stable operation id shared across retries, defaults to None + :param operation_telemetry: A caller-owned operation span outliving this call + (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. :return: Output of the calling func() """ @@ -2004,6 +2036,7 @@ async def _retry_internal( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + operation_telemetry=operation_telemetry, ).run() async def _retryable_read( @@ -2017,6 +2050,7 @@ async def _retryable_read( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> T: """Execute an operation with consecutive retries if possible @@ -2035,6 +2069,8 @@ async def _retryable_read( :param is_run_command: If this is a runCommand operation, defaults to False. :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False. :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. """ # Ensure that the client supports retrying on reads and there is no session in @@ -2055,7 +2091,61 @@ async def _retryable_read( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + operation_telemetry=operation_telemetry, + ) + + async def _retryable_read_cursor( + self, + func: _ReadCall[_CommandCursor], + read_pref: _ServerMode, + session: Optional[AsyncClientSession], + operation: str, + address: Optional[_Address] = None, + retryable: bool = True, + operation_id: Optional[int] = None, + is_run_command: bool = False, + is_aggregate_write: bool = False, + *, + dbname: str, + collection: Optional[str] = None, + ) -> _CommandCursor: + """Run a command-cursor read within its own operation span. + + Takes the same arguments as :meth:`_retryable_read`, plus the namespace + for the span. A command cursor's first batch is fetched inside that + 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. + """ + operation_telemetry = _operation_telemetry_or_none( + self.options.tracing, + operation, + session, + dbname=dbname, + collection=collection, + set_current=False, + ) + try: + cmd_cursor = await self._retryable_read( + func, + read_pref, + session, + operation, + address, + retryable, + operation_id, + is_run_command, + is_aggregate_write, + operation_telemetry=operation_telemetry, ) + except BaseException as exc: + if operation_telemetry is not None: + operation_telemetry.failed(exc) + raise + if operation_telemetry is not None: + operation_telemetry.succeeded() + return cmd_cursor async def _retryable_write( self, @@ -2193,9 +2283,15 @@ async def _kill_cursor_impl( conn: AsyncConnection, ) -> None: namespace = address.namespace - db, coll = namespace.split(".", 1) + db, coll = helpers_shared._split_namespace(namespace) spec = {"killCursors": coll, "cursors": cursor_ids} - await conn.command(db, spec, session=session, client=self) + # killCursors bypasses _retry_internal (it must never be retried), so + # start its span here. + telemetry = _operation_telemetry_or_none( + self.options.tracing, _Op.KILL_CURSORS, session, dbname=db, collection=coll + ) + with telemetry or contextlib.nullcontext(): + await conn.command(db, spec, session=session, client=self) async def _process_kill_cursors(self) -> None: """Process any pending kill cursors requests.""" @@ -2863,6 +2959,8 @@ class _ClientConnectionRetryable(Generic[T]): "_max_retries", "_operation", "_operation_id", + "_operation_telemetry", + "_owns_telemetry", "_read_pref", "_retry_policy", "_retryable", @@ -2886,6 +2984,7 @@ def __init__( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ): self._last_error: Optional[Exception] = None self._retrying = False @@ -2910,12 +3009,36 @@ def __init__( if operation_id is None: operation_id = _generate_op_id_or_none(self._client._event_listeners) self._operation_id = operation_id + # 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 + if self._owns_telemetry: + operation_telemetry = _operation_telemetry_or_none( + mongo_client.options.tracing, operation, session, is_run_command=is_run_command + ) + self._operation_telemetry = operation_telemetry self._attempt_number = 0 self._is_run_command = is_run_command self._is_aggregate_write = is_aggregate_write self._base_backoff_ms: Optional[float] = None async def run(self) -> T: + """Run the operation, retrying as allowed, within its operation span.""" + if self._operation_telemetry is None: + return await self._run() + if not self._owns_telemetry: + with self._operation_telemetry.use(): + return await self._run() + try: + result = await self._run() + except BaseException as exc: + self._operation_telemetry.failed(exc) + raise + else: + self._operation_telemetry.succeeded() + return result + + async def _run(self) -> T: """Runs the supplied func() and attempts a retry :raises: self._last_error: Last exception raised diff --git a/pymongo/client_options.py b/pymongo/client_options.py index d2d33f6805..f94cb23c77 100644 --- a/pymongo/client_options.py +++ b/pymongo/client_options.py @@ -24,6 +24,7 @@ from bson.codec_options import _parse_codec_options from pymongo import common +from pymongo._otel import _resolve_tracing_options from pymongo.compression_support import CompressionSettings from pymongo.errors import ConfigurationError from pymongo.monitoring import _EventListener, _EventListeners @@ -248,9 +249,13 @@ def __init__( if "enable_overload_retargeting" in options else options.get("enableoverloadretargeting", common.ENABLE_OVERLOAD_RETARGETING) ) - self.__tracing = cast( - "_otel.TracingOptions", - options.get("tracing") or {"enabled": False, "query_text_max_length": None}, + # Fold the OTEL_* environment variables in once, here. They are + # process-startup input, so nothing re-reads them per command. + self.__tracing = _resolve_tracing_options( + cast( + "_otel.TracingOptions", + options.get("tracing") or {"enabled": False, "query_text_max_length": None}, + ) ) @property diff --git a/pymongo/cursor_shared.py b/pymongo/cursor_shared.py index df0e1e2f58..5a26eef5cd 100644 --- a/pymongo/cursor_shared.py +++ b/pymongo/cursor_shared.py @@ -55,6 +55,7 @@ class _AgnosticCursorBase(Generic[_DocumentType], ABC): _sock_mgr: Any _session: Optional[Any] _killed: bool + _operation_telemetry: Optional[Any] = None @abstractmethod def _get_namespace(self) -> str: @@ -115,6 +116,24 @@ def _prepare_to_die(self, already_killed: bool) -> tuple[int, Optional[_CursorAd address = None return cursor_id, address + 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. + """ + telemetry = self._operation_telemetry + if telemetry is None: + return + self._operation_telemetry = None + if exc is None: + telemetry.succeeded() + else: + telemetry.failed(exc) + def _die_no_lock(self) -> None: """Closes this cursor without acquiring a lock.""" try: @@ -123,6 +142,7 @@ def _die_no_lock(self) -> None: # ___init__ did not run to completion (or at all). return + self._end_operation_telemetry() cursor_id, address = self._prepare_to_die(already_killed) self._collection.database.client._cleanup_cursor_no_lock( cursor_id, address, self._sock_mgr, self._session diff --git a/pymongo/helpers_shared.py b/pymongo/helpers_shared.py index ccd1b942e3..06f39acbd3 100644 --- a/pymongo/helpers_shared.py +++ b/pymongo/helpers_shared.py @@ -135,6 +135,17 @@ def format_timeout_details(details: Optional[dict[str, float]]) -> str: return result +def _split_namespace(namespace: str) -> tuple[str, str]: + """Split a ``"dbname.collname"`` namespace into its two parts. + + Collection names may contain dots, so only the first separator is + significant. If no dot is present, the entire string is treated as the + database name and the collection name is empty. + """ + dbname, _, collname = namespace.partition(".") + return dbname, collname + + def _gen_index_name(keys: _IndexList) -> str: """Generate an index name from the set of fields it is over.""" return "_".join(["{}_{}".format(*item) for item in keys]) diff --git a/pymongo/periodic_executor.py b/pymongo/periodic_executor.py index 4b979ca9f9..1ef988b409 100644 --- a/pymongo/periodic_executor.py +++ b/pymongo/periodic_executor.py @@ -23,7 +23,7 @@ import weakref from typing import Any, Optional -from pymongo import _csot, _op_id +from pymongo import _csot, _op_id, _otel from pymongo._asyncio_task import create_task from pymongo.lock import _create_lock @@ -94,9 +94,12 @@ def skip_sleep(self) -> None: self._skip_sleep = True async def _run(self) -> None: - # The CSOT and op id contextvars must be cleared inside the executor task before execution begins + # create_task froze a copy of the context this executor was opened in, + # which for the kill-cursors executor is the middle of the client's first + # operation. Clear it so every tick starts clean. _csot.reset_all() _op_id.reset() + _otel.reset_context() while not self._stopped: if self._task and self._task.cancelling(): # type: ignore[unused-ignore, attr-defined] raise asyncio.CancelledError @@ -232,6 +235,14 @@ def _should_stop(self) -> bool: return False def _run(self) -> None: + # Same reason as AsyncPeriodicExecutor._run. Where + # sys.flags.thread_inherit_context is set, which is the default on + # free-threaded builds, a thread runs its target in a copy of the + # creating thread's context, so this thread would otherwise inherit the + # context the executor was opened in for the rest of the process's life. + _csot.reset_all() + _op_id.reset() + _otel.reset_context() while not self._should_stop(): try: if not self._target(): diff --git a/pymongo/synchronous/aggregation.py b/pymongo/synchronous/aggregation.py index d540484fa8..dbc892c5e9 100644 --- a/pymongo/synchronous/aggregation.py +++ b/pymongo/synchronous/aggregation.py @@ -22,6 +22,7 @@ from pymongo import common from pymongo.collation import validate_collation_or_none from pymongo.errors import ConfigurationError +from pymongo.helpers_shared import _split_namespace from pymongo.read_preferences import ReadPreference, _AggWritePref if TYPE_CHECKING: @@ -251,5 +252,5 @@ def _cursor_collection(self, cursor: Mapping[str, Any]) -> Collection[Any]: # Collection level aggregate may not always return the "ns" field # according to our MockupDB tests. Let's handle that case for db level # aggregate too by defaulting to the .$cmd.aggregate namespace. - _, collname = cursor.get("ns", self._cursor_namespace).split(".", 1) + _, collname = _split_namespace(cursor.get("ns", self._cursor_namespace)) return self._database[collname] diff --git a/pymongo/synchronous/client_bulk.py b/pymongo/synchronous/client_bulk.py index 3dca2f7234..7f7181d826 100644 --- a/pymongo/synchronous/client_bulk.py +++ b/pymongo/synchronous/client_bulk.py @@ -32,7 +32,7 @@ from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument from pymongo import _csot, common -from pymongo._telemetry import _generate_op_id_or_none +from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.synchronous.client_session import ( ClientSession, _validate_session_write_concern, @@ -628,13 +628,28 @@ def execute( session = _validate_session_write_concern(session, self.write_concern) if not self.write_concern.acknowledged: - with self.client._conn_for_writes(session, operation) as connection: - if connection.max_wire_version < 25: - raise InvalidOperation( - "MongoClient.bulk_write requires MongoDB server version 8.0+." - ) - self.execute_no_results(connection) - return ClientBulkWriteResult(None, False, False) # type: ignore[arg-type] + # This path never reaches the command-span code that would otherwise + # fill in the namespace, so pass it here. A client bulk write always + # runs against admin and spans multiple namespaces, so it reports no + # collection. + operation_telemetry = _operation_telemetry_or_none( + self.client.options.tracing, operation, session, dbname="admin" + ) + try: + with self.client._conn_for_writes(session, operation) as connection: + if connection.max_wire_version < 25: + raise InvalidOperation( + "MongoClient.bulk_write requires MongoDB server version 8.0+." + ) + self.execute_no_results(connection) + except BaseException as exc: + if operation_telemetry is not None: + operation_telemetry.failed(exc) + raise + else: + if operation_telemetry is not None: + operation_telemetry.succeeded() + return ClientBulkWriteResult(None, False, False) # type: ignore[arg-type] result = self.execute_command(session, operation) return ClientBulkWriteResult( diff --git a/pymongo/synchronous/collection.py b/pymongo/synchronous/collection.py index 3051ac9839..0821eb4573 100644 --- a/pymongo/synchronous/collection.py +++ b/pymongo/synchronous/collection.py @@ -2592,8 +2592,13 @@ def _cmd( return cmd_cursor with self._database.client._tmp_session(session) as s: - return self._database.client._retryable_read( - _cmd, read_pref, s, operation=_Op.LIST_INDEXES + return self._database.client._retryable_read_cursor( + _cmd, + read_pref, + s, + operation=_Op.LIST_INDEXES, + dbname=self._database.name, + collection=self._name, ) def index_information( @@ -2691,12 +2696,14 @@ def list_search_indexes( user_fields={"cursor": {"firstBatch": 1}}, ) - return self._database.client._retryable_read( + return self._database.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(session), # type: ignore[arg-type] session, retryable=not cmd._performs_write, operation=_Op.LIST_SEARCH_INDEX, + dbname=self._database.name, + collection=self.name, ) def create_search_index( @@ -2938,13 +2945,15 @@ def _aggregate( user_fields={"cursor": {"firstBatch": 1}}, ) - return self._database.client._retryable_read( + return self._database.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(session), # type: ignore[arg-type] session, retryable=not cmd._performs_write, operation=_Op.AGGREGATE, is_aggregate_write=cmd._performs_write, + dbname=self._database.name, + collection=self._name, ) def aggregate( diff --git a/pymongo/synchronous/cursor.py b/pymongo/synchronous/cursor.py index fddbcc520f..6554c62a09 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._telemetry import _operation_telemetry_or_none from pymongo.collation import validate_collation_or_none from pymongo.common import ( validate_is_document_type, @@ -971,8 +972,14 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None: raise InvalidOperation("exhaust cursors do not support auto encryption") 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, + ) except OperationFailure as exc: + self._end_operation_telemetry(exc) if exc.code in _CURSOR_CLOSED_ERRORS or self._exhaust: # Don't send killCursors because the cursor is already closed. self._killed = True @@ -990,12 +997,14 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None: ): return raise - except ConnectionFailure: + except ConnectionFailure as exc: + self._end_operation_telemetry(exc) self._killed = True self.close() raise # Catch KeyboardInterrupt, CancelledError, etc. and cleanup. - except BaseException: + except BaseException as exc: + self._end_operation_telemetry(exc) self.close() raise self._address = response.address @@ -1013,7 +1022,7 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None: # Update the namespace used for future getMore commands. ns = cursor.get("ns") if ns: - self._dbname, self._collname = ns.split(".", 1) + self._dbname, self._collname = helpers_shared._split_namespace(ns) else: documents = cursor["nextBatch"] self._data = deque(documents) @@ -1068,7 +1077,16 @@ def _refresh(self) -> int: self._allow_disk_use, self._exhaust, ) - self._send_message(q) + client = self._collection.database.client + self._operation_telemetry = _operation_telemetry_or_none( + client.options.tracing, + q.name, + self._session, + dbname=self._collection.database.name, + collection=self._collection.name, + set_current=False, + ) + self._send_message_in_operation_span(q) elif self._id: # Get More if self._limit: limit = self._limit - self._retrieved @@ -1095,6 +1113,21 @@ def _refresh(self) -> int: 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. + + ``_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. + """ + try: + self._send_message(operation) + except BaseException as exc: + self._end_operation_telemetry(exc) + raise + self._end_operation_telemetry() + def rewind(self) -> Cursor[_DocumentType]: """Rewind this cursor to its unevaluated state. diff --git a/pymongo/synchronous/cursor_base.py b/pymongo/synchronous/cursor_base.py index 4cad4e0c09..03686e8051 100644 --- a/pymongo/synchronous/cursor_base.py +++ b/pymongo/synchronous/cursor_base.py @@ -175,6 +175,7 @@ def _die_lock(self) -> None: # ___init__ did not run to completion (or at all). return + self._end_operation_telemetry() cursor_id, address = self._prepare_to_die(already_killed) self._collection.database.client._cleanup_cursor_lock( cursor_id, diff --git a/pymongo/synchronous/database.py b/pymongo/synchronous/database.py index 65c5d2a7f0..2f9b1a7eb8 100644 --- a/pymongo/synchronous/database.py +++ b/pymongo/synchronous/database.py @@ -708,12 +708,13 @@ def aggregate( kwargs, user_fields={"cursor": {"firstBatch": 1}}, ) - return self.client._retryable_read( + return self.client._retryable_read_cursor( cmd.get_cursor, cmd.get_read_preference(s), # type: ignore[arg-type] s, retryable=not cmd._performs_write, operation=_Op.AGGREGATE, + dbname=self.name, ) @overload @@ -1051,8 +1052,8 @@ def inner( else: raise InvalidOperation("Command does not return a cursor.") - return self.client._retryable_read( - inner, read_preference, tmp_session, command_name, None, False + return self.client._retryable_read_cursor( + inner, read_preference, tmp_session, command_name, None, False, dbname=self.name ) def _retryable_read_command( @@ -1147,8 +1148,8 @@ def _cmd( ) -> CommandCursor[MutableMapping[str, Any]]: return self._list_collections(conn, session, read_preference=read_preference, **kwargs) - return self._client._retryable_read( - _cmd, read_pref, session, operation=_Op.LIST_COLLECTIONS + return self._client._retryable_read_cursor( + _cmd, read_pref, session, operation=_Op.LIST_COLLECTIONS, dbname=self.name ) def list_collections( diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py index 5fab9d533d..4e910dbb83 100644 --- a/pymongo/synchronous/mongo_client.py +++ b/pymongo/synchronous/mongo_client.py @@ -56,7 +56,12 @@ 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._telemetry import _generate_op_id_or_none, log_command_retry +from pymongo._telemetry import ( + _generate_op_id_or_none, + _operation_telemetry_or_none, + _OperationTelemetry, + log_command_retry, +) from pymongo.client_options import ClientOptions from pymongo.driver_info import DriverInfo from pymongo.errors import ( @@ -145,6 +150,7 @@ T = TypeVar("T") +_CommandCursor = TypeVar("_CommandCursor", bound=CommandCursor[Any]) _WriteCall = Callable[[Optional["ClientSession"], "Connection", bool], T] _ReadCall = Callable[ @@ -618,7 +624,8 @@ def __init__( | **OpenTelemetry options:** | (Requires the ``opentelemetry-api`` package; install with the ``pymongo[opentelemetry]`` extra.) - - `tracing`: (dict) Configuration for OpenTelemetry command spans, with keys: + - `tracing`: (dict) Configuration for OpenTelemetry command, operation, and + transaction spans, with keys: - ``enabled``: (boolean) Whether to create spans for server commands issued by this client. Defaults to ``False``. Also controlled by the @@ -634,7 +641,10 @@ def __init__( .. seealso:: The MongoDB documentation on `connections `_. .. versionchanged:: 4.18 - Added the ``tracing`` keyword argument. + Added the ``tracing`` keyword argument. Every public API call + produces an operation span, which contains one span per command + sent to the server. Inside a transaction, those operation spans + nest under a ``"transaction"`` span. .. versionchanged:: 4.17 Added the ``max_adaptive_retries`` and ``enable_overload_retargeting`` URI and keyword arguments. @@ -1732,7 +1742,13 @@ def _end_sessions(self, session_ids: list[_ServerSession]) -> None: for i in range(0, len(session_ids), common._MAX_END_SESSIONS): spec = {"endSessions": session_ids[i : i + common._MAX_END_SESSIONS]} - conn.command("admin", spec, read_preference=read_pref, client=self) + # endSessions bypasses _retry_internal (errors are ignored per + # spec, and it must not be retried), so start its span here. + telemetry = _operation_telemetry_or_none( + self.options.tracing, _Op.END_SESSIONS, None, dbname="admin" + ) + with telemetry or contextlib.nullcontext(): + conn.command("admin", spec, read_preference=read_pref, client=self) except PyMongoError: # Drivers MUST ignore any errors returned by the endSessions # command. @@ -1879,6 +1895,7 @@ def _run_operation( operation: Union[_Query, _GetMore], run_with_conn: Callable, # type: ignore[type-arg] address: Optional[_Address] = None, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> Response: """Run a _Query/_GetMore operation and return a Response. @@ -1887,6 +1904,7 @@ 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. """ if operation.conn_mgr: server = self._select_server( @@ -1903,9 +1921,17 @@ def _run_operation( operation.session, # type: ignore[arg-type] operation.conn_mgr.conn, ): - return run_with_conn( - operation.conn_mgr.conn, operation, operation.read_preference - ) + # Exhaust/pinned cursors bypass _retry_internal, so make the + # caller's span current here to keep their command spans + # nested under it. + with ( + operation_telemetry.use() + if operation_telemetry + else contextlib.nullcontext() + ): + return run_with_conn( + operation.conn_mgr.conn, operation, operation.read_preference + ) def _cmd( _session: Optional[ClientSession], @@ -1923,6 +1949,7 @@ def _cmd( address=address, retryable=isinstance(operation, _Query), operation=operation.name, + operation_telemetry=operation_telemetry, ) def _retry_with_session( @@ -1969,6 +1996,7 @@ def _retry_internal( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> T: """Internal retryable helper for all client transactions. @@ -1983,6 +2011,10 @@ def _retry_internal( :param is_run_command: If this is a runCommand operation, defaults to False :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False. :param operation_id: Stable operation id shared across retries, defaults to None + :param operation_telemetry: A caller-owned operation span outliving this call + (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. :return: Output of the calling func() """ @@ -1999,6 +2031,7 @@ def _retry_internal( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + operation_telemetry=operation_telemetry, ).run() def _retryable_read( @@ -2012,6 +2045,7 @@ def _retryable_read( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ) -> T: """Execute an operation with consecutive retries if possible @@ -2030,6 +2064,8 @@ def _retryable_read( :param is_run_command: If this is a runCommand operation, defaults to False. :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False. :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. """ # Ensure that the client supports retrying on reads and there is no session in @@ -2050,7 +2086,61 @@ def _retryable_read( operation_id=operation_id, is_run_command=is_run_command, is_aggregate_write=is_aggregate_write, + operation_telemetry=operation_telemetry, + ) + + def _retryable_read_cursor( + self, + func: _ReadCall[_CommandCursor], + read_pref: _ServerMode, + session: Optional[ClientSession], + operation: str, + address: Optional[_Address] = None, + retryable: bool = True, + operation_id: Optional[int] = None, + is_run_command: bool = False, + is_aggregate_write: bool = False, + *, + dbname: str, + collection: Optional[str] = None, + ) -> _CommandCursor: + """Run a command-cursor read within its own operation span. + + Takes the same arguments as :meth:`_retryable_read`, plus the namespace + for the span. A command cursor's first batch is fetched inside that + 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. + """ + operation_telemetry = _operation_telemetry_or_none( + self.options.tracing, + operation, + session, + dbname=dbname, + collection=collection, + set_current=False, + ) + try: + cmd_cursor = self._retryable_read( + func, + read_pref, + session, + operation, + address, + retryable, + operation_id, + is_run_command, + is_aggregate_write, + operation_telemetry=operation_telemetry, ) + except BaseException as exc: + if operation_telemetry is not None: + operation_telemetry.failed(exc) + raise + if operation_telemetry is not None: + operation_telemetry.succeeded() + return cmd_cursor def _retryable_write( self, @@ -2188,9 +2278,15 @@ def _kill_cursor_impl( conn: Connection, ) -> None: namespace = address.namespace - db, coll = namespace.split(".", 1) + db, coll = helpers_shared._split_namespace(namespace) spec = {"killCursors": coll, "cursors": cursor_ids} - conn.command(db, spec, session=session, client=self) + # killCursors bypasses _retry_internal (it must never be retried), so + # start its span here. + telemetry = _operation_telemetry_or_none( + self.options.tracing, _Op.KILL_CURSORS, session, dbname=db, collection=coll + ) + with telemetry or contextlib.nullcontext(): + conn.command(db, spec, session=session, client=self) def _process_kill_cursors(self) -> None: """Process any pending kill cursors requests.""" @@ -2852,6 +2948,8 @@ class _ClientConnectionRetryable(Generic[T]): "_max_retries", "_operation", "_operation_id", + "_operation_telemetry", + "_owns_telemetry", "_read_pref", "_retry_policy", "_retryable", @@ -2875,6 +2973,7 @@ def __init__( operation_id: Optional[int] = None, is_run_command: bool = False, is_aggregate_write: bool = False, + operation_telemetry: Optional[_OperationTelemetry] = None, ): self._last_error: Optional[Exception] = None self._retrying = False @@ -2899,12 +2998,36 @@ def __init__( if operation_id is None: operation_id = _generate_op_id_or_none(self._client._event_listeners) self._operation_id = operation_id + # 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 + if self._owns_telemetry: + operation_telemetry = _operation_telemetry_or_none( + mongo_client.options.tracing, operation, session, is_run_command=is_run_command + ) + self._operation_telemetry = operation_telemetry self._attempt_number = 0 self._is_run_command = is_run_command self._is_aggregate_write = is_aggregate_write self._base_backoff_ms: Optional[float] = None def run(self) -> T: + """Run the operation, retrying as allowed, within its operation span.""" + if self._operation_telemetry is None: + return self._run() + if not self._owns_telemetry: + with self._operation_telemetry.use(): + return self._run() + try: + result = self._run() + except BaseException as exc: + self._operation_telemetry.failed(exc) + raise + else: + self._operation_telemetry.succeeded() + return result + + def _run(self) -> T: """Runs the supplied func() and attempts a retry :raises: self._last_error: Last exception raised diff --git a/test/asynchronous/test_common.py b/test/asynchronous/test_common.py index 8440ca98dc..7190f393f3 100644 --- a/test/asynchronous/test_common.py +++ b/test/asynchronous/test_common.py @@ -25,12 +25,26 @@ from bson.codec_options import CodecOptions from bson.objectid import ObjectId from pymongo.errors import OperationFailure +from pymongo.helpers_shared import _split_namespace from pymongo.write_concern import WriteConcern from test.asynchronous import AsyncIntegrationTest, async_client_context, connected, unittest _IS_SYNC = False +class TestSplitNamespace(unittest.TestCase): + def test_plain_namespace(self): + self.assertEqual(_split_namespace("db.coll"), ("db", "coll")) + + def test_collection_name_with_dots(self): + self.assertEqual(_split_namespace("db.coll.with.dots"), ("db", "coll.with.dots")) + + def test_no_dot(self): + # No separator: the whole string is treated as the database name + # and the collection name is empty, matching str.partition. + self.assertEqual(_split_namespace("dbonly"), ("dbonly", "")) + + class TestCommon(AsyncIntegrationTest): async def test_uuid_representation(self): coll = self.db.uuid diff --git a/test/asynchronous/test_operation_id_retry.py b/test/asynchronous/test_operation_id_retry.py index eedd25547a..8d08d2b8d0 100644 --- a/test/asynchronous/test_operation_id_retry.py +++ b/test/asynchronous/test_operation_id_retry.py @@ -152,13 +152,16 @@ async def test_retry_without_telemetry_creates_no_operation_id(self): find_op_ids = [] original_init = _CommandTelemetry.__init__ + # Accept and forward any trailing arguments (e.g. tracing_options, + # speculative_hello) so this stays working as _CommandTelemetry gains + # parameters; only cmd and op_id are of interest here. def recording_init( - self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, name=None + self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, *args, **kwargs ): if next(iter(cmd)) == "find": find_op_ids.append(op_id) original_init( - self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, name=name + self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, *args, **kwargs ) fail_point = { diff --git a/test/asynchronous/test_otel.py b/test/asynchronous/test_otel.py index 7eaeafe734..4a46625b8f 100644 --- a/test/asynchronous/test_otel.py +++ b/test/asynchronous/test_otel.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Test OpenTelemetry command-span support.""" +"""Test OpenTelemetry operation spans, excluding cursor getMores.""" from __future__ import annotations +import gc import os import sys from typing import Optional @@ -26,18 +27,30 @@ import pytest import pymongo._otel as _otel -from pymongo import common -from pymongo.errors import ConfigurationError, OperationFailure +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, unittest +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 import TracerProvider 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: @@ -48,29 +61,261 @@ pytestmark = pytest.mark.otel -def _shared_test_provider() -> TracerProvider: - """Return a process-wide SDK TracerProvider for tests to attach exporters to. +def _tracing_opts() -> _otel.TracingOptions: + """Return tracing options with tracing on and ``db.query.text`` disabled.""" + return {"enabled": True, "query_text_max_length": None} + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelOperationSpanPrimitives(unittest.TestCase): + """Unit tests for the pymongo._otel operation-span primitives.""" + + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def _finished_span(self, operation: str = "find", **kwargs): + """Start an operation span, end it successfully, and return the one finished span.""" + handle = _otel.start_operation_span(_tracing_opts(), operation, None, **kwargs) + self.assertIsNotNone(handle) + _otel.end_operation_span_success(handle) + (span,) = self.exporter.get_finished_spans() + return span + + def test_start_operation_span_success_sets_provisional_attributes(self): + span = self._finished_span() + self.assertEqual(span.name, "find") + self.assertEqual(span.attributes["db.system.name"], "mongodb") + self.assertEqual(span.attributes["db.operation.name"], "find") + self.assertEqual(span.status.status_code, StatusCode.UNSET) + + def test_start_operation_span_failure_records_exception(self): + handle = _otel.start_operation_span(_tracing_opts(), "insert", None) + _otel.end_operation_span_failure(handle, ValueError("boom")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + self.assertEqual(len(span.events), 1) + self.assertEqual(span.events[0].name, "exception") + + def test_start_operation_span_with_parent(self): + parent_handle = _otel.start_operation_span(_tracing_opts(), "transaction", None) + handle = _otel.start_operation_span(_tracing_opts(), "insert", parent_handle.span) + _otel.end_operation_span_success(handle) + _otel.end_operation_span_success(parent_handle) + child, parent = self.exporter.get_finished_spans() + self.assertEqual(child.parent.span_id, parent.context.span_id) + + def test_current_operation_name_contextvar_scoped_correctly(self): + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + handle = _otel.start_operation_span(_tracing_opts(), "find", None) + self.assertEqual(_otel._CURRENT_OPERATION_NAME.get(), "find") + _otel.end_operation_span_success(handle) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + + def test_eager_dbname_and_collection_set_at_creation(self): + span = self._finished_span(dbname="mydb", collection="mycoll") + self.assertEqual(span.name, "find mydb.mycoll") + self.assertEqual(span.attributes["db.namespace"], "mydb") + self.assertEqual(span.attributes["db.collection.name"], "mycoll") + self.assertEqual(span.attributes["db.operation.summary"], "find mydb.mycoll") + self.assertEqual(span.attributes["db.operation.name"], "find") + + def test_eager_dbname_without_collection_omits_collection_attribute(self): + span = self._finished_span("listCollections", dbname="mydb") + self.assertEqual(span.name, "listCollections mydb") + self.assertEqual(span.attributes["db.operation.summary"], "listCollections mydb") + self.assertNotIn("db.collection.name", span.attributes) + + def test_no_eager_attributes_leaves_provisional_name(self): + span = self._finished_span() + self.assertEqual(span.name, "find") + self.assertNotIn("db.namespace", span.attributes) + # db.operation.summary is Required (unlike db.namespace, which is only + # "Required if available"), so it always falls back to the bare + # operation name when no dbname is given. + self.assertEqual(span.attributes["db.operation.summary"], "find") + + def test_detached_span_is_not_current_until_used(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + self.assertIsNotNone(handle) + # Not current, and the operation-name contextvar is untouched. + self.assertIsNot(trace.get_current_span(), handle.span) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + with _otel.use_operation_span(handle): + self.assertIs(trace.get_current_span(), handle.span) + self.assertEqual(_otel._CURRENT_OPERATION_NAME.get(), "find") + # Restored afterwards, and the span is still open (not ended). + self.assertIsNot(trace.get_current_span(), handle.span) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + self.assertEqual(self.exporter.get_finished_spans(), ()) + _otel.end_operation_span_success(handle) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "find") + + def test_detached_span_reused_across_multiple_use_blocks(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + for _ in range(3): + with _otel.use_operation_span(handle): + pass + self.assertEqual(self.exporter.get_finished_spans(), ()) + _otel.end_operation_span_success(handle) + self.assertEqual(len(self.exporter.get_finished_spans()), 1) + + def test_use_operation_span_with_none_handle_is_noop(self): + with _otel.use_operation_span(None): + pass + self.assertEqual(self.exporter.get_finished_spans(), ()) + + def test_detached_span_failure_inside_use_block_records_exception_once(self): + # Regression test: use_span's record_exception/set_status_on_exception + # default to True, so without disabling them an exception leaving the + # block is recorded twice, here and by end_operation_span_failure. + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + exc = ValueError("boom") + try: + with _otel.use_operation_span(handle): + raise exc + except ValueError: + pass + _otel.end_operation_span_failure(handle, exc) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + exception_events = [e for e in span.events if e.name == "exception"] + self.assertEqual(len(exception_events), 1) + + def test_detached_span_failure_without_use_block(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + _otel.end_operation_span_failure(handle, ValueError("boom")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + exception_events = [e for e in span.events if e.name == "exception"] + self.assertEqual(len(exception_events), 1) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetry(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_succeeded_with_no_session(self): + telemetry = _telemetry._OperationTelemetry(_tracing_opts(), "find", None) + telemetry.succeeded() + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "find") + self.assertIsNone(span.parent) + + def test_failed_records_exception(self): + telemetry = _telemetry._OperationTelemetry(_tracing_opts(), "insert", None) + telemetry.failed(RuntimeError("nope")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + + def test_disabled_is_a_no_op(self): + telemetry = _telemetry._OperationTelemetry(None, "find", None) + telemetry.succeeded() # must not raise + telemetry2 = _telemetry._OperationTelemetry(None, "find", None) + telemetry2.failed(RuntimeError("x")) # must not raise + self.assertEqual(self.exporter.get_finished_spans(), ()) + + def test_run_command_operation_name_override(self): + # Per the spec, Database.command() produces a "runCommand" operation + # span, not one named after the command actually sent. + telemetry = _telemetry._OperationTelemetry( + _tracing_opts(), "ping", None, is_run_command=True + ) + telemetry.succeeded() + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "runCommand") + self.assertEqual(span.attributes["db.operation.name"], "runCommand") - ``trace.set_tracer_provider`` only takes effect once per process (later calls - are silently ignored), so tests must share one provider and each register - their own span processor rather than trying to install a fresh provider. - """ - current = trace.get_tracer_provider() - if isinstance(current, TracerProvider): - return current - provider = TracerProvider() - trace.set_tracer_provider(provider) - return provider + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetryContextManager(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_context_manager_success_ends_span(self): + with _OperationTelemetry( + _tracing_opts(), "killCursors", None, dbname="mydb", collection="c" + ): + pass + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "killCursors mydb.c") + self.assertEqual(span.status.status_code, StatusCode.UNSET) + + def test_context_manager_failure_records_exception(self): + with self.assertRaises(ValueError): + with _OperationTelemetry(_tracing_opts(), "killCursors", None, dbname="mydb"): + raise ValueError("boom") + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + self.assertEqual(span.attributes["exception.type"], "ValueError") + + def test_detached_telemetry_use_makes_span_current(self): + telemetry = _OperationTelemetry( + _tracing_opts(), "find", None, dbname="mydb", collection="c", set_current=False + ) + self.assertIsNot(trace.get_current_span(), telemetry.handle.span) + with telemetry.use(): + self.assertIs(trace.get_current_span(), telemetry.handle.span) + self.assertEqual(self.exporter.get_finished_spans(), ()) + telemetry.succeeded() + self.assertEqual(len(self.exporter.get_finished_spans()), 1) @unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") class TestOTelSpans(AsyncIntegrationTest): + """Operation and command spans for a single round trip.""" + @classmethod def setUpClass(cls): super().setUpClass() cls.exporter = InMemorySpanExporter() _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + @classmethod + def tearDownClass(cls): + # A span processor cannot be removed from the shared process-wide + # TracerProvider, so without this shutdown() the exporter accumulates + # every span from every client for the rest of the run. + cls.exporter.shutdown() + super().tearDownClass() + async def asyncSetUp(self): await super().asyncSetUp() self.exporter.clear() @@ -81,52 +326,69 @@ def spans(self, name: str | None = None): return list(finished) return [s for s in finished if s.name == name] - # TODO(PYTHON-5947): once the unified test format runner supports - # expectTracingMessages/operation spans, this is superseded by the spec's - # find_without_query_text.yml and insert.yml. - async def test_span_created_for_insert_and_find(self): + @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 tests asserting tracing produced *nothing*. An empty-exporter + assertion would also catch unrelated spans, since a cursor abandoned + earlier ends its span from a finalizer that runs at an unpredictable + point on interpreters without reference counting. + """ + 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_operation_span_records_failure(self): client = await self.async_rs_or_single_client(tracing={"enabled": True}) - coll = client[self.db.name].test_otel - await coll.drop() - self.exporter.clear() - await coll.insert_one({"x": 1}) - - insert_spans = self.spans("insert") - self.assertEqual(len(insert_spans), 1) - attrs = insert_spans[0].attributes - self.assertEqual(attrs["db.system.name"], "mongodb") - self.assertEqual(attrs["db.namespace"], self.db.name) - self.assertEqual(attrs["db.collection.name"], "test_otel") - self.assertEqual(attrs["db.command.name"], "insert") - self.assertEqual(attrs["db.query.summary"], f"insert {self.db.name}.test_otel") - self.assertIn("server.address", attrs) - self.assertIn("server.port", attrs) - self.assertIn(attrs["network.transport"], ("tcp", "unix")) - self.assertIn("db.mongodb.driver_connection_id", attrs) - self.assertNotIn("db.query.text", attrs) - - self.exporter.clear() - docs = await coll.find({}).to_list() - self.assertEqual(len(docs), 1) - find_spans = self.spans("find") - self.assertEqual(len(find_spans), 1) - self.assertEqual(find_spans[0].attributes["db.command.name"], "find") - - 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)]) + coll = client[self.db.name].test 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") + with self.assertRaises(OperationFailure): + await coll.find_one({"$invalidOperator": 1}) + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "find" + ] + self.assertEqual(len(matching), 1) + op_span = matching[0] + self.assertEqual(op_span.status.status_code, StatusCode.ERROR) + # The spec requires exception.type/message/stacktrace as *attributes* + # on the operation span, not just the event record_exception attaches. + self.assertTrue(any(event.name == "exception" for event in op_span.events)) + self.assertIn("exception.type", op_span.attributes) + self.assertIn("exception.message", op_span.attributes) + self.assertIn("exception.stacktrace", op_span.attributes) async def test_explain_retains_collection_name(self): # explain wraps the real command ({"explain": {"find": "coll", ...}}), the @@ -171,8 +433,23 @@ async def test_sensitive_command_produces_no_span(self): with self.assertRaises(OperationFailure): await client.admin.command("saslStart", mechanism="SCRAM-SHA-256", payload=b"") - names = [s.name for s in self.spans()] - self.assertNotIn("saslStart", names) + # The inner command span must stay fully suppressed for sensitive commands. + command_span_names = [s.name for s in self.spans() if "db.command.name" in s.attributes] + self.assertNotIn("saslStart", command_span_names) + + # The sensitive name must not leak onto the operation span either. + # Database.command() runs with is_run_command=True, so that span reads + # "runCommand" and "saslStart" appears nowhere. It still carries the + # Required db.namespace/db.operation.summary, backfilled before + # start_command_span's sensitive-command return. + finished = self.exporter.get_finished_spans() + operation_names = [s.attributes.get("db.operation.name") for s in finished] + self.assertNotIn("saslStart", operation_names) + self.assertIn("runCommand", operation_names) + op_span = next(s for s in finished if s.attributes.get("db.operation.name") == "runCommand") + self.assertEqual(op_span.name, "runCommand admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + self.assertEqual(op_span.attributes["db.operation.summary"], "runCommand admin") async def test_admin_command_omits_collection_name(self): # usersInfo's command value is a username string, not a collection, and @@ -188,13 +465,35 @@ async def test_admin_command_omits_collection_name(self): self.assertNotIn("db.collection.name", attrs) self.assertEqual(attrs["db.query.summary"], "usersInfo admin") + async def test_database_command_produces_run_command_operation_span(self): + # The spec names anything reached through Database.command() + # "runCommand", so admin.command("ping") yields an operation span named + # "runCommand admin", not "ping admin". + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + await client.admin.command("ping") + + finished = self.exporter.get_finished_spans() + matching = [s for s in finished if s.attributes.get("db.operation.name") == "runCommand"] + self.assertEqual(len(matching), 1) + op_span = matching[0] + self.assertEqual(op_span.name, "runCommand admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + # The wire-level command span is unaffected: it's still named/attributed + # after the actual command sent. + cmd_spans = [s for s in finished if s.attributes.get("db.command.name") == "ping"] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].name, "ping") + async def test_failure_records_exception_and_status_code(self): client = await self.async_rs_or_single_client(tracing={"enabled": True}) self.exporter.clear() with self.assertRaises(OperationFailure): await client[self.db.name].command("thisCommandDoesNotExist") - spans = self.spans() + # This also produces an ERROR operation span, so narrow to the command + # span, which alone carries db.response.status_code. + spans = [s for s in self.spans() if "db.response.status_code" in s.attributes] self.assertEqual(len(spans), 1) span = spans[0] self.assertEqual(span.status.status_code, trace.StatusCode.ERROR) @@ -205,31 +504,65 @@ async def test_tracing_disabled_by_default(self): client = await self.async_rs_or_single_client() self.exporter.clear() await client.admin.command("ping") - self.assertEqual(self.spans(), []) + self.assertEqual(self.ping_spans(), []) - # TODO(PYTHON-5947): once operation spans exist, also assert that the - # "ping" *operation* span (not just the command span) is absent/present - # here, and that self.spans() counts both. async def test_prose_1_tracing_enable_disable_via_env_var(self): """Prose Test 1: Tracing Enable/Disable via Environment Variable.""" with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "false"}): client = await self.async_rs_or_single_client() self.exporter.clear() await client.admin.command("ping") - self.assertEqual(self.spans(), []) + # When tracing is disabled we must suppress both the operation span and + # the command span it wraps: db.command() routes through _retry_internal + # same as any CRUD call, so both would exist if tracing weren't fully off. + self.assertEqual(self.ping_spans(), []) with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true"}): client = await self.async_rs_or_single_client() self.exporter.clear() await client.admin.command("ping") - self.assertIn("ping", [s.name for s in self.spans()]) + finished = self.exporter.get_finished_spans() + # start_command_span renames the operation span in place, so span.name + # cannot tell the two apart; db.command.name and db.operation.name can. + # The operation span reads "runCommand" rather than "ping". + self.assertIn("ping", [s.attributes.get("db.command.name") for s in finished]) + self.assertIn("runCommand", [s.attributes.get("db.operation.name") for s in finished]) + + async def test_env_var_tracing_does_not_trace_monitor_commands(self): + # Monitor and handshake connections have no client, so their tracing + # options are None. Enablement must not fall back to the environment + # variable there, or every hello would get a span the spec excludes. + self.assertFalse(_otel._is_tracing_enabled(None)) + with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true"}): + self.assertFalse(_otel._is_tracing_enabled(None)) + client = await self.async_rs_or_single_client() + self.exporter.clear() + await client.admin.command("ping") + + finished = self.exporter.get_finished_spans() + # The env var still enables tracing for the client's own commands... + self.assertIn("ping", [s.attributes.get("db.command.name") for s in finished]) + # ...but nothing traces the monitors' hellos. + hello_spans = [ + s + for s in finished + if s.attributes.get("db.command.name") in _HELLO_COMMANDS or s.name in _HELLO_COMMANDS + ] + self.assertEqual(hello_spans, [], [s.name for s in finished]) - # TODO(PYTHON-5947): once operation spans exist, self.spans("find") will - # also match the outer find *operation* span; disambiguate (e.g. by - # db.command.name vs db.operation.name) so this only asserts on the - # command span's db.query.text attribute. async def test_prose_2_command_payload_emission_via_env_var(self): """Prose Test 2: Command Payload Emission via Environment Variable.""" + + def command_spans(): + # self.spans("find") would also match the outer operation span, so + # filter on db.command.name to isolate the command span that + # carries db.query.text. + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "find" + ] + env = { "OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true", "OTEL_PYTHON_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH": "1024", @@ -238,7 +571,7 @@ async def test_prose_2_command_payload_emission_via_env_var(self): client = await self.async_rs_or_single_client() self.exporter.clear() await client[self.db.name].test_otel.find({}).to_list() - spans = self.spans("find") + spans = command_spans() self.assertEqual(len(spans), 1) self.assertIn("db.query.text", spans[0].attributes) @@ -246,27 +579,10 @@ async def test_prose_2_command_payload_emission_via_env_var(self): client = await self.async_rs_or_single_client() self.exporter.clear() await client[self.db.name].test_otel.find({}).to_list() - spans = self.spans("find") + spans = command_spans() self.assertEqual(len(spans), 1) self.assertNotIn("db.query.text", spans[0].attributes) - # TODO(PYTHON-5947): once the unified test format runner supports - # expectTracingMessages/operation spans, this is superseded by the spec's - # find.yml (db.query.text assertion). - async def test_query_text_included_when_configured(self): - client = await self.async_rs_or_single_client( - tracing={"enabled": True, "query_text_max_length": 1000} - ) - coll = client[self.db.name].test_otel - await coll.drop() - self.exporter.clear() - await coll.insert_one({"x": 1}) - - spans = self.spans("insert") - self.assertEqual(len(spans), 1) - self.assertIn("db.query.text", spans[0].attributes) - self.assertNotIn("lsid", spans[0].attributes["db.query.text"]) - async def test_explicit_query_text_max_length_zero_overrides_env_var(self): # An explicit client-side 0 must win over the environment variable, unlike # unset (which defers to it) - otherwise an app can't reliably opt out. @@ -300,10 +616,162 @@ async def test_query_text_truncation_shrinks_oversized_field_values(self): self.assertLessEqual(len(query_text), 200) self.assertNotIn("a" * 500, query_text) + @async_client_context.require_version_min(8, 0, 0, -24) + async def test_bulk_write_unacknowledged_gets_operation_span(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}, w=0) + self.exporter.clear() + await client.bulk_write( + [InsertOne(namespace=f"{self.db.name}.test", document={"x": 1})], + ordered=False, + ) + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "bulkWrite" + ] + self.assertEqual(len(matching), 1) + self.assertEqual(matching[0].attributes["db.namespace"], "admin") + + async def test_operation_span_falls_back_to_bare_name_when_no_command_is_sent(self): + # Failing during server selection builds no command, so the backfill in + # start_command_span never runs, and insert_one threads no namespace + # eagerly. db.operation.summary (Required) falls back to the bare + # operation name; db.namespace/db.collection.name are absent. + client = await self.async_rs_or_single_client( + "mongodb://localhost:1/", + tracing={"enabled": True}, + serverSelectionTimeoutMS=10, + connect=False, + ) + self.exporter.clear() + with self.assertRaises(ServerSelectionTimeoutError): + await client.mydb.mycoll.insert_one({}) + (span,) = [s for s in self.exporter.get_finished_spans() if s.name == "insert"] + self.assertEqual(span.attributes["db.operation.name"], "insert") + self.assertEqual(span.attributes["db.operation.summary"], "insert") + self.assertNotIn("db.namespace", span.attributes) + self.assertNotIn("db.collection.name", span.attributes) + self.assertEqual(span.status.status_code, StatusCode.ERROR) + + async def test_operation_span_name_can_differ_from_command_name(self): + # count_documents' operation span is named "count" but sends an + # aggregate, so an operation span name is not the command beneath it. + # count.json covers estimated_document_count, where the two coincide. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + await db.mycoll.insert_one({"x": 1}) + self.exporter.clear() + await db.mycoll.count_documents({}) + + (op_span,) = self.spans("count pymongo_test.mycoll") + self.assertEqual(op_span.attributes["db.operation.name"], "count") + self.assertEqual(op_span.attributes["db.namespace"], "pymongo_test") + (cmd_span,) = self.spans("aggregate") + self.assertEqual(cmd_span.attributes["db.command.name"], "aggregate") + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + async def test_kill_cursors_gets_operation_span(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.kill_cursors_span + await coll.drop() + await coll.insert_many([{"i": i} for i in range(10)]) + cursor = coll.find({}, batch_size=2) + await cursor.next() + self.exporter.clear() + await cursor.close() # Sends killCursors, since batches remain. + + op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "killCursors" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(op_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + (op_span,) = op_spans + self.assertEqual(op_span.name, "killCursors pymongo_test.kill_cursors_span") + self.assertEqual(op_span.attributes["db.namespace"], "pymongo_test") + self.assertEqual(op_span.attributes["db.collection.name"], "kill_cursors_span") + + cmd_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "killCursors" + ] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].parent.span_id, op_span.context.span_id) + + async def test_background_kill_cursors_span_is_a_trace_root(self): + # Regression test for PYTHON-5947: create_task freezes the calling + # coroutine's context, so without the reset in + # AsyncPeriodicExecutor._run every killCursors span the background tick + # emits is parented under whatever operation opened the executor. + # + # connect=False keeps _get_topology() out of construction, so the + # executor opens inside coll.drop() below with that span current. Waking + # the existing task runs the tick in that frozen context; calling + # _process_kill_cursors() here would use this coroutine's clean one. + client = await self.async_rs_or_single_client(tracing={"enabled": True}, connect=False) + coll = client.pymongo_test.bg_kill_cursors + await coll.drop() + await coll.insert_many([{"i": i} for i in range(10)]) + + cursor = coll.find({}, batch_size=2) + await cursor.next() + del cursor + gc.collect() # Queues a deferred killCursors. + + self.exporter.clear() + + def _kill_op_spans(): + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "killCursors" + and "db.command.name" not in s.attributes + ] + + executor = client._kill_cursors_executor + executor.skip_sleep() + executor.wake() + await async_wait_until(_kill_op_spans, "background killCursors span emitted") + + kill_spans = _kill_op_spans() + self.assertEqual(len(kill_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + # The background tick must not inherit a parent from whatever span + # happened to be current when the executor task was created. + self.assertIsNone(kill_spans[0].parent) + + async def test_end_sessions_gets_operation_span(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + await client.pymongo_test.end_sessions_span.find_one({}) # Uses an implicit session. + self.exporter.clear() + await client.close() # Sends endSessions. + + op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "endSessions" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(op_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + (op_span,) = op_spans + self.assertEqual(op_span.name, "endSessions admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + self.assertNotIn("db.collection.name", op_span.attributes) + + cmd_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "endSessions" + ] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].parent.span_id, op_span.context.span_id) + + +# These unit tests cover the validator's edge cases: the rejection paths and the +# explicit-zero vs unset distinction for query_text_max_length. + -# TODO(PYTHON-5947): superseded once the unified test format's -# expectTracingMessages/observeTracingMessages tests exercise this validator -# indirectly through real client construction; remove this class then. class TestValidateTracingOrNone(unittest.TestCase): def test_none(self): self.assertIsNone(common.validate_tracing_or_none("tracing", None)) @@ -352,12 +820,11 @@ def test_rejects_negative_query_text_max_length(self): class TestOTelTracerCaching(unittest.TestCase): - """Regression test for the tracer-caching implementation in ``pymongo/_otel.py``. + """Regression test for the tracer caching in ``pymongo/_otel.py``. - ``opentelemetry.trace.get_tracer()`` must only be called once, at import - time (cached as module-level ``_otel._TRACER``). Calling it per command - allocates two objects, takes a process-wide lock, and mutates the global - ``warnings`` filter list on every call, even on a cache hit. + ``get_tracer()`` must be called once at import time, cached as + ``_otel._TRACER``. Per command it would allocate, take a process-wide lock, + and mutate the global ``warnings`` filter list even on a cache hit. """ @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") diff --git a/test/test_common.py b/test/test_common.py index b0c706b9f2..705dc5f78e 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -25,12 +25,26 @@ from bson.codec_options import CodecOptions from bson.objectid import ObjectId from pymongo.errors import OperationFailure +from pymongo.helpers_shared import _split_namespace from pymongo.write_concern import WriteConcern from test import IntegrationTest, client_context, connected, unittest _IS_SYNC = True +class TestSplitNamespace(unittest.TestCase): + def test_plain_namespace(self): + self.assertEqual(_split_namespace("db.coll"), ("db", "coll")) + + def test_collection_name_with_dots(self): + self.assertEqual(_split_namespace("db.coll.with.dots"), ("db", "coll.with.dots")) + + def test_no_dot(self): + # No separator: the whole string is treated as the database name + # and the collection name is empty, matching str.partition. + self.assertEqual(_split_namespace("dbonly"), ("dbonly", "")) + + class TestCommon(IntegrationTest): def test_uuid_representation(self): coll = self.db.uuid diff --git a/test/test_json_util.py b/test/test_json_util.py index ada04955c4..d049674208 100644 --- a/test/test_json_util.py +++ b/test/test_json_util.py @@ -657,6 +657,27 @@ class MyBinary(Binary): expected_json = json_util.dumps(Binary(b"bin", USER_DEFINED_SUBTYPE)) self.assertEqual(json_util.dumps(MyBinary(b"bin", USER_DEFINED_SUBTYPE)), expected_json) + def test_truncate_documents_retains_falsy_values(self): + # Regression test: _truncate_documents (used by pymongo/logger.py for + # structured command logging, and by pymongo/_otel.py for OTel's + # db.query.text) must not drop fields whose value is falsy-but-present + # (0, False, "", {}, []); only fields that genuinely don't fit within + # the remaining budget should be omitted. A prior implementation used + # `if truncated_v:` to decide whether to keep a field, which silently + # dropped legitimate falsy values along with truly-out-of-room ones. + doc = { + "a": 0, + "b": False, + "c": "", + "d": {}, + "e": [], + "f": None, + "g": [0, False, "", {}, [], None], + } + truncated, remaining = json_util._truncate_documents(doc, 1000) + self.assertEqual(truncated, doc) + self.assertGreater(remaining, 0) + if __name__ == "__main__": unittest.main() diff --git a/test/test_operation_id_retry.py b/test/test_operation_id_retry.py index 237447f994..e4e50c2b9c 100644 --- a/test/test_operation_id_retry.py +++ b/test/test_operation_id_retry.py @@ -150,13 +150,16 @@ def test_retry_without_telemetry_creates_no_operation_id(self): find_op_ids = [] original_init = _CommandTelemetry.__init__ + # Accept and forward any trailing arguments (e.g. tracing_options, + # speculative_hello) so this stays working as _CommandTelemetry gains + # parameters; only cmd and op_id are of interest here. def recording_init( - self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, name=None + self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, *args, **kwargs ): if next(iter(cmd)) == "find": find_op_ids.append(op_id) original_init( - self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, name=name + self, topology_id, conn, listeners, cmd, dbname, request_id, op_id, *args, **kwargs ) fail_point = { diff --git a/test/test_otel.py b/test/test_otel.py index d0e3a55fe5..dfe28db351 100644 --- a/test/test_otel.py +++ b/test/test_otel.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Test OpenTelemetry command-span support.""" +"""Test OpenTelemetry operation spans, excluding cursor getMores.""" from __future__ import annotations +import gc import os import sys from typing import Optional @@ -26,18 +27,30 @@ import pytest import pymongo._otel as _otel -from pymongo import common -from pymongo.errors import ConfigurationError, OperationFailure +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, unittest +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 import TracerProvider 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: @@ -48,29 +61,261 @@ pytestmark = pytest.mark.otel -def _shared_test_provider() -> TracerProvider: - """Return a process-wide SDK TracerProvider for tests to attach exporters to. +def _tracing_opts() -> _otel.TracingOptions: + """Return tracing options with tracing on and ``db.query.text`` disabled.""" + return {"enabled": True, "query_text_max_length": None} - ``trace.set_tracer_provider`` only takes effect once per process (later calls - are silently ignored), so tests must share one provider and each register - their own span processor rather than trying to install a fresh provider. - """ - current = trace.get_tracer_provider() - if isinstance(current, TracerProvider): - return current - provider = TracerProvider() - trace.set_tracer_provider(provider) - return provider + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelOperationSpanPrimitives(unittest.TestCase): + """Unit tests for the pymongo._otel operation-span primitives.""" + + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def _finished_span(self, operation: str = "find", **kwargs): + """Start an operation span, end it successfully, and return the one finished span.""" + handle = _otel.start_operation_span(_tracing_opts(), operation, None, **kwargs) + self.assertIsNotNone(handle) + _otel.end_operation_span_success(handle) + (span,) = self.exporter.get_finished_spans() + return span + + def test_start_operation_span_success_sets_provisional_attributes(self): + span = self._finished_span() + self.assertEqual(span.name, "find") + self.assertEqual(span.attributes["db.system.name"], "mongodb") + self.assertEqual(span.attributes["db.operation.name"], "find") + self.assertEqual(span.status.status_code, StatusCode.UNSET) + + def test_start_operation_span_failure_records_exception(self): + handle = _otel.start_operation_span(_tracing_opts(), "insert", None) + _otel.end_operation_span_failure(handle, ValueError("boom")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + self.assertEqual(len(span.events), 1) + self.assertEqual(span.events[0].name, "exception") + + def test_start_operation_span_with_parent(self): + parent_handle = _otel.start_operation_span(_tracing_opts(), "transaction", None) + handle = _otel.start_operation_span(_tracing_opts(), "insert", parent_handle.span) + _otel.end_operation_span_success(handle) + _otel.end_operation_span_success(parent_handle) + child, parent = self.exporter.get_finished_spans() + self.assertEqual(child.parent.span_id, parent.context.span_id) + + def test_current_operation_name_contextvar_scoped_correctly(self): + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + handle = _otel.start_operation_span(_tracing_opts(), "find", None) + self.assertEqual(_otel._CURRENT_OPERATION_NAME.get(), "find") + _otel.end_operation_span_success(handle) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + + def test_eager_dbname_and_collection_set_at_creation(self): + span = self._finished_span(dbname="mydb", collection="mycoll") + self.assertEqual(span.name, "find mydb.mycoll") + self.assertEqual(span.attributes["db.namespace"], "mydb") + self.assertEqual(span.attributes["db.collection.name"], "mycoll") + self.assertEqual(span.attributes["db.operation.summary"], "find mydb.mycoll") + self.assertEqual(span.attributes["db.operation.name"], "find") + + def test_eager_dbname_without_collection_omits_collection_attribute(self): + span = self._finished_span("listCollections", dbname="mydb") + self.assertEqual(span.name, "listCollections mydb") + self.assertEqual(span.attributes["db.operation.summary"], "listCollections mydb") + self.assertNotIn("db.collection.name", span.attributes) + + def test_no_eager_attributes_leaves_provisional_name(self): + span = self._finished_span() + self.assertEqual(span.name, "find") + self.assertNotIn("db.namespace", span.attributes) + # db.operation.summary is Required (unlike db.namespace, which is only + # "Required if available"), so it always falls back to the bare + # operation name when no dbname is given. + self.assertEqual(span.attributes["db.operation.summary"], "find") + + def test_detached_span_is_not_current_until_used(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + self.assertIsNotNone(handle) + # Not current, and the operation-name contextvar is untouched. + self.assertIsNot(trace.get_current_span(), handle.span) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + with _otel.use_operation_span(handle): + self.assertIs(trace.get_current_span(), handle.span) + self.assertEqual(_otel._CURRENT_OPERATION_NAME.get(), "find") + # Restored afterwards, and the span is still open (not ended). + self.assertIsNot(trace.get_current_span(), handle.span) + self.assertIsNone(_otel._CURRENT_OPERATION_NAME.get()) + self.assertEqual(self.exporter.get_finished_spans(), ()) + _otel.end_operation_span_success(handle) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "find") + + def test_detached_span_reused_across_multiple_use_blocks(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + for _ in range(3): + with _otel.use_operation_span(handle): + pass + self.assertEqual(self.exporter.get_finished_spans(), ()) + _otel.end_operation_span_success(handle) + self.assertEqual(len(self.exporter.get_finished_spans()), 1) + + def test_use_operation_span_with_none_handle_is_noop(self): + with _otel.use_operation_span(None): + pass + self.assertEqual(self.exporter.get_finished_spans(), ()) + + def test_detached_span_failure_inside_use_block_records_exception_once(self): + # Regression test: use_span's record_exception/set_status_on_exception + # default to True, so without disabling them an exception leaving the + # block is recorded twice, here and by end_operation_span_failure. + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + exc = ValueError("boom") + try: + with _otel.use_operation_span(handle): + raise exc + except ValueError: + pass + _otel.end_operation_span_failure(handle, exc) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + exception_events = [e for e in span.events if e.name == "exception"] + self.assertEqual(len(exception_events), 1) + + def test_detached_span_failure_without_use_block(self): + handle = _otel.start_operation_span(_tracing_opts(), "find", None, set_current=False) + _otel.end_operation_span_failure(handle, ValueError("boom")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + exception_events = [e for e in span.events if e.name == "exception"] + self.assertEqual(len(exception_events), 1) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetry(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_succeeded_with_no_session(self): + telemetry = _telemetry._OperationTelemetry(_tracing_opts(), "find", None) + telemetry.succeeded() + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "find") + self.assertIsNone(span.parent) + + def test_failed_records_exception(self): + telemetry = _telemetry._OperationTelemetry(_tracing_opts(), "insert", None) + telemetry.failed(RuntimeError("nope")) + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + + def test_disabled_is_a_no_op(self): + telemetry = _telemetry._OperationTelemetry(None, "find", None) + telemetry.succeeded() # must not raise + telemetry2 = _telemetry._OperationTelemetry(None, "find", None) + telemetry2.failed(RuntimeError("x")) # must not raise + self.assertEqual(self.exporter.get_finished_spans(), ()) + + def test_run_command_operation_name_override(self): + # Per the spec, Database.command() produces a "runCommand" operation + # span, not one named after the command actually sent. + telemetry = _telemetry._OperationTelemetry( + _tracing_opts(), "ping", None, is_run_command=True + ) + telemetry.succeeded() + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "runCommand") + self.assertEqual(span.attributes["db.operation.name"], "runCommand") + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetryContextManager(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_context_manager_success_ends_span(self): + with _OperationTelemetry( + _tracing_opts(), "killCursors", None, dbname="mydb", collection="c" + ): + pass + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.name, "killCursors mydb.c") + self.assertEqual(span.status.status_code, StatusCode.UNSET) + + def test_context_manager_failure_records_exception(self): + with self.assertRaises(ValueError): + with _OperationTelemetry(_tracing_opts(), "killCursors", None, dbname="mydb"): + raise ValueError("boom") + (span,) = self.exporter.get_finished_spans() + self.assertEqual(span.status.status_code, StatusCode.ERROR) + self.assertEqual(span.attributes["exception.type"], "ValueError") + + def test_detached_telemetry_use_makes_span_current(self): + telemetry = _OperationTelemetry( + _tracing_opts(), "find", None, dbname="mydb", collection="c", set_current=False + ) + self.assertIsNot(trace.get_current_span(), telemetry.handle.span) + with telemetry.use(): + self.assertIs(trace.get_current_span(), telemetry.handle.span) + self.assertEqual(self.exporter.get_finished_spans(), ()) + telemetry.succeeded() + self.assertEqual(len(self.exporter.get_finished_spans()), 1) @unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") class TestOTelSpans(IntegrationTest): + """Operation and command spans for a single round trip.""" + @classmethod def setUpClass(cls): super().setUpClass() cls.exporter = InMemorySpanExporter() _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + @classmethod + def tearDownClass(cls): + # A span processor cannot be removed from the shared process-wide + # TracerProvider, so without this shutdown() the exporter accumulates + # every span from every client for the rest of the run. + cls.exporter.shutdown() + super().tearDownClass() + def setUp(self): super().setUp() self.exporter.clear() @@ -81,52 +326,69 @@ def spans(self, name: str | None = None): return list(finished) return [s for s in finished if s.name == name] - # TODO(PYTHON-5947): once the unified test format runner supports - # expectTracingMessages/operation spans, this is superseded by the spec's - # find_without_query_text.yml and insert.yml. - def test_span_created_for_insert_and_find(self): - client = self.rs_or_single_client(tracing={"enabled": True}) - coll = client[self.db.name].test_otel - coll.drop() - self.exporter.clear() - coll.insert_one({"x": 1}) - - insert_spans = self.spans("insert") - self.assertEqual(len(insert_spans), 1) - attrs = insert_spans[0].attributes - self.assertEqual(attrs["db.system.name"], "mongodb") - self.assertEqual(attrs["db.namespace"], self.db.name) - self.assertEqual(attrs["db.collection.name"], "test_otel") - self.assertEqual(attrs["db.command.name"], "insert") - self.assertEqual(attrs["db.query.summary"], f"insert {self.db.name}.test_otel") - self.assertIn("server.address", attrs) - self.assertIn("server.port", attrs) - self.assertIn(attrs["network.transport"], ("tcp", "unix")) - self.assertIn("db.mongodb.driver_connection_id", attrs) - self.assertNotIn("db.query.text", attrs) - - self.exporter.clear() - docs = coll.find({}).to_list() - self.assertEqual(len(docs), 1) - find_spans = self.spans("find") - self.assertEqual(len(find_spans), 1) - self.assertEqual(find_spans[0].attributes["db.command.name"], "find") - - def test_span_created_for_get_more(self): + @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 tests asserting tracing produced *nothing*. An empty-exporter + assertion would also catch unrelated spans, since a cursor abandoned + earlier ends its span from a finalizer that runs at an unpredictable + point on interpreters without reference counting. + """ + 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_operation_span_records_failure(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)]) + coll = client[self.db.name].test 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") + with self.assertRaises(OperationFailure): + coll.find_one({"$invalidOperator": 1}) + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "find" + ] + self.assertEqual(len(matching), 1) + op_span = matching[0] + self.assertEqual(op_span.status.status_code, StatusCode.ERROR) + # The spec requires exception.type/message/stacktrace as *attributes* + # on the operation span, not just the event record_exception attaches. + self.assertTrue(any(event.name == "exception" for event in op_span.events)) + self.assertIn("exception.type", op_span.attributes) + self.assertIn("exception.message", op_span.attributes) + self.assertIn("exception.stacktrace", op_span.attributes) def test_explain_retains_collection_name(self): # explain wraps the real command ({"explain": {"find": "coll", ...}}), the @@ -171,8 +433,23 @@ def test_sensitive_command_produces_no_span(self): with self.assertRaises(OperationFailure): client.admin.command("saslStart", mechanism="SCRAM-SHA-256", payload=b"") - names = [s.name for s in self.spans()] - self.assertNotIn("saslStart", names) + # The inner command span must stay fully suppressed for sensitive commands. + command_span_names = [s.name for s in self.spans() if "db.command.name" in s.attributes] + self.assertNotIn("saslStart", command_span_names) + + # The sensitive name must not leak onto the operation span either. + # Database.command() runs with is_run_command=True, so that span reads + # "runCommand" and "saslStart" appears nowhere. It still carries the + # Required db.namespace/db.operation.summary, backfilled before + # start_command_span's sensitive-command return. + finished = self.exporter.get_finished_spans() + operation_names = [s.attributes.get("db.operation.name") for s in finished] + self.assertNotIn("saslStart", operation_names) + self.assertIn("runCommand", operation_names) + op_span = next(s for s in finished if s.attributes.get("db.operation.name") == "runCommand") + self.assertEqual(op_span.name, "runCommand admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + self.assertEqual(op_span.attributes["db.operation.summary"], "runCommand admin") def test_admin_command_omits_collection_name(self): # usersInfo's command value is a username string, not a collection, and @@ -188,13 +465,35 @@ def test_admin_command_omits_collection_name(self): self.assertNotIn("db.collection.name", attrs) self.assertEqual(attrs["db.query.summary"], "usersInfo admin") + def test_database_command_produces_run_command_operation_span(self): + # The spec names anything reached through Database.command() + # "runCommand", so admin.command("ping") yields an operation span named + # "runCommand admin", not "ping admin". + client = self.rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + client.admin.command("ping") + + finished = self.exporter.get_finished_spans() + matching = [s for s in finished if s.attributes.get("db.operation.name") == "runCommand"] + self.assertEqual(len(matching), 1) + op_span = matching[0] + self.assertEqual(op_span.name, "runCommand admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + # The wire-level command span is unaffected: it's still named/attributed + # after the actual command sent. + cmd_spans = [s for s in finished if s.attributes.get("db.command.name") == "ping"] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].name, "ping") + def test_failure_records_exception_and_status_code(self): client = self.rs_or_single_client(tracing={"enabled": True}) self.exporter.clear() with self.assertRaises(OperationFailure): client[self.db.name].command("thisCommandDoesNotExist") - spans = self.spans() + # This also produces an ERROR operation span, so narrow to the command + # span, which alone carries db.response.status_code. + spans = [s for s in self.spans() if "db.response.status_code" in s.attributes] self.assertEqual(len(spans), 1) span = spans[0] self.assertEqual(span.status.status_code, trace.StatusCode.ERROR) @@ -205,31 +504,65 @@ def test_tracing_disabled_by_default(self): client = self.rs_or_single_client() self.exporter.clear() client.admin.command("ping") - self.assertEqual(self.spans(), []) + self.assertEqual(self.ping_spans(), []) - # TODO(PYTHON-5947): once operation spans exist, also assert that the - # "ping" *operation* span (not just the command span) is absent/present - # here, and that self.spans() counts both. def test_prose_1_tracing_enable_disable_via_env_var(self): """Prose Test 1: Tracing Enable/Disable via Environment Variable.""" with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "false"}): client = self.rs_or_single_client() self.exporter.clear() client.admin.command("ping") - self.assertEqual(self.spans(), []) + # When tracing is disabled we must suppress both the operation span and + # the command span it wraps: db.command() routes through _retry_internal + # same as any CRUD call, so both would exist if tracing weren't fully off. + self.assertEqual(self.ping_spans(), []) with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true"}): client = self.rs_or_single_client() self.exporter.clear() client.admin.command("ping") - self.assertIn("ping", [s.name for s in self.spans()]) + finished = self.exporter.get_finished_spans() + # start_command_span renames the operation span in place, so span.name + # cannot tell the two apart; db.command.name and db.operation.name can. + # The operation span reads "runCommand" rather than "ping". + self.assertIn("ping", [s.attributes.get("db.command.name") for s in finished]) + self.assertIn("runCommand", [s.attributes.get("db.operation.name") for s in finished]) + + def test_env_var_tracing_does_not_trace_monitor_commands(self): + # Monitor and handshake connections have no client, so their tracing + # options are None. Enablement must not fall back to the environment + # variable there, or every hello would get a span the spec excludes. + self.assertFalse(_otel._is_tracing_enabled(None)) + with patch.dict(os.environ, {"OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true"}): + self.assertFalse(_otel._is_tracing_enabled(None)) + client = self.rs_or_single_client() + self.exporter.clear() + client.admin.command("ping") + + finished = self.exporter.get_finished_spans() + # The env var still enables tracing for the client's own commands... + self.assertIn("ping", [s.attributes.get("db.command.name") for s in finished]) + # ...but nothing traces the monitors' hellos. + hello_spans = [ + s + for s in finished + if s.attributes.get("db.command.name") in _HELLO_COMMANDS or s.name in _HELLO_COMMANDS + ] + self.assertEqual(hello_spans, [], [s.name for s in finished]) - # TODO(PYTHON-5947): once operation spans exist, self.spans("find") will - # also match the outer find *operation* span; disambiguate (e.g. by - # db.command.name vs db.operation.name) so this only asserts on the - # command span's db.query.text attribute. def test_prose_2_command_payload_emission_via_env_var(self): """Prose Test 2: Command Payload Emission via Environment Variable.""" + + def command_spans(): + # self.spans("find") would also match the outer operation span, so + # filter on db.command.name to isolate the command span that + # carries db.query.text. + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "find" + ] + env = { "OTEL_PYTHON_INSTRUMENTATION_MONGODB_ENABLED": "true", "OTEL_PYTHON_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH": "1024", @@ -238,7 +571,7 @@ def test_prose_2_command_payload_emission_via_env_var(self): client = self.rs_or_single_client() self.exporter.clear() client[self.db.name].test_otel.find({}).to_list() - spans = self.spans("find") + spans = command_spans() self.assertEqual(len(spans), 1) self.assertIn("db.query.text", spans[0].attributes) @@ -246,25 +579,10 @@ def test_prose_2_command_payload_emission_via_env_var(self): client = self.rs_or_single_client() self.exporter.clear() client[self.db.name].test_otel.find({}).to_list() - spans = self.spans("find") + spans = command_spans() self.assertEqual(len(spans), 1) self.assertNotIn("db.query.text", spans[0].attributes) - # TODO(PYTHON-5947): once the unified test format runner supports - # expectTracingMessages/operation spans, this is superseded by the spec's - # find.yml (db.query.text assertion). - def test_query_text_included_when_configured(self): - client = self.rs_or_single_client(tracing={"enabled": True, "query_text_max_length": 1000}) - coll = client[self.db.name].test_otel - coll.drop() - self.exporter.clear() - coll.insert_one({"x": 1}) - - spans = self.spans("insert") - self.assertEqual(len(spans), 1) - self.assertIn("db.query.text", spans[0].attributes) - self.assertNotIn("lsid", spans[0].attributes["db.query.text"]) - def test_explicit_query_text_max_length_zero_overrides_env_var(self): # An explicit client-side 0 must win over the environment variable, unlike # unset (which defers to it) - otherwise an app can't reliably opt out. @@ -294,10 +612,162 @@ def test_query_text_truncation_shrinks_oversized_field_values(self): self.assertLessEqual(len(query_text), 200) self.assertNotIn("a" * 500, query_text) + @client_context.require_version_min(8, 0, 0, -24) + def test_bulk_write_unacknowledged_gets_operation_span(self): + client = self.rs_or_single_client(tracing={"enabled": True}, w=0) + self.exporter.clear() + client.bulk_write( + [InsertOne(namespace=f"{self.db.name}.test", document={"x": 1})], + ordered=False, + ) + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "bulkWrite" + ] + self.assertEqual(len(matching), 1) + self.assertEqual(matching[0].attributes["db.namespace"], "admin") + + def test_operation_span_falls_back_to_bare_name_when_no_command_is_sent(self): + # Failing during server selection builds no command, so the backfill in + # start_command_span never runs, and insert_one threads no namespace + # eagerly. db.operation.summary (Required) falls back to the bare + # operation name; db.namespace/db.collection.name are absent. + client = self.rs_or_single_client( + "mongodb://localhost:1/", + tracing={"enabled": True}, + serverSelectionTimeoutMS=10, + connect=False, + ) + self.exporter.clear() + with self.assertRaises(ServerSelectionTimeoutError): + client.mydb.mycoll.insert_one({}) + (span,) = [s for s in self.exporter.get_finished_spans() if s.name == "insert"] + self.assertEqual(span.attributes["db.operation.name"], "insert") + self.assertEqual(span.attributes["db.operation.summary"], "insert") + self.assertNotIn("db.namespace", span.attributes) + self.assertNotIn("db.collection.name", span.attributes) + self.assertEqual(span.status.status_code, StatusCode.ERROR) + + def test_operation_span_name_can_differ_from_command_name(self): + # count_documents' operation span is named "count" but sends an + # aggregate, so an operation span name is not the command beneath it. + # count.json covers estimated_document_count, where the two coincide. + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + db.mycoll.insert_one({"x": 1}) + self.exporter.clear() + db.mycoll.count_documents({}) + + (op_span,) = self.spans("count pymongo_test.mycoll") + self.assertEqual(op_span.attributes["db.operation.name"], "count") + self.assertEqual(op_span.attributes["db.namespace"], "pymongo_test") + (cmd_span,) = self.spans("aggregate") + self.assertEqual(cmd_span.attributes["db.command.name"], "aggregate") + self.assertEqual(cmd_span.parent.span_id, op_span.context.span_id) + + def test_kill_cursors_gets_operation_span(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.kill_cursors_span + coll.drop() + coll.insert_many([{"i": i} for i in range(10)]) + cursor = coll.find({}, batch_size=2) + cursor.next() + self.exporter.clear() + cursor.close() # Sends killCursors, since batches remain. + + op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "killCursors" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(op_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + (op_span,) = op_spans + self.assertEqual(op_span.name, "killCursors pymongo_test.kill_cursors_span") + self.assertEqual(op_span.attributes["db.namespace"], "pymongo_test") + self.assertEqual(op_span.attributes["db.collection.name"], "kill_cursors_span") + + cmd_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "killCursors" + ] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].parent.span_id, op_span.context.span_id) + + def test_background_kill_cursors_span_is_a_trace_root(self): + # Regression test for PYTHON-5947: create_task freezes the calling + # coroutine's context, so without the reset in + # PeriodicExecutor._run every killCursors span the background tick + # emits is parented under whatever operation opened the executor. + # + # connect=False keeps _get_topology() out of construction, so the + # executor opens inside coll.drop() below with that span current. Waking + # the existing task runs the tick in that frozen context; calling + # _process_kill_cursors() here would use this coroutine's clean one. + client = self.rs_or_single_client(tracing={"enabled": True}, connect=False) + coll = client.pymongo_test.bg_kill_cursors + coll.drop() + coll.insert_many([{"i": i} for i in range(10)]) + + cursor = coll.find({}, batch_size=2) + cursor.next() + del cursor + gc.collect() # Queues a deferred killCursors. + + self.exporter.clear() + + def _kill_op_spans(): + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "killCursors" + and "db.command.name" not in s.attributes + ] + + executor = client._kill_cursors_executor + executor.skip_sleep() + executor.wake() + wait_until(_kill_op_spans, "background killCursors span emitted") + + kill_spans = _kill_op_spans() + self.assertEqual(len(kill_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + # The background tick must not inherit a parent from whatever span + # happened to be current when the executor task was created. + self.assertIsNone(kill_spans[0].parent) + + def test_end_sessions_gets_operation_span(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + client.pymongo_test.end_sessions_span.find_one({}) # Uses an implicit session. + self.exporter.clear() + client.close() # Sends endSessions. + + op_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "endSessions" + and "db.command.name" not in s.attributes + ] + self.assertEqual(len(op_spans), 1, [s.name for s in self.exporter.get_finished_spans()]) + (op_span,) = op_spans + self.assertEqual(op_span.name, "endSessions admin") + self.assertEqual(op_span.attributes["db.namespace"], "admin") + self.assertNotIn("db.collection.name", op_span.attributes) + + cmd_spans = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "endSessions" + ] + self.assertEqual(len(cmd_spans), 1) + self.assertEqual(cmd_spans[0].parent.span_id, op_span.context.span_id) + + +# These unit tests cover the validator's edge cases: the rejection paths and the +# explicit-zero vs unset distinction for query_text_max_length. + -# TODO(PYTHON-5947): superseded once the unified test format's -# expectTracingMessages/observeTracingMessages tests exercise this validator -# indirectly through real client construction; remove this class then. class TestValidateTracingOrNone(unittest.TestCase): def test_none(self): self.assertIsNone(common.validate_tracing_or_none("tracing", None)) @@ -346,12 +816,11 @@ def test_rejects_negative_query_text_max_length(self): class TestOTelTracerCaching(unittest.TestCase): - """Regression test for the tracer-caching implementation in ``pymongo/_otel.py``. + """Regression test for the tracer caching in ``pymongo/_otel.py``. - ``opentelemetry.trace.get_tracer()`` must only be called once, at import - time (cached as module-level ``_otel._TRACER``). Calling it per command - allocates two objects, takes a process-wide lock, and mutates the global - ``warnings`` filter list on every call, even on a cache hit. + ``get_tracer()`` must be called once at import time, cached as + ``_otel._TRACER``. Per command it would allocate, take a process-wide lock, + and mutate the global ``warnings`` filter list even on a cache hit. """ @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") diff --git a/test/unified_format_shared.py b/test/unified_format_shared.py index 8a0a3cd46b..77eea4d753 100644 --- a/test/unified_format_shared.py +++ b/test/unified_format_shared.py @@ -29,6 +29,7 @@ from collections.abc import MutableMapping from typing import Any, Union +import pymongo._otel as _otel from bson import ( RE_TYPE, Binary, @@ -252,6 +253,29 @@ def parse_client_bulk_write_error_result(error): return parse_client_bulk_write_result(write_result) +if _otel._HAS_OPENTELEMETRY: + try: + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + except ImportError: + pass + + +def _shared_test_provider() -> TracerProvider: + """Return a process-wide SDK TracerProvider for tests to attach exporters to. + + ``trace.set_tracer_provider`` only takes effect once per process (later calls + are silently ignored), so tests must share one provider and each register + their own span processor rather than trying to install a fresh provider. + """ + current = trace.get_tracer_provider() + if isinstance(current, TracerProvider): + return current + provider = TracerProvider() + trace.set_tracer_provider(provider) + return provider + + class EventListenerUtil( CMAPListener, CommandListener, ServerListener, ServerHeartbeatListener, TopologyListener ):