PYTHON-5855 Propagate the W3C traceparent to the server

.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

.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

.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]
 
 

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

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