From 7b6b45463c99ca0b8b86f6f207c58faff61bce69 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 07:01:47 -0500 Subject: [PATCH] PYTHON-6035 Add OpenTelemetry transaction spans Wrap a transaction's operation spans in a "transaction" pseudo-span, per the OpenTelemetry driver specification. The span is stored on the session's _Transaction and passed as the explicit parent when an operation span starts, rather than read from ambient context, so a concurrently running unrelated session cannot pick up this transaction by accident. with_transaction() pins one span across all of its retries, so a retried call still yields a single span rather than one per attempt. Its retry loop moves into a helper to keep the span bookkeeping readable. A nested with_transaction() call on the same session now raises instead of clobbering and leaking the outer call's span, and a direct-API commit retry starts a fresh span, the previous attempt having already ended its own. --- pymongo/_otel.py | 24 + pymongo/_telemetry.py | 5 +- pymongo/asynchronous/client_session.py | 79 +++- pymongo/synchronous/client_session.py | 79 +++- test/asynchronous/test_otel_transactions.py | 461 ++++++++++++++++++++ test/test_otel_transactions.py | 461 ++++++++++++++++++++ 6 files changed, 1106 insertions(+), 3 deletions(-) create mode 100644 test/asynchronous/test_otel_transactions.py create mode 100644 test/test_otel_transactions.py diff --git a/pymongo/_otel.py b/pymongo/_otel.py index 66d5a29153..64180a1d5b 100644 --- a/pymongo/_otel.py +++ b/pymongo/_otel.py @@ -523,3 +523,27 @@ def end_operation_span_failure(handle: Optional[_OperationSpanHandle], exc: Base return _CURRENT_OPERATION_NAME.reset(handle._name_token) handle._cm.__exit__(None, None, None) + + +def start_transaction_span(tracing_options: Optional[TracingOptions]) -> Optional[Span]: + """Start (but do not make current) the ``"transaction"`` pseudo-span, or None. + + Not pushed as ambient/current context; it's stored explicitly on + ``session._transaction.span`` and passed as the explicit ``parent_span`` + wherever an operation span is started under this transaction (see + :func:`start_operation_span`). Per the OTel driver spec, this span has + exactly one attribute. + """ + if not _is_tracing_enabled(tracing_options): + return None + assert _TRACER is not None + return _TRACER.start_span( + "transaction", kind=SpanKind.CLIENT, attributes={"db.system.name": "mongodb"} + ) + + +def end_transaction_span(span: Optional[Span]) -> None: + """End the transaction span, if any.""" + if span is None: + return + span.end() diff --git a/pymongo/_telemetry.py b/pymongo/_telemetry.py index f9dde5520b..f4aea9d3b6 100644 --- a/pymongo/_telemetry.py +++ b/pymongo/_telemetry.py @@ -290,10 +290,13 @@ def __init__( collection: Optional[str] = None, set_current: bool = True, ) -> None: + parent_span = None + if session is not None and session.in_transaction: + parent_span = session._transaction.span self.handle = _otel.start_operation_span( tracing_options, _otel._build_operation_name(operation, is_run_command), - None, + parent_span, dbname=dbname, collection=collection, set_current=set_current, diff --git a/pymongo/asynchronous/client_session.py b/pymongo/asynchronous/client_session.py index 72b1ac100e..5e10096dde 100644 --- a/pymongo/asynchronous/client_session.py +++ b/pymongo/asynchronous/client_session.py @@ -156,7 +156,7 @@ from bson.binary import Binary from bson.int64 import Int64 from bson.timestamp import Timestamp -from pymongo import _csot +from pymongo import _csot, _otel from pymongo.asynchronous.cursor_base import _ConnectionManager from pymongo.errors import ( ConfigurationError, @@ -427,6 +427,7 @@ def __init__(self, opts: Optional[TransactionOptions], client: AsyncMongoClient[ self.attempt = 0 self.client = client self.has_completed_command = False + self.span: Optional[Any] = None def active(self) -> bool: return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS) @@ -467,6 +468,7 @@ async def reset(self) -> None: self.recovery_token = None self.attempt = 0 self.has_completed_command = False + self.span = None def __del__(self) -> None: if self.conn_mgr: @@ -562,6 +564,10 @@ def __init__( # Is this an implicitly created session? self._implicit = implicit self._transaction = _Transaction(None, client) + # The one "transaction" span shared across every retry of a single + # with_transaction() call, or None outside of it, where + # start/commit/abort_transaction each manage their own span. + self._with_transaction_span: Optional[Any] = None # Is this session attached to a cursor? self._attached_to_cursor = False # Should we leave the session alive when the cursor is closed? @@ -769,6 +775,43 @@ async def callback(session, custom_arg, custom_kwarg=None): .. _transactions specification: https://github.com/mongodb/specifications/blob/master/source/transactions-convenient-api/transactions-convenient-api.md#handling-errors-inside-the-callback """ + if self._with_transaction_span is not None: + # Raise before any span bookkeeping, so a nested call cannot + # clobber and leak the outer call's span. + raise InvalidOperation( + "Cannot call with_transaction() while a previous with_transaction() " + "call on this session has not returned; sessions do not support " + "nested or concurrent with_transaction() calls" + ) + # One span for the whole call: start_transaction reuses it and + # commit/abort leave it open, so a retried with_transaction() yields a + # single span. Skipped when a direct-API transaction is already active, + # since start_transaction() raises below and the span would be empty. + tracing_options = self._client.options.tracing + if _otel._is_tracing_enabled(tracing_options) and not self.in_transaction: + self._with_transaction_span = _otel.start_transaction_span(tracing_options) + try: + return await self._with_transaction_retry_loop( + callback, read_concern, write_concern, read_preference, max_commit_time_ms + ) + finally: + if self._with_transaction_span is not None: + _otel.end_transaction_span(self._with_transaction_span) + # Only clear the span this call owns; a concurrent direct-API + # transaction's span belongs to that transaction. + if self._transaction.span is self._with_transaction_span: + self._transaction.span = None + self._with_transaction_span = None + + async def _with_transaction_retry_loop( + self, + callback: Callable[[AsyncClientSession], Awaitable[_T]], + read_concern: Optional[ReadConcern], + write_concern: Optional[WriteConcern], + read_preference: Optional[_ServerMode], + max_commit_time_ms: Optional[int], + ) -> _T: + """Run with_transaction's retry loop; see with_transaction.""" start_time = time.monotonic() retry = 0 last_error: Optional[BaseException] = None @@ -864,9 +907,30 @@ async def start_transaction( ) await self._transaction.reset() self._transaction.state = _TxnState.STARTING + if self._with_transaction_span is not None: + # Reuse with_transaction's shared span so a retried call still + # produces exactly one "transaction" span. + self._transaction.span = self._with_transaction_span + elif _otel._is_tracing_enabled(self._transaction.client.options.tracing): + self._transaction.span = _otel.start_transaction_span( + self._transaction.client.options.tracing + ) self._start_retryable_write() return _TransactionContext(self) + def _end_own_transaction_span(self) -> None: + """End and clear the transaction span, unless with_transaction() owns it. + + with_transaction() pins one shared span across all of its retries in + ``self._with_transaction_span`` (see its comments); while that's set, + the span must survive until with_transaction() itself ends it, so this + is a no-op here. Otherwise a retried with_transaction() would end the + shared span prematurely on the first failed attempt. + """ + if self._transaction.span is not None and self._with_transaction_span is None: + _otel.end_transaction_span(self._transaction.span) + self._transaction.span = None + async def commit_transaction(self) -> None: """Commit a multi-statement transaction. @@ -879,6 +943,7 @@ async def commit_transaction(self) -> None: elif state in (_TxnState.STARTING, _TxnState.COMMITTED_EMPTY): # Server transaction was never started, no need to send a command. self._transaction.state = _TxnState.COMMITTED_EMPTY + self._end_own_transaction_span() return elif state is _TxnState.ABORTED: raise InvalidOperation("Cannot call commitTransaction after calling abortTransaction") @@ -886,6 +951,15 @@ async def commit_transaction(self) -> None: # We're explicitly retrying the commit, move the state back to # "in progress" so that in_transaction returns true. self._transaction.state = _TxnState.IN_PROGRESS + # A direct-API retry needs a fresh span: the prior attempt's + # finally block already ended and cleared it. with_transaction + # pins its shared span instead, see _end_own_transaction_span. + if self._transaction.span is None and _otel._is_tracing_enabled( + self._transaction.client.options.tracing + ): + self._transaction.span = _otel.start_transaction_span( + self._transaction.client.options.tracing + ) try: await self._finish_transaction_with_retry("commitTransaction") @@ -909,6 +983,7 @@ async def commit_transaction(self) -> None: _reraise_with_unknown_commit(exc) finally: self._transaction.state = _TxnState.COMMITTED + self._end_own_transaction_span() async def abort_transaction(self) -> None: """Abort a multi-statement transaction. @@ -923,6 +998,7 @@ async def abort_transaction(self) -> None: elif state is _TxnState.STARTING: # Server transaction was never started, no need to send a command. self._transaction.state = _TxnState.ABORTED + self._end_own_transaction_span() return elif state is _TxnState.ABORTED: raise InvalidOperation("Cannot call abortTransaction twice") @@ -936,6 +1012,7 @@ async def abort_transaction(self) -> None: pass finally: self._transaction.state = _TxnState.ABORTED + self._end_own_transaction_span() await self._unpin() async def _finish_transaction_with_retry(self, command_name: str) -> dict[str, Any]: diff --git a/pymongo/synchronous/client_session.py b/pymongo/synchronous/client_session.py index 0774c182de..8254320e40 100644 --- a/pymongo/synchronous/client_session.py +++ b/pymongo/synchronous/client_session.py @@ -155,7 +155,7 @@ from bson.binary import Binary from bson.int64 import Int64 from bson.timestamp import Timestamp -from pymongo import _csot +from pymongo import _csot, _otel from pymongo.errors import ( ConfigurationError, ConnectionFailure, @@ -426,6 +426,7 @@ def __init__(self, opts: Optional[TransactionOptions], client: MongoClient[Any]) self.attempt = 0 self.client = client self.has_completed_command = False + self.span: Optional[Any] = None def active(self) -> bool: return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS) @@ -466,6 +467,7 @@ def reset(self) -> None: self.recovery_token = None self.attempt = 0 self.has_completed_command = False + self.span = None def __del__(self) -> None: if self.conn_mgr: @@ -561,6 +563,10 @@ def __init__( # Is this an implicitly created session? self._implicit = implicit self._transaction = _Transaction(None, client) + # The one "transaction" span shared across every retry of a single + # with_transaction() call, or None outside of it, where + # start/commit/abort_transaction each manage their own span. + self._with_transaction_span: Optional[Any] = None # Is this session attached to a cursor? self._attached_to_cursor = False # Should we leave the session alive when the cursor is closed? @@ -768,6 +774,43 @@ def callback(session, custom_arg, custom_kwarg=None): .. _transactions specification: https://github.com/mongodb/specifications/blob/master/source/transactions-convenient-api/transactions-convenient-api.md#handling-errors-inside-the-callback """ + if self._with_transaction_span is not None: + # Raise before any span bookkeeping, so a nested call cannot + # clobber and leak the outer call's span. + raise InvalidOperation( + "Cannot call with_transaction() while a previous with_transaction() " + "call on this session has not returned; sessions do not support " + "nested or concurrent with_transaction() calls" + ) + # One span for the whole call: start_transaction reuses it and + # commit/abort leave it open, so a retried with_transaction() yields a + # single span. Skipped when a direct-API transaction is already active, + # since start_transaction() raises below and the span would be empty. + tracing_options = self._client.options.tracing + if _otel._is_tracing_enabled(tracing_options) and not self.in_transaction: + self._with_transaction_span = _otel.start_transaction_span(tracing_options) + try: + return self._with_transaction_retry_loop( + callback, read_concern, write_concern, read_preference, max_commit_time_ms + ) + finally: + if self._with_transaction_span is not None: + _otel.end_transaction_span(self._with_transaction_span) + # Only clear the span this call owns; a concurrent direct-API + # transaction's span belongs to that transaction. + if self._transaction.span is self._with_transaction_span: + self._transaction.span = None + self._with_transaction_span = None + + def _with_transaction_retry_loop( + self, + callback: Callable[[ClientSession], _T], + read_concern: Optional[ReadConcern], + write_concern: Optional[WriteConcern], + read_preference: Optional[_ServerMode], + max_commit_time_ms: Optional[int], + ) -> _T: + """Run with_transaction's retry loop; see with_transaction.""" start_time = time.monotonic() retry = 0 last_error: Optional[BaseException] = None @@ -861,9 +904,30 @@ def start_transaction( ) self._transaction.reset() self._transaction.state = _TxnState.STARTING + if self._with_transaction_span is not None: + # Reuse with_transaction's shared span so a retried call still + # produces exactly one "transaction" span. + self._transaction.span = self._with_transaction_span + elif _otel._is_tracing_enabled(self._transaction.client.options.tracing): + self._transaction.span = _otel.start_transaction_span( + self._transaction.client.options.tracing + ) self._start_retryable_write() return _TransactionContext(self) + def _end_own_transaction_span(self) -> None: + """End and clear the transaction span, unless with_transaction() owns it. + + with_transaction() pins one shared span across all of its retries in + ``self._with_transaction_span`` (see its comments); while that's set, + the span must survive until with_transaction() itself ends it, so this + is a no-op here. Otherwise a retried with_transaction() would end the + shared span prematurely on the first failed attempt. + """ + if self._transaction.span is not None and self._with_transaction_span is None: + _otel.end_transaction_span(self._transaction.span) + self._transaction.span = None + def commit_transaction(self) -> None: """Commit a multi-statement transaction. @@ -876,6 +940,7 @@ def commit_transaction(self) -> None: elif state in (_TxnState.STARTING, _TxnState.COMMITTED_EMPTY): # Server transaction was never started, no need to send a command. self._transaction.state = _TxnState.COMMITTED_EMPTY + self._end_own_transaction_span() return elif state is _TxnState.ABORTED: raise InvalidOperation("Cannot call commitTransaction after calling abortTransaction") @@ -883,6 +948,15 @@ def commit_transaction(self) -> None: # We're explicitly retrying the commit, move the state back to # "in progress" so that in_transaction returns true. self._transaction.state = _TxnState.IN_PROGRESS + # A direct-API retry needs a fresh span: the prior attempt's + # finally block already ended and cleared it. with_transaction + # pins its shared span instead, see _end_own_transaction_span. + if self._transaction.span is None and _otel._is_tracing_enabled( + self._transaction.client.options.tracing + ): + self._transaction.span = _otel.start_transaction_span( + self._transaction.client.options.tracing + ) try: self._finish_transaction_with_retry("commitTransaction") @@ -906,6 +980,7 @@ def commit_transaction(self) -> None: _reraise_with_unknown_commit(exc) finally: self._transaction.state = _TxnState.COMMITTED + self._end_own_transaction_span() def abort_transaction(self) -> None: """Abort a multi-statement transaction. @@ -920,6 +995,7 @@ def abort_transaction(self) -> None: elif state is _TxnState.STARTING: # Server transaction was never started, no need to send a command. self._transaction.state = _TxnState.ABORTED + self._end_own_transaction_span() return elif state is _TxnState.ABORTED: raise InvalidOperation("Cannot call abortTransaction twice") @@ -933,6 +1009,7 @@ def abort_transaction(self) -> None: pass finally: self._transaction.state = _TxnState.ABORTED + self._end_own_transaction_span() self._unpin() def _finish_transaction_with_retry(self, command_name: str) -> dict[str, Any]: diff --git a/test/asynchronous/test_otel_transactions.py b/test/asynchronous/test_otel_transactions.py new file mode 100644 index 0000000000..ca7a15a945 --- /dev/null +++ b/test/asynchronous/test_otel_transactions.py @@ -0,0 +1,461 @@ +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test the OpenTelemetry transaction pseudo-span.""" + +from __future__ import annotations + +import os +import sys +from typing import Optional +from unittest.mock import patch + +sys.path[0:0] = [""] + +import pytest + +import pymongo._otel as _otel +from pymongo import _telemetry, common +from pymongo._telemetry import _OperationTelemetry +from pymongo.errors import ( + ClientBulkWriteException, + ConfigurationError, + InvalidOperation, + OperationFailure, + ServerSelectionTimeoutError, +) +from pymongo.logger import _HELLO_COMMANDS +from pymongo.operations import InsertOne +from pymongo.read_preferences import ReadPreference +from pymongo.typings import _Address +from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest +from test.asynchronous.utils import async_wait_until +from test.unified_format_shared import _shared_test_provider + +_HAS_OTEL_TEST_DEPS = False +if _otel._HAS_OPENTELEMETRY: + try: + from opentelemetry import trace + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + _HAS_OTEL_TEST_DEPS = True + except ImportError: + pass + +_IS_SYNC = False + +pytestmark = pytest.mark.otel + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelTransactionSpanPrimitives(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_start_transaction_span_has_only_one_attribute(self): + opts: _otel.TracingOptions = {"enabled": True, "query_text_max_length": None} + span = _otel.start_transaction_span(opts) + _otel.end_transaction_span(span) + (finished,) = self.exporter.get_finished_spans() + self.assertEqual(finished.name, "transaction") + self.assertEqual(dict(finished.attributes), {"db.system.name": "mongodb"}) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetryInTransaction(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_nests_under_active_transaction_span(self): + opts: _otel.TracingOptions = {"enabled": True, "query_text_max_length": None} + txn_span = _otel.start_transaction_span(opts) + + class _FakeTransaction: + span = txn_span + + class _FakeSession: + in_transaction = True + _transaction = _FakeTransaction() + + telemetry = _telemetry._OperationTelemetry(opts, "insert", _FakeSession()) + telemetry.succeeded() + _otel.end_transaction_span(txn_span) + child, parent = self.exporter.get_finished_spans() + self.assertEqual(child.parent.span_id, parent.context.span_id) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelTransactionSpans(AsyncIntegrationTest): + """Transaction spans and the operations nested under them.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # See the matching comment in test/asynchronous/unified_format.py's + # UnifiedSpecTestMixinV1.tearDownClass: the span processor can never + # be removed from the shared process-wide TracerProvider, so without + # this shutdown() the exporter keeps accumulating every span from + # every client for the rest of the test run. + 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 operation_spans(finished, operation: str): + """Return the operation spans for ``operation``, excluding command spans. + + Only command spans carry db.command.name, so its absence is what tells + the two kinds apart when both name the same operation. + """ + return [ + s + for s in finished + if s.attributes.get("db.operation.name") == operation + and "db.command.name" not in s.attributes + ] + + @staticmethod + def command_spans(finished, command: str): + """Return the command spans for ``command``.""" + return [s for s in finished if s.attributes.get("db.command.name") == command] + + def ping_spans(self): + """Return the spans belonging to a ``ping`` run through ``db.command()``. + + For the tests that assert tracing produced *nothing*. Asserting the + exporter is empty would also catch spans no test asked for: a cursor + abandoned earlier in the class ends its operation span from a + finalizer, and on an interpreter that does not reference count, that + finalizer runs at an unpredictable point and lands in whichever test + happens to be running. Naming the ping's own spans keeps the assertion + about this client while staying immune to that. + """ + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "ping" + or s.attributes.get("db.operation.name") == "runCommand" + ] + + def _aggregate_operation_span(self): + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "aggregate" + ] + self.assertEqual(len(matching), 1) + return matching[0] + + @async_client_context.require_transactions + async def test_committing_empty_transaction_ends_span(self): + # No operation is ever run against the server, so commit_transaction + # takes the STARTING/COMMITTED_EMPTY early-return path rather than + # actually sending a commitTransaction command. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + + async with client.start_session() as session: + await session.start_transaction() + await session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_span = next(s for s in finished if s.name == "transaction") + self.assertTrue(txn_span.end_time is not None) + + @async_client_context.require_transactions + async def test_aborting_empty_transaction_ends_span(self): + # No operation is ever run against the server, so abort_transaction + # takes the STARTING early-return path rather than actually sending + # an abortTransaction command. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + + async with client.start_session() as session: + await session.start_transaction() + await session.abort_transaction() + + finished = self.exporter.get_finished_spans() + txn_span = next(s for s in finished if s.name == "transaction") + self.assertTrue(txn_span.end_time is not None) + + @async_client_context.require_transactions + async def test_direct_commit_retry_gives_each_span_its_own_end(self): + # Explicitly retrying a successful commit moves the transaction state + # COMMITTED -> IN_PROGRESS -> (back through the try/finally) -> + # COMMITTED again. The prior attempt's span was already ended and + # cleared, so the retry gets a fresh "transaction" span of its own + # (this is the direct-API path, not with_transaction; see + # test_with_transaction_retry_reuses_one_transaction_span for the + # with_transaction case, which shares a single span across retries + # instead); each span's ending finally block must run exactly once + # for its own span, never double-ending the same span and never + # leaving one unended. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name].test + await coll.drop() + await client[self.db.name].create_collection("test") + self.exporter.clear() + + async with client.start_session() as session: + async with await session.start_transaction(): + await coll.insert_one({"x": 5}, session=session) + # The transaction context manager already committed on clean + # exit; retry the commit explicitly. + await session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 2) + self.assertNotEqual(txn_spans[0].context.span_id, txn_spans[1].context.span_id) + for txn_span in txn_spans: + self.assertTrue(txn_span.end_time is not None) + + @async_client_context.require_transactions + async def test_with_transaction_retry_reuses_one_transaction_span(self): + # A retried with_transaction() call must still produce exactly one + # "transaction" span for the whole logical call, not one sibling + # span per full-transaction retry, and no separately-named wrapper + # span either (the vendored transaction/convenient.json fixture + # pins "transaction" itself as the trace root for withTransaction). + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.with_txn_spans + await coll.drop() + await client.pymongo_test.create_collection("with_txn_spans") + + attempts = [] + + async def callback(session): + attempts.append(1) + await coll.insert_one({"n": len(attempts)}, session=session) + if len(attempts) == 1: + exc = OperationFailure("transient", 251) + exc._add_error_label("TransientTransactionError") + raise exc + + self.exporter.clear() + async with client.start_session() as session: + await session.with_transaction(callback) + + self.assertEqual(len(attempts), 2) + finished = self.exporter.get_finished_spans() + self.assertFalse( + [s.name for s in finished if s.name.startswith("withTransaction")], + [s.name for s in finished], + ) + + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + self.assertTrue(txn_spans[0].end_time is not None) + + insert_op_spans = [s for s in finished if s.attributes.get("db.operation.name") == "insert"] + self.assertEqual(len(insert_op_spans), 2) + for op_span in insert_op_spans: + self.assertEqual(op_span.parent.span_id, txn_spans[0].context.span_id) + + @async_client_context.require_transactions + async def test_reentrant_with_transaction_raises_and_does_not_leak_span(self): + # A callback that illegally re-enters with_transaction() on the same + # session must be rejected with a clear InvalidOperation, and the + # outer call's "transaction" span must still end exactly once, + # never leaked (created but never ended) and never double-ended. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.reentrant_with_txn + await coll.drop() + await client.pymongo_test.create_collection("reentrant_with_txn") + + async def inner_callback(session): + await coll.insert_one({"x": 1}, session=session) + + async def outer_callback(session): + await coll.insert_one({"x": 2}, session=session) + # Illegal: with_transaction() is not reentrant on one session. + await session.with_transaction(inner_callback) + + self.exporter.clear() + async with client.start_session() as session: + with self.assertRaises(InvalidOperation): + await session.with_transaction(outer_callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + # Only the outer call ever gets far enough to create a span; the + # guard rejects the inner call before it creates one of its own. + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + for txn_span in txn_spans: + self.assertIsNotNone(txn_span.end_time) + + @async_client_context.require_transactions + async def test_nested_with_transaction_on_another_session_keeps_spans_separate(self): + # Nesting with_transaction() is legal on a *different* session, unlike + # the same-session case above. Each session's operations must parent to + # its own transaction span, which holds because an operation span takes + # its parent explicitly from session._transaction.span instead of from + # ambient context. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + outer_coll = db.two_session_outer + inner_coll = db.two_session_inner + # Create both up front: creating a collection inside a transaction is + # illegal before server 4.4. + await outer_coll.drop() + await inner_coll.drop() + await db.create_collection("two_session_outer") + await db.create_collection("two_session_inner") + + async def inner_callback(inner_session): + await inner_coll.insert_one({"x": 1}, session=inner_session) + + async def outer_callback(outer_session): + await outer_coll.insert_one({"x": 1}, session=outer_session) + async with client.start_session() as inner_session: + await inner_session.with_transaction(inner_callback) + + self.exporter.clear() + async with client.start_session() as outer_session: + await outer_session.with_transaction(outer_callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 2, [s.name for s in finished]) + for txn_span in txn_spans: + self.assertIsNotNone(txn_span.end_time) + # Transaction spans are never made current, so neither ends up + # nested under the other. + self.assertIsNone(txn_span.parent) + + def insert_parent_id(collname: str) -> int: + (span,) = [ + s + for s in finished + if s.attributes.get("db.operation.name") == "insert" + and s.attributes.get("db.collection.name") == collname + ] + return span.parent.span_id + + outer_parent = insert_parent_id("two_session_outer") + inner_parent = insert_parent_id("two_session_inner") + self.assertNotEqual(outer_parent, inner_parent) + self.assertEqual({outer_parent, inner_parent}, {s.context.span_id for s in txn_spans}) + + @async_client_context.require_transactions + async def test_with_transaction_while_direct_api_transaction_active_does_not_corrupt_span( + self, + ): + # Calling with_transaction() while a transaction started with the + # DIRECT API is already active on the same session is illegal: + # start_transaction() inside with_transaction() raises "Transaction + # already in progress", but the direct-API transaction's own + # "transaction" span must survive that failure: with_transaction()'s + # finally must not end/null it out from under the still-active + # transaction (Important #1). Operations run on the session + # afterwards must still parent to that span rather than becoming + # trace roots, and the failed call must not leave behind a second, + # spurious "transaction" span of its own. + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.direct_api_with_txn_conflict + await coll.drop() + await client.pymongo_test.create_collection("direct_api_with_txn_conflict") + + async def callback(session): + raise AssertionError("never reached; start_transaction() raises first") + + self.exporter.clear() + async with client.start_session() as session: + await session.start_transaction() + await coll.insert_one({"x": 1}, session=session) + + with self.assertRaises(InvalidOperation): + await session.with_transaction(callback) + + # The original transaction is still active; this must still + # nest under its span, not become a trace root. + await coll.insert_one({"x": 2}, session=session) + await session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + txn_span = txn_spans[0] + self.assertIsNotNone(txn_span.end_time) + + insert_op_spans = [s for s in finished if s.attributes.get("db.operation.name") == "insert"] + self.assertEqual(len(insert_op_spans), 2) + for op_span in insert_op_spans: + self.assertEqual(op_span.parent.span_id, txn_span.context.span_id) + + @async_client_context.require_transactions + async def test_retried_commit_has_a_transaction_span(self): + client = await self.async_rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.retried_commit_spans + await coll.drop() + await client.pymongo_test.create_collection("retried_commit_spans") + + async with client.start_session() as session: + await session.start_transaction() + await coll.insert_one({"x": 1}, session=session) + await session.commit_transaction() + self.exporter.clear() + # An explicit second commit re-enters the COMMITTED -> IN_PROGRESS + # branch, which previously ran with no transaction span at all. + await session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + commit_cmd_spans = [ + s for s in finished if s.attributes.get("db.command.name") == "commitTransaction" + ] + self.assertGreaterEqual(len(commit_cmd_spans), 1) + for cmd_span in commit_cmd_spans: + self.assertIsNotNone(cmd_span.parent) diff --git a/test/test_otel_transactions.py b/test/test_otel_transactions.py new file mode 100644 index 0000000000..6321542c6f --- /dev/null +++ b/test/test_otel_transactions.py @@ -0,0 +1,461 @@ +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test the OpenTelemetry transaction pseudo-span.""" + +from __future__ import annotations + +import os +import sys +from typing import Optional +from unittest.mock import patch + +sys.path[0:0] = [""] + +import pytest + +import pymongo._otel as _otel +from pymongo import _telemetry, common +from pymongo._telemetry import _OperationTelemetry +from pymongo.errors import ( + ClientBulkWriteException, + ConfigurationError, + InvalidOperation, + OperationFailure, + ServerSelectionTimeoutError, +) +from pymongo.logger import _HELLO_COMMANDS +from pymongo.operations import InsertOne +from pymongo.read_preferences import ReadPreference +from pymongo.typings import _Address +from test import IntegrationTest, client_context, unittest +from test.unified_format_shared import _shared_test_provider +from test.utils import wait_until + +_HAS_OTEL_TEST_DEPS = False +if _otel._HAS_OPENTELEMETRY: + try: + from opentelemetry import trace + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + _HAS_OTEL_TEST_DEPS = True + except ImportError: + pass + +_IS_SYNC = True + +pytestmark = pytest.mark.otel + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelTransactionSpanPrimitives(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_start_transaction_span_has_only_one_attribute(self): + opts: _otel.TracingOptions = {"enabled": True, "query_text_max_length": None} + span = _otel.start_transaction_span(opts) + _otel.end_transaction_span(span) + (finished,) = self.exporter.get_finished_spans() + self.assertEqual(finished.name, "transaction") + self.assertEqual(dict(finished.attributes), {"db.system.name": "mongodb"}) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOperationTelemetryInTransaction(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # The span processor can never be removed from the shared process-wide + # TracerProvider, so without this the exporter keeps accumulating every + # span from every client for the rest of the test run. + cls.exporter.shutdown() + + def setUp(self): + self.exporter.clear() + + def test_nests_under_active_transaction_span(self): + opts: _otel.TracingOptions = {"enabled": True, "query_text_max_length": None} + txn_span = _otel.start_transaction_span(opts) + + class _FakeTransaction: + span = txn_span + + class _FakeSession: + in_transaction = True + _transaction = _FakeTransaction() + + telemetry = _telemetry._OperationTelemetry(opts, "insert", _FakeSession()) + telemetry.succeeded() + _otel.end_transaction_span(txn_span) + child, parent = self.exporter.get_finished_spans() + self.assertEqual(child.parent.span_id, parent.context.span_id) + + +@unittest.skipUnless(_HAS_OTEL_TEST_DEPS, "opentelemetry-sdk is not installed") +class TestOTelTransactionSpans(IntegrationTest): + """Transaction spans and the operations nested under them.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.exporter = InMemorySpanExporter() + _shared_test_provider().add_span_processor(SimpleSpanProcessor(cls.exporter)) + + @classmethod + def tearDownClass(cls): + # See the matching comment in test/synchronous/unified_format.py's + # UnifiedSpecTestMixinV1.tearDownClass: the span processor can never + # be removed from the shared process-wide TracerProvider, so without + # this shutdown() the exporter keeps accumulating every span from + # every client for the rest of the test run. + cls.exporter.shutdown() + super().tearDownClass() + + def setUp(self): + super().setUp() + 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 operation_spans(finished, operation: str): + """Return the operation spans for ``operation``, excluding command spans. + + Only command spans carry db.command.name, so its absence is what tells + the two kinds apart when both name the same operation. + """ + return [ + s + for s in finished + if s.attributes.get("db.operation.name") == operation + and "db.command.name" not in s.attributes + ] + + @staticmethod + def command_spans(finished, command: str): + """Return the command spans for ``command``.""" + return [s for s in finished if s.attributes.get("db.command.name") == command] + + def ping_spans(self): + """Return the spans belonging to a ``ping`` run through ``db.command()``. + + For the tests that assert tracing produced *nothing*. Asserting the + exporter is empty would also catch spans no test asked for: a cursor + abandoned earlier in the class ends its operation span from a + finalizer, and on an interpreter that does not reference count, that + finalizer runs at an unpredictable point and lands in whichever test + happens to be running. Naming the ping's own spans keeps the assertion + about this client while staying immune to that. + """ + return [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.command.name") == "ping" + or s.attributes.get("db.operation.name") == "runCommand" + ] + + def _aggregate_operation_span(self): + matching = [ + s + for s in self.exporter.get_finished_spans() + if s.attributes.get("db.operation.name") == "aggregate" + ] + self.assertEqual(len(matching), 1) + return matching[0] + + @client_context.require_transactions + def test_committing_empty_transaction_ends_span(self): + # No operation is ever run against the server, so commit_transaction + # takes the STARTING/COMMITTED_EMPTY early-return path rather than + # actually sending a commitTransaction command. + client = self.rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + + with client.start_session() as session: + session.start_transaction() + session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_span = next(s for s in finished if s.name == "transaction") + self.assertTrue(txn_span.end_time is not None) + + @client_context.require_transactions + def test_aborting_empty_transaction_ends_span(self): + # No operation is ever run against the server, so abort_transaction + # takes the STARTING early-return path rather than actually sending + # an abortTransaction command. + client = self.rs_or_single_client(tracing={"enabled": True}) + self.exporter.clear() + + with client.start_session() as session: + session.start_transaction() + session.abort_transaction() + + finished = self.exporter.get_finished_spans() + txn_span = next(s for s in finished if s.name == "transaction") + self.assertTrue(txn_span.end_time is not None) + + @client_context.require_transactions + def test_direct_commit_retry_gives_each_span_its_own_end(self): + # Explicitly retrying a successful commit moves the transaction state + # COMMITTED -> IN_PROGRESS -> (back through the try/finally) -> + # COMMITTED again. The prior attempt's span was already ended and + # cleared, so the retry gets a fresh "transaction" span of its own + # (this is the direct-API path, not with_transaction; see + # test_with_transaction_retry_reuses_one_transaction_span for the + # with_transaction case, which shares a single span across retries + # instead); each span's ending finally block must run exactly once + # for its own span, never double-ending the same span and never + # leaving one unended. + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client[self.db.name].test + coll.drop() + client[self.db.name].create_collection("test") + self.exporter.clear() + + with client.start_session() as session: + with session.start_transaction(): + coll.insert_one({"x": 5}, session=session) + # The transaction context manager already committed on clean + # exit; retry the commit explicitly. + session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 2) + self.assertNotEqual(txn_spans[0].context.span_id, txn_spans[1].context.span_id) + for txn_span in txn_spans: + self.assertTrue(txn_span.end_time is not None) + + @client_context.require_transactions + def test_with_transaction_retry_reuses_one_transaction_span(self): + # A retried with_transaction() call must still produce exactly one + # "transaction" span for the whole logical call, not one sibling + # span per full-transaction retry, and no separately-named wrapper + # span either (the vendored transaction/convenient.json fixture + # pins "transaction" itself as the trace root for withTransaction). + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.with_txn_spans + coll.drop() + client.pymongo_test.create_collection("with_txn_spans") + + attempts = [] + + def callback(session): + attempts.append(1) + coll.insert_one({"n": len(attempts)}, session=session) + if len(attempts) == 1: + exc = OperationFailure("transient", 251) + exc._add_error_label("TransientTransactionError") + raise exc + + self.exporter.clear() + with client.start_session() as session: + session.with_transaction(callback) + + self.assertEqual(len(attempts), 2) + finished = self.exporter.get_finished_spans() + self.assertFalse( + [s.name for s in finished if s.name.startswith("withTransaction")], + [s.name for s in finished], + ) + + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + self.assertTrue(txn_spans[0].end_time is not None) + + insert_op_spans = [s for s in finished if s.attributes.get("db.operation.name") == "insert"] + self.assertEqual(len(insert_op_spans), 2) + for op_span in insert_op_spans: + self.assertEqual(op_span.parent.span_id, txn_spans[0].context.span_id) + + @client_context.require_transactions + def test_reentrant_with_transaction_raises_and_does_not_leak_span(self): + # A callback that illegally re-enters with_transaction() on the same + # session must be rejected with a clear InvalidOperation, and the + # outer call's "transaction" span must still end exactly once, + # never leaked (created but never ended) and never double-ended. + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.reentrant_with_txn + coll.drop() + client.pymongo_test.create_collection("reentrant_with_txn") + + def inner_callback(session): + coll.insert_one({"x": 1}, session=session) + + def outer_callback(session): + coll.insert_one({"x": 2}, session=session) + # Illegal: with_transaction() is not reentrant on one session. + session.with_transaction(inner_callback) + + self.exporter.clear() + with client.start_session() as session: + with self.assertRaises(InvalidOperation): + session.with_transaction(outer_callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + # Only the outer call ever gets far enough to create a span; the + # guard rejects the inner call before it creates one of its own. + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + for txn_span in txn_spans: + self.assertIsNotNone(txn_span.end_time) + + @client_context.require_transactions + def test_nested_with_transaction_on_another_session_keeps_spans_separate(self): + # Nesting with_transaction() is legal on a *different* session, unlike + # the same-session case above. Each session's operations must parent to + # its own transaction span, which holds because an operation span takes + # its parent explicitly from session._transaction.span instead of from + # ambient context. + client = self.rs_or_single_client(tracing={"enabled": True}) + db = client.pymongo_test + outer_coll = db.two_session_outer + inner_coll = db.two_session_inner + # Create both up front: creating a collection inside a transaction is + # illegal before server 4.4. + outer_coll.drop() + inner_coll.drop() + db.create_collection("two_session_outer") + db.create_collection("two_session_inner") + + def inner_callback(inner_session): + inner_coll.insert_one({"x": 1}, session=inner_session) + + def outer_callback(outer_session): + outer_coll.insert_one({"x": 1}, session=outer_session) + with client.start_session() as inner_session: + inner_session.with_transaction(inner_callback) + + self.exporter.clear() + with client.start_session() as outer_session: + outer_session.with_transaction(outer_callback) + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 2, [s.name for s in finished]) + for txn_span in txn_spans: + self.assertIsNotNone(txn_span.end_time) + # Transaction spans are never made current, so neither ends up + # nested under the other. + self.assertIsNone(txn_span.parent) + + def insert_parent_id(collname: str) -> int: + (span,) = [ + s + for s in finished + if s.attributes.get("db.operation.name") == "insert" + and s.attributes.get("db.collection.name") == collname + ] + return span.parent.span_id + + outer_parent = insert_parent_id("two_session_outer") + inner_parent = insert_parent_id("two_session_inner") + self.assertNotEqual(outer_parent, inner_parent) + self.assertEqual({outer_parent, inner_parent}, {s.context.span_id for s in txn_spans}) + + @client_context.require_transactions + def test_with_transaction_while_direct_api_transaction_active_does_not_corrupt_span( + self, + ): + # Calling with_transaction() while a transaction started with the + # DIRECT API is already active on the same session is illegal: + # start_transaction() inside with_transaction() raises "Transaction + # already in progress", but the direct-API transaction's own + # "transaction" span must survive that failure: with_transaction()'s + # finally must not end/null it out from under the still-active + # transaction (Important #1). Operations run on the session + # afterwards must still parent to that span rather than becoming + # trace roots, and the failed call must not leave behind a second, + # spurious "transaction" span of its own. + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.direct_api_with_txn_conflict + coll.drop() + client.pymongo_test.create_collection("direct_api_with_txn_conflict") + + def callback(session): + raise AssertionError("never reached; start_transaction() raises first") + + self.exporter.clear() + with client.start_session() as session: + session.start_transaction() + coll.insert_one({"x": 1}, session=session) + + with self.assertRaises(InvalidOperation): + session.with_transaction(callback) + + # The original transaction is still active; this must still + # nest under its span, not become a trace root. + coll.insert_one({"x": 2}, session=session) + session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + txn_span = txn_spans[0] + self.assertIsNotNone(txn_span.end_time) + + insert_op_spans = [s for s in finished if s.attributes.get("db.operation.name") == "insert"] + self.assertEqual(len(insert_op_spans), 2) + for op_span in insert_op_spans: + self.assertEqual(op_span.parent.span_id, txn_span.context.span_id) + + @client_context.require_transactions + def test_retried_commit_has_a_transaction_span(self): + client = self.rs_or_single_client(tracing={"enabled": True}) + coll = client.pymongo_test.retried_commit_spans + coll.drop() + client.pymongo_test.create_collection("retried_commit_spans") + + with client.start_session() as session: + session.start_transaction() + coll.insert_one({"x": 1}, session=session) + session.commit_transaction() + self.exporter.clear() + # An explicit second commit re-enters the COMMITTED -> IN_PROGRESS + # branch, which previously ran with no transaction span at all. + session.commit_transaction() + + finished = self.exporter.get_finished_spans() + txn_spans = [s for s in finished if s.name == "transaction"] + self.assertEqual(len(txn_spans), 1, [s.name for s in finished]) + commit_cmd_spans = [ + s for s in finished if s.attributes.get("db.command.name") == "commitTransaction" + ] + self.assertGreaterEqual(len(commit_cmd_spans), 1) + for cmd_span in commit_cmd_spans: + self.assertIsNotNone(cmd_span.parent)