diff --git a/docs/mysql.md b/docs/mysql.md new file mode 100644 index 00000000..d42ac834 --- /dev/null +++ b/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---'; +``` + +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. diff --git a/src/splunk_otel/mysql.py b/src/splunk_otel/mysql.py new file mode 100644 index 00000000..866ccd2c --- /dev/null +++ b/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__) diff --git a/tests/integration/mysql/test_correlation.py b/tests/integration/mysql/test_correlation.py new file mode 100644 index 00000000..1efea1bd --- /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 + +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 + +mysql = pytest.importorskip("mysql.connector") +instrumentation = pytest.importorskip("opentelemetry.instrumentation.mysql") + +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() diff --git a/tests/test_mysql.py b/tests/test_mysql.py new file mode 100644 index 00000000..aba94d06 --- /dev/null +++ b/tests/test_mysql.py @@ -0,0 +1,162 @@ +# 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): + message = "internal failure" + raise RuntimeError(message) + + +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, match="query failure") 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, match="query failure") 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()