Add experimental MySQL query-to-trace correlation

docs/mysql.md

@@ -0,0 +1,145 @@
+# 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
+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-<trace-id>-<span-id>-<flags>';
+```
+
+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.
+
+## 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
+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.
+
+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
+
+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.
+
+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.

src/splunk_otel/mysql.py

@@ -0,0 +1,175 @@
+# 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 inspect
+import logging
+import re
+import threading
+from functools import partial
+from importlib import import_module
+
+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
+_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  # noqa: PLW0603 - Process-lifetime activation is serialized by _enable_lock.
+
+    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 = 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(
+            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) < _CURSOR_AND_QUERY_METHOD_ARGUMENTS:
+        return wrapped(*args, **kwargs)
+    cursor, query_method, *query_args = args
+    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
+        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()  # noqa: SLF001 - Connector cursors have no public connection accessor.
+    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):
+    """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__)