From 0fbef917a327b38fc918d53b54e0e3cad50b4cdc Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Fri, 23 Oct 2026 15:51:27 -0400 Subject: [PATCH 1/3] Add experimental MySQL trace context execution helper --- docs/mysql.md | 66 ++++++++++++++++ src/splunk_otel/mysql.py | 52 +++++++++++++ tests/test_mysql.py | 161 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 docs/mysql.md create mode 100644 src/splunk_otel/mysql.py create mode 100644 tests/test_mysql.py diff --git a/docs/mysql.md b/docs/mysql.md new file mode 100644 index 00000000..20d05ee3 --- /dev/null +++ b/docs/mysql.md @@ -0,0 +1,66 @@ +# MySQL query correlation: first implementation step + +The goal is to link a query in Database Monitoring (DBMon) to the application +trace that caused it. A trace describes an application request. A span describes +one operation within that request, such as a database call. + +## How the link works + +Before running a query, we store the trace and database span IDs in a MySQL +session variable named `@traceparent`. A session variable belongs to one database +connection and keeps its value until it is changed or the session is cleared. + +`traceparent` is an ordinary variable name to MySQL. Our code and the Collector +agree to use that name. The value follows the standard W3C `traceparent` format: + +```sql +SET @traceparent = '00---'; +``` + +The Collector reads this variable when it samples queries. DBMon can then use +the IDs to link a sampled query to its application trace. We still need to test +that full path with this implementation. + +## What is implemented + +`splunk_otel.mysql.execute_with_trace_context` does three things: + +1. Writes the supplied database span's context to `@traceparent`. +2. Calls the supplied query method once. +3. Attempts to clear `@traceparent` with `SET @traceparent = NULL`, even if the + query fails. + +Clearing the variable prevents a later operation from inheriting the completed +query's IDs. The helper uses a separate cursor for its SET commands and disables +tracing around those commands so they do not create extra spans. + +The helper uses only the span context supplied by its caller. Valid IDs are +written even when the trace is not selected for export. If the supplied context +is invalid, the helper clears the variable before the query. It does not use the +surrounding application's span instead. + +## What the caller must provide + +The helper requires the actual driver connection, without an instrumentation +wrapper, and the context of the span for this database call. It does not create +that span itself. + +The connection must have no pending results from earlier queries. The query must +return at most one result set and use a buffered cursor: the driver downloads all +rows before the query method returns. MySQL may reject the cleanup SET if rows +are still waiting to be downloaded. + +The helper does not check these restrictions yet. Driver integration must enforce +them before the helper can be used automatically. + +## Failures and remaining work + +If an internal SET command fails, the helper logs a warning. It still returns the +query's result or raises the query's original exception. A failed SET can leave +old IDs on the connection, so later queries could receive an incorrect trace +link. + +This is an experimental building block. It is not enabled at startup and has no +environment setting. The remaining work includes driver integration, transaction +and connection-pool handling, streaming results, multiple result sets, custom +cursors, asynchronous drivers, and Collector/DBMon testing. diff --git a/src/splunk_otel/mysql.py b/src/splunk_otel/mysql.py new file mode 100644 index 00000000..195afb49 --- /dev/null +++ b/src/splunk_otel/mysql.py @@ -0,0 +1,52 @@ +# Copyright Splunk 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. + +import logging + +from opentelemetry.instrumentation.utils import suppress_instrumentation +from opentelemetry.trace import NonRecordingSpan, SpanContext, set_span_in_context +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +_logger = logging.getLogger(__name__) + + +def execute_with_trace_context(connection, query_method, span_context: SpanContext, *args, **kwargs): + """Experimental helper for one buffered, single-result MySQL execution. + + The caller supplies the unwrapped connection and the database CLIENT span's + context. All earlier results must be consumed. Driver hooks must enforce + these constraints before using this helper with customer queries. + """ + carrier = {} + # An explicit context prevents an invalid database span from injecting a parent. + context = set_span_in_context(NonRecordingSpan(span_context)) + TraceContextTextMapPropagator().inject(carrier, context=context) + try: + _set_traceparent(connection, carrier.get("traceparent"), phase="setup") + return query_method(*args, **kwargs) + finally: + # Attempt cleanup even if setup failed after the server accepted its SET. + _set_traceparent(connection, None, phase="cleanup") + + +def _set_traceparent(connection, value, *, phase): + try: + with suppress_instrumentation(): + cursor = connection.cursor() + try: + cursor.execute("SET @traceparent = %s", (value,)) + finally: + cursor.close() + except Exception as error: # noqa: BLE001 - Internal commands must not replace the query result or error. + _logger.warning("MySQL trace context %s failed (%s)", phase, type(error).__name__) diff --git a/tests/test_mysql.py b/tests/test_mysql.py new file mode 100644 index 00000000..58975705 --- /dev/null +++ b/tests/test_mysql.py @@ -0,0 +1,161 @@ +# Copyright Splunk 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. + +import pytest + +from opentelemetry.instrumentation.utils import is_instrumentation_enabled +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import INVALID_SPAN_CONTEXT, NonRecordingSpan, SpanContext, SpanKind, TraceFlags, use_span + +from splunk_otel.mysql import execute_with_trace_context + + +class _FakeConnection: + def __init__(self, *, failure_phase=None, failure_operation=None): + self.traceparent = "previous-context" + self.failure_phase = failure_phase + self.failure_operation = failure_operation + self.cursors = [] + + def cursor(self): + phase = "setup" if not self.cursors else "cleanup" + cursor = _FakeInternalCursor(self, phase) + self.cursors.append(cursor) + self.fail_if_requested(phase, "cursor") + return cursor + + def fail_if_requested(self, phase, operation): + if (phase, operation) == (self.failure_phase, self.failure_operation): + raise RuntimeError("internal failure") + + +class _FakeInternalCursor: + def __init__(self, connection, phase): + self.connection = connection + self.phase = phase + self.closed = False + + def execute(self, statement, parameters): + assert not is_instrumentation_enabled() + assert statement == "SET @traceparent = %s" + self.connection.traceparent = parameters[0] + self.connection.fail_if_requested(self.phase, "execute") + + def close(self): + assert not is_instrumentation_enabled() + self.closed = True + self.connection.fail_if_requested(self.phase, "close") + + +class _FakeQuery: + def __init__(self, connection, *, error=None): + self.connection = connection + self.error = error + self.result = object() + self.calls = [] + + def execute(self, *args, **kwargs): + self.calls.append((args, kwargs, self.connection.traceparent)) + assert is_instrumentation_enabled() + if self.error is not None: + raise self.error + return self.result + + +@pytest.mark.parametrize("flags", [TraceFlags.DEFAULT, TraceFlags.SAMPLED]) +def test_execution_uses_supplied_context_and_clears_it(flags): + connection = _FakeConnection() + query = _FakeQuery(connection) + database_context = SpanContext(0x123, 0x456, False, TraceFlags(flags)) + parent_context = SpanContext(0xABC, 0xDEF, False) + + with use_span(NonRecordingSpan(parent_context)): + result = execute_with_trace_context( + connection, query.execute, database_context, "SELECT %s", (17,), driver_option=True + ) + + assert result is query.result + assert query.calls == [(("SELECT %s", (17,)), {"driver_option": True}, f"00-{0x123:032x}-{0x456:016x}-{flags:02x}")] + assert connection.traceparent is None + assert len(connection.cursors) == 2 + assert all(cursor.closed for cursor in connection.cursors) + + +def test_invalid_database_context_clears_previous_value_without_injecting_parent(): + connection = _FakeConnection() + query = _FakeQuery(connection) + parent_context = SpanContext(0xABC, 0xDEF, False, TraceFlags(TraceFlags.SAMPLED)) + + with use_span(NonRecordingSpan(parent_context)): + execute_with_trace_context(connection, query.execute, INVALID_SPAN_CONTEXT, "SELECT 1") + + assert query.calls == [(("SELECT 1",), {}, None)] + assert connection.traceparent is None + + +@pytest.mark.parametrize("failure_phase", ["setup", "cleanup"]) +@pytest.mark.parametrize("failure_operation", ["cursor", "execute", "close"]) +@pytest.mark.parametrize("query_fails", [False, True]) +def test_internal_failure_does_not_replace_query_result_or_error(failure_phase, failure_operation, query_fails, caplog): + connection = _FakeConnection(failure_phase=failure_phase, failure_operation=failure_operation) + query_error = ValueError("query failure") if query_fails else None + query = _FakeQuery(connection, error=query_error) + database_context = SpanContext(0x123, 0x456, False) + + if query_fails: + with pytest.raises(ValueError) as caught: + execute_with_trace_context(connection, query.execute, database_context, "SELECT 1") + assert caught.value is query_error + else: + assert execute_with_trace_context(connection, query.execute, database_context, "SELECT 1") is query.result + + assert len(query.calls) == 1 + assert len(connection.cursors) == 2 + assert f"MySQL trace context {failure_phase} failed (RuntimeError)" in caplog.text + if failure_phase == "setup": + assert connection.traceparent is None + + +def test_query_error_clears_context_and_closes_internal_cursors(): + connection = _FakeConnection() + error = ValueError("query failure") + query = _FakeQuery(connection, error=error) + + with pytest.raises(ValueError) as caught: + execute_with_trace_context(connection, query.execute, SpanContext(0x123, 0x456, False), "SELECT 1") + + assert caught.value is error + assert connection.traceparent is None + assert all(cursor.closed for cursor in connection.cursors) + + +def test_execution_uses_existing_database_span_without_creating_more_spans(): + connection = _FakeConnection() + query = _FakeQuery(connection) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + tracer = provider.get_tracer(__name__) + with tracer.start_as_current_span("query", kind=SpanKind.CLIENT) as span: + execute_with_trace_context(connection, query.execute, span.get_span_context(), "SELECT 1") + context = span.get_span_context() + assert query.calls[0][2] == ( + f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags):02x}" + ) + assert len(exporter.get_finished_spans()) == 1 + finally: + provider.shutdown() From 2a76dc34a08d18dd20f258e17976a458d0f6f36c Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Fri, 23 Oct 2026 17:01:04 -0400 Subject: [PATCH 2/3] Integrate MySQL trace context with buffered Connector/Python queries --- docs/mysql.md | 95 +++++- src/splunk_otel/mysql.py | 127 +++++++- tests/integration/mysql/test_correlation.py | 324 ++++++++++++++++++++ 3 files changed, 536 insertions(+), 10 deletions(-) create mode 100644 tests/integration/mysql/test_correlation.py diff --git a/docs/mysql.md b/docs/mysql.md index 20d05ee3..d42ac834 100644 --- a/docs/mysql.md +++ b/docs/mysql.md @@ -1,4 +1,4 @@ -# MySQL query correlation: first implementation step +# Experimental MySQL query correlation The goal is to link a query in Database Monitoring (DBMon) to the application trace that caused it. A trace describes an application request. A span describes @@ -39,7 +39,58 @@ written even when the trace is not selected for export. If the supplied context is invalid, the helper clears the variable before the query. It does not use the surrounding application's span instead. -## What the caller must provide +## Connector/Python integration + +Call `enable_mysql_trace_context_propagation()` explicitly to install the +process-lifetime hook. It returns `False` and logs a warning if Connector/Python +or compatible OpenTelemetry DBAPI instrumentation is unavailable. Repeated calls +do not install additional wrappers. There is no startup activation or environment +setting yet. + +```python +import mysql.connector +from opentelemetry.instrumentation.mysql import MySQLInstrumentor +from splunk_otel.mysql import enable_mysql_trace_context_propagation + +assert enable_mysql_trace_context_propagation() +MySQLInstrumentor().instrument() + +with mysql.connector.connect(database="example", buffered=True) as connection: + with connection.cursor() as cursor: + cursor.execute("SELECT 1") + rows = cursor.fetchall() +``` + +Configure an OpenTelemetry tracer provider before running this example. The hook +also works with connections wrapped by `MySQLInstrumentor.instrument_connection()`. +It uses the database span created by OpenTelemetry and adds no spans for its +internal SET commands. Suppressed calls and no-op tracers clear the variable for +eligible queries without copying an application parent span. + +The first integration supports Connector/Python's built-in buffered tuple, +raw, and dictionary cursors, in pure Python and C extension modes. Before adding +session commands, the hook requires all of the following: + +- The cursor uses the built-in `execute` method and a built-in physical connection. +- The connection reports no unread rows or additional result sets. +- The SQL is a string beginning with `SELECT`, `INSERT`, `UPDATE`, `DELETE`, + `REPLACE`, `WITH`, `SHOW`, `EXPLAIN`, `DESCRIBE`, or `DESC` followed by whitespace. +- The SQL contains no semicolons except an optional final terminator, no comment + markers (`/*`, `--`, or `#`), no NUL character, and no `DELIMITER` word. + +The SQL check is deliberately conservative. It also skips markers inside quoted +strings. SQLCommenter adds comments to queries, so those queries are skipped. +Stored-procedure calls, scripts, batches, prepared cursors, unbuffered cursors, +custom cursor and connection subclasses, bytes SQL, other drivers, and asynchronous +calls run through their existing instrumentation without correlation commands. +Skipped calls do not clear an existing session value. + +Connections must not be shared by concurrent calls. The hook depends on the +DBAPI execution callback and Connector/Python cursor connection reference, which +are private interfaces. The tested versions are DBAPI instrumentation 0.66b1 and +Connector/Python 26.7.0; other versions need separate validation. + +## Direct helper calls The helper requires the actual driver connection, without an instrumentation wrapper, and the context of the span for this database call. It does not create @@ -50,8 +101,33 @@ return at most one result set and use a buffered cursor: the driver downloads al rows before the query method returns. MySQL may reject the cleanup SET if rows are still waiting to be downloaded. -The helper does not check these restrictions yet. Driver integration must enforce -them before the helper can be used automatically. +Direct helper calls must meet these restrictions themselves. The Connector/Python +hook checks its supported cursor types, connection state, and SQL before calling +the helper. + +## Validation + +The driver tests use a disposable MySQL database. Install the project, pytest, +`mysql-connector-python==26.7.0`, and +`opentelemetry-instrumentation-mysql==0.66b1` into a virtual environment. Start a +test server, then run the tests from the repository root: + +```sh +docker run --detach --rm --name mysql-correlation-tests \ + --env MYSQL_ROOT_PASSWORD=prototype-only --env MYSQL_DATABASE=correlation \ + --publish 127.0.0.1:33060:3306 mysql:8.4 +# Wait until mysqladmin reports that the server is alive. +docker exec --env MYSQL_PWD=prototype-only mysql-correlation-tests \ + mysqladmin --protocol=tcp --host=127.0.0.1 --user=root ping --silent +MYSQL_CORRELATION_TEST_PORT=33060 python -m pytest \ + tests/test_mysql.py tests/integration/mysql/test_correlation.py -q +docker stop mysql-correlation-tests +``` + +The driver tests skip when the port variable or optional packages are absent. +They check exact child-span IDs, buffered rows after cleanup, error recovery, +suppression, unsampled spans, transaction boundaries, and unsupported query paths +against MySQL 8.4.11 in both driver modes. They do not validate Collector or DBMon. ## Failures and remaining work @@ -60,7 +136,10 @@ query's result or raises the query's original exception. A failed SET can leave old IDs on the connection, so later queries could receive an incorrect trace link. -This is an experimental building block. It is not enabled at startup and has no -environment setting. The remaining work includes driver integration, transaction -and connection-pool handling, streaming results, multiple result sets, custom -cursors, asynchronous drivers, and Collector/DBMon testing. +Every eligible query adds two SET commands. These commands also change MySQL +session diagnostics, such as `ROW_COUNT()` and the warning list. Applications +that depend on the previous command's diagnostics need further design work. + +The remaining work includes startup configuration, connection-pool validation, +streaming results, multiple result sets, custom cursors, additional drivers, +asynchronous drivers, overhead measurement, and Collector/DBMon testing. diff --git a/src/splunk_otel/mysql.py b/src/splunk_otel/mysql.py index 195afb49..3b656cf6 100644 --- a/src/splunk_otel/mysql.py +++ b/src/splunk_otel/mysql.py @@ -12,13 +12,136 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect import logging +import re +import threading +from functools import partial +from importlib import import_module -from opentelemetry.instrumentation.utils import suppress_instrumentation -from opentelemetry.trace import NonRecordingSpan, SpanContext, set_span_in_context +import wrapt +from opentelemetry import trace +from opentelemetry.instrumentation.utils import is_instrumentation_enabled, suppress_instrumentation +from opentelemetry.trace import INVALID_SPAN_CONTEXT, NonRecordingSpan, SpanContext, set_span_in_context from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator _logger = logging.getLogger(__name__) +_enable_lock = threading.Lock() +_enabled = False +_QUERY_START = re.compile(r"^(SELECT|INSERT|UPDATE|DELETE|REPLACE|WITH|SHOW|EXPLAIN|DESCRIBE|DESC)\s", re.IGNORECASE) + + +def enable_mysql_trace_context_propagation() -> bool: + """Enable experimental correlation for built-in Connector/Python buffered cursors.""" + global _enabled + + with _enable_lock: + if _enabled: + return True + try: + dbapi = import_module("opentelemetry.instrumentation.dbapi") + connector = import_module("mysql.connector") + cursor_types = _load_connector_cursor_types(connector) + except ImportError: + _logger.warning("MySQL trace context propagation requires Connector/Python and DBAPI instrumentation") + return False + + cursor_tracer = getattr(dbapi, "CursorTracer", None) + method = getattr(cursor_tracer, "traced_execution", None) + try: + parameters = tuple(inspect.signature(method).parameters) + except (ValueError, TypeError): + parameters = () + if parameters[:3] != ("self", "cursor", "query_method"): + _logger.warning("MySQL trace context propagation requires compatible DBAPI execution methods") + return False + + wrapper = partial(_execute_with_mysql_trace_context, connector=connector, cursor_types=cursor_types) + wrapt.wrap_function_wrapper(cursor_tracer, "traced_execution", wrapper) + _enabled = True + return True + + +def _load_connector_cursor_types(connector): + cursors = import_module("mysql.connector.cursor") + cursor_types = { + cursor_type: (connector.MySQLConnection, "have_next_result") + for cursor_type in ( + cursors.MySQLCursorBuffered, + cursors.MySQLCursorBufferedRaw, + cursors.MySQLCursorBufferedDict, + ) + } + if connector.HAVE_CEXT: + cext = import_module("mysql.connector.cursor_cext") + cursor_types.update( + { + cursor_type: (connector.CMySQLConnection, "more_results") + for cursor_type in ( + cext.CMySQLCursorBuffered, + cext.CMySQLCursorBufferedRaw, + cext.CMySQLCursorBufferedDict, + ) + } + ) + return cursor_types + + +def _execute_with_mysql_trace_context(wrapped, cursor_tracer, args, kwargs, *, connector, cursor_types): + if len(args) < 2: + return wrapped(*args, **kwargs) + cursor, query_method, *query_args = args + # The pinned DBAPI callback has no public extension hook. + integration = cursor_tracer._db_api_integration + if ( + integration.database_system != "mysql" + or integration.connect_module is not connector + or type(cursor) not in cursor_types + or getattr(query_method, "__self__", None) is not cursor + or getattr(query_method, "__func__", None) is not type(cursor).execute + ): + return wrapped(*args, **kwargs) + + connection_type, pending_attribute = cursor_types[type(cursor)] + try: + # Connector stores a weak proxy; get_self() returns the physical connection. + connection = cursor._connection.get_self() + except (AttributeError, ReferenceError): + return wrapped(*args, **kwargs) + if type(connection) is not connection_type: + return wrapped(*args, **kwargs) + parent_context = trace.get_current_span().get_span_context() + instrumentation_enabled = is_instrumentation_enabled() + + def _execute_query(*execution_args, **execution_kwargs): + operation = execution_args[0] if execution_args else execution_kwargs.get("operation") + if not _is_single_result_query(operation) or not _connection_is_ready(connection, pending_attribute): + return query_method(*execution_args, **execution_kwargs) + context = trace.get_current_span().get_span_context() + if not instrumentation_enabled or context == parent_context: + context = INVALID_SPAN_CONTEXT + return execute_with_trace_context(connection, query_method, context, *execution_args, **execution_kwargs) + + return wrapped(cursor, _execute_query, *query_args, **kwargs) + + +def _is_single_result_query(operation): + if type(operation) is not str: + return False + statement = operation.strip().removesuffix(";").rstrip() + # Skip ambiguous SQL rather than parse strings, comments, or custom delimiters. + return ( + _QUERY_START.match(statement) is not None + and not any(token in statement for token in (";", "/*", "--", "#", "\x00")) + and re.search(r"\bDELIMITER\b", statement, re.IGNORECASE) is None + ) + + +def _connection_is_ready(connection, pending_attribute): + try: + return connection.unread_result is False and getattr(connection, pending_attribute) is False + except Exception: # noqa: BLE001 - A closed connection must fail through the driver's query method. + return False def execute_with_trace_context(connection, query_method, span_context: SpanContext, *args, **kwargs): diff --git a/tests/integration/mysql/test_correlation.py b/tests/integration/mysql/test_correlation.py new file mode 100644 index 00000000..4524c7d3 --- /dev/null +++ b/tests/integration/mysql/test_correlation.py @@ -0,0 +1,324 @@ +# Copyright Splunk 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. + +import os + +import pytest + +mysql = pytest.importorskip("mysql.connector") +instrumentation = pytest.importorskip("opentelemetry.instrumentation.mysql") + +from opentelemetry.instrumentation.utils import suppress_instrumentation +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import ALWAYS_OFF +from opentelemetry.trace import NoOpTracerProvider, SpanKind + +from splunk_otel.mysql import enable_mysql_trace_context_propagation + +pytestmark = pytest.mark.manual + + +@pytest.fixture +def span_exporter(): + return InMemorySpanExporter() + + +@pytest.fixture +def tracer_provider(span_exporter): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + yield provider + provider.shutdown() + + +@pytest.fixture(params=[True, False], ids=["python", "cext"]) +def raw_connection(request): + port = os.environ.get("MYSQL_CORRELATION_TEST_PORT") + if not port: + pytest.skip("Set MYSQL_CORRELATION_TEST_PORT to run against a disposable MySQL server") + if not request.param and not mysql.HAVE_CEXT: + pytest.skip("Connector/Python C extension is not installed") + connection = mysql.connect( + host="127.0.0.1", + port=int(port), + user="root", + password="prototype-only", + database="correlation", + use_pure=request.param, + ) + yield connection + connection.close() + + +@pytest.fixture +def connection(raw_connection, tracer_provider): + assert enable_mysql_trace_context_propagation() + assert enable_mysql_trace_context_propagation() + return instrumentation.MySQLInstrumentor().instrument_connection(raw_connection, tracer_provider=tracer_provider) + + +def _read_traceparent(raw_connection): + with raw_connection.cursor(buffered=True) as cursor: + cursor.execute("SELECT @traceparent") + return cursor.fetchone()[0] + + +def _set_previous_context(raw_connection): + with raw_connection.cursor() as cursor: + cursor.execute("SET @traceparent = 'previous-context'") + + +@pytest.mark.parametrize("cursor_options", [{}, {"dictionary": True}, {"raw": True}]) +def test_buffered_query_uses_database_child_span_and_keeps_rows( + connection, raw_connection, tracer_provider, span_exporter, cursor_options +): + with ( + tracer_provider.get_tracer(__name__).start_as_current_span("parent") as parent, + connection.cursor(buffered=True, **cursor_options) as cursor, + ): + cursor.execute("SELECT @traceparent AS context, %s AS value;", (17,)) + assert _read_traceparent(raw_connection) is None + row = cursor.fetchone() + spans = span_exporter.get_finished_spans() + database_span = next(span for span in spans if span.kind == SpanKind.CLIENT) + context = database_span.context + expected = f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags):02x}" + if cursor_options.get("dictionary"): + assert row == {"context": expected, "value": 17} + elif cursor_options.get("raw"): + assert row == (expected.encode(), b"17") + else: + assert row == (expected, 17) + assert database_span.parent == parent.get_span_context() + assert len(spans) == 2 + + +@pytest.mark.parametrize( + "operation", + [ + "SELECT @traceparent; SELECT 2", + "SELECT @traceparent /* comment */", + "SELECT @traceparent -- comment\n", + "SELECT @traceparent # comment\n", + "SELECT /*! @traceparent */", + "SELECT @traceparent, ';'", + b"SELECT @traceparent", + ], +) +def test_unsupported_sql_runs_without_session_commands(connection, raw_connection, operation): + _set_previous_context(raw_connection) + with connection.cursor(buffered=True) as cursor: + cursor.execute(operation) + assert cursor.fetchone()[0] == "previous-context" + while cursor.nextset(): + if cursor.with_rows: + cursor.fetchall() + assert _read_traceparent(raw_connection) == "previous-context" + + +def test_unbuffered_query_keeps_unread_rows_and_session_value(connection, raw_connection): + _set_previous_context(raw_connection) + with connection.cursor() as cursor: + cursor.execute("SELECT @traceparent UNION ALL SELECT 'second'") + assert cursor.fetchall() == [("previous-context",), ("second",)] + assert _read_traceparent(raw_connection) == "previous-context" + + +def test_prepared_cursor_runs_without_session_commands(connection, raw_connection): + _set_previous_context(raw_connection) + with connection.cursor(prepared=True) as cursor: + cursor.execute("SELECT @traceparent") + assert cursor.fetchall() == [("previous-context",)] + assert _read_traceparent(raw_connection) == "previous-context" + + +def test_batch_runs_without_session_commands(connection, raw_connection, span_exporter): + _set_previous_context(raw_connection) + with raw_connection.cursor() as cursor: + cursor.execute("CREATE TEMPORARY TABLE correlation_batch (context TEXT, value INT)") + with connection.cursor(buffered=True) as cursor: + cursor.executemany("INSERT INTO correlation_batch VALUES (@traceparent, %s)", [(17,), (23,)]) + with raw_connection.cursor(buffered=True) as cursor: + cursor.execute("SELECT context, value FROM correlation_batch ORDER BY value") + assert cursor.fetchall() == [("previous-context", 17), ("previous-context", 23)] + assert _read_traceparent(raw_connection) == "previous-context" + assert len(span_exporter.get_finished_spans()) == 1 + + +def test_sqlcommenter_query_runs_without_session_commands(connection, raw_connection, tracer_provider, span_exporter): + connection = instrumentation.MySQLInstrumentor().instrument_connection( + raw_connection, tracer_provider=tracer_provider, enable_commenter=True + ) + _set_previous_context(raw_connection) + with connection.cursor(buffered=True) as cursor: + cursor.execute("SELECT @traceparent") + assert cursor.fetchone() == ("previous-context",) + assert _read_traceparent(raw_connection) == "previous-context" + assert len(span_exporter.get_finished_spans()) == 1 + + +def test_pending_rows_are_left_to_the_driver(connection, raw_connection): + _set_previous_context(raw_connection) + with raw_connection.cursor() as pending, connection.cursor(buffered=True) as cursor: + pending.execute("SELECT 17") + with pytest.raises(mysql.InternalError): + cursor.execute("SELECT @traceparent") + assert pending.fetchall() == [(17,)] + assert _read_traceparent(raw_connection) == "previous-context" + + +def test_stored_procedure_results_are_not_consumed(connection, raw_connection): + _set_previous_context(raw_connection) + with raw_connection.cursor() as cursor: + cursor.execute("DROP PROCEDURE IF EXISTS correlation_results") + cursor.execute("CREATE PROCEDURE correlation_results() BEGIN SELECT @traceparent; SELECT 2; END") + try: + with connection.cursor(buffered=True) as cursor: + cursor.execute("CALL correlation_results()") + assert cursor.fetchone() == ("previous-context",) + assert cursor.nextset() + assert cursor.fetchone() == (2,) + while cursor.nextset(): + if cursor.with_rows: + cursor.fetchall() + assert _read_traceparent(raw_connection) == "previous-context" + finally: + with raw_connection.cursor() as cursor: + cursor.execute("DROP PROCEDURE correlation_results") + + +def test_suppressed_query_clears_context_without_injecting_parent( + connection, raw_connection, tracer_provider, span_exporter +): + _set_previous_context(raw_connection) + with ( + tracer_provider.get_tracer(__name__).start_as_current_span("parent"), + suppress_instrumentation(), + connection.cursor(buffered=True) as cursor, + ): + cursor.execute("SELECT @traceparent") + assert cursor.fetchone() == (None,) + assert _read_traceparent(raw_connection) is None + assert [span.name for span in span_exporter.get_finished_spans()] == ["parent"] + + +def test_noop_tracer_does_not_inject_parent(connection, raw_connection, tracer_provider): + connection = instrumentation.MySQLInstrumentor().instrument_connection( + raw_connection, tracer_provider=NoOpTracerProvider() + ) + _set_previous_context(raw_connection) + with ( + tracer_provider.get_tracer(__name__).start_as_current_span("parent"), + connection.cursor(buffered=True) as cursor, + ): + cursor.execute("SELECT @traceparent") + assert cursor.fetchone() == (None,) + assert _read_traceparent(raw_connection) is None + + +def test_unsampled_database_context_is_propagated(connection, raw_connection): + provider = TracerProvider(sampler=ALWAYS_OFF) + connection = instrumentation.MySQLInstrumentor().instrument_connection(raw_connection, tracer_provider=provider) + try: + with provider.get_tracer(__name__).start_as_current_span("parent") as parent: + with connection.cursor(buffered=True) as cursor: + cursor.execute("SELECT @traceparent") + value = cursor.fetchone()[0] + context = parent.get_span_context() + assert value.startswith(f"00-{context.trace_id:032x}-") + assert value.split("-")[2] != f"{context.span_id:016x}" + assert int(value.split("-")[3], 16) & 1 == 0 + assert _read_traceparent(raw_connection) is None + finally: + provider.shutdown() + + +def test_query_failure_clears_context_and_keeps_driver_error(connection, raw_connection, span_exporter): + with connection.cursor(buffered=True) as cursor: + with pytest.raises(mysql.ProgrammingError) as caught: + cursor.execute("SELECT * FROM correlation_missing_table") + assert caught.value.errno == 1146 + assert _read_traceparent(raw_connection) is None + cursor.execute("SELECT @traceparent") + assert cursor.fetchone()[0] is not None + assert _read_traceparent(raw_connection) is None + assert len(span_exporter.get_finished_spans()) == 2 + + +@pytest.mark.parametrize("boundary", ["commit", "rollback"]) +def test_transaction_boundary_has_no_completed_context(connection, raw_connection, span_exporter, boundary): + with raw_connection.cursor() as cursor: + cursor.execute( + "CREATE TEMPORARY TABLE correlation_transaction (id INT AUTO_INCREMENT PRIMARY KEY, context TEXT)" + ) + with connection.cursor(buffered=True) as cursor: + cursor.execute("INSERT INTO correlation_transaction (context) VALUES (@traceparent)") + assert cursor.rowcount == 1 + assert cursor.lastrowid == 1 + assert raw_connection.in_transaction + assert _read_traceparent(raw_connection) is None + getattr(connection, boundary)() + assert _read_traceparent(raw_connection) is None + with raw_connection.cursor(buffered=True) as cursor: + cursor.execute("SELECT context FROM correlation_transaction") + rows = cursor.fetchall() + span = span_exporter.get_finished_spans()[0] + context = span.context + expected = f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags):02x}" + assert rows == ([(expected,)] if boundary == "commit" else []) + assert len(span_exporter.get_finished_spans()) == 1 + + +def test_custom_cursor_is_not_correlated(connection, raw_connection): + base = ( + mysql.cursor.MySQLCursorBuffered + if isinstance(raw_connection, mysql.MySQLConnection) + else mysql.cursor_cext.CMySQLCursorBuffered + ) + + class CustomCursor(base): + pass + + _set_previous_context(raw_connection) + with connection.cursor(cursor_class=CustomCursor) as cursor: + cursor.execute("SELECT @traceparent") + assert cursor.fetchone() == ("previous-context",) + assert _read_traceparent(raw_connection) == "previous-context" + + +def test_automatic_instrumentation_uses_same_hook(raw_connection, tracer_provider, span_exporter): + assert enable_mysql_trace_context_propagation() + instrumentor = instrumentation.MySQLInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider) + try: + with ( + mysql.connect( + host="127.0.0.1", + port=int(os.environ["MYSQL_CORRELATION_TEST_PORT"]), + user="root", + password="prototype-only", + database="correlation", + use_pure=isinstance(raw_connection, mysql.MySQLConnection), + ) as connection, + connection.cursor(buffered=True) as cursor, + ): + cursor.execute("SELECT @traceparent") + value = cursor.fetchone()[0] + context = span_exporter.get_finished_spans()[0].context + assert value == f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags):02x}" + assert len(span_exporter.get_finished_spans()) == 1 + finally: + instrumentor.uninstrument() From 1951203015006b905de1b446f9f126942317247b Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Fri, 23 Oct 2026 17:08:20 -0400 Subject: [PATCH 3/3] Fix MySQL correlation lint checks under Hatch --- src/splunk_otel/mysql.py | 30 ++++++++++----------- tests/integration/mysql/test_correlation.py | 6 ++--- tests/test_mysql.py | 7 ++--- 3 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/splunk_otel/mysql.py b/src/splunk_otel/mysql.py index 3b656cf6..866ccd2c 100644 --- a/src/splunk_otel/mysql.py +++ b/src/splunk_otel/mysql.py @@ -28,12 +28,13 @@ _logger = logging.getLogger(__name__) _enable_lock = threading.Lock() _enabled = False +_CURSOR_AND_QUERY_METHOD_ARGUMENTS = 2 _QUERY_START = re.compile(r"^(SELECT|INSERT|UPDATE|DELETE|REPLACE|WITH|SHOW|EXPLAIN|DESCRIBE|DESC)\s", re.IGNORECASE) def enable_mysql_trace_context_propagation() -> bool: """Enable experimental correlation for built-in Connector/Python buffered cursors.""" - global _enabled + global _enabled # noqa: PLW0603 - Process-lifetime activation is serialized by _enable_lock. with _enable_lock: if _enabled: @@ -64,35 +65,34 @@ def enable_mysql_trace_context_propagation() -> bool: def _load_connector_cursor_types(connector): cursors = import_module("mysql.connector.cursor") - cursor_types = { - cursor_type: (connector.MySQLConnection, "have_next_result") - for cursor_type in ( + cursor_types = dict.fromkeys( + ( cursors.MySQLCursorBuffered, cursors.MySQLCursorBufferedRaw, cursors.MySQLCursorBufferedDict, - ) - } + ), + (connector.MySQLConnection, "have_next_result"), + ) if connector.HAVE_CEXT: cext = import_module("mysql.connector.cursor_cext") cursor_types.update( - { - cursor_type: (connector.CMySQLConnection, "more_results") - for cursor_type in ( + dict.fromkeys( + ( cext.CMySQLCursorBuffered, cext.CMySQLCursorBufferedRaw, cext.CMySQLCursorBufferedDict, - ) - } + ), + (connector.CMySQLConnection, "more_results"), + ) ) return cursor_types def _execute_with_mysql_trace_context(wrapped, cursor_tracer, args, kwargs, *, connector, cursor_types): - if len(args) < 2: + if len(args) < _CURSOR_AND_QUERY_METHOD_ARGUMENTS: return wrapped(*args, **kwargs) cursor, query_method, *query_args = args - # The pinned DBAPI callback has no public extension hook. - integration = cursor_tracer._db_api_integration + integration = cursor_tracer._db_api_integration # noqa: SLF001 - The DBAPI callback has no public integration accessor. if ( integration.database_system != "mysql" or integration.connect_module is not connector @@ -105,7 +105,7 @@ def _execute_with_mysql_trace_context(wrapped, cursor_tracer, args, kwargs, *, c connection_type, pending_attribute = cursor_types[type(cursor)] try: # Connector stores a weak proxy; get_self() returns the physical connection. - connection = cursor._connection.get_self() + connection = cursor._connection.get_self() # noqa: SLF001 - Connector cursors have no public connection accessor. except (AttributeError, ReferenceError): return wrapped(*args, **kwargs) if type(connection) is not connection_type: diff --git a/tests/integration/mysql/test_correlation.py b/tests/integration/mysql/test_correlation.py index 4524c7d3..1efea1bd 100644 --- a/tests/integration/mysql/test_correlation.py +++ b/tests/integration/mysql/test_correlation.py @@ -16,9 +16,6 @@ import pytest -mysql = pytest.importorskip("mysql.connector") -instrumentation = pytest.importorskip("opentelemetry.instrumentation.mysql") - from opentelemetry.instrumentation.utils import suppress_instrumentation from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor @@ -28,6 +25,9 @@ from splunk_otel.mysql import enable_mysql_trace_context_propagation +mysql = pytest.importorskip("mysql.connector") +instrumentation = pytest.importorskip("opentelemetry.instrumentation.mysql") + pytestmark = pytest.mark.manual diff --git a/tests/test_mysql.py b/tests/test_mysql.py index 58975705..aba94d06 100644 --- a/tests/test_mysql.py +++ b/tests/test_mysql.py @@ -39,7 +39,8 @@ def cursor(self): def fail_if_requested(self, phase, operation): if (phase, operation) == (self.failure_phase, self.failure_operation): - raise RuntimeError("internal failure") + message = "internal failure" + raise RuntimeError(message) class _FakeInternalCursor: @@ -116,7 +117,7 @@ def test_internal_failure_does_not_replace_query_result_or_error(failure_phase, database_context = SpanContext(0x123, 0x456, False) if query_fails: - with pytest.raises(ValueError) as caught: + with pytest.raises(ValueError, match="query failure") as caught: execute_with_trace_context(connection, query.execute, database_context, "SELECT 1") assert caught.value is query_error else: @@ -134,7 +135,7 @@ def test_query_error_clears_context_and_closes_internal_cursors(): error = ValueError("query failure") query = _FakeQuery(connection, error=error) - with pytest.raises(ValueError) as caught: + with pytest.raises(ValueError, match="query failure") as caught: execute_with_trace_context(connection, query.execute, SpanContext(0x123, 0x456, False), "SELECT 1") assert caught.value is error