diff --git a/.evergreen/generated_configs/functions.yml b/.evergreen/generated_configs/functions.yml index 9fbc1f8500..2b77d51d70 100644 --- a/.evergreen/generated_configs/functions.yml +++ b/.evergreen/generated_configs/functions.yml @@ -111,6 +111,7 @@ functions: - LOAD_BALANCER - LOCAL_ATLAS - NO_EXT + - OTEL type: test - command: expansions.update params: @@ -152,6 +153,7 @@ functions: - IS_WIN32 - REQUIRE_FIPS - TEST_MIN_DEPS + - OTEL_TRACE_DIR type: test - command: subprocess.exec params: @@ -160,6 +162,8 @@ functions: - .evergreen/just.sh - run-tests working_dir: src + include_expansions_in_env: + - OTEL_TRACE_DIR type: test # Send dashboard data diff --git a/.evergreen/generated_configs/variants.yml b/.evergreen/generated_configs/variants.yml index bde14deada..7a370d4faf 100644 --- a/.evergreen/generated_configs/variants.yml +++ b/.evergreen/generated_configs/variants.yml @@ -445,16 +445,16 @@ buildvariants: # Otel tests - name: otel-rhel8 tasks: - - 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 + - name: .test-non-standard .replica_set-noauth-ssl .server-latest + - name: .test-non-standard .sharded_cluster-auth-ssl .server-latest .python-pypy3.11 + - name: .test-non-standard .standalone-noauth-nossl .server-latest .python-3.10 display_name: OTel RHEL8 run_on: - rhel87-small expansions: TEST_NAME: otel COVERAGE: "1" + OTEL: "1" tags: [pr] # Perf tests diff --git a/.evergreen/scripts/generate_config.py b/.evergreen/scripts/generate_config.py index a46a989ba6..686ced7054 100644 --- a/.evergreen/scripts/generate_config.py +++ b/.evergreen/scripts/generate_config.py @@ -455,25 +455,28 @@ def create_doctests_variants(): def create_otel_variants(): host = DEFAULT_HOST # Merge otel's coverage into the combined report; see setup_tests.py's COVERAGE handling. - expansions = dict(TEST_NAME="otel", COVERAGE="1") + # OTEL=1 makes drivers-evergreen-tools enable the server's OpenTelemetry file exporter + # and export OTEL_TRACE_DIR, which TestServerTraceContext requires. + expansions = dict(TEST_NAME="otel", COVERAGE="1", OTEL="1") return [ create_variant( [ - # All three topologies, subset to keep the task count at 22. + # All three topologies, one task each to keep the task count small. # - # Replica set keeps every task: 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", + # OTEL=1 enables the server's OpenTelemetry file exporter, which + # requires MongoDB 9.0+ and a binary that accepts every OTel + # setParameter; only the latest nightly qualifies (the v9.0 + # nightly rejects openTelemetryTracingFileFlushCount), so only + # latest tasks are selected. + ".test-non-standard .replica_set-noauth-ssl .server-latest", # 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 + # sensitive-command redaction and prose 9. PyPy covers the + # alternate implementation. + ".test-non-standard .sharded_cluster-auth-ssl .server-latest .python-pypy3.11", + # Standalone for its min-deps task, which resolves # opentelemetry-api down to the floor in requirements/. - ".test-non-standard .standalone-noauth-nossl .python-3.10", + ".test-non-standard .standalone-noauth-nossl .server-latest .python-3.10", ], get_variant_name("OTel", host), host=host, @@ -1276,6 +1279,9 @@ def create_run_server_func(): "LOAD_BALANCER", "LOCAL_ATLAS", "NO_EXT", + # Enables the server's OpenTelemetry file exporter; run-mongodb.sh exports + # OTEL_TRACE_DIR through mo-expansion.yml when it is set. + "OTEL", ] args = [".evergreen/just.sh", "run-server", "${TEST_NAME}"] sub_cmd = get_subprocess_exec(include_expansions_in_env=includes, args=args) @@ -1309,10 +1315,13 @@ def create_run_tests_func(): "IS_WIN32", "REQUIRE_FIPS", "TEST_MIN_DEPS", + "OTEL_TRACE_DIR", ] args = [".evergreen/just.sh", "setup-tests", "${TEST_NAME}", "${SUB_TEST_NAME}"] setup_cmd = get_subprocess_exec(include_expansions_in_env=includes, args=args) - test_cmd = get_subprocess_exec(args=[".evergreen/just.sh", "run-tests"]) + test_cmd = get_subprocess_exec( + include_expansions_in_env=["OTEL_TRACE_DIR"], args=[".evergreen/just.sh", "run-tests"] + ) return "run tests", [setup_cmd, test_cmd] diff --git a/.evergreen/scripts/setup_tests.py b/.evergreen/scripts/setup_tests.py index 7a46603105..da185aeba8 100644 --- a/.evergreen/scripts/setup_tests.py +++ b/.evergreen/scripts/setup_tests.py @@ -504,6 +504,17 @@ def handle_test_env() -> None: TEST_SUITE = TEST_SUITE_MAP.get(test_name) if TEST_SUITE: TEST_ARGS = f"-m {TEST_SUITE} {TEST_ARGS}" + if test_name == "otel": + # Collect only the otel test files: sweeping the whole tree imports + # unrelated modules, whose import-time skips (test_ocsp_support without + # the ocsp extra) end up in the results. + TEST_ARGS += ( + " test/asynchronous/test_otel.py test/asynchronous/test_otel_getmore.py" + " test/asynchronous/test_otel_transactions.py" + " test/asynchronous/test_open_telemetry_unified.py" + " test/test_otel.py test/test_otel_getmore.py" + " test/test_otel_transactions.py test/test_open_telemetry_unified.py" + ) write_env("TEST_ARGS", TEST_ARGS) write_env("UV_ARGS", " ".join(UV_ARGS)) diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index 1597652d69..82e7a0696d 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -212,6 +212,57 @@ static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { return result; } +static int +_write_telemetry_section(struct module_state *state, buffer_t buffer, + const char* traceparent, codec_options_t* options) { + PyObject* otel_dict = NULL; + PyObject* parent_dict = NULL; + PyObject* tp_str = NULL; + int tp_size; + + if (!traceparent) { + return 1; + } + otel_dict = PyDict_New(); + parent_dict = PyDict_New(); + if (!otel_dict || !parent_dict) { + Py_XDECREF(otel_dict); + Py_XDECREF(parent_dict); + return 0; + } + tp_str = PyUnicode_FromString(traceparent); + if (!tp_str) { + Py_DECREF(otel_dict); + Py_DECREF(parent_dict); + return 0; + } + if (PyDict_SetItemString(parent_dict, "traceparent", tp_str) < 0) { + Py_DECREF(tp_str); + Py_DECREF(otel_dict); + Py_DECREF(parent_dict); + return 0; + } + Py_DECREF(tp_str); + if (PyDict_SetItemString(otel_dict, "otel", parent_dict) < 0) { + Py_DECREF(otel_dict); + Py_DECREF(parent_dict); + return 0; + } + Py_DECREF(parent_dict); + + /* Payload type 3 section */ + if (!buffer_write_bytes(buffer, "\x03", 1)) { + Py_DECREF(otel_dict); + return 0; + } + tp_size = write_dict(state->_cbson, buffer, otel_dict, 0, options, 1); + Py_DECREF(otel_dict); + if (!tp_size) { + return 0; + } + return 1; +} + /* * NOTE this method handles multiple documents in a type one payload but * it does not perform batch splitting and the total message size is @@ -234,20 +285,22 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { int max_doc_size = 0; PyObject* result = NULL; PyObject* iterator = NULL; + const char* traceparent = NULL; struct module_state *state = GETSTATE(self); if (!state) { return NULL; } - /*flags, command, identifier, docs, opts*/ - if (!(PyArg_ParseTuple(args, "IOet#OO", + /*flags, command, identifier, docs, opts, [traceparent]*/ + if (!(PyArg_ParseTuple(args, "IOet#OO|z", &flags, &command, "utf-8", &identifier, &identifier_length, &docs, - &options_obj) && + &options_obj, + &traceparent) && convert_codec_options(state->_cbson, options_obj, &options))) { return NULL; } @@ -313,6 +366,14 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) { total_size += payload_length; } + if (traceparent) { + int before = pymongo_buffer_get_position(buffer); + if (!_write_telemetry_section(state, buffer, traceparent, &options)) { + goto fail; + } + total_size += pymongo_buffer_get_position(buffer) - before; + } + message_length = pymongo_buffer_get_position(buffer) - length_location; buffer_write_int32_at_position( buffer, length_location, (int32_t)message_length); @@ -358,7 +419,8 @@ _batched_op_msg( unsigned char op, unsigned char ack, PyObject* command, PyObject* docs, PyObject* ctx, PyObject* to_publish, codec_options_t options, - buffer_t buffer, struct module_state *state) { + buffer_t buffer, struct module_state *state, + const char* traceparent) { long max_bson_size; long max_write_batch_size; @@ -395,6 +457,14 @@ _batched_op_msg( return 0; } + if (traceparent) { + /* The Payload Type 3 section is appended after the batch is split, so + * reserve its encoded size now to keep the message within + * max_message_size. The section is the traceparent plus 35 bytes of + * BSON overhead; see _write_telemetry_section. */ + max_message_size -= (long)strlen(traceparent) + 35; + } + if (!buffer_write_bytes(buffer, flags, 4)) { return 0; } @@ -523,6 +593,11 @@ _batched_op_msg( position = pymongo_buffer_get_position(buffer); length = position - size_location; buffer_write_int32_at_position(buffer, size_location, (int32_t)length); + if (traceparent) { + if (!_write_telemetry_section(state, buffer, traceparent, &options)) { + goto fail; + } + } return 1; fail: @@ -543,14 +618,15 @@ _cbson_encode_batched_op_msg(PyObject* self, PyObject* args) { PyObject* options_obj = NULL; codec_options_t options; buffer_t buffer; + const char* traceparent = NULL; struct module_state *state = GETSTATE(self); if (!state) { return NULL; } - if (!(PyArg_ParseTuple(args, "bOObOO", + if (!(PyArg_ParseTuple(args, "bOObOO|z", &op, &command, &docs, &ack, - &options_obj, &ctx) && + &options_obj, &ctx, &traceparent) && convert_codec_options(state->_cbson, options_obj, &options))) { return NULL; } @@ -571,7 +647,8 @@ _cbson_encode_batched_op_msg(PyObject* self, PyObject* args) { to_publish, options, buffer, - state)) { + state, + traceparent)) { goto fail; } @@ -600,14 +677,15 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { PyObject* options_obj = NULL; codec_options_t options; buffer_t buffer; + const char* traceparent = NULL; struct module_state *state = GETSTATE(self); if (!state) { return NULL; } - if (!(PyArg_ParseTuple(args, "bOObOO", + if (!(PyArg_ParseTuple(args, "bOObOO|z", &op, &command, &docs, &ack, - &options_obj, &ctx) && + &options_obj, &ctx, &traceparent) && convert_codec_options(state->_cbson, options_obj, &options))) { return NULL; } @@ -638,7 +716,8 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) { to_publish, options, buffer, - state)) { + state, + traceparent)) { goto fail; } diff --git a/pymongo/_otel.py b/pymongo/_otel.py index a54ff907bd..8b040055e7 100644 --- a/pymongo/_otel.py +++ b/pymongo/_otel.py @@ -33,6 +33,7 @@ from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Optional, TypedDict +from bson import encode as _bson_encode from bson import json_util from bson.json_util import _truncate_documents from pymongo._version import __version__ @@ -396,6 +397,63 @@ def _set_operation_cursor_id(cursor_id: int) -> None: current_span.set_attribute("db.mongodb.cursor_id", cursor_id) +def _set_command_span_query_text( + span: Optional[Span], + tracing_options: Optional[TracingOptions], + cmd: Mapping[str, Any], +) -> None: + """Rebuild db.query.text from a command that gained fields after span creation. + + Client-level bulkWrite adds its ops and nsInfo only at send time, after the + command span exists, so its text is refreshed once they are known. + """ + if span is None: + return + max_query_text_length = _get_query_text_max_length(tracing_options) + if max_query_text_length > 0: + span.set_attribute("db.query.text", _build_query_text(cmd, max_query_text_length)) + + +_TELEMETRY_TRACEPARENT_KEY = "traceparent" +_TELEMETRY_OTEL_KEY = "otel" +# MongoDB 9.0, the first server that accepts the Payload Type 3 telemetry section. +_TELEMETRY_MIN_WIRE_VERSION = 29 +# Marks "tracing is on but this command gets no span" (a sensitive command). +# Distinct from None, which means the caller did not pre-create a span. +_NO_COMMAND_SPAN = object() + + +def _traceparent_from_span(span: Optional[Span]) -> Optional[str]: + """Return the W3C traceparent string for ``span``, or ``None``. + + The traceparent carries the command span's own span id as the parent-id + so server spans join the trace as children of that command. The span + need not be recording: an unsampled context (trace-flags ``00``) is + still propagated. Returns ``None`` when the span or its context is + invalid (e.g. all-zero trace-id or span-id). + """ + if span is None: + return None + ctx = span.get_span_context() + if not ctx.is_valid: + return None + # :02x only pads to a minimum width, so flags wider than one byte would + # push the traceparent past 55 characters. + if not 0 <= ctx.trace_flags <= 0xFF: + return None + return f"00-{ctx.trace_id:032x}-{ctx.span_id:016x}-{ctx.trace_flags:02x}" + + +def _telemetry_section(traceparent: str) -> bytes: + """Encode the OP_MSG Payload Type 3 telemetry section. + + Returns ``\\x03`` (payload type) followed by the BSON document + ``{"otel": {"traceparent": traceparent}}``. The caller appends this + to the OP_MSG body after the Type 0 and Type 1 sections. + """ + return b"\x03" + _bson_encode({_TELEMETRY_OTEL_KEY: {_TELEMETRY_TRACEPARENT_KEY: traceparent}}) + + def end_command_span_success(span: Optional[Span], reply: _DocumentOut) -> None: """Set the cursor id (if any open cursor) and end the span.""" if span is None: @@ -457,6 +515,38 @@ def end_command_span_failure( span.end() +@contextlib.contextmanager +def _command_span_for_encoding( + tracing_options: Optional[TracingOptions], + conn: _ConnectionTelemetryInfo, + cmd: MutableMapping[str, Any], + dbname: str, + command_name: str, + speculative_hello: bool = False, +) -> Iterator[tuple[Optional[Span], Any, Optional[str]]]: + """Start the command span for a command that is about to be encoded. + + Yields ``(span, precreated_span, traceparent)``. An exception raised while + encoding ends the span, so an encode failure that never reaches the send + cannot leak it. Otherwise the caller passes ``precreated_span`` to + ``_CommandTelemetry``, which ends the span after the send. + """ + span = None + traceparent = None + if _is_tracing_enabled(tracing_options): + span = start_command_span( + tracing_options, conn, cmd, dbname, command_name, speculative_hello + ) + if conn.max_wire_version >= _TELEMETRY_MIN_WIRE_VERSION: + traceparent = _traceparent_from_span(span) + precreated_span = span if span is not None else _NO_COMMAND_SPAN + try: + yield span, precreated_span, traceparent + except Exception as exc: + end_command_span_failure(span, {}, exc) + raise + + class _OperationSpanHandle: """Bundles an operation span with what is needed to end it later. diff --git a/pymongo/_telemetry.py b/pymongo/_telemetry.py index 4f21f25b0f..55f6f1fd39 100644 --- a/pymongo/_telemetry.py +++ b/pymongo/_telemetry.py @@ -92,6 +92,7 @@ class _CommandTelemetry: "_listeners", "_name", "_op_id", + "_precreated_span", "_publish", "_request_id", "_should_log", @@ -115,6 +116,7 @@ def __init__( tracing_options: Optional[_otel.TracingOptions] = None, speculative_hello: bool = False, name: Optional[str] = None, + precreated_span: Optional[Any] = None, ) -> None: # NOTE: the _run_command fast path in command_runner.py inline this gate for performance # They must be kept in sync with any gating changes @@ -137,6 +139,7 @@ def __init__( self._request_id = request_id self._op_id = op_id if op_id is not None else _op_id.OP_ID.get() self._speculative_hello = speculative_hello + self._precreated_span = precreated_span def _emit_log(self, message: _CommandStatusMessage, **extra: Any) -> None: _debug_log( @@ -175,7 +178,12 @@ def started(self, orig: MutableMapping[str, Any], ensure_db: bool) -> None: self._op_id, service_id=self._conn.service_id, ) - if self._tracing_enabled: + if self._precreated_span is _otel._NO_COMMAND_SPAN: + # Tracing is on but the command is sensitive, so it has no span. + self._span = None + elif self._precreated_span is not None: + self._span = self._precreated_span + elif self._tracing_enabled: self._span = _otel.start_command_span( self._tracing_options, self._conn, diff --git a/pymongo/asynchronous/bulk.py b/pymongo/asynchronous/bulk.py index 8028b04a79..8096cfd230 100644 --- a/pymongo/asynchronous/bulk.py +++ b/pymongo/asynchronous/bulk.py @@ -28,11 +28,12 @@ Any, Optional, Union, + cast, ) from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument -from pymongo import _csot, common +from pymongo import _csot, _otel, common from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.asynchronous.client_session import AsyncClientSession, _validate_session_write_concern from pymongo.asynchronous.command_runner import ( @@ -236,7 +237,6 @@ def gen_unordered(self) -> Iterator[_Run]: if run.ops: yield run - @_handle_reauth async def write_command( self, bwc: _BulkWriteContext, @@ -245,6 +245,8 @@ async def write_command( msg: bytes, docs: list[Mapping[str, Any]], client: AsyncMongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> dict[str, Any]: """Run a batch write command, returning the response as a dict.""" cmd[bwc.field] = docs @@ -254,6 +256,7 @@ async def write_command( request_id, msg, client=client, + precreated_span=precreated_span, ) return result_docs[0] @@ -266,6 +269,8 @@ async def unack_write( max_doc_size: int, docs: list[Mapping[str, Any]], client: AsyncMongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> Optional[Mapping[str, Any]]: """Send an unacknowledged batch write command.""" # Historically the STARTED log omits the documents while the published @@ -282,6 +287,7 @@ async def unack_write( orig=published, max_doc_size=max_doc_size, unacknowledged=True, + precreated_span=precreated_span, ) return None @@ -302,13 +308,23 @@ async def _execute_batch_unack( client=client, # type: ignore[arg-type] ) else: - request_id, msg, to_send = bwc.batch_command(cmd, ops) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send = bwc.batch_command(cmd, ops, traceparent=traceparent) + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, bwc.field: to_send} + ) + msg = cast(bytes, msg) # Though this isn't strictly a "legacy" write, the helper # handles publishing commands and sending our message # without receiving a result. Send 0 for max_doc_size # to disable size checking. Size checking is handled while # the documents are encoded to BSON. - await self.unack_write(bwc, cmd, request_id, msg, 0, to_send, client) # type: ignore[arg-type] + await self.unack_write( + bwc, cmd, request_id, msg, 0, to_send, client, precreated_span=precreated_span + ) # type: ignore[arg-type] return to_send @@ -316,10 +332,12 @@ async def _execute_batch( self, bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext], cmd: dict[str, Any], - ops: list[Mapping[str, Any]], + run_ops: list[Mapping[str, Any]], + idx_offset: int, client: AsyncMongoClient[Any], ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]: if self.is_encrypted: + ops = cast("list[Mapping[str, Any]]", islice(run_ops, idx_offset, None)) _, batched_cmd, to_send = bwc.batch_command(cmd, ops) result = await bwc.conn.command( # type: ignore[misc] bwc.db_name, @@ -328,11 +346,31 @@ async def _execute_batch( session=bwc.session, # type: ignore[arg-type] client=client, # type: ignore[arg-type] ) - else: - request_id, msg, to_send = bwc.batch_command(cmd, ops) - result = await self.write_command(bwc, cmd, request_id, msg, to_send, client) # type: ignore[arg-type] + return result, to_send # type: ignore[return-value] + + async def run( + ctx: _BulkWriteContext, + ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]: + # Reauthentication re-runs this closure, so undo the field that + # write_command added and re-encode against a fresh span and + # traceparent. + cmd.pop(ctx.field, None) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, ctx.conn, cmd, ctx.db_name, ctx.name + ) as (span, precreated_span, traceparent): + ops = cast("list[Mapping[str, Any]]", islice(run_ops, idx_offset, None)) + request_id, msg, to_send = ctx.batch_command(cmd, ops, traceparent=traceparent) + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, ctx.field: to_send} + ) + msg = cast(bytes, msg) + result = await self.write_command( + ctx, cmd, request_id, msg, to_send, client, precreated_span=precreated_span + ) + return result, to_send - return result, to_send # type: ignore[return-value] + return await _handle_reauth(run)(bwc) async def _execute_command( self, @@ -406,7 +444,9 @@ async def _execute_command( # Run as many ops as possible in one command. if write_concern.acknowledged: - result, to_send = await self._execute_batch(bwc, cmd, ops, client) + result, to_send = await self._execute_batch( + bwc, cmd, run.ops, run.idx_offset, client + ) # Retryable writeConcernErrors halt the execution of this run. wce = result.get("writeConcernError", {}) diff --git a/pymongo/asynchronous/client_bulk.py b/pymongo/asynchronous/client_bulk.py index 09a0c368c6..a26c90a9f5 100644 --- a/pymongo/asynchronous/client_bulk.py +++ b/pymongo/asynchronous/client_bulk.py @@ -28,11 +28,12 @@ Any, Optional, Union, + cast, ) from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument -from pymongo import _csot, common +from pymongo import _csot, _otel, common from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.asynchronous.client_session import ( AsyncClientSession, @@ -236,6 +237,8 @@ async def write_command( op_docs: list[Mapping[str, Any]], ns_docs: list[Mapping[str, Any]], client: AsyncMongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> dict[str, Any]: """Run a client-level batch write command, returning the response as a dict.""" cmd["ops"] = op_docs @@ -247,6 +250,7 @@ async def write_command( request_id, msg, # type: ignore[arg-type] client=client, + precreated_span=precreated_span, ) reply = result_docs[0] except Exception as exc: @@ -263,6 +267,8 @@ async def unack_write( op_docs: list[Mapping[str, Any]], ns_docs: list[Mapping[str, Any]], client: AsyncMongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> Optional[Mapping[str, Any]]: """Send an unacknowledged client-level batch write command.""" # Historically the STARTED log omits the ops/nsInfo while the published @@ -281,6 +287,7 @@ async def unack_write( orig=published, max_doc_size=bwc.max_bson_size, unacknowledged=True, + precreated_span=precreated_span, ) reply: Mapping[str, Any] = result_docs[0] except Exception as exc: @@ -296,8 +303,29 @@ async def _execute_batch_unack( namespaces: list[str], ) -> tuple[list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Executes a batch of bulkWrite server commands (unack).""" - request_id, msg, to_send_ops, to_send_ns = bwc.batch_command(cmd, ops, namespaces) - await self.unack_write(bwc, cmd, request_id, msg, to_send_ops, to_send_ns, self.client) # type: ignore[arg-type] + tracing_options = self.client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send_ops, to_send_ns = bwc.batch_command( + cmd, ops, namespaces, traceparent=traceparent + ) + # The ops and nsInfo are only known after encoding, so the query text + # is refreshed from the full published command. + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, "ops": to_send_ops, "nsInfo": to_send_ns} + ) + msg = cast(bytes, msg) + await self.unack_write( + bwc, + cmd, + request_id, + msg, + to_send_ops, + to_send_ns, + self.client, + precreated_span=precreated_span, + ) # type: ignore[arg-type] return to_send_ops, to_send_ns async def _execute_batch( @@ -308,9 +336,28 @@ async def _execute_batch( namespaces: list[str], ) -> tuple[dict[str, Any], list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Executes a batch of bulkWrite server commands (ack).""" - request_id, msg, to_send_ops, to_send_ns = bwc.batch_command(cmd, ops, namespaces) + tracing_options = self.client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send_ops, to_send_ns = bwc.batch_command( + cmd, ops, namespaces, traceparent=traceparent + ) + # The ops and nsInfo are only known after encoding, so the query text + # is refreshed from the full published command. + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, "ops": to_send_ops, "nsInfo": to_send_ns} + ) + msg = cast(bytes, msg) result = await self.write_command( - bwc, cmd, request_id, msg, to_send_ops, to_send_ns, self.client + bwc, + cmd, + request_id, + msg, + to_send_ops, + to_send_ns, + self.client, + precreated_span=precreated_span, ) # type: ignore[arg-type] return result, to_send_ops, to_send_ns # type: ignore[return-value] diff --git a/pymongo/asynchronous/command_runner.py b/pymongo/asynchronous/command_runner.py index f13d893dbf..f73d0700d8 100644 --- a/pymongo/asynchronous/command_runner.py +++ b/pymongo/asynchronous/command_runner.py @@ -102,6 +102,7 @@ async def _run_command( set_conn_more_to_come: bool = False, unpack_res: Optional[Callable[..., Any]] = None, cursor_id: Optional[int] = None, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], float]: """Send ``msg`` over ``conn`` and return ``(docs, reply, duration_s)``, where ``duration_s`` is the round-trip duration in seconds. @@ -185,6 +186,7 @@ async def _run_command( tracing_options=tracing_options, speculative_hello=speculative_hello, name=name, + precreated_span=precreated_span, ) telemetry.started(orig, ensure_db) start = 0.0 @@ -266,6 +268,7 @@ async def run_bulk_write_command( orig: Optional[MutableMapping[str, Any]] = None, max_doc_size: int = 0, unacknowledged: bool = False, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], float]: """Send a bulk write batch and return ``(docs, reply, duration_s)``. @@ -297,6 +300,7 @@ async def run_bulk_write_command( max_doc_size=max_doc_size, unacknowledged=unacknowledged, decrypt_reply=False, + precreated_span=precreated_span, ) @@ -318,6 +322,7 @@ async def run_cursor_command( more_to_come: bool = False, unpack_res: Optional[Callable[..., Any]] = None, cursor_id: Optional[int] = None, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], datetime.timedelta]: """Run a cursor ``find``/``getMore`` operation over ``conn``. @@ -359,6 +364,7 @@ async def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, + precreated_span=precreated_span, ) # The cursor path stores the duration on Response, which expects a timedelta. return docs, reply, datetime.timedelta(seconds=duration_s) @@ -441,16 +447,26 @@ async def run_command( flags = _OpMsg.MORE_TO_COME if unacknowledged else 0 flags |= _OpMsg.EXHAUST_ALLOWED if exhaust_allowed else 0 - request_id, msg, size, max_doc_size = message._op_msg( - flags, spec, dbname, read_preference, codec_options, ctx=compression_ctx - ) - # If this is an unacknowledged write then make sure the encoded doc(s) - # are small enough, otherwise rely on the server to return an error. - if unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size: - message._raise_document_too_large(name, size, max_bson_size) + tracing_options = client.options.tracing if client is not None else None + with _otel._command_span_for_encoding( + tracing_options, conn, spec, dbname, name, speculative_hello + ) as (_, precreated_span, traceparent): + request_id, msg, size, max_doc_size = message._op_msg( + flags, + spec, + dbname, + read_preference, + codec_options, + ctx=compression_ctx, + traceparent=traceparent, + ) + # If this is an unacknowledged write then make sure the encoded doc(s) + # are small enough, otherwise rely on the server to return an error. + if unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size: + message._raise_document_too_large(name, size, max_bson_size) - if max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD: - message._raise_document_too_large(name, size, max_bson_size + message._COMMAND_OVERHEAD) + if max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD: + message._raise_document_too_large(name, size, max_bson_size + message._COMMAND_OVERHEAD) docs, _, _ = await _run_command( conn, spec, @@ -470,5 +486,6 @@ async def run_command( speculative_hello=speculative_hello, unacknowledged=unacknowledged, set_conn_more_to_come=True, + precreated_span=precreated_span, ) return docs[0] # type: ignore[return-value] diff --git a/pymongo/asynchronous/cursor_base.py b/pymongo/asynchronous/cursor_base.py index bcd8125cd3..d1d47c741a 100644 --- a/pymongo/asynchronous/cursor_base.py +++ b/pymongo/asynchronous/cursor_base.py @@ -20,7 +20,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Optional, Union -from pymongo import _csot +from pymongo import _csot, _otel from pymongo.asynchronous.command_runner import run_cursor_command from pymongo.asynchronous.helpers import _handle_reauth from pymongo.cursor_shared import _CURSOR_DOC_FIELDS, _AgnosticCursorBase, _split_message @@ -113,11 +113,19 @@ async def _run_with_conn( use_cmd = operation.use_command(conn) more_to_come = bool(operation.conn_mgr and operation.conn_mgr.more_to_come) cmd, dbn = await _operation_to_command(operation, conn, use_cmd) - if more_to_come: - request_id, data, max_doc_size = 0, b"", 0 - else: - message = operation.get_message(read_preference, conn, use_cmd) - request_id, data, max_doc_size = _split_message(message) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding(tracing_options, conn, cmd, dbn, operation.name) as ( + _, + precreated_span, + traceparent, + ): + if more_to_come: + request_id, data, max_doc_size = 0, b"", 0 + else: + message = operation.get_message( + read_preference, conn, use_cmd, traceparent=traceparent + ) + request_id, data, max_doc_size = _split_message(message) user_fields = _CURSOR_DOC_FIELDS if use_cmd else None docs, reply, duration = await run_cursor_command( conn, @@ -136,6 +144,7 @@ async def _run_with_conn( more_to_come=more_to_come, unpack_res=self._unpack_response, cursor_id=operation.cursor_id, + precreated_span=precreated_span, ) assert reply is not None if client._should_pin_cursor(operation.session) or operation.exhaust: # type: ignore[arg-type] diff --git a/pymongo/message.py b/pymongo/message.py index b696ba2302..c4a97695ce 100644 --- a/pymongo/message.py +++ b/pymongo/message.py @@ -43,6 +43,7 @@ RawBSONDocument, _inflate_bson, ) +from pymongo import _otel from pymongo.common import MONGOS_EXHAUST_WIRE_VERSION from pymongo.monitoring import _EventListeners @@ -289,6 +290,7 @@ def _op_msg_no_header( identifier: str, docs: Optional[list[Mapping[str, Any]]], opts: CodecOptions[Any], + traceparent: Optional[str] = None, ) -> tuple[bytes, int, int]: """Get a OP_MSG message. @@ -301,6 +303,7 @@ def _op_msg_no_header( flags_type = _pack_op_msg_flags_type(flags, 0) total_size = len(encoded) max_doc_size = 0 + data: list[bytes] = [flags_type, encoded] if identifier and docs is not None: type_one = _pack_byte(1) cstring = _make_c_string(identifier) @@ -309,9 +312,11 @@ def _op_msg_no_header( encoded_size = _pack_int(size) total_size += size max_doc_size = max(len(doc) for doc in encoded_docs) - data = [flags_type, encoded, type_one, encoded_size, cstring, *encoded_docs] - else: - data = [flags_type, encoded] + data.extend([type_one, encoded_size, cstring, *encoded_docs]) + if traceparent is not None: + section = _otel._telemetry_section(traceparent) + data.append(section) + total_size += len(section) return b"".join(data), total_size, max_doc_size @@ -322,9 +327,12 @@ def _op_msg_compressed( docs: Optional[list[Mapping[str, Any]]], opts: CodecOptions[Any], ctx: Union[SnappyContext, ZlibContext, ZstdContext], + traceparent: Optional[str] = None, ) -> tuple[int, bytes, int, int]: """Internal OP_MSG message helper.""" - msg, total_size, max_bson_size = _op_msg_no_header(flags, command, identifier, docs, opts) + msg, total_size, max_bson_size = _op_msg_no_header( + flags, command, identifier, docs, opts, traceparent + ) rid, msg = _compress(2013, msg, ctx) return rid, msg, total_size, max_bson_size @@ -335,9 +343,12 @@ def _op_msg_uncompressed( identifier: str, docs: Optional[list[Mapping[str, Any]]], opts: CodecOptions[Any], + traceparent: Optional[str] = None, ) -> tuple[int, bytes, int, int]: """Internal compressed OP_MSG message helper.""" - data, total_size, max_bson_size = _op_msg_no_header(flags, command, identifier, docs, opts) + data, total_size, max_bson_size = _op_msg_no_header( + flags, command, identifier, docs, opts, traceparent + ) request_id, op_message = __pack_message(2013, data) return request_id, op_message, total_size, max_bson_size @@ -353,6 +364,7 @@ def _op_msg( read_preference: Optional[_ServerMode], opts: CodecOptions[Any], ctx: Union[SnappyContext, ZlibContext, ZstdContext, None] = None, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, int, int]: """Get a OP_MSG message.""" command["$db"] = dbname @@ -370,8 +382,8 @@ def _op_msg( docs = None try: if ctx: - return _op_msg_compressed(flags, command, identifier, docs, opts, ctx) - return _op_msg_uncompressed(flags, command, identifier, docs, opts) + return _op_msg_compressed(flags, command, identifier, docs, opts, ctx, traceparent) + return _op_msg_uncompressed(flags, command, identifier, docs, opts, traceparent) finally: # Add the field back to the command. if identifier: @@ -526,11 +538,14 @@ def __init__( ) def batch_command( - self, cmd: MutableMapping[str, Any], docs: list[Mapping[str, Any]] + self, + cmd: MutableMapping[str, Any], + docs: list[Mapping[str, Any]], + traceparent: Optional[str] = None, ) -> tuple[int, Union[bytes, dict[str, Any]], list[Mapping[str, Any]]]: namespace = self.db_name + ".$cmd" request_id, msg, to_send = _do_batched_op_msg( - namespace, self.op_type, cmd, docs, self.codec, self + namespace, self.op_type, cmd, docs, self.codec, self, traceparent ) if not to_send: raise InvalidOperation("cannot do an empty bulk write") @@ -541,8 +556,13 @@ class _EncryptedBulkWriteContext(_BulkWriteContext): __slots__ = () def batch_command( - self, cmd: MutableMapping[str, Any], docs: list[Mapping[str, Any]] + self, + cmd: MutableMapping[str, Any], + docs: list[Mapping[str, Any]], + traceparent: Optional[str] = None, ) -> tuple[int, dict[str, Any], list[Mapping[str, Any]]]: + # traceparent is unused: the batched command is sent through + # conn.command(), whose own command span carries the traceparent. namespace = self.db_name + ".$cmd" msg, to_send = _encode_batched_write_command( namespace, self.op_type, cmd, docs, self.codec, self @@ -591,11 +611,16 @@ def _batched_op_msg_impl( opts: CodecOptions[Any], ctx: _BulkWriteContext, buf: _BytesIO, + traceparent: Optional[str] = None, ) -> tuple[list[Mapping[str, Any]], int]: """Create a batched OP_MSG write.""" max_bson_size = ctx.max_bson_size max_write_batch_size = ctx.max_write_batch_size max_message_size = ctx.max_message_size + section = _otel._telemetry_section(traceparent) if traceparent is not None else b"" + # The telemetry section is appended after the batch is split, so reserve its + # bytes now to keep the final message within max_message_size. + max_message_size -= len(section) flags = b"\x00\x00\x00\x00" if ack else b"\x02\x00\x00\x00" # Flags @@ -643,10 +668,15 @@ def _batched_op_msg_impl( # Write type 1 section size length = buf.tell() + if section: + buf.write(section) + total_length = buf.tell() + else: + total_length = length buf.seek(size_location) buf.write(_pack_int(length - size_location)) - return to_send, length + return to_send, total_length def _encode_batched_op_msg( @@ -656,13 +686,14 @@ def _encode_batched_op_msg( ack: bool, opts: CodecOptions[Any], ctx: _BulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[bytes, list[Mapping[str, Any]]]: """Encode the next batched insert, update, or delete operation as OP_MSG. """ buf = _BytesIO() - to_send, _ = _batched_op_msg_impl(operation, command, docs, ack, opts, ctx, buf) + to_send, _ = _batched_op_msg_impl(operation, command, docs, ack, opts, ctx, buf, traceparent) return buf.getvalue(), to_send @@ -677,11 +708,12 @@ def _batched_op_msg_compressed( ack: bool, opts: CodecOptions[Any], ctx: _BulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]]]: """Create the next batched insert, update, or delete operation with OP_MSG, compressed. """ - data, to_send = _encode_batched_op_msg(operation, command, docs, ack, opts, ctx) + data, to_send = _encode_batched_op_msg(operation, command, docs, ack, opts, ctx, traceparent) assert ctx.conn.compression_context is not None request_id, msg = _compress(2013, data, ctx.conn.compression_context) @@ -695,6 +727,7 @@ def _batched_op_msg( ack: bool, opts: CodecOptions[Any], ctx: _BulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]]]: """OP_MSG implementation entry point.""" buf = _BytesIO() @@ -704,7 +737,9 @@ def _batched_op_msg( # responseTo, opCode buf.write(b"\x00\x00\x00\x00\xdd\x07\x00\x00") - to_send, length = _batched_op_msg_impl(operation, command, docs, ack, opts, ctx, buf) + to_send, length = _batched_op_msg_impl( + operation, command, docs, ack, opts, ctx, buf, traceparent + ) # Header - request id and message length buf.seek(4) @@ -727,6 +762,7 @@ def _do_batched_op_msg( docs: list[Mapping[str, Any]], opts: CodecOptions[Any], ctx: _BulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]]]: """Create the next batched insert, update, or delete operation using OP_MSG. @@ -737,8 +773,8 @@ def _do_batched_op_msg( else: ack = True if ctx.conn.compression_context: - return _batched_op_msg_compressed(operation, command, docs, ack, opts, ctx) - return _batched_op_msg(operation, command, docs, ack, opts, ctx) + return _batched_op_msg_compressed(operation, command, docs, ack, opts, ctx, traceparent) + return _batched_op_msg(operation, command, docs, ack, opts, ctx, traceparent) class _ClientBulkWriteContext(_BulkWriteContextBase): @@ -772,9 +808,10 @@ def batch_command( cmd: MutableMapping[str, Any], operations: list[tuple[str, Mapping[str, Any]]], namespaces: list[str], + traceparent: Optional[str] = None, ) -> tuple[int, Union[bytes, dict[str, Any]], list[Mapping[str, Any]], list[Mapping[str, Any]]]: request_id, msg, to_send_ops, to_send_ns = _client_do_batched_op_msg( - cmd, operations, namespaces, self.codec, self + cmd, operations, namespaces, self.codec, self, traceparent ) if not to_send_ops: raise InvalidOperation("cannot do an empty bulk write") @@ -790,6 +827,7 @@ def _client_construct_op_msg( to_send_ns_encoded: list[bytes], ack: bool, buf: _BytesIO, + traceparent: Optional[str] = None, ) -> int: # Write flags flags = b"\x00\x00\x00\x00" if ack else b"\x02\x00\x00\x00" @@ -829,6 +867,11 @@ def _client_construct_op_msg( buf.seek(size_location) buf.write(_pack_int(length - size_location)) + if traceparent is not None: + buf.seek(0, 2) + buf.write(_otel._telemetry_section(traceparent)) + length = buf.tell() + return length @@ -840,6 +883,7 @@ def _client_batched_op_msg_impl( opts: CodecOptions[Any], ctx: _ClientBulkWriteContext, buf: _BytesIO, + traceparent: Optional[str] = None, ) -> tuple[list[Mapping[str, Any]], list[Mapping[str, Any]], int]: """Create a batched OP_MSG write for client-level bulk write.""" @@ -854,6 +898,7 @@ def _check_doc_size_limits( max_bson_size = ctx.max_bson_size max_write_batch_size = ctx.max_write_batch_size max_message_size = ctx.max_message_size + section = _otel._telemetry_section(traceparent) if traceparent is not None else b"" command_encoded = _dict_to_bson(command, False, opts) # When OP_MSG is used unacknowledged we have to check command @@ -872,8 +917,11 @@ def _check_doc_size_limits( command_abridged = {key: command[key] for key in abridged_keys} command_len_abridged = len(_dict_to_bson(command_abridged, False, opts)) - # Maximum combined size of the ops and nsInfo document sequences. - max_doc_sequences_bytes = max_message_size - (_OP_MSG_OVERHEAD + command_len_abridged) + # Maximum combined size of the ops and nsInfo document sequences. Reserve + # the telemetry section, which _client_construct_op_msg appends last. + max_doc_sequences_bytes = max_message_size - ( + _OP_MSG_OVERHEAD + command_len_abridged + len(section) + ) ns_info = {} to_send_ops: list[Mapping[str, Any]] = [] @@ -941,7 +989,7 @@ def _check_doc_size_limits( # Construct the entire OP_MSG. length = _client_construct_op_msg( - command_encoded, to_send_ops_encoded, to_send_ns_encoded, ack, buf + command_encoded, to_send_ops_encoded, to_send_ns_encoded, ack, buf, traceparent ) return to_send_ops, to_send_ns, length @@ -954,6 +1002,7 @@ def _client_encode_batched_op_msg( ack: bool, opts: CodecOptions[Any], ctx: _ClientBulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[bytes, list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Encode the next batched client-level bulkWrite operation as OP_MSG. @@ -961,7 +1010,7 @@ def _client_encode_batched_op_msg( buf = _BytesIO() to_send_ops, to_send_ns, _ = _client_batched_op_msg_impl( - command, operations, namespaces, ack, opts, ctx, buf + command, operations, namespaces, ack, opts, ctx, buf, traceparent ) return buf.getvalue(), to_send_ops, to_send_ns @@ -973,12 +1022,13 @@ def _client_batched_op_msg_compressed( ack: bool, opts: CodecOptions[Any], ctx: _ClientBulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Create the next batched client-level bulkWrite operation with OP_MSG, compressed. """ data, to_send_ops, to_send_ns = _client_encode_batched_op_msg( - command, operations, namespaces, ack, opts, ctx + command, operations, namespaces, ack, opts, ctx, traceparent ) assert ctx.conn.compression_context is not None @@ -993,6 +1043,7 @@ def _client_batched_op_msg( ack: bool, opts: CodecOptions[Any], ctx: _ClientBulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]], list[Mapping[str, Any]]]: """OP_MSG implementation entry point for client-level bulkWrite.""" buf = _BytesIO() @@ -1003,7 +1054,7 @@ def _client_batched_op_msg( buf.write(b"\x00\x00\x00\x00\xdd\x07\x00\x00") to_send_ops, to_send_ns, length = _client_batched_op_msg_impl( - command, operations, namespaces, ack, opts, ctx, buf + command, operations, namespaces, ack, opts, ctx, buf, traceparent ) # Header - request id and message length @@ -1022,6 +1073,7 @@ def _client_do_batched_op_msg( namespaces: list[str], opts: CodecOptions[Any], ctx: _ClientBulkWriteContext, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Create the next batched client-level bulkWrite operation using OP_MSG. @@ -1032,8 +1084,10 @@ def _client_do_batched_op_msg( else: ack = True if ctx.conn.compression_context: - return _client_batched_op_msg_compressed(command, operations, namespaces, ack, opts, ctx) - return _client_batched_op_msg(command, operations, namespaces, ack, opts, ctx) + return _client_batched_op_msg_compressed( + command, operations, namespaces, ack, opts, ctx, traceparent + ) + return _client_batched_op_msg(command, operations, namespaces, ack, opts, ctx, traceparent) # End OP_MSG ----------------------------------------------------- @@ -1353,7 +1407,11 @@ def as_command( return self._as_command def get_message( - self, read_preference: _ServerMode, conn: _AgnosticConnection, use_cmd: bool = False + self, + read_preference: _ServerMode, + conn: _AgnosticConnection, + use_cmd: bool = False, + traceparent: Optional[str] = None, ) -> tuple[int, bytes, int]: """Get a query message""" # Use the read_preference decided by _socket_from_server. @@ -1367,6 +1425,7 @@ def get_message( read_preference, self.codec_options, ctx=conn.compression_context, + traceparent=traceparent, ) return request_id, msg, size @@ -1464,7 +1523,11 @@ def as_command( return self._as_command def get_message( - self, dummy0: Any, conn: _AgnosticConnection, use_cmd: bool = False + self, + dummy0: Any, + conn: _AgnosticConnection, + use_cmd: bool = False, + traceparent: Optional[str] = None, ) -> Union[tuple[int, bytes, int], tuple[int, bytes]]: """Get a getmore message.""" ns = self.namespace() @@ -1477,7 +1540,13 @@ def get_message( else: flags = 0 request_id, msg, size, _ = _op_msg( - flags, spec, self.db, None, self.codec_options, ctx=conn.compression_context + flags, + spec, + self.db, + None, + self.codec_options, + ctx=conn.compression_context, + traceparent=traceparent, ) return request_id, msg, size diff --git a/pymongo/pool_shared.py b/pymongo/pool_shared.py index ca7bf9370b..41b7368b78 100644 --- a/pymongo/pool_shared.py +++ b/pymongo/pool_shared.py @@ -57,6 +57,7 @@ class _ConnectionTelemetryInfo(Protocol): server_connection_id: Optional[int] address: _Address service_id: Optional[ObjectId] + max_wire_version: int def _get_ssl_session(ssl_sock: Any) -> Optional[Any]: diff --git a/pymongo/synchronous/bulk.py b/pymongo/synchronous/bulk.py index 93c78869f4..cbdc0dd5dc 100644 --- a/pymongo/synchronous/bulk.py +++ b/pymongo/synchronous/bulk.py @@ -28,11 +28,12 @@ Any, Optional, Union, + cast, ) from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument -from pymongo import _csot, common +from pymongo import _csot, _otel, common from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.bulk_shared import ( _COMMANDS, @@ -236,7 +237,6 @@ def gen_unordered(self) -> Iterator[_Run]: if run.ops: yield run - @_handle_reauth def write_command( self, bwc: _BulkWriteContext, @@ -245,6 +245,8 @@ def write_command( msg: bytes, docs: list[Mapping[str, Any]], client: MongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> dict[str, Any]: """Run a batch write command, returning the response as a dict.""" cmd[bwc.field] = docs @@ -254,6 +256,7 @@ def write_command( request_id, msg, client=client, + precreated_span=precreated_span, ) return result_docs[0] @@ -266,6 +269,8 @@ def unack_write( max_doc_size: int, docs: list[Mapping[str, Any]], client: MongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> Optional[Mapping[str, Any]]: """Send an unacknowledged batch write command.""" # Historically the STARTED log omits the documents while the published @@ -282,6 +287,7 @@ def unack_write( orig=published, max_doc_size=max_doc_size, unacknowledged=True, + precreated_span=precreated_span, ) return None @@ -302,13 +308,23 @@ def _execute_batch_unack( client=client, # type: ignore[arg-type] ) else: - request_id, msg, to_send = bwc.batch_command(cmd, ops) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send = bwc.batch_command(cmd, ops, traceparent=traceparent) + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, bwc.field: to_send} + ) + msg = cast(bytes, msg) # Though this isn't strictly a "legacy" write, the helper # handles publishing commands and sending our message # without receiving a result. Send 0 for max_doc_size # to disable size checking. Size checking is handled while # the documents are encoded to BSON. - self.unack_write(bwc, cmd, request_id, msg, 0, to_send, client) # type: ignore[arg-type] + self.unack_write( + bwc, cmd, request_id, msg, 0, to_send, client, precreated_span=precreated_span + ) # type: ignore[arg-type] return to_send @@ -316,10 +332,12 @@ def _execute_batch( self, bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext], cmd: dict[str, Any], - ops: list[Mapping[str, Any]], + run_ops: list[Mapping[str, Any]], + idx_offset: int, client: MongoClient[Any], ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]: if self.is_encrypted: + ops = cast("list[Mapping[str, Any]]", islice(run_ops, idx_offset, None)) _, batched_cmd, to_send = bwc.batch_command(cmd, ops) result = bwc.conn.command( # type: ignore[misc] bwc.db_name, @@ -328,11 +346,31 @@ def _execute_batch( session=bwc.session, # type: ignore[arg-type] client=client, # type: ignore[arg-type] ) - else: - request_id, msg, to_send = bwc.batch_command(cmd, ops) - result = self.write_command(bwc, cmd, request_id, msg, to_send, client) # type: ignore[arg-type] + return result, to_send # type: ignore[return-value] + + def run( + ctx: _BulkWriteContext, + ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]: + # Reauthentication re-runs this closure, so undo the field that + # write_command added and re-encode against a fresh span and + # traceparent. + cmd.pop(ctx.field, None) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, ctx.conn, cmd, ctx.db_name, ctx.name + ) as (span, precreated_span, traceparent): + ops = cast("list[Mapping[str, Any]]", islice(run_ops, idx_offset, None)) + request_id, msg, to_send = ctx.batch_command(cmd, ops, traceparent=traceparent) + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, ctx.field: to_send} + ) + msg = cast(bytes, msg) + result = self.write_command( + ctx, cmd, request_id, msg, to_send, client, precreated_span=precreated_span + ) + return result, to_send - return result, to_send # type: ignore[return-value] + return _handle_reauth(run)(bwc) def _execute_command( self, @@ -406,7 +444,7 @@ def _execute_command( # Run as many ops as possible in one command. if write_concern.acknowledged: - result, to_send = self._execute_batch(bwc, cmd, ops, client) + result, to_send = self._execute_batch(bwc, cmd, run.ops, run.idx_offset, client) # Retryable writeConcernErrors halt the execution of this run. wce = result.get("writeConcernError", {}) diff --git a/pymongo/synchronous/client_bulk.py b/pymongo/synchronous/client_bulk.py index 4093e9ab39..278a1e5f64 100644 --- a/pymongo/synchronous/client_bulk.py +++ b/pymongo/synchronous/client_bulk.py @@ -28,11 +28,12 @@ Any, Optional, Union, + cast, ) from bson.objectid import ObjectId from bson.raw_bson import RawBSONDocument -from pymongo import _csot, common +from pymongo import _csot, _otel, common from pymongo._telemetry import _generate_op_id_or_none, _operation_telemetry_or_none from pymongo.synchronous.client_session import ( ClientSession, @@ -236,6 +237,8 @@ def write_command( op_docs: list[Mapping[str, Any]], ns_docs: list[Mapping[str, Any]], client: MongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> dict[str, Any]: """Run a client-level batch write command, returning the response as a dict.""" cmd["ops"] = op_docs @@ -247,6 +250,7 @@ def write_command( request_id, msg, # type: ignore[arg-type] client=client, + precreated_span=precreated_span, ) reply = result_docs[0] except Exception as exc: @@ -263,6 +267,8 @@ def unack_write( op_docs: list[Mapping[str, Any]], ns_docs: list[Mapping[str, Any]], client: MongoClient[Any], + *, + precreated_span: Optional[Any] = None, ) -> Optional[Mapping[str, Any]]: """Send an unacknowledged client-level batch write command.""" # Historically the STARTED log omits the ops/nsInfo while the published @@ -281,6 +287,7 @@ def unack_write( orig=published, max_doc_size=bwc.max_bson_size, unacknowledged=True, + precreated_span=precreated_span, ) reply: Mapping[str, Any] = result_docs[0] except Exception as exc: @@ -296,8 +303,29 @@ def _execute_batch_unack( namespaces: list[str], ) -> tuple[list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Executes a batch of bulkWrite server commands (unack).""" - request_id, msg, to_send_ops, to_send_ns = bwc.batch_command(cmd, ops, namespaces) - self.unack_write(bwc, cmd, request_id, msg, to_send_ops, to_send_ns, self.client) # type: ignore[arg-type] + tracing_options = self.client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send_ops, to_send_ns = bwc.batch_command( + cmd, ops, namespaces, traceparent=traceparent + ) + # The ops and nsInfo are only known after encoding, so the query text + # is refreshed from the full published command. + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, "ops": to_send_ops, "nsInfo": to_send_ns} + ) + msg = cast(bytes, msg) + self.unack_write( + bwc, + cmd, + request_id, + msg, + to_send_ops, + to_send_ns, + self.client, + precreated_span=precreated_span, + ) # type: ignore[arg-type] return to_send_ops, to_send_ns def _execute_batch( @@ -308,8 +336,29 @@ def _execute_batch( namespaces: list[str], ) -> tuple[dict[str, Any], list[Mapping[str, Any]], list[Mapping[str, Any]]]: """Executes a batch of bulkWrite server commands (ack).""" - request_id, msg, to_send_ops, to_send_ns = bwc.batch_command(cmd, ops, namespaces) - result = self.write_command(bwc, cmd, request_id, msg, to_send_ops, to_send_ns, self.client) # type: ignore[arg-type] + tracing_options = self.client.options.tracing + with _otel._command_span_for_encoding( + tracing_options, bwc.conn, cmd, bwc.db_name, bwc.name + ) as (span, precreated_span, traceparent): + request_id, msg, to_send_ops, to_send_ns = bwc.batch_command( + cmd, ops, namespaces, traceparent=traceparent + ) + # The ops and nsInfo are only known after encoding, so the query text + # is refreshed from the full published command. + _otel._set_command_span_query_text( + span, tracing_options, {**cmd, "ops": to_send_ops, "nsInfo": to_send_ns} + ) + msg = cast(bytes, msg) + result = self.write_command( + bwc, + cmd, + request_id, + msg, + to_send_ops, + to_send_ns, + self.client, + precreated_span=precreated_span, + ) # type: ignore[arg-type] return result, to_send_ops, to_send_ns # type: ignore[return-value] def _process_results_cursor( diff --git a/pymongo/synchronous/command_runner.py b/pymongo/synchronous/command_runner.py index 27380eb068..3b034caced 100644 --- a/pymongo/synchronous/command_runner.py +++ b/pymongo/synchronous/command_runner.py @@ -102,6 +102,7 @@ def _run_command( set_conn_more_to_come: bool = False, unpack_res: Optional[Callable[..., Any]] = None, cursor_id: Optional[int] = None, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], float]: """Send ``msg`` over ``conn`` and return ``(docs, reply, duration_s)``, where ``duration_s`` is the round-trip duration in seconds. @@ -185,6 +186,7 @@ def _run_command( tracing_options=tracing_options, speculative_hello=speculative_hello, name=name, + precreated_span=precreated_span, ) telemetry.started(orig, ensure_db) start = 0.0 @@ -266,6 +268,7 @@ def run_bulk_write_command( orig: Optional[MutableMapping[str, Any]] = None, max_doc_size: int = 0, unacknowledged: bool = False, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], float]: """Send a bulk write batch and return ``(docs, reply, duration_s)``. @@ -297,6 +300,7 @@ def run_bulk_write_command( max_doc_size=max_doc_size, unacknowledged=unacknowledged, decrypt_reply=False, + precreated_span=precreated_span, ) @@ -318,6 +322,7 @@ def run_cursor_command( more_to_come: bool = False, unpack_res: Optional[Callable[..., Any]] = None, cursor_id: Optional[int] = None, + precreated_span: Optional[Any] = None, ) -> tuple[list[dict[str, Any]], Optional[_OpMsg], datetime.timedelta]: """Run a cursor ``find``/``getMore`` operation over ``conn``. @@ -359,6 +364,7 @@ def run_cursor_command( more_to_come=more_to_come, unpack_res=unpack_res, cursor_id=cursor_id, + precreated_span=precreated_span, ) # The cursor path stores the duration on Response, which expects a timedelta. return docs, reply, datetime.timedelta(seconds=duration_s) @@ -441,16 +447,26 @@ def run_command( flags = _OpMsg.MORE_TO_COME if unacknowledged else 0 flags |= _OpMsg.EXHAUST_ALLOWED if exhaust_allowed else 0 - request_id, msg, size, max_doc_size = message._op_msg( - flags, spec, dbname, read_preference, codec_options, ctx=compression_ctx - ) - # If this is an unacknowledged write then make sure the encoded doc(s) - # are small enough, otherwise rely on the server to return an error. - if unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size: - message._raise_document_too_large(name, size, max_bson_size) + tracing_options = client.options.tracing if client is not None else None + with _otel._command_span_for_encoding( + tracing_options, conn, spec, dbname, name, speculative_hello + ) as (_, precreated_span, traceparent): + request_id, msg, size, max_doc_size = message._op_msg( + flags, + spec, + dbname, + read_preference, + codec_options, + ctx=compression_ctx, + traceparent=traceparent, + ) + # If this is an unacknowledged write then make sure the encoded doc(s) + # are small enough, otherwise rely on the server to return an error. + if unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size: + message._raise_document_too_large(name, size, max_bson_size) - if max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD: - message._raise_document_too_large(name, size, max_bson_size + message._COMMAND_OVERHEAD) + if max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD: + message._raise_document_too_large(name, size, max_bson_size + message._COMMAND_OVERHEAD) docs, _, _ = _run_command( conn, spec, @@ -470,5 +486,6 @@ def run_command( speculative_hello=speculative_hello, unacknowledged=unacknowledged, set_conn_more_to_come=True, + precreated_span=precreated_span, ) return docs[0] # type: ignore[return-value] diff --git a/pymongo/synchronous/cursor_base.py b/pymongo/synchronous/cursor_base.py index c4376e8842..4bf2a7daed 100644 --- a/pymongo/synchronous/cursor_base.py +++ b/pymongo/synchronous/cursor_base.py @@ -20,7 +20,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Optional, Union -from pymongo import _csot +from pymongo import _csot, _otel from pymongo.cursor_shared import _CURSOR_DOC_FIELDS, _AgnosticCursorBase, _split_message from pymongo.lock import _create_lock from pymongo.message import _GetMore, _OpMsg, _Query @@ -113,11 +113,19 @@ def _run_with_conn( use_cmd = operation.use_command(conn) more_to_come = bool(operation.conn_mgr and operation.conn_mgr.more_to_come) cmd, dbn = _operation_to_command(operation, conn, use_cmd) - if more_to_come: - request_id, data, max_doc_size = 0, b"", 0 - else: - message = operation.get_message(read_preference, conn, use_cmd) - request_id, data, max_doc_size = _split_message(message) + tracing_options = client.options.tracing + with _otel._command_span_for_encoding(tracing_options, conn, cmd, dbn, operation.name) as ( + _, + precreated_span, + traceparent, + ): + if more_to_come: + request_id, data, max_doc_size = 0, b"", 0 + else: + message = operation.get_message( + read_preference, conn, use_cmd, traceparent=traceparent + ) + request_id, data, max_doc_size = _split_message(message) user_fields = _CURSOR_DOC_FIELDS if use_cmd else None docs, reply, duration = run_cursor_command( conn, @@ -136,6 +144,7 @@ def _run_with_conn( more_to_come=more_to_come, unpack_res=self._unpack_response, cursor_id=operation.cursor_id, + precreated_span=precreated_span, ) assert reply is not None if client._should_pin_cursor(operation.session) or operation.exhaust: # type: ignore[arg-type] diff --git a/test/asynchronous/test_otel.py b/test/asynchronous/test_otel.py index e122bb7ea5..8c2605923c 100644 --- a/test/asynchronous/test_otel.py +++ b/test/asynchronous/test_otel.py @@ -16,12 +16,13 @@ from __future__ import annotations +import asyncio import gc import os import subprocess import sys -from typing import Optional -from unittest.mock import patch +from typing import Callable, Optional +from unittest.mock import MagicMock, patch sys.path[0:0] = [""] @@ -35,6 +36,7 @@ ClientBulkWriteException, ConfigurationError, ConnectionFailure, + InvalidDocument, InvalidOperation, NetworkTimeout, OperationFailure, @@ -269,6 +271,45 @@ def test_run_command_operation_name_override(self): self.assertEqual(span.attributes["db.operation.name"], "runCommand") +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestCommandTelemetryPrecreatedSpan(unittest.TestCase): + """A pre-created command span is used as-is, including the "no span" sentinel.""" + + def _telemetry(self, **kwargs): + return _telemetry._CommandTelemetry( + None, + MagicMock(), + None, + {"insert": "coll"}, + "db", + 1, + None, + tracing_options=_tracing_opts(), + name="insert", + **kwargs, + ) + + def test_omitted_span_is_created_in_started(self): + telemetry = self._telemetry() + with patch.object(_otel, "start_command_span", return_value=None) as start: + telemetry.started({"insert": "coll"}, False) + start.assert_called_once() + + def test_sensitive_sentinel_skips_creation(self): + telemetry = self._telemetry(precreated_span=_otel._NO_COMMAND_SPAN) + with patch.object(_otel, "start_command_span", return_value=None) as start: + telemetry.started({"insert": "coll"}, False) + start.assert_not_called() + + def test_precreated_span_is_used(self): + span = MagicMock() + telemetry = self._telemetry(precreated_span=span) + with patch.object(_otel, "start_command_span", return_value=None) as start: + telemetry.started({"insert": "coll"}, False) + start.assert_not_called() + self.assertIs(telemetry._span, span) + + @unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") class TestOperationTelemetryContextManager(unittest.TestCase): @classmethod @@ -469,6 +510,7 @@ class _FakeUnixConn: server_connection_id: Optional[int] = None address: _Address = ("/tmp/fake-otel-test.sock", None) service_id = None + max_wire_version = 30 self.exporter.clear() span = _otel.start_command_span( @@ -560,7 +602,9 @@ async def test_failure_records_exception_and_status_code(self): self.assertTrue(any(event.name == "exception" for event in span.events)) @async_client_context.require_failCommand_fail_point - async def test_operation_span_error_type_is_exception_class_name_for_server_error(self): + async def test_prose_6_error_type_on_operation_span_is_exception_class_name_for_server_error( + self, + ): # A non-retryable server error names the exception class on the operation # span, not the server error code. client = await self.async_rs_or_single_client(tracing={"enabled": True}, retryReads=False) @@ -583,7 +627,7 @@ async def test_operation_span_error_type_is_exception_class_name_for_server_erro ) @async_client_context.require_failCommand_fail_point - async def test_error_type_is_exception_class_name_for_connection_failure(self): + async def test_prose_5_error_type_is_exception_class_name_for_non_server_error(self): # A closed connection produces no server reply, so error.type uses the class name. client = await self.async_rs_or_single_client(tracing={"enabled": True}, retryReads=False) fail_command = { @@ -778,6 +822,55 @@ async def test_collection_bulk_write_unacknowledged_gets_operation_span(self): self.assertEqual(matching[0].attributes["db.namespace"], self.db.name) self.assertEqual(matching[0].attributes["db.collection.name"], "test") + async def test_collection_bulk_write_query_text_includes_documents(self): + # The batch is encoded before write_command adds the documents sequence, + # so the pre-created span's query text must be refreshed. + client = await self.async_rs_or_single_client( + tracing={"enabled": True, "query_text_max_length": 1024} + ) + coll = client[self.db.name]["test_bulk_query_text"] + self.exporter.clear() + await coll.bulk_write([InsertOne({"x": 1})], ordered=True) + + (span,) = self.command_spans(self.exporter.get_finished_spans(), "insert") + self.assertIn("documents", span.attributes["db.query.text"]) + + async def test_command_span_ended_when_encoding_fails(self): + # The command span is created before encoding, so an encoding error that + # never reaches the server must still end it rather than leak it. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + with self.assertRaises(InvalidDocument): + await client[self.db.name]["test_encoding_failure"].insert_one({"x": object()}) + + (span,) = self.command_spans(self.exporter.get_finished_spans(), "insert") + self.assertEqual(span.status.status_code, StatusCode.ERROR) + + @async_client_context.require_auth + @async_client_context.require_failCommand_fail_point + async def test_bulk_write_reauth_creates_a_new_command_span(self): + # Reauthentication resends the batch, so it must re-create the span and + # re-encode the traceparent rather than reuse the ended one. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_reauth"] + await coll.delete_many({}) + await client.admin.command( + "configureFailPoint", + "failCommand", + mode={"times": 1}, + data={"failCommands": ["insert"], "errorCode": 391}, + ) + try: + self.exporter.clear() + await coll.bulk_write([InsertOne({"x": 1})], ordered=True) + finally: + await client.admin.command("configureFailPoint", "failCommand", mode="off") + + spans = self.command_spans(self.exporter.get_finished_spans(), "insert") + self.assertEqual(len(spans), 2) + self.assertEqual(spans[0].status.status_code, StatusCode.ERROR) + self.assertNotEqual(spans[0].context.span_id, spans[1].context.span_id) + 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 @@ -935,6 +1028,326 @@ async def test_end_sessions_gets_operation_span(self): self.assertEqual(len(cmd_spans), 1) self.assertEqual(cmd_spans[0].parent.span_id, op_span.context.span_id) + @async_client_context.require_version_min(9, 0) + async def test_traceparent_injected_when_tracing_enabled(self): + """Verify the telemetry section is present in the OP_MSG when tracing is on.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + await client.admin.command("ping") + cmd_spans = self.command_spans(self.exporter.get_finished_spans(), "ping") + self.assertEqual(len(cmd_spans), 1) + tp = _otel._traceparent_from_span(cmd_spans[0]) + self.assertIsNotNone(tp) + assert tp is not None + self.assertEqual(len(tp), 55) + + async def test_no_traceparent_when_tracing_disabled(self): + """Verify no telemetry section when tracing is off.""" + client = await self.async_rs_or_single_client() + self.exporter.clear() + await client.admin.command("ping") + # No spans at all when tracing is disabled + self.assertEqual(self.ping_spans(), []) + + async def test_no_traceparent_for_sensitive_commands(self): + """Verify no command span (and thus no traceparent) for saslStart.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + with self.assertRaises(OperationFailure): + await client.admin.command("saslStart", mechanism="SCRAM-SHA-256", payload=b"") + command_spans = [s for s in self.spans() if "db.command.name" in s.attributes] + self.assertNotIn("saslStart", [s.name for s in command_spans]) + + +class TestTraceparent(unittest.TestCase): + """Unit tests for traceparent extraction and telemetry section encoding.""" + + @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") + def test_traceparent_from_span_valid(self): + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + + ctx = SpanContext( + trace_id=int("0123456789abcdef0123456789abcdef", 16), + span_id=int("0123456789abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + + tp = _otel._traceparent_from_span(NonRecordingSpan(context=ctx)) + self.assertIsNotNone(tp) + assert tp is not None + self.assertEqual(len(tp), 55) + self.assertEqual(tp, "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01") + + @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") + def test_traceparent_from_span_unsampled(self): + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + + ctx = SpanContext( + trace_id=int("0123456789abcdef0123456789abcdef", 16), + span_id=int("0123456789abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.DEFAULT), + ) + + tp = _otel._traceparent_from_span(NonRecordingSpan(context=ctx)) + self.assertIsNotNone(tp) + self.assertEqual(tp, "00-0123456789abcdef0123456789abcdef-0123456789abcdef-00") + + @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") + def test_traceparent_from_span_none(self): + self.assertIsNone(_otel._traceparent_from_span(None)) + + @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") + def test_traceparent_from_span_invalid_context(self): + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + + ctx = SpanContext( + trace_id=0, + span_id=0, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.DEFAULT), + ) + + self.assertIsNone(_otel._traceparent_from_span(NonRecordingSpan(context=ctx))) + + @unittest.skipUnless(_otel._HAS_OPENTELEMETRY, "opentelemetry is not installed") + def test_traceparent_from_span_flags_above_one_byte(self): + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + + ctx = SpanContext( + trace_id=int("0123456789abcdef0123456789abcdef", 16), + span_id=int("0123456789abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(0x1234), + ) + + self.assertIsNone(_otel._traceparent_from_span(NonRecordingSpan(context=ctx))) + + def test_telemetry_section_format(self): + tp = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01" + section = _otel._telemetry_section(tp) + # First byte is payload type 3 + self.assertEqual(section[0:1], b"\x03") + # Rest is a valid BSON document + import bson as _bson + + doc = _bson.decode(section[1:]) + self.assertEqual(doc, {"otel": {"traceparent": tp}}) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestServerTraceContext(AsyncIntegrationTest): + """Prose tests 5-7: server trace context propagation. + + These tests require a MongoDB 9.0+ deployment with the OpenTelemetry + file exporter enabled. The OTEL_TRACE_DIR environment variable + (set by drivers-evergreen-tools with OTEL=1) points to the directory + where the server writes OTLP JSON span files. Skip when unset. + """ + + OTEL_TRACE_DIR = os.environ.get("OTEL_TRACE_DIR", "") + + @classmethod + def setUpClass(cls): + if not cls.OTEL_TRACE_DIR: + raise unittest.SkipTest("OTEL_TRACE_DIR not set") + super().setUpClass() + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "exporter"): + cls.exporter.shutdown() + super().tearDownClass() + + async def asyncSetUp(self): + await super().asyncSetUp() + self.exporter.clear() + + def spans(self, name: str | None = None): + finished = self.exporter.get_finished_spans() + if name is None: + return list(finished) + return [s for s in finished if s.name == name] + + @staticmethod + def 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] + + @staticmethod + async def _read_server_spans(trace_dir: str) -> list[dict]: + """Read all OTLP JSON span files under trace_dir.""" + import json + from pathlib import Path + + def read() -> list[dict]: + spans: list[dict] = [] + for p in Path(trace_dir).rglob("*"): + if not p.is_file(): + continue + for line in p.read_text().splitlines(): + line = line.strip() + if not line: + continue + try: + batch = json.loads(line) + except json.JSONDecodeError: + continue + for rs in batch.get("resourceSpans", []): + for ss in rs.get("scopeSpans", []): + spans.extend(ss.get("spans", [])) + return spans + + # File reads block, so keep them off the event loop in the async driver. + if not _IS_SYNC: + return await asyncio.get_running_loop().run_in_executor(None, read) # type: ignore[return-value] + return read() + + @staticmethod + async def _poll_server_spans( + trace_dir: str, + trace_id: str, + predicate: Optional[Callable[[list[dict]], bool]] = None, + timeout: float = 30.0, + ) -> list[dict]: + """Poll trace_dir for server spans matching trace_id. + + Server spans are batched, so return as soon as ``predicate`` holds for + the matching spans (or any matching span exists when ``predicate`` is + ``None``), or when the timeout elapses. + """ + import time + + deadline = time.monotonic() + timeout + matching: list[dict] = [] + while True: + spans = await TestServerTraceContext._read_server_spans(trace_dir) + matching = [s for s in spans if s.get("traceId") == trace_id] + if predicate is not None: + if predicate(matching): + return matching + elif matching: + return matching + if time.monotonic() >= deadline: + return matching + await asyncio.sleep(0.5) + + @async_client_context.require_version_min(9, 0) + async def test_prose_7_server_spans_join_driver_trace(self): + """Prose Test 5: Server spans join the driver's trace.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_5"] + await coll.delete_many({}) + await coll.insert_one({"x": 1}) + cmd_spans = self.command_spans(self.exporter.get_finished_spans(), "insert") + self.assertEqual(len(cmd_spans), 1) + cmd_span = cmd_spans[0] + trace_id = f"{cmd_span.context.trace_id:032x}" + span_id = f"{cmd_span.context.span_id:016x}" + + # On a sharded cluster the command fans out through mongos to the shard, + # and every hop exports its own span named for the command. The ingress + # span is the one whose parent is the driver's command span. Poll for it + # explicitly: the deeper hops end (and flush) first, so a poll that + # stops at the first span of the trace would miss it. + def has_ingress_span(spans: list[dict]) -> bool: + ingress = [s for s in spans if s.get("parentSpanId") == span_id] + return len(ingress) == 1 and ingress[0].get("name") == "insert" + + server_spans = await self._poll_server_spans( + self.OTEL_TRACE_DIR, trace_id, has_ingress_span + ) + ingress_spans = [s for s in server_spans if s.get("parentSpanId") == span_id] + self.assertEqual(len(ingress_spans), 1) + self.assertEqual(ingress_spans[0].get("name"), "insert") + + @async_client_context.require_version_min(9, 0) + @async_client_context.require_failCommand_fail_point + async def test_prose_8_one_server_span_per_retry_attempt(self): + """Prose Test 6: One server span per retry attempt.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_6"] + await coll.delete_many({}) + await coll.insert_one({"x": 1}) + # Configure failpoint + await client.admin.command( + "configureFailPoint", + "failCommand", + mode={"times": 1}, + data={"failCommands": ["find"], "errorCode": 91}, + ) + try: + self.exporter.clear() + cursor = coll.find({}) + docs = await cursor.to_list() + self.assertEqual(len(docs), 1) + finally: + await client.admin.command( + "configureFailPoint", + "failCommand", + mode="off", + ) + find_cmd_spans = self.command_spans(self.exporter.get_finished_spans(), "find") + self.assertEqual(len(find_cmd_spans), 2) + self.assertEqual( + f"{find_cmd_spans[0].context.trace_id:032x}", + f"{find_cmd_spans[1].context.trace_id:032x}", + ) + span_ids = {f"{s.context.span_id:016x}" for s in find_cmd_spans} + self.assertEqual(len(span_ids), 2) + trace_id = f"{find_cmd_spans[0].context.trace_id:032x}" + + def _both_attempts_have_child(spans: list[dict]) -> bool: + return span_ids <= {s.get("parentSpanId") for s in spans} + + server_spans = await self._poll_server_spans( + self.OTEL_TRACE_DIR, trace_id, _both_attempts_have_child + ) + parented = [s for s in server_spans if s.get("parentSpanId") in span_ids] + self.assertEqual(len(parented), 2) + for span_id in span_ids: + children = [s for s in parented if s.get("parentSpanId") == span_id] + self.assertEqual(len(children), 1) + + @async_client_context.require_version_min(9, 0) + @async_client_context.require_auth + async def test_prose_9_no_trace_context_for_auth_monitoring(self): + """Prose Test 7: No trace context for authentication and monitoring commands.""" + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_7"] + # Handshakes and authentication run lazily on the first operation, so + # their spans are part of the driver's trace ids and must not be cleared. + await coll.delete_many({}) + await coll.insert_one({"x": 1}) + await coll.find_one({}) + finished = self.exporter.get_finished_spans() + driver_trace_ids = {f"{s.context.trace_id:032x}" for s in finished} + # Wait for the find server span to appear + find_cmd_spans = self.command_spans(finished, "find") + self.assertTrue(find_cmd_spans) + trace_id = f"{find_cmd_spans[0].context.trace_id:032x}" + await self._poll_server_spans(self.OTEL_TRACE_DIR, trace_id) + # Collect all server spans that match any driver trace id + all_server_spans = await self._read_server_spans(self.OTEL_TRACE_DIR) + driver_server_spans = [s for s in all_server_spans if s.get("traceId") in driver_trace_ids] + auth_monitor_names = { + "hello", + "ismaster", + "isMaster", + "saslStart", + "saslContinue", + "authenticate", + } + for s in driver_server_spans: + self.assertNotIn( + s.get("name"), + auth_monitor_names, + f"Server span for {s.get('name')} joined a driver trace", + ) + # These unit tests cover the validator's edge cases: the rejection paths and the # explicit-zero vs unset distinction for query_text_max_length. @@ -1006,6 +1419,7 @@ class _FakeConn: server_connection_id: Optional[int] = None address: _Address = ("localhost", 27017) service_id = None + max_wire_version = 30 with patch.object(_otel, "trace") as mock_trace: for _ in range(3): diff --git a/test/test_message.py b/test/test_message.py index 9c11ffef5b..ccb1cf90ce 100644 --- a/test/test_message.py +++ b/test/test_message.py @@ -18,26 +18,31 @@ import struct import sys +from io import BytesIO from typing import Any from unittest.mock import MagicMock sys.path[0:0] = [""] -from bson import CodecOptions, encode +from bson import CodecOptions, decode, encode from bson.objectid import ObjectId from pymongo.common import MIN_SUPPORTED_WIRE_VERSION, MONGOS_EXHAUST_WIRE_VERSION from pymongo.compression_support import ZlibContext, _have_zlib from pymongo.errors import DocumentTooLarge, InvalidOperation, OperationFailure from pymongo.hello import _get_server_type from pymongo.message import ( + _batched_op_msg_impl, _check_exhaust_supported, + _client_do_batched_op_msg, _convert_client_bulk_exception, _convert_exception, + _do_batched_op_msg, _gen_find_command, _gen_get_more_command, _GetMore, _maybe_add_read_preference, _op_msg, + _op_msg_no_header, _Query, _raise_document_too_large, ) @@ -484,5 +489,119 @@ def test_get_more_comment_not_added_on_low_wire_version(self): self.assertNotIn("comment", cmd) +class TestTelemetrySection(unittest.TestCase): + """Wire-format tests for the OP_MSG Payload Type 3 telemetry section. + + ``_op_msg`` and ``_do_batched_op_msg`` dispatch to the C extension when it + is built and to the pure-Python encoders otherwise. ``_op_msg_no_header``, + ``_batched_op_msg_impl``, and ``_client_do_batched_op_msg`` are always the + pure-Python encoders. + """ + + TRACEPARENT = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01" + MAX_BSON_SIZE = 16 * 1024 * 1024 + + @staticmethod + def _sections(payload: bytes) -> list[tuple[int, bytes]]: + """Split an OP_MSG body (flag bits first) into (kind, bytes) sections.""" + sections = [] + pos = 4 # Skip flagBits. + while pos < len(payload): + kind = payload[pos] + pos += 1 + size = struct.unpack(" None: + self.assertEqual(sections[-1][0], 3) + self.assertEqual(decode(sections[-1][1]), {"otel": {"traceparent": self.TRACEPARENT}}) + + def _ctx(self, max_message_size: int = MAX_BSON_SIZE) -> MagicMock: + ctx = MagicMock() + ctx.max_bson_size = self.MAX_BSON_SIZE + ctx.max_write_batch_size = 1000 + ctx.max_message_size = max_message_size + ctx.conn.compression_context = None + return ctx + + def test_op_msg_telemetry_section_is_last(self): + _, msg, _, _ = _op_msg( + 0, {"insert": "coll"}, "db", None, _OPTS, traceparent=self.TRACEPARENT + ) + self.assertEqual(struct.unpack(" list[dict]: + """Read all OTLP JSON span files under trace_dir.""" + import json + from pathlib import Path + + def read() -> list[dict]: + spans: list[dict] = [] + for p in Path(trace_dir).rglob("*"): + if not p.is_file(): + continue + for line in p.read_text().splitlines(): + line = line.strip() + if not line: + continue + try: + batch = json.loads(line) + except json.JSONDecodeError: + continue + for rs in batch.get("resourceSpans", []): + for ss in rs.get("scopeSpans", []): + spans.extend(ss.get("spans", [])) + return spans + + # File reads block, so keep them off the event loop in the async driver. + if not _IS_SYNC: + return asyncio.get_running_loop().run_in_executor(None, read) # type: ignore[return-value] + return read() + + @staticmethod + def _poll_server_spans( + trace_dir: str, + trace_id: str, + predicate: Optional[Callable[[list[dict]], bool]] = None, + timeout: float = 30.0, + ) -> list[dict]: + """Poll trace_dir for server spans matching trace_id. + + Server spans are batched, so return as soon as ``predicate`` holds for + the matching spans (or any matching span exists when ``predicate`` is + ``None``), or when the timeout elapses. + """ + import time + + deadline = time.monotonic() + timeout + matching: list[dict] = [] + while True: + spans = TestServerTraceContext._read_server_spans(trace_dir) + matching = [s for s in spans if s.get("traceId") == trace_id] + if predicate is not None: + if predicate(matching): + return matching + elif matching: + return matching + if time.monotonic() >= deadline: + return matching + time.sleep(0.5) + + @client_context.require_version_min(9, 0) + def test_prose_7_server_spans_join_driver_trace(self): + """Prose Test 5: Server spans join the driver's trace.""" + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_5"] + coll.delete_many({}) + coll.insert_one({"x": 1}) + cmd_spans = self.command_spans(self.exporter.get_finished_spans(), "insert") + self.assertEqual(len(cmd_spans), 1) + cmd_span = cmd_spans[0] + trace_id = f"{cmd_span.context.trace_id:032x}" + span_id = f"{cmd_span.context.span_id:016x}" + + # On a sharded cluster the command fans out through mongos to the shard, + # and every hop exports its own span named for the command. The ingress + # span is the one whose parent is the driver's command span. Poll for it + # explicitly: the deeper hops end (and flush) first, so a poll that + # stops at the first span of the trace would miss it. + def has_ingress_span(spans: list[dict]) -> bool: + ingress = [s for s in spans if s.get("parentSpanId") == span_id] + return len(ingress) == 1 and ingress[0].get("name") == "insert" + + server_spans = self._poll_server_spans(self.OTEL_TRACE_DIR, trace_id, has_ingress_span) + ingress_spans = [s for s in server_spans if s.get("parentSpanId") == span_id] + self.assertEqual(len(ingress_spans), 1) + self.assertEqual(ingress_spans[0].get("name"), "insert") + + @client_context.require_version_min(9, 0) + @client_context.require_failCommand_fail_point + def test_prose_8_one_server_span_per_retry_attempt(self): + """Prose Test 6: One server span per retry attempt.""" + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_6"] + coll.delete_many({}) + coll.insert_one({"x": 1}) + # Configure failpoint + client.admin.command( + "configureFailPoint", + "failCommand", + mode={"times": 1}, + data={"failCommands": ["find"], "errorCode": 91}, + ) + try: + self.exporter.clear() + cursor = coll.find({}) + docs = cursor.to_list() + self.assertEqual(len(docs), 1) + finally: + client.admin.command( + "configureFailPoint", + "failCommand", + mode="off", + ) + find_cmd_spans = self.command_spans(self.exporter.get_finished_spans(), "find") + self.assertEqual(len(find_cmd_spans), 2) + self.assertEqual( + f"{find_cmd_spans[0].context.trace_id:032x}", + f"{find_cmd_spans[1].context.trace_id:032x}", + ) + span_ids = {f"{s.context.span_id:016x}" for s in find_cmd_spans} + self.assertEqual(len(span_ids), 2) + trace_id = f"{find_cmd_spans[0].context.trace_id:032x}" + + def _both_attempts_have_child(spans: list[dict]) -> bool: + return span_ids <= {s.get("parentSpanId") for s in spans} + + server_spans = self._poll_server_spans( + self.OTEL_TRACE_DIR, trace_id, _both_attempts_have_child + ) + parented = [s for s in server_spans if s.get("parentSpanId") in span_ids] + self.assertEqual(len(parented), 2) + for span_id in span_ids: + children = [s for s in parented if s.get("parentSpanId") == span_id] + self.assertEqual(len(children), 1) + + @client_context.require_version_min(9, 0) + @client_context.require_auth + def test_prose_9_no_trace_context_for_auth_monitoring(self): + """Prose Test 7: No trace context for authentication and monitoring commands.""" + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name]["test_prose_7"] + # Handshakes and authentication run lazily on the first operation, so + # their spans are part of the driver's trace ids and must not be cleared. + coll.delete_many({}) + coll.insert_one({"x": 1}) + coll.find_one({}) + finished = self.exporter.get_finished_spans() + driver_trace_ids = {f"{s.context.trace_id:032x}" for s in finished} + # Wait for the find server span to appear + find_cmd_spans = self.command_spans(finished, "find") + self.assertTrue(find_cmd_spans) + trace_id = f"{find_cmd_spans[0].context.trace_id:032x}" + self._poll_server_spans(self.OTEL_TRACE_DIR, trace_id) + # Collect all server spans that match any driver trace id + all_server_spans = self._read_server_spans(self.OTEL_TRACE_DIR) + driver_server_spans = [s for s in all_server_spans if s.get("traceId") in driver_trace_ids] + auth_monitor_names = { + "hello", + "ismaster", + "isMaster", + "saslStart", + "saslContinue", + "authenticate", + } + for s in driver_server_spans: + self.assertNotIn( + s.get("name"), + auth_monitor_names, + f"Server span for {s.get('name')} joined a driver trace", + ) + # These unit tests cover the validator's edge cases: the rejection paths and the # explicit-zero vs unset distinction for query_text_max_length. @@ -1002,6 +1411,7 @@ class _FakeConn: server_connection_id: Optional[int] = None address: _Address = ("localhost", 27017) service_id = None + max_wire_version = 30 with patch.object(_otel, "trace") as mock_trace: for _ in range(3):