From a899611a9e69b80d08166ebfaf5a21244305804c Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Tue, 13 Oct 2026 14:44:22 -0400 Subject: [PATCH 01/10] Add experimental PostgreSQL trace correlation Assisted-by: Codex (GPT-6) --- .github/workflows/ci-main.yml | 43 ++ README.md | 35 ++ src/splunk_otel/configurator.py | 2 + src/splunk_otel/env.py | 1 + src/splunk_otel/postgresql.py | 234 +++++++++ src/splunk_otel/runtime.py | 8 +- tests/smoke/postgresql_trace_context.py | 334 +++++++++++++ tests/test_configurator.py | 27 ++ tests/test_postgresql.py | 604 ++++++++++++++++++++++++ tests/test_runtime.py | 16 +- 10 files changed, 1302 insertions(+), 2 deletions(-) create mode 100644 src/splunk_otel/postgresql.py create mode 100644 tests/smoke/postgresql_trace_context.py create mode 100644 tests/test_postgresql.py diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 4997936e..115df8ab 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -49,3 +49,46 @@ jobs: - name: Run oteltests run: oteltest tests/integration/*.py + + postgresql-smoke: + name: PostgreSQL trace-context smoke test (Python ${{ matrix.python-version }}, ${{ matrix.mode }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + mode: [auto, manual, disabled] + python-version: ['3.10', '3.14'] + env: + POSTGRESQL_SMOKE_INSTRUMENTATION: ${{ matrix.mode }} + SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED: ${{ matrix.mode != 'disabled' }} + services: + postgres: + image: postgres:16-alpine + env: + POSTGRES_DB: otel + POSTGRES_USER: otel + POSTGRES_PASSWORD: otel + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U otel -d otel" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + + - name: Set up Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Install smoke-test dependencies + run: >- + pip install . "psycopg[binary]" psycopg2-binary + opentelemetry-instrumentation-psycopg + opentelemetry-instrumentation-psycopg2 + + - name: Run PostgreSQL smoke test + run: python tests/smoke/postgresql_trace_context.py diff --git a/README.md b/README.md index 9e46c0d7..32fd8cd9 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,41 @@ Then run your application with `opentelemetry-instrument`: OTEL_SERVICE_NAME=my-python-service opentelemetry-instrument python app.py ``` +### Experimental PostgreSQL trace correlation + +Set `SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true` to place each +instrumented psycopg or psycopg2 database span's W3C `traceparent` in +PostgreSQL's `application_name` immediately before the query. This sends one +additional `SET application_name` command per traced query. +The feature targets Psycopg 3 synchronous and asynchronous connections and +Psycopg2 synchronous connections. Psycopg2 asynchronous connections pass through +without trace propagation. + +This opt-in replaces the connection's `application_name` and leaves the value +set after the query. If instrumentation creates no valid new database span, the +feature resets the setting to the connection's configured default. It does not +restore a value that the application previously set through SQL. Queries that +bypass instrumentation, including suppressed queries and cursors created before +instrumentation, can inherit an earlier query's trace context. +Psycopg 3 named cursors use a separate server-cursor factory that the current +upstream instrumentor does not wrap. Those queries, cursor fetch operations, +and COPY operations are outside this feature's automatic coverage. + +The feature logs internal command failures and attempts the customer query. A +failed command can leave a transaction aborted, so the customer query can also +fail. The feature does not roll back application transactions. Database span +duration and database operation duration metrics include lock waiting and the +additional command. + +Manual instrumentation must initialize the Splunk distro first with +`init_splunk_otel()`, or call `enable_postgresql_trace_context_propagation()` +before using `PsycopgInstrumentor` or `Psycopg2Instrumentor`. + +Install the corresponding OpenTelemetry driver instrumentation package. The +environment switch applies when `OTEL_CONFIG_FILE` is unset. Activation lasts +for the process lifetime. To turn it off, unset the switch and restart. Only +queries using the instrumented driver participate in its connection lock. + By default, OTLP exporters send telemetry to a locally running [Splunk OpenTelemetry Collector](https://github.com/signalfx/splunk-otel-collector) at `http://localhost:4317`. To send telemetry directly to Splunk Observability diff --git a/src/splunk_otel/configurator.py b/src/splunk_otel/configurator.py index 2cbe354a..6930e916 100644 --- a/src/splunk_otel/configurator.py +++ b/src/splunk_otel/configurator.py @@ -36,6 +36,7 @@ from splunk_otel.runtime import ( _SERVER_TIMING_DEFAULT_ENABLED, configure_logging_instrumentation, + configure_postgresql_trace_context_propagation, ) _logger = logging.getLogger(__name__) @@ -68,6 +69,7 @@ def _configure(self, **kwargs): # A later change will configure profiling and callgraphs from OTEL_CONFIG_FILE. # For now, start them from environment variables only when no config file is set. if not config_file: + configure_postgresql_trace_context_propagation(env) _start_profiling_if_enabled() _configure_callgraphs_if_enabled() diff --git a/src/splunk_otel/env.py b/src/splunk_otel/env.py index 8d23edcb..b9dc0fda 100644 --- a/src/splunk_otel/env.py +++ b/src/splunk_otel/env.py @@ -47,6 +47,7 @@ SPLUNK_OPAMP_ENABLED = "SPLUNK_OPAMP_ENABLED" SPLUNK_OPAMP_ENDPOINT = "SPLUNK_OPAMP_ENDPOINT" SPLUNK_OPAMP_POLLING_INTERVAL = "SPLUNK_OPAMP_POLLING_INTERVAL" +SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED = "SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED" SPLUNK_PROFILER_ENABLED = "SPLUNK_PROFILER_ENABLED" SPLUNK_PROFILER_CALL_STACK_INTERVAL = "SPLUNK_PROFILER_CALL_STACK_INTERVAL" SPLUNK_PROFILER_LOGS_ENDPOINT = "SPLUNK_PROFILER_LOGS_ENDPOINT" diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py new file mode 100644 index 00000000..61e19f15 --- /dev/null +++ b/src/splunk_otel/postgresql.py @@ -0,0 +1,234 @@ +# 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. + +from __future__ import annotations + +import asyncio +import inspect +import logging +import threading +import weakref +from importlib import import_module +from typing import TYPE_CHECKING, Any + +import wrapt +from opentelemetry import trace +from opentelemetry.instrumentation.utils import is_instrumentation_enabled, suppress_instrumentation + +if TYPE_CHECKING: + from collections.abc import Callable + + from opentelemetry.trace import SpanContext + +logger = logging.getLogger(__name__) + +_SET_APPLICATION_NAME_DEFAULT = "SET application_name = DEFAULT" +_DRIVER_MODULE_NAMES = ("psycopg", "psycopg2") +_CURSOR_AND_QUERY_METHOD_ARGUMENTS = 2 + +_enabled = False +_enable_lock = threading.Lock() +_connection_locks_lock = threading.Lock() +# Only queries using this adapter participate; raw or suppressed queries can bypass it. +_sync_connection_locks: weakref.WeakKeyDictionary[Any, threading.Lock] = weakref.WeakKeyDictionary() +_async_connection_locks: weakref.WeakKeyDictionary[Any, asyncio.Lock] = weakref.WeakKeyDictionary() + + +def enable_postgresql_trace_context_propagation() -> bool: + """Enable process-lifetime PostgreSQL correlation for installed DBAPI instrumentation.""" + global _enabled # noqa: PLW0603 + + with _enable_lock: + if _enabled: + return True + + try: + dbapi = import_module("opentelemetry.instrumentation.dbapi") + except ImportError: + logger.warning("PostgreSQL trace context propagation requires DBAPI instrumentation") + return False + + cursor_tracer = getattr(dbapi, "CursorTracer", None) + for name in ("traced_execution", "traced_execution_async"): + method = getattr(cursor_tracer, name, None) + if not _has_execution_signature(method): + logger.warning("PostgreSQL trace context propagation requires compatible DBAPI execution methods") + return False + + wrapt.wrap_function_wrapper(cursor_tracer, "traced_execution", _execute_with_postgresql_trace_context) + wrapt.wrap_function_wrapper( + cursor_tracer, + "traced_execution_async", + _execute_with_postgresql_trace_context_async, + ) + _enabled = True + return True + + +def _has_execution_signature(method: Any) -> bool: + if not callable(method): + return False + try: + parameters = tuple(inspect.signature(method).parameters.values())[:3] + except (TypeError, ValueError): + return False + return tuple(parameter.name for parameter in parameters) == ("self", "cursor", "query_method") and all( + parameter.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + for parameter in parameters + ) + + +def _execute_with_postgresql_trace_context(traced_execution, _cursor_tracer, args, kwargs): + if not is_instrumentation_enabled() or len(args) < _CURSOR_AND_QUERY_METHOD_ARGUMENTS: + return traced_execution(*args, **kwargs) + + cursor, query_method, *query_args = args + if not _is_postgresql_cursor(cursor) or getattr(cursor.connection, "async_", False): + return traced_execution(*args, **kwargs) + + span_context_before = trace.get_current_span().get_span_context() + + def correlated_query_method(*method_args, **method_kwargs): + return _execute_correlated( + cursor, + query_method, + span_context_before, + method_args, + method_kwargs, + ) + + return traced_execution(cursor, correlated_query_method, *query_args, **kwargs) + + +async def _execute_with_postgresql_trace_context_async(traced_execution_async, _cursor_tracer, args, kwargs): + if not is_instrumentation_enabled() or len(args) < _CURSOR_AND_QUERY_METHOD_ARGUMENTS: + return await traced_execution_async(*args, **kwargs) + + cursor, query_method, *query_args = args + if not _is_postgresql_cursor(cursor) or getattr(cursor.connection, "async_", False): + return await traced_execution_async(*args, **kwargs) + + span_context_before = trace.get_current_span().get_span_context() + + async def correlated_query_method(*method_args, **method_kwargs): + return await _execute_correlated_async( + cursor, + query_method, + span_context_before, + method_args, + method_kwargs, + ) + + return await traced_execution_async(cursor, correlated_query_method, *query_args, **kwargs) + + +def _execute_correlated( + cursor, + query_method: Callable[..., Any], + span_context_before: SpanContext, + args, + kwargs, +): + lock = _get_sync_connection_lock(cursor.connection) + with lock: + traceparent = _current_database_traceparent(span_context_before) + _update_application_name(cursor.connection, traceparent) + return query_method(*args, **kwargs) + + +async def _execute_correlated_async( + cursor, + query_method: Callable[..., Any], + span_context_before: SpanContext, + args, + kwargs, +): + lock = _get_async_connection_lock(cursor.connection) + async with lock: + traceparent = _current_database_traceparent(span_context_before) + await _update_application_name_async(cursor.connection, traceparent) + return await query_method(*args, **kwargs) + + +def _get_sync_connection_lock(connection) -> threading.Lock: + with _connection_locks_lock: + lock = _sync_connection_locks.get(connection) + if lock is None: + lock = threading.Lock() + _sync_connection_locks[connection] = lock + return lock + + +def _get_async_connection_lock(connection) -> asyncio.Lock: + with _connection_locks_lock: + lock = _async_connection_locks.get(connection) + if lock is None: + lock = asyncio.Lock() + _async_connection_locks[connection] = lock + return lock + + +def _current_database_traceparent(span_context_before: SpanContext) -> str | None: + current = trace.get_current_span().get_span_context() + if not current.is_valid or current == span_context_before: + return None + return f"00-{current.trace_id:032x}-{current.span_id:016x}-{int(current.trace_flags) & 0xFF:02x}" + + +def _update_application_name(connection, traceparent: str | None) -> None: + statement = _application_name_statement(traceparent) + try: + cursor = connection.cursor() + try: + with suppress_instrumentation(): + result = cursor.execute(statement) + if inspect.iscoroutine(result): + result.close() + if inspect.isawaitable(result): + logger.error("Synchronous PostgreSQL cursor returned an awaitable") + finally: + cursor.close() + except Exception: + logger.exception("Failed to set PostgreSQL application_name") + + +async def _update_application_name_async(connection, traceparent: str | None) -> None: + statement = _application_name_statement(traceparent) + try: + cursor = connection.cursor() + try: + with suppress_instrumentation(): + result = cursor.execute(statement) + if inspect.isawaitable(result): + await result + finally: + close_result = cursor.close() + if inspect.isawaitable(close_result): + await close_result + except Exception: + logger.exception("Failed to set PostgreSQL application_name") + + +def _application_name_statement(traceparent: str | None) -> str: + if traceparent is None: + return _SET_APPLICATION_NAME_DEFAULT + return f"SET application_name = '{traceparent}'" + + +def _is_postgresql_cursor(cursor) -> bool: + return any(_is_postgresql_module(cursor_class.__module__) for cursor_class in type(cursor).__mro__) + + +def _is_postgresql_module(module_name: str) -> bool: + return any(module_name == name or module_name.startswith(f"{name}.") for name in _DRIVER_MODULE_NAMES) diff --git a/src/splunk_otel/runtime.py b/src/splunk_otel/runtime.py index fe19a54a..41963f2a 100644 --- a/src/splunk_otel/runtime.py +++ b/src/splunk_otel/runtime.py @@ -17,13 +17,19 @@ from opentelemetry.instrumentation.environment_variables import OTEL_PYTHON_DISABLED_INSTRUMENTATIONS from opentelemetry.instrumentation.logging import LoggingInstrumentor -from splunk_otel.env import Env +from splunk_otel.env import SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED, Env +from splunk_otel.postgresql import enable_postgresql_trace_context_propagation _DISABLED_INSTRUMENTATIONS_WILDCARD = "*" _LOGGING_INSTRUMENTATION_NAME = "logging" _SERVER_TIMING_DEFAULT_ENABLED = True +def configure_postgresql_trace_context_propagation(env: Env) -> None: + if env.is_true(SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED, "false"): + enable_postgresql_trace_context_propagation() + + class _LoggingInstrumentor(Protocol): def instrument(self) -> None: ... diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py new file mode 100644 index 00000000..34c702c3 --- /dev/null +++ b/tests/smoke/postgresql_trace_context.py @@ -0,0 +1,334 @@ +import asyncio +import os +from concurrent.futures import ThreadPoolExecutor +from importlib.metadata import version + + +os.environ.setdefault("OTEL_LOGS_EXPORTER", "none") +os.environ.setdefault("OTEL_METRICS_EXPORTER", "none") +os.environ.setdefault("OTEL_PYTHON_DISABLED_INSTRUMENTATIONS", "system_metrics") +os.environ.setdefault("OTEL_SERVICE_NAME", "postgresql-trace-context-smoke-test") +os.environ.setdefault("OTEL_TRACES_EXPORTER", "none") +os.environ.setdefault("OTEL_TRACES_SAMPLER", "parentbased_always_on") +os.environ.setdefault("SPLUNK_OPAMP_ENABLED", "false") +os.environ.setdefault("SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") +os.environ.setdefault("PGHOST", "127.0.0.1") +os.environ.setdefault("PGPORT", "5432") +os.environ.setdefault("PGDATABASE", "otel") +os.environ.setdefault("PGUSER", "otel") +os.environ.setdefault("PGPASSWORD", "otel") + +from opentelemetry import trace +from opentelemetry.instrumentation.utils import suppress_instrumentation +from opentelemetry.trace import NonRecordingSpan, SpanContext, SpanKind, TraceFlags + + +_QUERY = "SELECT current_setting('application_name'), pg_sleep(0.01)" +_DEFAULT_APPLICATION_NAME = "startup-app" + + +def _traceparent(span) -> str: + context = span.get_span_context() + return f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags) & 0xFF:02x}" + + +def _assert_database_span(exporter, application_name, parent_context=None): + matches = [span for span in exporter.get_finished_spans() if _traceparent(span) == application_name] + assert len(matches) == 1, f"PostgreSQL value {application_name!r} did not identify one exported span" + span = matches[0] + assert span.kind == SpanKind.CLIENT + if parent_context is not None: + assert span.parent.span_id == parent_context.span_id + assert span.context.span_id != parent_context.span_id + return span + + +def _sync_queries(driver, exporter, tracer): + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + for prepared in [False, True, True] if driver.__name__ == "psycopg" else [False]: + exporter.clear() + with tracer.start_as_current_span("caller") as parent, connection.cursor() as cursor: + cursor.execute(_QUERY, **({"prepare": True} if prepared else {})) + application_name = cursor.fetchone()[0] + _assert_database_span(exporter, application_name, parent.get_span_context()) + assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 1 + + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("CREATE TEMP TABLE correlation (application_name text)") + exporter.clear() + with connection.cursor() as cursor: + cursor.executemany("INSERT INTO correlation VALUES (current_setting('application_name'))", [(), ()]) + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SELECT application_name FROM correlation") + values = [row[0] for row in cursor.fetchall()] + assert len(values) == 2 + assert values[0] == values[1] + application_name = values[0] + _assert_database_span(exporter, values[0]) + assert len(exporter.get_finished_spans()) == 1 + + if driver.__name__ == "psycopg2": + exporter.clear() + with connection.cursor(name="correlation_cursor") as cursor: + cursor.execute(_QUERY) + application_name = cursor.fetchone()[0] + _assert_database_span(exporter, application_name) + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute( + "CREATE FUNCTION pg_temp.correlation_name() RETURNS text LANGUAGE sql " + "AS $$ SELECT current_setting('application_name') $$" + ) + exporter.clear() + with connection.cursor() as cursor: + cursor.callproc("pg_temp.correlation_name") + application_name = cursor.fetchone()[0] + _assert_database_span(exporter, application_name) + + before = len(exporter.get_finished_spans()) + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute(_QUERY) + assert cursor.fetchone()[0] == application_name + assert len(exporter.get_finished_spans()) == before + + +def _sync_concurrency(driver, exporter, tracer): + with driver.connect() as connection: + connection.autocommit = True + exporter.clear() + + def query(_index): + with tracer.start_as_current_span("concurrent caller") as parent, connection.cursor() as cursor: + cursor.execute(_QUERY) + return cursor.fetchone()[0], parent.get_span_context() + + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(query, range(12))) + for application_name, parent_context in results: + _assert_database_span(exporter, application_name, parent_context) + assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 + + +async def _async_queries(exporter, tracer): + import psycopg + + async with await psycopg.AsyncConnection.connect() as connection: + exporter.clear() + async with connection.cursor() as cursor: + await cursor.execute(_QUERY, prepare=True) + application_name = (await cursor.fetchone())[0] + _assert_database_span(exporter, application_name) + assert len(exporter.get_finished_spans()) == 1 + + with suppress_instrumentation(): + async with connection.cursor() as cursor: + await cursor.execute("CREATE TEMP TABLE correlation (application_name text)") + exporter.clear() + async with connection.cursor() as cursor: + await cursor.executemany("INSERT INTO correlation VALUES (current_setting('application_name'))", [(), ()]) + with suppress_instrumentation(): + async with connection.cursor() as cursor: + await cursor.execute("SELECT application_name FROM correlation") + values = [row[0] for row in await cursor.fetchall()] + assert len(values) == 2 + assert values[0] == values[1] + _assert_database_span(exporter, values[0]) + assert len(exporter.get_finished_spans()) == 1 + + async with await psycopg.AsyncConnection.connect(autocommit=True) as connection: + exporter.clear() + + async def query(): + with tracer.start_as_current_span("async caller") as parent: + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + return (await cursor.fetchone())[0], parent.get_span_context() + + results = await asyncio.wait_for(asyncio.gather(*(query() for _ in range(12))), timeout=10) + for application_name, parent_context in results: + _assert_database_span(exporter, application_name, parent_context) + assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 + + with suppress_instrumentation(): + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + assert (await cursor.fetchone())[0] == results[-1][0] + assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 + + class RejectingSetCursor(psycopg.AsyncCursor): + async def execute(self, query, *args, **kwargs): + if query.startswith("SET application_name"): + query = "SET nonexistent_correlation_parameter = 'failure'" + return await super().execute(query, *args, **kwargs) + + connection = await psycopg.AsyncConnection.connect(cursor_factory=RejectingSetCursor) + try: + async with connection.cursor() as cursor: + failed = False + try: + await cursor.execute("SELECT 1") + except psycopg.errors.InFailedSqlTransaction: + failed = True + assert failed + assert connection.info.transaction_status == psycopg.pq.TransactionStatus.INERROR + finally: + await connection.close() + + +def _unsampled_queries(driver, exporter): + with driver.connect() as connection: + for flags in (0, 2): + parent = SpanContext( + trace_id=0x11111111111111111111111111111111, + span_id=0x1111111111111111, + is_remote=True, + trace_flags=TraceFlags(flags), + ) + exporter.clear() + with trace.use_span(NonRecordingSpan(parent)), connection.cursor() as cursor: + cursor.execute(_QUERY) + value = cursor.fetchone()[0] + fields = value.split("-") + assert fields[1] == f"{parent.trace_id:032x}" + assert int(fields[2], 16) not in (0, parent.span_id) + assert int(fields[3], 16) == flags + assert not exporter.get_finished_spans() + + +def _reset_to_default(driver): + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + from opentelemetry.trace import NoOpTracerProvider + + instrumentor = PsycopgInstrumentor if driver.__name__ == "psycopg" else Psycopg2Instrumentor + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + with connection.cursor() as cursor: + cursor.execute(_QUERY) + assert cursor.fetchone()[0].startswith("00-") + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SET application_name = 'sql-app'") + instrumentor.uninstrument_connection(connection) + instrumentor.instrument_connection(connection, tracer_provider=NoOpTracerProvider()) + with connection.cursor() as cursor: + cursor.execute(_QUERY) + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +def _failed_internal_command(driver): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + transaction_error = driver.errors.InFailedSqlTransaction + + class RejectingSetCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query.startswith("SET application_name"): + query = "SET nonexistent_correlation_parameter = 'failure'" + return super().execute(query, *args, **kwargs) + + connection = driver.connect(cursor_factory=RejectingSetCursor) + try: + with connection.cursor() as cursor: + try: + cursor.execute("SELECT 1") + except transaction_error: + pass + else: + message = "Customer query should report the aborted transaction after failed SET" + raise AssertionError(message) + status = ( + connection.info.transaction_status if driver.__name__ == "psycopg" else connection.get_transaction_status() + ) + assert status == 3 + finally: + connection.close() + + connection = driver.connect(cursor_factory=RejectingSetCursor) + try: + connection.autocommit = True + with connection.cursor() as cursor: + cursor.execute("SELECT 1") + assert cursor.fetchone()[0] == 1 + finally: + connection.close() + + +def _custom_cursor(exporter): + import psycopg + import psycopg2 + from psycopg2.extras import RealDictCursor + + for driver, cursor_class in ((psycopg, psycopg.ClientCursor), (psycopg2, RealDictCursor)): + with driver.connect(cursor_factory=cursor_class) as connection: + exporter.clear() + with connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name') AS application_name") + row = cursor.fetchone() + application_name = row["application_name"] if isinstance(row, dict) else row[0] + _assert_database_span(exporter, application_name) + + +def _psycopg_server_cursor_boundary(exporter): + import psycopg + + with psycopg.connect() as connection: + with connection.cursor() as cursor: + cursor.execute(_QUERY) + earlier_context = cursor.fetchone()[0] + exporter.clear() + with connection.cursor(name="raw_server_cursor") as cursor: + cursor.execute(_QUERY) + assert cursor.fetchone()[0] == earlier_context + assert not exporter.get_finished_spans() + + +def main(): + from opentelemetry.instrumentation.auto_instrumentation import initialize + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from splunk_otel import init_splunk_otel + from splunk_otel.postgresql import enable_postgresql_trace_context_propagation + + mode = os.environ.get("POSTGRESQL_SMOKE_INSTRUMENTATION", "auto") + if mode == "manual": + init_splunk_otel() + PsycopgInstrumentor().instrument() + Psycopg2Instrumentor().instrument() + else: + initialize(swallow_exceptions=False) + + exporter = InMemorySpanExporter() + trace.get_tracer_provider().add_span_processor(SimpleSpanProcessor(exporter)) + tracer = trace.get_tracer(__name__) + + import psycopg + import psycopg2 + + if mode == "disabled": + for driver in (psycopg, psycopg2): + with ( + driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, + connection.cursor() as cursor, + ): + cursor.execute(_QUERY) + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + else: + assert enable_postgresql_trace_context_propagation() + assert enable_postgresql_trace_context_propagation() + for driver in (psycopg, psycopg2): + _sync_queries(driver, exporter, tracer) + _sync_concurrency(driver, exporter, tracer) + _unsampled_queries(driver, exporter) + _reset_to_default(driver) + _failed_internal_command(driver) + asyncio.run(_async_queries(exporter, tracer)) + _custom_cursor(exporter) + _psycopg_server_cursor_boundary(exporter) + + print( # noqa: T201 + f"POSTGRESQL_CORRELATION_OK mode={mode} psycopg={version('psycopg')} " + f"psycopg2={version('psycopg2-binary')} instrumentation={version('opentelemetry-instrumentation')}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_configurator.py b/tests/test_configurator.py index cf11db99..089f1ca9 100644 --- a/tests/test_configurator.py +++ b/tests/test_configurator.py @@ -26,6 +26,7 @@ _configure_declarative_sdk, _create_declarative_resource, _declarative_server_timing_enabled, + SplunkConfigurator, ) @@ -87,3 +88,29 @@ def test_declarative_server_timing_reads_splunk_http_configuration(): ) assert not _declarative_server_timing_enabled(config) + + +@pytest.mark.parametrize("use_config_file", [False, True]) +def test_configurator_activates_postgresql_only_without_config_file(monkeypatch, use_config_file): + enable_calls = [] + monkeypatch.setenv("SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") + if use_config_file: + monkeypatch.setenv("OTEL_CONFIG_FILE", "config.yaml") + else: + monkeypatch.delenv("OTEL_CONFIG_FILE", raising=False) + monkeypatch.setattr("opentelemetry.sdk._configuration._OTelSDKConfigurator._configure", lambda *_a, **_k: None) + monkeypatch.setattr( + "splunk_otel.configurator._configure_declarative_sdk", + lambda *_a, **_k: OpenTelemetryConfiguration(file_format="1.0"), + ) + monkeypatch.setattr("splunk_otel.configurator.configure_logging_instrumentation", lambda _env: None) + monkeypatch.setattr("splunk_otel.configurator._start_profiling_if_enabled", lambda: None) + monkeypatch.setattr("splunk_otel.configurator._configure_callgraphs_if_enabled", lambda: None) + monkeypatch.setattr( + "splunk_otel.runtime.enable_postgresql_trace_context_propagation", + lambda: enable_calls.append(True), + ) + + SplunkConfigurator().configure() + + assert enable_calls == ([] if use_config_file else [True]) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py new file mode 100644 index 00000000..986329e1 --- /dev/null +++ b/tests/test_postgresql.py @@ -0,0 +1,604 @@ +# 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 asyncio +import threading +from types import SimpleNamespace + +import pytest + +from opentelemetry import trace +from opentelemetry.instrumentation.utils import suppress_instrumentation +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + +from splunk_otel.postgresql import ( + _execute_with_postgresql_trace_context, + _execute_with_postgresql_trace_context_async, +) +from splunk_otel import postgresql + + +PARENT_CONTEXT = SpanContext( + trace_id=0x11111111111111111111111111111111, + span_id=0x1111111111111111, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), +) +DATABASE_CONTEXT = SpanContext( + trace_id=PARENT_CONTEXT.trace_id, + span_id=0x2222222222222222, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), +) +SECOND_DATABASE_CONTEXT = SpanContext( + trace_id=PARENT_CONTEXT.trace_id, + span_id=0x3333333333333333, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), +) +TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01" +SECOND_TRACEPARENT = "00-11111111111111111111111111111111-3333333333333333-01" + + +class FakeConnection: + def __init__(self): + self.default_application_name = "startup-app" + self.application_name = "customer-app" + self.calls = [] + self.cursors = [] + self.before_execute = lambda _query: None + + def cursor(self): + cursor = FakePsycopgCursor(self) + self.cursors.append(cursor) + return cursor + + +class FakeAsyncConnection(FakeConnection): + def cursor(self): + cursor = FakeAsyncPsycopgCursor(self) + self.cursors.append(cursor) + return cursor + + +class FakePsycopgCursor: + __module__ = "psycopg" + + def __init__(self, connection): + self.connection = connection + self.calls = [] + self.closed = False + + def execute(self, query): + self.calls.append(query) + self.connection.calls.append(query) + self.connection.before_execute(query) + if query == "SET application_name = DEFAULT": + self.connection.application_name = self.connection.default_application_name + elif query.startswith("SET application_name = '"): + self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") + return self.connection.application_name + + def close(self): + self.closed = True + + +class FakeAsyncPsycopgCursor: + __module__ = "psycopg.fake" + + def __init__(self, connection): + self.connection = connection + self.calls = [] + self.closed = False + + async def execute(self, query): + self.calls.append(query) + self.connection.calls.append(query) + if query == "SET application_name = DEFAULT": + self.connection.application_name = self.connection.default_application_name + elif query.startswith("SET application_name = '"): + self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") + return self.connection.application_name + + async def close(self): + self.closed = True + + +def _fake_traced_execution(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(DATABASE_CONTEXT)): + return query_method(*args, **kwargs) + + +async def _fake_traced_execution_async(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(DATABASE_CONTEXT)): + return await query_method(*args, **kwargs) + + +def _fake_untraced_execution(_cursor, query_method, *args, **kwargs): + return query_method(*args, **kwargs) + + +def test_query_sets_database_span_traceparent(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + with trace.use_span(NonRecordingSpan(PARENT_CONTEXT)): + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT application_name"), + {}, + ) + + assert result == TRACEPARENT + assert cursor.calls == ["SELECT application_name"] + assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] + assert connection.cursors[0].closed + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT application_name", + ] + + +def test_query_without_database_span_clears_stale_traceparent(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + result = _execute_with_postgresql_trace_context( + _fake_untraced_execution, + None, + (cursor, cursor.execute, "SELECT 2"), + {}, + ) + + assert result == "startup-app" + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT 1", + "SET application_name = DEFAULT", + "SELECT 2", + ] + + +def test_non_postgresql_cursor_is_unchanged(): + class OtherCursor: + pass + + cursor = OtherCursor() + calls = [] + + def wrapped(*args, **kwargs): + calls.append((args, kwargs)) + return "result" + + assert _execute_with_postgresql_trace_context(wrapped, None, (cursor, "query", "SELECT 1"), {}) == "result" + assert calls == [((cursor, "query", "SELECT 1"), {})] + + +def test_suppressed_instrumentation_is_unchanged(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + with suppress_instrumentation(): + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert result == "customer-app" + assert cursor.calls == ["SELECT 1"] + assert connection.cursors == [] + + +def test_psycopg2_asynchronous_connection_sends_only_customer_query(): + class FakePsycopg2Cursor(FakePsycopgCursor): + __module__ = "psycopg2.extensions" + + connection = FakeConnection() + connection.async_ = 1 + cursor = FakePsycopg2Cursor(connection) + + _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert connection.calls == ["SELECT 1"] + assert connection.cursors == [] + + +def test_psycopg2_named_cursor_is_not_used_for_application_name(): + class FakePsycopg2NamedCursor(FakePsycopgCursor): + __module__ = "psycopg2.extensions" + + def execute(self, query): + if self.calls: + raise RuntimeError + return super().execute(query) + + connection = FakeConnection() + cursor = FakePsycopg2NamedCursor(connection) + + with trace.use_span(NonRecordingSpan(PARENT_CONTEXT)): + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert result == TRACEPARENT + assert cursor.calls == ["SELECT 1"] + assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] + assert connection.cursors[0].closed + + +def test_async_query_sets_database_span_traceparent(): + async def run(): + connection = FakeAsyncConnection() + cursor = FakeAsyncPsycopgCursor(connection) + + with trace.use_span(NonRecordingSpan(PARENT_CONTEXT)): + result = await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, + None, + (cursor, cursor.execute, "SELECT application_name"), + {}, + ) + + return result, connection, cursor + + result, connection, cursor = asyncio.run(run()) + + assert result == TRACEPARENT + assert cursor.calls == ["SELECT application_name"] + assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] + assert connection.cursors[0].closed + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT application_name", + ] + + +def test_shared_connection_keeps_set_and_query_together(monkeypatch): + connection = FakeConnection() + first_cursor = FakePsycopgCursor(connection) + second_cursor = FakePsycopgCursor(connection) + first_query_started = threading.Event() + release_first_query = threading.Event() + second_query_waiting = threading.Event() + + class FakeConnectionLock: + def __init__(self): + self.lock = threading.Lock() + + def __enter__(self): + if self.lock.locked(): + second_query_waiting.set() + self.lock.acquire() + + def __exit__(self, *_args): + self.lock.release() + + connection_lock = FakeConnectionLock() + monkeypatch.setattr(postgresql, "_get_sync_connection_lock", lambda _connection: connection_lock) + + def before_execute(query): + if query == "SELECT first": + first_query_started.set() + release_first_query.wait(timeout=1) + + connection.before_execute = before_execute + + def run(cursor, query, span_context): + def wrapped(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(span_context)): + return query_method(*args, **kwargs) + + _execute_with_postgresql_trace_context(wrapped, None, (cursor, cursor.execute, query), {}) + + first_thread = threading.Thread(target=run, args=(first_cursor, "SELECT first", DATABASE_CONTEXT)) + second_thread = threading.Thread(target=run, args=(second_cursor, "SELECT second", SECOND_DATABASE_CONTEXT)) + + try: + first_thread.start() + assert first_query_started.wait(timeout=1) + second_thread.start() + assert second_query_waiting.wait(timeout=1) + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT first", + ] + finally: + release_first_query.set() + first_thread.join(timeout=1) + second_thread.join(timeout=1) + + assert not first_thread.is_alive() + assert not second_thread.is_alive() + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT first", + f"SET application_name = '{SECOND_TRACEPARENT}'", + "SELECT second", + ] + + +class FakeCursorTracer: + def traced_execution(self, cursor, query_method, *args, **kwargs): + return _fake_traced_execution(cursor, query_method, *args, **kwargs) + + async def traced_execution_async(self, cursor, query_method, *args, **kwargs): + return await _fake_traced_execution_async(cursor, query_method, *args, **kwargs) + + +def test_public_enable_is_idempotent(monkeypatch): + class CursorTracer(FakeCursorTracer): + pass + + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", lambda _name: SimpleNamespace(CursorTracer=CursorTracer)) + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + assert postgresql.enable_postgresql_trace_context_propagation() + assert postgresql.enable_postgresql_trace_context_propagation() + result = CursorTracer().traced_execution(cursor, cursor.execute, "SELECT 1") + + assert result == TRACEPARENT + assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT 1"] + + +def test_public_enable_without_dbapi_reports_missing_dependency(monkeypatch, caplog): + def _fake_dbapi_import(_name): + raise ImportError + + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", _fake_dbapi_import) + + assert not postgresql.enable_postgresql_trace_context_propagation() + assert "requires DBAPI instrumentation" in caplog.text + + +def test_public_enable_rejects_incompatible_dbapi_before_wrapping(monkeypatch, caplog): + class CursorTracer: + def traced_execution(self, cursor, query_method, *args, **kwargs): + return query_method(*args, **kwargs) + + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", lambda _name: SimpleNamespace(CursorTracer=CursorTracer)) + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + assert not postgresql.enable_postgresql_trace_context_propagation() + CursorTracer().traced_execution(cursor, cursor.execute, "SELECT 1") + + assert connection.calls == ["SELECT 1"] + assert "compatible DBAPI execution methods" in caplog.text + + +def test_public_enable_rejects_keyword_only_execution_methods(monkeypatch): + class CursorTracer: + def traced_execution(self, *, cursor, query_method): + return query_method(cursor) + + async def traced_execution_async(self, *, cursor, query_method): + return await query_method(cursor) + + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", lambda _name: SimpleNamespace(CursorTracer=CursorTracer)) + + assert not postgresql.enable_postgresql_trace_context_propagation() + + +@pytest.mark.parametrize("flags", [0, 1, 2, 3]) +def test_query_propagates_valid_context_flags(flags): + context = SpanContext( + trace_id=DATABASE_CONTEXT.trace_id, + span_id=DATABASE_CONTEXT.span_id, + is_remote=False, + trace_flags=TraceFlags(flags), + ) + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + def _fake_traced_execution_with_flags(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(context)): + return query_method(*args, **kwargs) + + result = _execute_with_postgresql_trace_context( + _fake_traced_execution_with_flags, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert result == f"00-{context.trace_id:032x}-{context.span_id:016x}-{flags:02x}" + + +def test_failed_internal_command_attempts_customer_query_and_closes_cursor(caplog): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + def _fake_internal_command_failure(query): + if query.startswith("SET "): + message = "SET rejected" + raise RuntimeError(message) + + connection.before_execute = _fake_internal_command_failure + + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert result == "customer-app" + assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT 1"] + assert connection.cursors[0].closed + assert "Failed to set PostgreSQL application_name" in caplog.text + + +def test_customer_query_exception_is_returned_to_caller(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + error = RuntimeError("customer query failed") + + def _fake_customer_query_failure(query): + if query == "SELECT 1": + raise error + + connection.before_execute = _fake_customer_query_failure + + with pytest.raises(RuntimeError) as raised: + _execute_with_postgresql_trace_context( + _fake_traced_execution, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + + assert raised.value is error + assert connection.cursors[0].closed + + +def test_async_shared_connection_keeps_set_and_query_together(): + async def run(): + connection = FakeAsyncConnection() + first_query_started = asyncio.Event() + second_query_started = asyncio.Event() + release_first_query = asyncio.Event() + + class BlockingQueryCursor(FakeAsyncPsycopgCursor): + async def execute(self, query): + result = await super().execute(query) + first_query_started.set() + await release_first_query.wait() + return result + + async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(SECOND_DATABASE_CONTEXT)): + second_query_started.set() + return await query_method(*args, **kwargs) + + first_cursor = BlockingQueryCursor(connection) + second_cursor = FakeAsyncPsycopgCursor(connection) + first = asyncio.create_task( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, + None, + (first_cursor, first_cursor.execute, "SELECT first"), + {}, + ) + ) + second = None + try: + await asyncio.wait_for(first_query_started.wait(), timeout=1) + second = asyncio.create_task( + _execute_with_postgresql_trace_context_async( + _fake_second_traced_execution, + None, + (second_cursor, second_cursor.execute, "SELECT second"), + {}, + ) + ) + await asyncio.wait_for(second_query_started.wait(), timeout=1) + assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT first"] + finally: + release_first_query.set() + await asyncio.wait_for(asyncio.gather(*([first, second] if second is not None else [first])), timeout=1) + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT first", + f"SET application_name = '{SECOND_TRACEPARENT}'", + "SELECT second", + ] + + asyncio.run(run()) + + +def test_async_suppression_sends_only_customer_query(): + async def run(): + connection = FakeAsyncConnection() + cursor = FakeAsyncPsycopgCursor(connection) + with suppress_instrumentation(): + await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, + None, + (cursor, cursor.execute, "SELECT 1"), + {}, + ) + assert connection.calls == ["SELECT 1"] + assert connection.cursors == [] + + asyncio.run(run()) + + +def test_async_cancellation_closes_internal_cursor_and_releases_connection(): + async def run(): + internal_command_started = asyncio.Event() + + class BlockingInternalCursor(FakeAsyncPsycopgCursor): + async def execute(self, _query): + internal_command_started.set() + await asyncio.Event().wait() + + class Connection(FakeAsyncConnection): + def cursor(self): + if not self.cursors: + cursor = BlockingInternalCursor(self) + self.cursors.append(cursor) + return cursor + return super().cursor() + + connection = Connection() + customer_cursor = FakeAsyncPsycopgCursor(connection) + query = asyncio.create_task( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, + None, + (customer_cursor, customer_cursor.execute, "SELECT 1"), + {}, + ) + ) + await asyncio.wait_for(internal_command_started.wait(), timeout=1) + query.cancel() + with pytest.raises(asyncio.CancelledError): + await query + assert connection.cursors[0].closed + assert customer_cursor.calls == [] + + result = await asyncio.wait_for( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, + None, + (customer_cursor, customer_cursor.execute, "SELECT 2"), + {}, + ), + timeout=1, + ) + assert result == TRACEPARENT + + asyncio.run(run()) diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 234d6662..867c735e 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -15,7 +15,7 @@ import pytest from splunk_otel.env import Env -from splunk_otel.runtime import configure_logging_instrumentation +from splunk_otel.runtime import configure_logging_instrumentation, configure_postgresql_trace_context_propagation class _FakeLoggingInstrumentor: @@ -45,6 +45,20 @@ def test_logging_instrumentation_is_enabled_when_not_disabled(): assert logging_instrumentor.instrument_calls == 1 +@pytest.mark.parametrize(("value", "expected_calls"), [(None, []), ("false", []), ("true", [True])]) +def test_postgresql_trace_context_propagation_requires_opt_in(monkeypatch, value, expected_calls): + enable_calls = [] + monkeypatch.setattr( + "splunk_otel.runtime.enable_postgresql_trace_context_propagation", + lambda: enable_calls.append(True), + ) + env_store = {} if value is None else {"SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED": value} + + configure_postgresql_trace_context_propagation(Env(env_store)) + + assert enable_calls == expected_calls + + def _configure_logging(env_store): logging_instrumentor = _FakeLoggingInstrumentor() configure_logging_instrumentation(Env(env_store), logging_instrumentor) From efb4b7f199652b07ed5f32d35df21405dfc9a866 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Wed, 14 Oct 2026 16:53:48 -0400 Subject: [PATCH 02/10] Reset PostgreSQL trace context after query execution Keep SET, query execution, and reset under the connection lock. Skip named cursors, cover cleanup failures and cancellation, and move setup and limits into dedicated documentation. Assisted-by: Codex (GPT-6) --- README.md | 42 +----- docs/postgresql.md | 96 +++++++++++++ src/splunk_otel/postgresql.py | 25 +++- tests/smoke/postgresql_trace_context.py | 113 ++++++++++++++- tests/test_postgresql.py | 184 +++++++++++++++++++++--- 5 files changed, 392 insertions(+), 68 deletions(-) create mode 100644 docs/postgresql.md diff --git a/README.md b/README.md index 32fd8cd9..2f5ea863 100644 --- a/README.md +++ b/README.md @@ -51,41 +51,6 @@ Then run your application with `opentelemetry-instrument`: OTEL_SERVICE_NAME=my-python-service opentelemetry-instrument python app.py ``` -### Experimental PostgreSQL trace correlation - -Set `SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true` to place each -instrumented psycopg or psycopg2 database span's W3C `traceparent` in -PostgreSQL's `application_name` immediately before the query. This sends one -additional `SET application_name` command per traced query. -The feature targets Psycopg 3 synchronous and asynchronous connections and -Psycopg2 synchronous connections. Psycopg2 asynchronous connections pass through -without trace propagation. - -This opt-in replaces the connection's `application_name` and leaves the value -set after the query. If instrumentation creates no valid new database span, the -feature resets the setting to the connection's configured default. It does not -restore a value that the application previously set through SQL. Queries that -bypass instrumentation, including suppressed queries and cursors created before -instrumentation, can inherit an earlier query's trace context. -Psycopg 3 named cursors use a separate server-cursor factory that the current -upstream instrumentor does not wrap. Those queries, cursor fetch operations, -and COPY operations are outside this feature's automatic coverage. - -The feature logs internal command failures and attempts the customer query. A -failed command can leave a transaction aborted, so the customer query can also -fail. The feature does not roll back application transactions. Database span -duration and database operation duration metrics include lock waiting and the -additional command. - -Manual instrumentation must initialize the Splunk distro first with -`init_splunk_otel()`, or call `enable_postgresql_trace_context_propagation()` -before using `PsycopgInstrumentor` or `Psycopg2Instrumentor`. - -Install the corresponding OpenTelemetry driver instrumentation package. The -environment switch applies when `OTEL_CONFIG_FILE` is unset. Activation lasts -for the process lifetime. To turn it off, unset the switch and restart. Only -queries using the instrumented driver participate in its connection lock. - By default, OTLP exporters send telemetry to a locally running [Splunk OpenTelemetry Collector](https://github.com/signalfx/splunk-otel-collector) at `http://localhost:4317`. To send telemetry directly to Splunk Observability @@ -93,6 +58,13 @@ Cloud instead, set `SPLUNK_REALM` and `SPLUNK_ACCESS_TOKEN` before starting your application. `SPLUNK_REALM` configures direct ingest endpoints for traces and metrics. +### Experimental PostgreSQL trace correlation + +Set `SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true` to link PostgreSQL +query samples to APM traces and database client spans. This feature is disabled +by default. See [PostgreSQL trace correlation](docs/postgresql.md) for setup, +supported drivers, and limitations. + ## Declarative configuration The distribution includes experimental support for the OpenTelemetry declarative diff --git a/docs/postgresql.md b/docs/postgresql.md new file mode 100644 index 00000000..9da00de5 --- /dev/null +++ b/docs/postgresql.md @@ -0,0 +1,96 @@ +# PostgreSQL trace correlation + +This experimental feature connects sampled PostgreSQL queries in Splunk Database +Monitoring to the APM trace and database client span that executed them. It places +the database client span's W3C `traceparent` in PostgreSQL's `application_name`. +It sends only `traceparent`, without baggage or `tracestate`. Valid unsampled +span contexts also work. + +## Database Monitoring prerequisites + +Configure a Collector with the PostgreSQL receiver, database monitoring +permissions, and both `db.server.query_sample` and `db.server.top_query` enabled. +Export database metrics and query events to Splunk Database Monitoring. Query +events use `/v3/event` with the `X-splunk-instrumentation-library: dbmon` header. +See the [PostgreSQL receiver setup instructions](https://help.splunk.com/en/splunk-observability-cloud/monitor-databases/get-data-in/configure-receivers/postgresql-receiver). + +The receiver must support extracting trace and span IDs from `application_name`. +End-to-end correlation was validated with Splunk Collector 0.162.0. The minimum +compatible Collector version has not yet been established for this feature. +The application must also export traces to APM in the same organization; those +traces can use a separate export path. A query sample can link to a trace only +when that trace is available in APM. + +## Enablement + +The feature is disabled by default. In the same Python environment as your +application, install `splunk-opentelemetry` and your application's PostgreSQL +driver first. Then install instrumentation for the installed libraries: + +```sh +opentelemetry-bootstrap -a install +``` + +This command selects instrumentation versions that match the installed distro. +For PostgreSQL, it installs: + +| Application driver | Instrumentation package | +|---|---| +| `psycopg` (Psycopg 3) | [opentelemetry-instrumentation-psycopg](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg/psycopg.html) | +| `psycopg2` or `psycopg2-binary` | [opentelemetry-instrumentation-psycopg2](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html) | + +Start the application with: + +```sh +SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true \ +opentelemetry-instrument python app.py +``` + +The environment switch applies when `OTEL_CONFIG_FILE` is unset. Activation lasts +for the process lifetime. To turn it off, unset the switch and restart. + +Manual instrumentation must initialize the Splunk distro first with +`init_splunk_otel()`, or call +`splunk_otel.postgresql.enable_postgresql_trace_context_propagation()` before using +`PsycopgInstrumentor` or `Psycopg2Instrumentor`. + +## Driver coverage + +The feature targets Psycopg 3 synchronous and asynchronous connections and +Psycopg2 synchronous connections. Psycopg2 asynchronous connections pass through +without trace propagation. + +Named/server cursors, cursor fetch operations, and COPY operations are outside +this feature's automatic coverage. Named cursors can execute SQL later during +fetch, after the execute span has ended, so this feature skips propagation for +them in both drivers. + +## Connection behavior + +The feature uses a separate cursor on the same connection for internal commands +and suppresses instrumentation for those commands. A per-connection lock covers +SET, the customer query, and cleanup. Each traced query sends two additional +commands: SET before execution and SET DEFAULT afterward. + +This opt-in owns the connection's `application_name` during execution. Cleanup +resets the setting to the connection's configured default before releasing the +instrumentation lock, including when the query raises an exception. If +instrumentation creates no valid new database span, the feature sends SET DEFAULT +before execution. It does not restore a value that the application previously +set through SQL. + +The lock only coordinates wrapped queries. Raw cursors and suppressed queries +can bypass it and observe another query's context during execution. Failed +cleanup can also leave trace context set until the application recovers the +connection or rolls back its transaction. + +## Failure behavior and query duration + +The feature logs internal command failures and attempts the customer query. A +failed setup command can leave a transaction aborted, so the customer query can +also fail. Cleanup failures do not replace the customer query's result or error; +they can still leave an otherwise successful transaction aborted. The feature +does not roll back application transactions. + +The extra commands add latency. Database span duration and database operation +duration metrics include lock waiting and the additional commands. diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index 61e19f15..c7039547 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -94,7 +94,7 @@ def _execute_with_postgresql_trace_context(traced_execution, _cursor_tracer, arg return traced_execution(*args, **kwargs) cursor, query_method, *query_args = args - if not _is_postgresql_cursor(cursor) or getattr(cursor.connection, "async_", False): + if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return traced_execution(*args, **kwargs) span_context_before = trace.get_current_span().get_span_context() @@ -116,7 +116,7 @@ async def _execute_with_postgresql_trace_context_async(traced_execution_async, _ return await traced_execution_async(*args, **kwargs) cursor, query_method, *query_args = args - if not _is_postgresql_cursor(cursor) or getattr(cursor.connection, "async_", False): + if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return await traced_execution_async(*args, **kwargs) span_context_before = trace.get_current_span().get_span_context() @@ -143,8 +143,12 @@ def _execute_correlated( lock = _get_sync_connection_lock(cursor.connection) with lock: traceparent = _current_database_traceparent(span_context_before) - _update_application_name(cursor.connection, traceparent) - return query_method(*args, **kwargs) + try: + _update_application_name(cursor.connection, traceparent) + return query_method(*args, **kwargs) + finally: + if traceparent is not None: + _update_application_name(cursor.connection, None) async def _execute_correlated_async( @@ -157,8 +161,12 @@ async def _execute_correlated_async( lock = _get_async_connection_lock(cursor.connection) async with lock: traceparent = _current_database_traceparent(span_context_before) - await _update_application_name_async(cursor.connection, traceparent) - return await query_method(*args, **kwargs) + try: + await _update_application_name_async(cursor.connection, traceparent) + return await query_method(*args, **kwargs) + finally: + if traceparent is not None: + await _update_application_name_async(cursor.connection, None) def _get_sync_connection_lock(connection) -> threading.Lock: @@ -226,7 +234,10 @@ def _application_name_statement(traceparent: str | None) -> str: return f"SET application_name = '{traceparent}'" -def _is_postgresql_cursor(cursor) -> bool: +def _is_postgresql_client_cursor(cursor) -> bool: + # Named cursors execute SQL during later fetches, outside this query's span. + if getattr(cursor, "name", None): + return False return any(_is_postgresql_module(cursor_class.__module__) for cursor_class in type(cursor).__mro__) diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py index 34c702c3..799bac50 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/smoke/postgresql_trace_context.py @@ -72,7 +72,7 @@ def _sync_queries(driver, exporter, tracer): with connection.cursor(name="correlation_cursor") as cursor: cursor.execute(_QUERY) application_name = cursor.fetchone()[0] - _assert_database_span(exporter, application_name) + assert application_name == _DEFAULT_APPLICATION_NAME with suppress_instrumentation(), connection.cursor() as cursor: cursor.execute( "CREATE FUNCTION pg_temp.correlation_name() RETURNS text LANGUAGE sql " @@ -87,7 +87,7 @@ def _sync_queries(driver, exporter, tracer): before = len(exporter.get_finished_spans()) with suppress_instrumentation(), connection.cursor() as cursor: cursor.execute(_QUERY) - assert cursor.fetchone()[0] == application_name + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME assert len(exporter.get_finished_spans()) == before @@ -111,7 +111,7 @@ def query(_index): async def _async_queries(exporter, tracer): import psycopg - async with await psycopg.AsyncConnection.connect() as connection: + async with await psycopg.AsyncConnection.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: exporter.clear() async with connection.cursor() as cursor: await cursor.execute(_QUERY, prepare=True) @@ -134,7 +134,9 @@ async def _async_queries(exporter, tracer): _assert_database_span(exporter, values[0]) assert len(exporter.get_finished_spans()) == 1 - async with await psycopg.AsyncConnection.connect(autocommit=True) as connection: + async with await psycopg.AsyncConnection.connect( + autocommit=True, application_name=_DEFAULT_APPLICATION_NAME + ) as connection: exporter.clear() async def query(): @@ -151,7 +153,7 @@ async def query(): with suppress_instrumentation(): async with connection.cursor() as cursor: await cursor.execute(_QUERY) - assert (await cursor.fetchone())[0] == results[-1][0] + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 class RejectingSetCursor(psycopg.AsyncCursor): @@ -265,20 +267,114 @@ def _custom_cursor(exporter): _assert_database_span(exporter, application_name) +def _failed_cleanup_command(driver): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + + class RejectingCleanupCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query == "SET application_name = DEFAULT": + query = "SET nonexistent_correlation_parameter = 'failure'" + return super().execute(query, *args, **kwargs) + + connection = driver.connect(cursor_factory=RejectingCleanupCursor) + try: + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("CREATE TEMP TABLE cleanup_failure (value int)") + connection.commit() + with connection.cursor() as cursor: + cursor.execute("INSERT INTO cleanup_failure VALUES (1) RETURNING value") + assert cursor.fetchone()[0] == 1 + status = ( + connection.info.transaction_status if driver.__name__ == "psycopg" else connection.get_transaction_status() + ) + assert status == 3 + connection.rollback() + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SELECT count(*) FROM cleanup_failure") + assert cursor.fetchone()[0] == 0 + finally: + connection.close() + + def _psycopg_server_cursor_boundary(exporter): import psycopg - with psycopg.connect() as connection: + with psycopg.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: with connection.cursor() as cursor: cursor.execute(_QUERY) earlier_context = cursor.fetchone()[0] exporter.clear() with connection.cursor(name="raw_server_cursor") as cursor: cursor.execute(_QUERY) - assert cursor.fetchone()[0] == earlier_context + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + assert earlier_context != _DEFAULT_APPLICATION_NAME assert not exporter.get_finished_spans() +def _psycopg2_commit_boundary(exporter): + import psycopg2 + + connection = psycopg2.connect(application_name=_DEFAULT_APPLICATION_NAME) + try: + exporter.clear() + with connection, connection.cursor() as cursor: + cursor.execute(_QUERY) + _assert_database_span(exporter, cursor.fetchone()[0]) + assert connection.get_parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + + assert not connection.closed + assert connection.get_parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + connection.rollback() + assert connection.get_parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + finally: + connection.close() + + +async def _async_query_cancellation(): + import psycopg + + async with await psycopg.AsyncConnection.connect(autocommit=True) as observer: + for autocommit in (False, True): + async with ( + await psycopg.AsyncConnection.connect( + autocommit=autocommit, application_name=_DEFAULT_APPLICATION_NAME + ) as connection, + connection.cursor() as cursor, + ): + query = asyncio.create_task(cursor.execute("SELECT pg_sleep(10)")) + try: + with suppress_instrumentation(): + async with observer.cursor() as observer_cursor: + for _ in range(20): + await observer_cursor.execute( + "SELECT state, query FROM pg_stat_activity WHERE pid = %s", + (connection.info.backend_pid,), + ) + state, statement = await observer_cursor.fetchone() + if state == "active" and statement == "SELECT pg_sleep(10)": + break + await asyncio.sleep(0.05) + else: + raise AssertionError("Customer query did not start before cancellation") + finally: + query.cancel() + try: + await query + except asyncio.CancelledError: + pass + + if not autocommit: + assert connection.info.transaction_status == psycopg.pq.TransactionStatus.INERROR + await connection.rollback() + assert connection.info.parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + with suppress_instrumentation(): + await cursor.execute("SELECT 1") + assert (await cursor.fetchone())[0] == 1 + + def main(): from opentelemetry.instrumentation.auto_instrumentation import initialize from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor @@ -320,9 +416,12 @@ def main(): _unsampled_queries(driver, exporter) _reset_to_default(driver) _failed_internal_command(driver) + _failed_cleanup_command(driver) asyncio.run(_async_queries(exporter, tracer)) _custom_cursor(exporter) _psycopg_server_cursor_boundary(exporter) + _psycopg2_commit_boundary(exporter) + asyncio.run(_async_query_cancellation()) print( # noqa: T201 f"POSTGRESQL_CORRELATION_OK mode={mode} psycopg={version('psycopg')} " diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 986329e1..fea5ce2e 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -64,6 +64,10 @@ def cursor(self): self.cursors.append(cursor) return cursor + def commit(self): + self.calls.append("COMMIT") + return self.application_name + class FakeAsyncConnection(FakeConnection): def cursor(self): @@ -105,6 +109,7 @@ def __init__(self, connection): async def execute(self, query): self.calls.append(query) self.connection.calls.append(query) + self.connection.before_execute(query) if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name elif query.startswith("SET application_name = '"): @@ -144,23 +149,19 @@ def test_query_sets_database_span_traceparent(): assert result == TRACEPARENT assert cursor.calls == ["SELECT application_name"] assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] - assert connection.cursors[0].closed + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + assert connection.application_name == "startup-app" assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT application_name", + "SET application_name = DEFAULT", ] def test_query_without_database_span_clears_stale_traceparent(): connection = FakeConnection() + connection.application_name = TRACEPARENT cursor = FakePsycopgCursor(connection) - - _execute_with_postgresql_trace_context( - _fake_traced_execution, - None, - (cursor, cursor.execute, "SELECT 1"), - {}, - ) result = _execute_with_postgresql_trace_context( _fake_untraced_execution, None, @@ -170,8 +171,6 @@ def test_query_without_database_span_clears_stale_traceparent(): assert result == "startup-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", - "SELECT 1", "SET application_name = DEFAULT", "SELECT 2", ] @@ -228,7 +227,7 @@ class FakePsycopg2Cursor(FakePsycopgCursor): assert connection.cursors == [] -def test_psycopg2_named_cursor_is_not_used_for_application_name(): +def test_psycopg2_named_cursor_sends_only_customer_query(): class FakePsycopg2NamedCursor(FakePsycopgCursor): __module__ = "psycopg2.extensions" @@ -239,6 +238,7 @@ def execute(self, query): connection = FakeConnection() cursor = FakePsycopg2NamedCursor(connection) + cursor.name = "server_cursor" with trace.use_span(NonRecordingSpan(PARENT_CONTEXT)): result = _execute_with_postgresql_trace_context( @@ -248,10 +248,9 @@ def execute(self, query): {}, ) - assert result == TRACEPARENT + assert result == "customer-app" assert cursor.calls == ["SELECT 1"] - assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] - assert connection.cursors[0].closed + assert connection.cursors == [] def test_async_query_sets_database_span_traceparent(): @@ -274,10 +273,12 @@ async def run(): assert result == TRACEPARENT assert cursor.calls == ["SELECT application_name"] assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] - assert connection.cursors[0].closed + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + assert connection.application_name == "startup-app" assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT application_name", + "SET application_name = DEFAULT", ] @@ -340,8 +341,10 @@ def wrapped(_cursor, query_method, *args, **kwargs): assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT first", + "SET application_name = DEFAULT", f"SET application_name = '{SECOND_TRACEPARENT}'", "SELECT second", + "SET application_name = DEFAULT", ] @@ -367,7 +370,11 @@ class CursorTracer(FakeCursorTracer): result = CursorTracer().traced_execution(cursor, cursor.execute, "SELECT 1") assert result == TRACEPARENT - assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT 1"] + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT 1", + "SET application_name = DEFAULT", + ] def test_public_enable_without_dbapi_reports_missing_dependency(monkeypatch, caplog): @@ -381,6 +388,28 @@ def _fake_dbapi_import(_name): assert "requires DBAPI instrumentation" in caplog.text +def test_public_enable_clears_context_before_commit_and_raw_queries(monkeypatch): + class CursorTracer(FakeCursorTracer): + pass + + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", lambda _name: SimpleNamespace(CursorTracer=CursorTracer)) + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + + assert postgresql.enable_postgresql_trace_context_propagation() + assert CursorTracer().traced_execution(cursor, cursor.execute, "SELECT 1") == TRACEPARENT + assert connection.commit() == "startup-app" + assert cursor.execute("SELECT untraced") == "startup-app" + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT 1", + "SET application_name = DEFAULT", + "COMMIT", + "SELECT untraced", + ] + + def test_public_enable_rejects_incompatible_dbapi_before_wrapping(monkeypatch, caplog): class CursorTracer: def traced_execution(self, cursor, query_method, *args, **kwargs): @@ -456,8 +485,12 @@ def _fake_internal_command_failure(query): ) assert result == "customer-app" - assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT 1"] - assert connection.cursors[0].closed + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT 1", + "SET application_name = DEFAULT", + ] + assert all(internal_cursor.closed for internal_cursor in connection.cursors) assert "Failed to set PostgreSQL application_name" in caplog.text @@ -481,7 +514,41 @@ def _fake_customer_query_failure(query): ) assert raised.value is error - assert connection.cursors[0].closed + assert connection.calls[-1] == "SET application_name = DEFAULT" + assert connection.application_name == "startup-app" + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + + +@pytest.mark.parametrize("query_fails", [False, True]) +def test_failed_cleanup_does_not_replace_customer_result_or_exception(query_fails, caplog): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + error = RuntimeError("customer query failed") + + def _fake_query_or_cleanup_failure(query): + if query == "SET application_name = DEFAULT": + raise RuntimeError("cleanup failed") + if query_fails and query == "SELECT 1": + raise error + + connection.before_execute = _fake_query_or_cleanup_failure + if query_fails: + with pytest.raises(RuntimeError) as raised: + _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + assert raised.value is error + else: + assert ( + _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + == TRACEPARENT + ) + + assert connection.calls[-1] == "SET application_name = DEFAULT" + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + assert "Failed to set PostgreSQL application_name" in caplog.text def test_async_shared_connection_keeps_set_and_query_together(): @@ -532,8 +599,10 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT first", + "SET application_name = DEFAULT", f"SET application_name = '{SECOND_TRACEPARENT}'", "SELECT second", + "SET application_name = DEFAULT", ] asyncio.run(run()) @@ -556,6 +625,83 @@ async def run(): asyncio.run(run()) +def test_async_named_cursor_sends_only_customer_query(): + async def run(): + connection = FakeAsyncConnection() + cursor = FakeAsyncPsycopgCursor(connection) + cursor.name = "server_cursor" + result = await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + assert result == "customer-app" + assert connection.calls == ["SELECT 1"] + assert connection.cursors == [] + + asyncio.run(run()) + + +def test_async_customer_cancellation_clears_context_and_releases_connection(): + async def run(): + query_started = asyncio.Event() + connection = FakeAsyncConnection() + + class BlockingQueryCursor(FakeAsyncPsycopgCursor): + async def execute(self, query): + await super().execute(query) + query_started.set() + await asyncio.Event().wait() + + cursor = BlockingQueryCursor(connection) + query = asyncio.create_task( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + ) + await asyncio.wait_for(query_started.wait(), timeout=1) + query.cancel() + with pytest.raises(asyncio.CancelledError): + await query + + assert connection.calls[-1] == "SET application_name = DEFAULT" + assert connection.application_name == "startup-app" + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + assert not postgresql._get_async_connection_lock(connection).locked() + + asyncio.run(run()) + + +@pytest.mark.parametrize("query_fails", [False, True]) +def test_async_failed_cleanup_does_not_replace_customer_result_or_exception(query_fails, caplog): + async def run(): + connection = FakeAsyncConnection() + cursor = FakeAsyncPsycopgCursor(connection) + error = RuntimeError("customer query failed") + + def _fake_query_or_cleanup_failure(query): + if query == "SET application_name = DEFAULT": + raise RuntimeError("cleanup failed") + if query_fails and query == "SELECT 1": + raise error + + connection.before_execute = _fake_query_or_cleanup_failure + if query_fails: + with pytest.raises(RuntimeError) as raised: + await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + assert raised.value is error + else: + result = await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + assert result == TRACEPARENT + assert connection.calls[-1] == "SET application_name = DEFAULT" + assert all(internal_cursor.closed for internal_cursor in connection.cursors) + + asyncio.run(run()) + assert "Failed to set PostgreSQL application_name" in caplog.text + + def test_async_cancellation_closes_internal_cursor_and_releases_connection(): async def run(): internal_command_started = asyncio.Event() From 5a08f1e73b6fca036339be157b87aed24d6efc57 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Wed, 14 Oct 2026 17:16:04 -0400 Subject: [PATCH 03/10] Fix PostgreSQL regression test lint checks Use existing exception instances and contextlib.suppress. Check cancellation releases the connection by executing another query instead of reading the private lock. Assisted-by: Codex (GPT-6) --- tests/smoke/postgresql_trace_context.py | 9 ++++----- tests/test_postgresql.py | 15 ++++++++++++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py index 799bac50..afaff870 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/smoke/postgresql_trace_context.py @@ -1,6 +1,7 @@ import asyncio import os from concurrent.futures import ThreadPoolExecutor +from contextlib import suppress from importlib.metadata import version @@ -357,14 +358,12 @@ async def _async_query_cancellation(): if state == "active" and statement == "SELECT pg_sleep(10)": break await asyncio.sleep(0.05) - else: - raise AssertionError("Customer query did not start before cancellation") + assert state == "active", "Customer query did not start before cancellation" + assert statement == "SELECT pg_sleep(10)" finally: query.cancel() - try: + with suppress(asyncio.CancelledError): await query - except asyncio.CancelledError: - pass if not autocommit: assert connection.info.transaction_status == psycopg.pq.TransactionStatus.INERROR diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index fea5ce2e..1e4f07c8 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -524,10 +524,11 @@ def test_failed_cleanup_does_not_replace_customer_result_or_exception(query_fail connection = FakeConnection() cursor = FakePsycopgCursor(connection) error = RuntimeError("customer query failed") + cleanup_error = RuntimeError("cleanup failed") def _fake_query_or_cleanup_failure(query): if query == "SET application_name = DEFAULT": - raise RuntimeError("cleanup failed") + raise cleanup_error if query_fails and query == "SELECT 1": raise error @@ -665,7 +666,14 @@ async def execute(self, query): assert connection.calls[-1] == "SET application_name = DEFAULT" assert connection.application_name == "startup-app" assert all(internal_cursor.closed for internal_cursor in connection.cursors) - assert not postgresql._get_async_connection_lock(connection).locked() + next_cursor = FakeAsyncPsycopgCursor(connection) + result = await asyncio.wait_for( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (next_cursor, next_cursor.execute, "SELECT 2"), {} + ), + timeout=1, + ) + assert result == TRACEPARENT asyncio.run(run()) @@ -676,10 +684,11 @@ async def run(): connection = FakeAsyncConnection() cursor = FakeAsyncPsycopgCursor(connection) error = RuntimeError("customer query failed") + cleanup_error = RuntimeError("cleanup failed") def _fake_query_or_cleanup_failure(query): if query == "SET application_name = DEFAULT": - raise RuntimeError("cleanup failed") + raise cleanup_error if query_fails and query == "SELECT 1": raise error From 936110b4ddd0e3b1084386ce122e4b7c83f4a31b Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Wed, 14 Oct 2026 17:50:20 -0400 Subject: [PATCH 04/10] Handle nested PostgreSQL queries and SQL transaction boundaries Reuse connection lock ownership only in the same thread or task, and restore the outer span context after nested queries. Serialize transaction-control SQL without propagation commands so COMMIT stays idle and rollback cannot restore a BEGIN span. Add unit and real PostgreSQL regression coverage and document the conservative SQL exclusions. Assisted-by: Codex (GPT-6) --- .github/workflows/ci-main.yml | 1 + docs/postgresql.md | 25 ++- src/splunk_otel/postgresql.py | 106 ++++++++++- tests/smoke/postgresql_trace_context.py | 153 +++++++++++++++ tests/test_postgresql.py | 237 ++++++++++++++++++++++-- 5 files changed, 493 insertions(+), 29 deletions(-) diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 115df8ab..76c18e3c 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -91,4 +91,5 @@ jobs: opentelemetry-instrumentation-psycopg2 - name: Run PostgreSQL smoke test + timeout-minutes: 5 run: python tests/smoke/postgresql_trace_context.py diff --git a/docs/postgresql.md b/docs/postgresql.md index 9da00de5..c9c742dc 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -65,19 +65,30 @@ this feature's automatic coverage. Named cursors can execute SQL later during fetch, after the execute span has ended, so this feature skips propagation for them in both drivers. +SQL transaction-control statements (including BEGIN, COMMIT, ROLLBACK, and +savepoint commands), SQL CALL, and DO run without propagation commands. Those +statements can change transaction boundaries, making SET and cleanup unsafe. +SQL batches also run without propagation. The adapter conservatively skips any +statement with a semicolon other than a single trailing semicolon, including +semicolons inside strings or comments. These calls still participate in the +connection lock. + ## Connection behavior The feature uses a separate cursor on the same connection for internal commands and suppresses instrumentation for those commands. A per-connection lock covers -SET, the customer query, and cleanup. Each traced query sends two additional -commands: SET before execution and SET DEFAULT afterward. +SET, the customer query, and cleanup. Nested queries from a custom cursor can +reuse the lock in the same thread or async task. After a nested query, cleanup +restores the outer query's context. Other threads and tasks must acquire the +lock. Each correlated query sends two additional commands: SET before execution +and context restoration afterward. This opt-in owns the connection's `application_name` during execution. Cleanup -resets the setting to the connection's configured default before releasing the -instrumentation lock, including when the query raises an exception. If -instrumentation creates no valid new database span, the feature sends SET DEFAULT -before execution. It does not restore a value that the application previously -set through SQL. +of the outermost query resets the setting to the connection's configured default +before releasing the instrumentation lock, including when the query raises an +exception. If instrumentation creates no valid new database span for an eligible +statement, the feature sends SET DEFAULT before execution. It does not restore +a value that the application previously set through SQL. The lock only coordinates wrapped queries. Raw cursors and suppressed queries can bypass it and observe another query's context during execution. Failed diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index c7039547..4ed28ad7 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -17,8 +17,12 @@ import asyncio import inspect import logging +import re import threading import weakref +from contextlib import nullcontext +from contextvars import ContextVar +from dataclasses import dataclass from importlib import import_module from typing import TYPE_CHECKING, Any @@ -36,6 +40,19 @@ _SET_APPLICATION_NAME_DEFAULT = "SET application_name = DEFAULT" _DRIVER_MODULE_NAMES = ("psycopg", "psycopg2") _CURSOR_AND_QUERY_METHOD_ARGUMENTS = 2 +_TRANSACTION_OPERATIONS = frozenset( + {"BEGIN", "START", "COMMIT", "END", "ROLLBACK", "ABORT", "SAVEPOINT", "RELEASE", "PREPARE", "CALL", "DO"} +) + + +@dataclass(frozen=True) +class _QueryContext: + connection: Any + owner: object + traceparent: str | None + + +_active_queries: ContextVar[tuple[_QueryContext, ...]] = ContextVar("postgresql_active_queries", default=()) _enabled = False _enable_lock = threading.Lock() @@ -97,6 +114,7 @@ def _execute_with_postgresql_trace_context(traced_execution, _cursor_tracer, arg if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return traced_execution(*args, **kwargs) + propagate = _can_propagate_statement(cursor, query_args, kwargs) span_context_before = trace.get_current_span().get_span_context() def correlated_query_method(*method_args, **method_kwargs): @@ -106,6 +124,7 @@ def correlated_query_method(*method_args, **method_kwargs): span_context_before, method_args, method_kwargs, + propagate=propagate, ) return traced_execution(cursor, correlated_query_method, *query_args, **kwargs) @@ -119,6 +138,7 @@ async def _execute_with_postgresql_trace_context_async(traced_execution_async, _ if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return await traced_execution_async(*args, **kwargs) + propagate = _can_propagate_statement(cursor, query_args, kwargs) span_context_before = trace.get_current_span().get_span_context() async def correlated_query_method(*method_args, **method_kwargs): @@ -128,6 +148,7 @@ async def correlated_query_method(*method_args, **method_kwargs): span_context_before, method_args, method_kwargs, + propagate=propagate, ) return await traced_execution_async(cursor, correlated_query_method, *query_args, **kwargs) @@ -139,16 +160,25 @@ def _execute_correlated( span_context_before: SpanContext, args, kwargs, + *, + propagate: bool, ): - lock = _get_sync_connection_lock(cursor.connection) + owner = threading.current_thread() + parent = _parent_query(cursor.connection, owner) + lock = nullcontext() if parent is not None else _get_sync_connection_lock(cursor.connection) with lock: - traceparent = _current_database_traceparent(span_context_before) + traceparent = _current_database_traceparent(span_context_before) if propagate else None + token = _active_queries.set((*_active_queries.get(), _QueryContext(cursor.connection, owner, traceparent))) try: - _update_application_name(cursor.connection, traceparent) + if propagate: + _update_application_name(cursor.connection, traceparent) return query_method(*args, **kwargs) finally: - if traceparent is not None: - _update_application_name(cursor.connection, None) + try: + if propagate and (traceparent is not None or parent is not None): + _update_application_name(cursor.connection, parent.traceparent if parent is not None else None) + finally: + _active_queries.reset(token) async def _execute_correlated_async( @@ -157,16 +187,72 @@ async def _execute_correlated_async( span_context_before: SpanContext, args, kwargs, + *, + propagate: bool, ): - lock = _get_async_connection_lock(cursor.connection) + owner = asyncio.current_task() + parent = _parent_query(cursor.connection, owner) + lock = nullcontext() if parent is not None else _get_async_connection_lock(cursor.connection) async with lock: - traceparent = _current_database_traceparent(span_context_before) + traceparent = _current_database_traceparent(span_context_before) if propagate else None + token = _active_queries.set((*_active_queries.get(), _QueryContext(cursor.connection, owner, traceparent))) try: - await _update_application_name_async(cursor.connection, traceparent) + if propagate: + await _update_application_name_async(cursor.connection, traceparent) return await query_method(*args, **kwargs) finally: - if traceparent is not None: - await _update_application_name_async(cursor.connection, None) + try: + if propagate and (traceparent is not None or parent is not None): + await _update_application_name_async( + cursor.connection, parent.traceparent if parent is not None else None + ) + finally: + _active_queries.reset(token) + + +def _parent_query(connection, owner: object) -> _QueryContext | None: + # Context can be copied into another thread/task; that caller must acquire the lock. + return next( + (query for query in reversed(_active_queries.get()) if query.connection is connection and query.owner is owner), + None, + ) + + +def _can_propagate_statement(cursor, args, kwargs) -> bool: + statement = args[0] if args else kwargs.get("query") + if isinstance(statement, bytes): + statement = statement.decode("ascii", errors="replace") + elif any(base.__module__ in ("psycopg.sql", "psycopg2.sql") for base in type(statement).__mro__): + statement = statement.as_string(cursor.connection) + if not isinstance(statement, str): + return False + + # SET and reset cannot safely straddle SQL transaction boundaries. Avoid + # batches too: even a SELECT-first batch can commit or begin a transaction. + statement = statement.strip().removesuffix(";") + if ";" in statement: + return False + statement = _strip_leading_comments(statement) + operation = re.match(r"[a-z]+", statement, re.IGNORECASE) + return operation is not None and operation[0].upper() not in _TRANSACTION_OPERATIONS + + +def _strip_leading_comments(statement: str) -> str: + statement = statement.lstrip() + while statement.startswith(("--", "/*")): + if statement.startswith("--"): + _, _, statement = statement.partition("\n") + else: + depth = 0 + for delimiter in re.finditer(r"/\*|\*/", statement): + depth += 1 if delimiter[0] == "/*" else -1 + if depth == 0: + statement = statement[delimiter.end() :] + break + else: + return "" + statement = statement.lstrip() + return statement def _get_sync_connection_lock(connection) -> threading.Lock: diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py index afaff870..662d1155 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/smoke/postgresql_trace_context.py @@ -2,6 +2,7 @@ import os from concurrent.futures import ThreadPoolExecutor from contextlib import suppress +from importlib import import_module from importlib.metadata import version @@ -268,6 +269,154 @@ def _custom_cursor(exporter): _assert_database_span(exporter, application_name) +def _nested_cursor_queries(driver, exporter): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + + class CustomerCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query == _QUERY: + with self.connection.cursor() as nested: + nested.execute("SELECT current_setting('application_name')") + self.nested_context = nested.fetchone()[0] + return super().execute(query, *args, **kwargs) + + for autocommit in (False, True): + with driver.connect(cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.autocommit = autocommit + exporter.clear() + with connection.cursor() as cursor: + cursor.execute(_QUERY) + outer = _assert_database_span(exporter, cursor.fetchone()[0]) + nested = _assert_database_span(exporter, cursor.nested_context) + assert nested.parent.span_id == outer.context.span_id + assert nested.context.span_id != outer.context.span_id + assert len(exporter.get_finished_spans()) == 2 + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +async def _async_nested_cursor_queries(exporter): + import psycopg + + class CustomerCursor(psycopg.AsyncCursor): + async def execute(self, query, *args, **kwargs): + if query == _QUERY: + async with self.connection.cursor() as nested: + await nested.execute("SELECT current_setting('application_name')") + self.nested_context = (await nested.fetchone())[0] + return await super().execute(query, *args, **kwargs) + + for autocommit in (False, True): + async with await psycopg.AsyncConnection.connect( + autocommit=autocommit, cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME + ) as connection: + exporter.clear() + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + outer = _assert_database_span(exporter, (await cursor.fetchone())[0]) + nested = _assert_database_span(exporter, cursor.nested_context) + assert nested.parent.span_id == outer.context.span_id + assert nested.context.span_id != outer.context.span_id + assert len(exporter.get_finished_spans()) == 2 + with suppress_instrumentation(): + async with connection.cursor() as cursor: + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + + +def _transaction_sql(driver, exporter): + sql = import_module(f"{driver.__name__}.sql") + statements = ( + ("BEGIN", "COMMIT"), + ("START TRANSACTION", b"END"), + (sql.SQL("/* outer /* nested */ comment */ BEGIN"), sql.SQL("COMMIT")), + ("BEGIN", "ROLLBACK"), + ("BEGIN", "ABORT"), + ) + for autocommit in (False, True): + for begin, end in statements: + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.autocommit = autocommit + with connection.cursor() as cursor: + cursor.execute(begin) + exporter.clear() + cursor.execute(_QUERY) + _assert_database_span(exporter, cursor.fetchone()[0]) + connection.rollback() + with suppress_instrumentation(): + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + connection.rollback() + + cursor.execute(begin) + cursor.execute(end) + status = ( + connection.info.transaction_status + if driver.__name__ == "psycopg" + else connection.get_transaction_status() + ) + assert status == 0 + if driver.__name__ == "psycopg": + connection.autocommit = True + with suppress_instrumentation(): + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + connection.rollback() + connection.autocommit = True + cursor.execute("BEGIN; SELECT 1") + connection.rollback() + cursor.execute("SELECT 1; COMMIT") + with suppress_instrumentation(): + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +async def _async_transaction_sql(exporter): + import psycopg + + statements = ( + ("BEGIN", "COMMIT"), + ("START TRANSACTION", b"END"), + (psycopg.sql.SQL("/* outer /* nested */ comment */ BEGIN"), psycopg.sql.SQL("COMMIT")), + ("BEGIN", "ROLLBACK"), + ("BEGIN", "ABORT"), + ) + for autocommit in (False, True): + for begin, end in statements: + async with ( + await psycopg.AsyncConnection.connect( + autocommit=autocommit, application_name=_DEFAULT_APPLICATION_NAME + ) as connection, + connection.cursor() as cursor, + ): + await cursor.execute(begin) + exporter.clear() + await cursor.execute(_QUERY) + _assert_database_span(exporter, (await cursor.fetchone())[0]) + await connection.rollback() + with suppress_instrumentation(): + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + await connection.rollback() + + await cursor.execute(begin) + await cursor.execute(end) + assert connection.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + await connection.set_autocommit(True) + with suppress_instrumentation(): + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + + await cursor.execute("BEGIN; SELECT 1") + await connection.rollback() + await cursor.execute("SELECT 1; COMMIT") + with suppress_instrumentation(): + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + + def _failed_cleanup_command(driver): cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor @@ -416,7 +565,11 @@ def main(): _reset_to_default(driver) _failed_internal_command(driver) _failed_cleanup_command(driver) + _nested_cursor_queries(driver, exporter) + _transaction_sql(driver, exporter) asyncio.run(_async_queries(exporter, tracer)) + asyncio.run(_async_nested_cursor_queries(exporter)) + asyncio.run(_async_transaction_sql(exporter)) _custom_cursor(exporter) _psycopg_server_cursor_boundary(exporter) _psycopg2_commit_boundary(exporter) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 1e4f07c8..04f73367 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -14,6 +14,7 @@ import asyncio import threading +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -90,7 +91,7 @@ def execute(self, query): self.connection.before_execute(query) if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name - elif query.startswith("SET application_name = '"): + elif isinstance(query, str) and query.startswith("SET application_name = '"): self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") return self.connection.application_name @@ -112,7 +113,7 @@ async def execute(self, query): self.connection.before_execute(query) if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name - elif query.startswith("SET application_name = '"): + elif isinstance(query, str) and query.startswith("SET application_name = '"): self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") return self.connection.application_name @@ -130,10 +131,214 @@ async def _fake_traced_execution_async(_cursor, query_method, *args, **kwargs): return await query_method(*args, **kwargs) +def _fake_nested_traced_execution(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(SECOND_DATABASE_CONTEXT)): + return query_method(*args, **kwargs) + + +async def _fake_nested_traced_execution_async(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(SECOND_DATABASE_CONTEXT)): + return await query_method(*args, **kwargs) + + def _fake_untraced_execution(_cursor, query_method, *args, **kwargs): return query_method(*args, **kwargs) +@pytest.mark.parametrize("nested_fails", [False, True]) +def test_nested_customer_query_restores_outer_context_without_deadlock(nested_fails): + connection = FakeConnection() + error = RuntimeError("nested query failed") + results = [] + + def _fake_query_failure(query): + if nested_fails and query == "SELECT nested": + raise error + + connection.before_execute = _fake_query_failure + + class CustomerCursor(FakePsycopgCursor): + def execute(self, query): + nested = FakePsycopgCursor(connection) + with pytest.raises(RuntimeError) if nested_fails else nullcontext() as raised: + result = _execute_with_postgresql_trace_context( + _fake_nested_traced_execution, None, (nested, nested.execute, "SELECT nested"), {} + ) + assert result == SECOND_TRACEPARENT + if nested_fails: + assert raised.value is error + assert connection.application_name == TRACEPARENT + return super().execute(query) + + cursor = CustomerCursor(connection) + + def _execute_customer_query(): + results.append( + _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT outer"), {} + ) + ) + + thread = threading.Thread(target=_execute_customer_query, daemon=True) + thread.start() + thread.join(timeout=2) + assert not thread.is_alive(), "Nested customer query deadlocked" + assert results == [TRACEPARENT] + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + f"SET application_name = '{SECOND_TRACEPARENT}'", + "SELECT nested", + f"SET application_name = '{TRACEPARENT}'", + "SELECT outer", + "SET application_name = DEFAULT", + ] + + +@pytest.mark.parametrize("nested_fails", [False, True]) +def test_async_nested_customer_query_restores_outer_context(nested_fails): + async def run(): + connection = FakeAsyncConnection() + error = RuntimeError("nested query failed") + + def _fake_query_failure(query): + if nested_fails and query == "SELECT nested": + raise error + + connection.before_execute = _fake_query_failure + + class CustomerCursor(FakeAsyncPsycopgCursor): + async def execute(self, query): + nested = FakeAsyncPsycopgCursor(connection) + with pytest.raises(RuntimeError) if nested_fails else nullcontext() as raised: + result = await _execute_with_postgresql_trace_context_async( + _fake_nested_traced_execution_async, None, (nested, nested.execute, "SELECT nested"), {} + ) + assert result == SECOND_TRACEPARENT + if nested_fails: + assert raised.value is error + assert connection.application_name == TRACEPARENT + return await super().execute(query) + + cursor = CustomerCursor(connection) + result = await asyncio.wait_for( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT outer"), {} + ), + timeout=2, + ) + assert result == TRACEPARENT + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + f"SET application_name = '{SECOND_TRACEPARENT}'", + "SELECT nested", + f"SET application_name = '{TRACEPARENT}'", + "SELECT outer", + "SET application_name = DEFAULT", + ] + + asyncio.run(run()) + + +def test_async_child_task_waits_for_outer_connection_lock(): + async def run(): + connection = FakeAsyncConnection() + nested_started = asyncio.Event() + children = [] + + async def _execute_nested_customer_query(): + nested_started.set() + cursor = FakeAsyncPsycopgCursor(connection) + return await _execute_with_postgresql_trace_context_async( + _fake_nested_traced_execution_async, None, (cursor, cursor.execute, "SELECT child"), {} + ) + + class CustomerCursor(FakeAsyncPsycopgCursor): + async def execute(self, query): + children.append(asyncio.create_task(_execute_nested_customer_query())) + await nested_started.wait() + assert not children[0].done() + assert connection.application_name == TRACEPARENT + return await super().execute(query) + + cursor = CustomerCursor(connection) + result = await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT outer"), {} + ) + assert result == TRACEPARENT + assert await children[0] == SECOND_TRACEPARENT + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT outer", + "SET application_name = DEFAULT", + f"SET application_name = '{SECOND_TRACEPARENT}'", + "SELECT child", + "SET application_name = DEFAULT", + ] + + asyncio.run(asyncio.wait_for(run(), timeout=2)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "statement", + [ + "BEGIN", + "START TRANSACTION", + "COMMIT", + "END;", + "ROLLBACK", + "ABORT", + "SAVEPOINT customer_savepoint", + "RELEASE SAVEPOINT customer_savepoint", + "ROLLBACK TO SAVEPOINT customer_savepoint", + "PREPARE TRANSACTION 'customer_transaction'", + "COMMIT PREPARED 'customer_transaction'", + "CALL customer_procedure()", + "DO $$ BEGIN COMMIT; END $$", + "SELECT 1; COMMIT", + "SELECT 1; BEGIN", + "-- customer comment\n COMMIT", + "/* outer /* nested */ comment */ BEGIN", + b"COMMIT", + ], +) +def test_transaction_control_and_batches_send_only_customer_statement(statement, asynchronous): + if asynchronous: + + async def run(): + connection = FakeAsyncConnection() + cursor = FakeAsyncPsycopgCursor(connection) + result = await _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, statement), {} + ) + assert result == "customer-app" + assert connection.calls == [statement] + assert connection.cursors == [] + + asyncio.run(run()) + else: + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, statement), {} + ) + assert result == "customer-app" + assert connection.calls == [statement] + assert connection.cursors == [] + + +@pytest.mark.parametrize( + "statement", ["SELECT 1;", "-- comment\nSELECT 1", "/* outer /* nested */ comment */ SELECT 1"] +) +def test_data_statement_after_comments_propagates_context(statement): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + assert ( + _execute_with_postgresql_trace_context(_fake_traced_execution, None, (cursor, cursor.execute, statement), {}) + == TRACEPARENT + ) + + def test_query_sets_database_span_traceparent(): connection = FakeConnection() cursor = FakePsycopgCursor(connection) @@ -282,7 +487,8 @@ async def run(): ] -def test_shared_connection_keeps_set_and_query_together(monkeypatch): +@pytest.mark.parametrize("second_query", ["SELECT second", "COMMIT"]) +def test_shared_connection_keeps_set_and_query_together(monkeypatch, second_query): connection = FakeConnection() first_cursor = FakePsycopgCursor(connection) second_cursor = FakePsycopgCursor(connection) @@ -320,7 +526,7 @@ def wrapped(_cursor, query_method, *args, **kwargs): _execute_with_postgresql_trace_context(wrapped, None, (cursor, cursor.execute, query), {}) first_thread = threading.Thread(target=run, args=(first_cursor, "SELECT first", DATABASE_CONTEXT)) - second_thread = threading.Thread(target=run, args=(second_cursor, "SELECT second", SECOND_DATABASE_CONTEXT)) + second_thread = threading.Thread(target=run, args=(second_cursor, second_query, SECOND_DATABASE_CONTEXT)) try: first_thread.start() @@ -338,13 +544,16 @@ def wrapped(_cursor, query_method, *args, **kwargs): assert not first_thread.is_alive() assert not second_thread.is_alive() + second_calls = ( + [f"SET application_name = '{SECOND_TRACEPARENT}'", second_query, "SET application_name = DEFAULT"] + if second_query == "SELECT second" + else [second_query] + ) assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT first", "SET application_name = DEFAULT", - f"SET application_name = '{SECOND_TRACEPARENT}'", - "SELECT second", - "SET application_name = DEFAULT", + *second_calls, ] @@ -552,7 +761,8 @@ def _fake_query_or_cleanup_failure(query): assert "Failed to set PostgreSQL application_name" in caplog.text -def test_async_shared_connection_keeps_set_and_query_together(): +@pytest.mark.parametrize("second_query", ["SELECT second", "COMMIT"]) +def test_async_shared_connection_keeps_set_and_query_together(second_query): async def run(): connection = FakeAsyncConnection() first_query_started = asyncio.Event() @@ -588,7 +798,7 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): _execute_with_postgresql_trace_context_async( _fake_second_traced_execution, None, - (second_cursor, second_cursor.execute, "SELECT second"), + (second_cursor, second_cursor.execute, second_query), {}, ) ) @@ -597,13 +807,16 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): finally: release_first_query.set() await asyncio.wait_for(asyncio.gather(*([first, second] if second is not None else [first])), timeout=1) + second_calls = ( + [f"SET application_name = '{SECOND_TRACEPARENT}'", second_query, "SET application_name = DEFAULT"] + if second_query == "SELECT second" + else [second_query] + ) assert connection.calls == [ f"SET application_name = '{TRACEPARENT}'", "SELECT first", "SET application_name = DEFAULT", - f"SET application_name = '{SECOND_TRACEPARENT}'", - "SELECT second", - "SET application_name = DEFAULT", + *second_calls, ] asyncio.run(run()) From 44b07de03973a16c782f772da174a3e14e3d9e14 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Wed, 14 Oct 2026 18:38:40 -0400 Subject: [PATCH 05/10] Apply PostgreSQL trace context at driver execution Capture the upstream database span in a bounded cursor lease, then apply application_name only when customer super() calls reach the native driver. Keep auxiliary custom SQL outside the propagation lock, expire copied execution contexts, and carry the active setting through skipped nested calls. Accept parenthesized queries and classify callproc by method role. Use native internal cursors and retain custom factories in Psycopg 3 manual instrumentation. Add unit and PostgreSQL 16 regressions for all four review findings and document the private instrumentation hooks. Assisted-by: Codex (GPT-6) --- docs/postgresql.md | 38 ++- src/splunk_otel/postgresql.py | 297 ++++++++++++++++++++---- tests/smoke/postgresql_trace_context.py | 261 +++++++++++++++++++-- tests/test_postgresql.py | 274 ++++++++++++++++++++-- 4 files changed, 784 insertions(+), 86 deletions(-) diff --git a/docs/postgresql.md b/docs/postgresql.md index c9c742dc..e525fb87 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -52,7 +52,8 @@ for the process lifetime. To turn it off, unset the switch and restart. Manual instrumentation must initialize the Splunk distro first with `init_splunk_otel()`, or call `splunk_otel.postgresql.enable_postgresql_trace_context_propagation()` before using -`PsycopgInstrumentor` or `Psycopg2Instrumentor`. +`PsycopgInstrumentor` or `Psycopg2Instrumentor`. Enable the feature before creating +instrumented connections and cursors. ## Driver coverage @@ -60,6 +61,10 @@ The feature targets Psycopg 3 synchronous and asynchronous connections and Psycopg2 synchronous connections. Psycopg2 asynchronous connections pass through without trace propagation. +Custom cursor methods must delegate through `super()` to the driver method. +The feature sets context after those custom hooks, immediately before the driver +executes the query. Direct calls to native driver methods can bypass this adapter. + Named/server cursors, cursor fetch operations, and COPY operations are outside this feature's automatic coverage. Named cursors can execute SQL later during fetch, after the execute span has ended, so this feature skips propagation for @@ -75,12 +80,17 @@ connection lock. ## Connection behavior -The feature uses a separate cursor on the same connection for internal commands -and suppresses instrumentation for those commands. A per-connection lock covers -SET, the customer query, and cleanup. Nested queries from a custom cursor can -reuse the lock in the same thread or async task. After a nested query, cleanup -restores the outer query's context. Other threads and tasks must acquire the -lock. Each correlated query sends two additional commands: SET before execution +The feature uses a native driver cursor on the same connection for internal +commands. It suppresses instrumentation and bypasses customer cursor hooks for +those commands. A per-connection lock covers SET, driver execution, and cleanup. +Customer hooks run outside this lock. Auxiliary queries from those hooks finish +their own propagation before the outer query starts driver execution. + +Nested driver executions in the same thread or async task can reuse the lock. +Cleanup restores the active outer driver's setting, including through skipped +intermediate statements. Other threads and tasks must acquire the lock. Context +copied into a delayed callback cannot grant ownership after its query has finished. +Each correlated query sends two additional commands: SET before driver execution and context restoration afterward. This opt-in owns the connection's `application_name` during execution. Cleanup @@ -105,3 +115,17 @@ does not roll back application transactions. The extra commands add latency. Database span duration and database operation duration metrics include lock waiting and the additional commands. + +## Instrumentation compatibility + +The Splunk adapter uses the installed OpenTelemetry driver instrumentors to +create spans. It wraps shared DBAPI execution callbacks to capture the database +span context, then adapts the Psycopg cursor factories to apply that context at +driver execution. It also adapts Psycopg 3 manual connection instrumentation to +include the application's custom cursor factory. + +These hooks use private instrumentation APIs. Activation checks their signatures +and declines incompatible versions before installing wrappers. Regression tests +must also check behavior when instrumentation dependencies change. The current +validation covers instrumentation 0.66b1, Psycopg 3.3.6, Psycopg2 2.9.13, and +PostgreSQL 16; it does not establish a broader supported version range. diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index 4ed28ad7..ff2d3cd3 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -22,7 +22,8 @@ import weakref from contextlib import nullcontext from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field +from functools import partial from importlib import import_module from typing import TYPE_CHECKING, Any @@ -45,13 +46,43 @@ ) +class _Lifetime: + def __init__(self): + self._active = True + + def is_active(self) -> bool: + return self._active + + def finish(self) -> None: + self._active = False + + +@dataclass(frozen=True) +class _SpanLease: + cursor: Any + owner: object + traceparent: str | None + lifetime: _Lifetime = field(default_factory=_Lifetime) + + @dataclass(frozen=True) class _QueryContext: connection: Any owner: object traceparent: str | None + lifetime: _Lifetime = field(default_factory=_Lifetime) + + +@dataclass(frozen=True) +class _CursorAdapter: + module: Any + factory_name: str + base_cursor: type + new_internal_cursor: Callable[[Any], Any] + asynchronous: bool = False +_span_leases: ContextVar[tuple[_SpanLease, ...]] = ContextVar("postgresql_span_leases", default=()) _active_queries: ContextVar[tuple[_QueryContext, ...]] = ContextVar("postgresql_active_queries", default=()) _enabled = False @@ -63,7 +94,7 @@ class _QueryContext: def enable_postgresql_trace_context_propagation() -> bool: - """Enable process-lifetime PostgreSQL correlation for installed DBAPI instrumentation.""" + """Enable process-lifetime PostgreSQL correlation at driver execution.""" global _enabled # noqa: PLW0603 with _enable_lock: @@ -83,6 +114,31 @@ def enable_postgresql_trace_context_propagation() -> bool: logger.warning("PostgreSQL trace context propagation requires compatible DBAPI execution methods") return False + adapters = _load_cursor_adapters() + if not adapters: + logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") + return False + for adapter in adapters: + factory = getattr(adapter.module, adapter.factory_name, None) + if not _has_factory_signature(factory): + logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") + return False + + psycopg_adapter = next((adapter for adapter in adapters if hasattr(adapter.module, "psycopg")), None) + if psycopg_adapter is not None and not _has_parameter_signature( + getattr(psycopg_adapter.module.PsycopgInstrumentor, "instrument_connection", None), ("connection",) + ): + logger.warning("PostgreSQL trace context propagation requires compatible manual instrumentation") + return False + for adapter in adapters: + wrapt.wrap_function_wrapper(adapter.module, adapter.factory_name, _cursor_factory_wrapper(adapter)) + if psycopg_adapter is not None: + wrapt.wrap_function_wrapper( + psycopg_adapter.module.PsycopgInstrumentor, + "instrument_connection", + _manual_connection_wrapper(psycopg_adapter.module), + ) + wrapt.wrap_function_wrapper(cursor_tracer, "traced_execution", _execute_with_postgresql_trace_context) wrapt.wrap_function_wrapper( cursor_tracer, @@ -93,14 +149,132 @@ def enable_postgresql_trace_context_propagation() -> bool: return True +def _load_cursor_adapters() -> list[_CursorAdapter]: + adapters = [] + for name in _DRIVER_MODULE_NAMES: + try: + module = import_module(f"opentelemetry.instrumentation.{name}") + except ImportError: + continue + if name == "psycopg": + driver = module.psycopg + adapters.extend( + _CursorAdapter( + module=module, + factory_name=factory_name, + base_cursor=base_cursor, + new_internal_cursor=partial(base_cursor, row_factory=driver.rows.tuple_row), + asynchronous=asynchronous, + ) + for factory_name, base_cursor, asynchronous in ( + ("_new_cursor_factory", driver.Cursor, False), + ("_new_cursor_async_factory", driver.AsyncCursor, True), + ) + ) + else: + adapters.append( + _CursorAdapter( + module=module, + factory_name="_new_cursor_factory", + base_cursor=module.psycopg2.extensions.cursor, + new_internal_cursor=module.psycopg2.extensions.cursor, + ) + ) + return adapters + + +def _has_factory_signature(factory) -> bool: + return _has_parameter_signature(factory, ("db_api", "base_factory")) + + +def _cursor_factory_wrapper(adapter: _CursorAdapter): + mixin = _AsyncDriverCursor if adapter.asynchronous else _SyncDriverCursor + driver_cursor = type( + "SplunkPostgreSQLCursor", + (mixin, adapter.base_cursor), + {"__slots__": (), "_new_internal_cursor": staticmethod(adapter.new_internal_cursor)}, + ) + + def _wrap_cursor_factory(wrapped, _instance, args, kwargs): + base_factory = (args[1] if len(args) > 1 else kwargs.get("base_factory")) or adapter.base_cursor + if issubclass(base_factory, mixin): + factory = base_factory + elif base_factory is adapter.base_cursor: + factory = driver_cursor + else: + # Customer super() calls reach this class after the customer's hooks. + factory = type("SplunkPostgreSQLCustomCursor", (base_factory, driver_cursor), {"__slots__": ()}) + if len(args) > 1: + args = (args[0], factory, *args[2:]) + else: + kwargs = {**kwargs, "base_factory": factory} + return wrapped(*args, **kwargs) + + return _wrap_cursor_factory + + +def _manual_connection_wrapper(module): + def _instrument_connection(wrapped, _instance, args, kwargs): + arguments = inspect.signature(wrapped).bind(*args, **kwargs).arguments + connection = arguments["connection"] + original_factory = connection.cursor_factory + result = wrapped(*args, **kwargs) + asynchronous = isinstance(connection, module.psycopg.AsyncConnection) + default_factory = module.psycopg.AsyncCursor if asynchronous else module.psycopg.Cursor + if original_factory is not default_factory and connection.cursor_factory is not original_factory: + # Upstream's manual path omits base_factory; retain the customer's hooks. + factory = getattr(module, "_new_cursor_async_factory" if asynchronous else "_new_cursor_factory") + options = {name: value for name, value in arguments.items() if name != "connection"} + connection.cursor_factory = factory(base_factory=original_factory, **options) + return result + + return _instrument_connection + + +class _SyncDriverCursor: + __slots__ = () + + def execute(self, *args, **kwargs): + return _execute_driver(self, super().execute, args, kwargs, new_internal_cursor=self._new_internal_cursor) + + def executemany(self, *args, **kwargs): + return _execute_driver(self, super().executemany, args, kwargs, new_internal_cursor=self._new_internal_cursor) + + def callproc(self, *args, **kwargs): + return _execute_driver(self, super().callproc, args, kwargs, new_internal_cursor=self._new_internal_cursor) + + +class _AsyncDriverCursor: + __slots__ = () + + async def execute(self, *args, **kwargs): + return await _execute_driver_async( + self, super().execute, args, kwargs, new_internal_cursor=self._new_internal_cursor + ) + + async def executemany(self, *args, **kwargs): + return await _execute_driver_async( + self, super().executemany, args, kwargs, new_internal_cursor=self._new_internal_cursor + ) + + async def callproc(self, *args, **kwargs): + return await _execute_driver_async( + self, super().callproc, args, kwargs, new_internal_cursor=self._new_internal_cursor + ) + + def _has_execution_signature(method: Any) -> bool: + return _has_parameter_signature(method, ("self", "cursor", "query_method")) + + +def _has_parameter_signature(method: Any, names: tuple[str, ...]) -> bool: if not callable(method): return False try: - parameters = tuple(inspect.signature(method).parameters.values())[:3] + parameters = tuple(inspect.signature(method).parameters.values())[: len(names)] except (TypeError, ValueError): return False - return tuple(parameter.name for parameter in parameters) == ("self", "cursor", "query_method") and all( + return tuple(parameter.name for parameter in parameters) == names and all( parameter.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in parameters ) @@ -114,18 +288,16 @@ def _execute_with_postgresql_trace_context(traced_execution, _cursor_tracer, arg if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return traced_execution(*args, **kwargs) - propagate = _can_propagate_statement(cursor, query_args, kwargs) span_context_before = trace.get_current_span().get_span_context() def correlated_query_method(*method_args, **method_kwargs): - return _execute_correlated( - cursor, - query_method, - span_context_before, - method_args, - method_kwargs, - propagate=propagate, - ) + lease = _SpanLease(cursor, threading.current_thread(), _current_database_traceparent(span_context_before)) + token = _span_leases.set((*_span_leases.get(), lease)) + try: + return query_method(*method_args, **method_kwargs) + finally: + lease.lifetime.finish() + _span_leases.reset(token) return traced_execution(cursor, correlated_query_method, *query_args, **kwargs) @@ -138,87 +310,120 @@ async def _execute_with_postgresql_trace_context_async(traced_execution_async, _ if not _is_postgresql_client_cursor(cursor) or getattr(cursor.connection, "async_", False): return await traced_execution_async(*args, **kwargs) - propagate = _can_propagate_statement(cursor, query_args, kwargs) span_context_before = trace.get_current_span().get_span_context() async def correlated_query_method(*method_args, **method_kwargs): - return await _execute_correlated_async( - cursor, - query_method, - span_context_before, - method_args, - method_kwargs, - propagate=propagate, - ) + lease = _SpanLease(cursor, asyncio.current_task(), _current_database_traceparent(span_context_before)) + token = _span_leases.set((*_span_leases.get(), lease)) + try: + return await query_method(*method_args, **method_kwargs) + finally: + lease.lifetime.finish() + _span_leases.reset(token) return await traced_execution_async(cursor, correlated_query_method, *query_args, **kwargs) -def _execute_correlated( +def _execute_driver( cursor, query_method: Callable[..., Any], - span_context_before: SpanContext, args, kwargs, *, - propagate: bool, + new_internal_cursor: Callable[[Any], Any], ): owner = threading.current_thread() + lease = _live_span_lease(cursor, owner) + if not is_instrumentation_enabled() or lease is None: + return query_method(*args, **kwargs) + propagate = _can_propagate_statement(cursor, query_method, args, kwargs) parent = _parent_query(cursor.connection, owner) lock = nullcontext() if parent is not None else _get_sync_connection_lock(cursor.connection) with lock: - traceparent = _current_database_traceparent(span_context_before) if propagate else None - token = _active_queries.set((*_active_queries.get(), _QueryContext(cursor.connection, owner, traceparent))) + traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None + execution = _QueryContext(cursor.connection, owner, traceparent) + token = _active_queries.set((*_active_queries.get(), execution)) try: if propagate: - _update_application_name(cursor.connection, traceparent) + _update_application_name(cursor.connection, traceparent, new_internal_cursor=new_internal_cursor) return query_method(*args, **kwargs) finally: try: if propagate and (traceparent is not None or parent is not None): - _update_application_name(cursor.connection, parent.traceparent if parent is not None else None) + _update_application_name( + cursor.connection, + parent.traceparent if parent is not None else None, + new_internal_cursor=new_internal_cursor, + ) finally: + execution.lifetime.finish() _active_queries.reset(token) -async def _execute_correlated_async( +async def _execute_driver_async( cursor, query_method: Callable[..., Any], - span_context_before: SpanContext, args, kwargs, *, - propagate: bool, + new_internal_cursor: Callable[[Any], Any], ): owner = asyncio.current_task() + lease = _live_span_lease(cursor, owner) + if not is_instrumentation_enabled() or lease is None: + return await query_method(*args, **kwargs) + propagate = _can_propagate_statement(cursor, query_method, args, kwargs) parent = _parent_query(cursor.connection, owner) lock = nullcontext() if parent is not None else _get_async_connection_lock(cursor.connection) async with lock: - traceparent = _current_database_traceparent(span_context_before) if propagate else None - token = _active_queries.set((*_active_queries.get(), _QueryContext(cursor.connection, owner, traceparent))) + traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None + execution = _QueryContext(cursor.connection, owner, traceparent) + token = _active_queries.set((*_active_queries.get(), execution)) try: if propagate: - await _update_application_name_async(cursor.connection, traceparent) + await _update_application_name_async( + cursor.connection, traceparent, new_internal_cursor=new_internal_cursor + ) return await query_method(*args, **kwargs) finally: try: if propagate and (traceparent is not None or parent is not None): await _update_application_name_async( - cursor.connection, parent.traceparent if parent is not None else None + cursor.connection, + parent.traceparent if parent is not None else None, + new_internal_cursor=new_internal_cursor, ) finally: + execution.lifetime.finish() _active_queries.reset(token) +def _live_span_lease(cursor, owner: object) -> _SpanLease | None: + return next( + ( + lease + for lease in reversed(_span_leases.get()) + if lease.cursor is cursor and lease.owner is owner and lease.lifetime.is_active() + ), + None, + ) + + def _parent_query(connection, owner: object) -> _QueryContext | None: - # Context can be copied into another thread/task; that caller must acquire the lock. + # Context copies outlive their calls; neither copied nor expired ownership grants the lock. return next( - (query for query in reversed(_active_queries.get()) if query.connection is connection and query.owner is owner), + ( + query + for query in reversed(_active_queries.get()) + if query.connection is connection and query.owner is owner and query.lifetime.is_active() + ), None, ) -def _can_propagate_statement(cursor, args, kwargs) -> bool: +def _can_propagate_statement(cursor, query_method, args, kwargs) -> bool: + if query_method.__name__ == "callproc": + return True statement = args[0] if args else kwargs.get("query") if isinstance(statement, bytes): statement = statement.decode("ascii", errors="replace") @@ -233,7 +438,9 @@ def _can_propagate_statement(cursor, args, kwargs) -> bool: if ";" in statement: return False statement = _strip_leading_comments(statement) - operation = re.match(r"[a-z]+", statement, re.IGNORECASE) + if statement.startswith("("): + return True + operation = re.match(r"[a-z_][a-z_0-9$]*", statement, re.IGNORECASE) return operation is not None and operation[0].upper() not in _TRANSACTION_OPERATIONS @@ -280,10 +487,10 @@ def _current_database_traceparent(span_context_before: SpanContext) -> str | Non return f"00-{current.trace_id:032x}-{current.span_id:016x}-{int(current.trace_flags) & 0xFF:02x}" -def _update_application_name(connection, traceparent: str | None) -> None: +def _update_application_name(connection, traceparent: str | None, *, new_internal_cursor: Callable[[Any], Any]) -> None: statement = _application_name_statement(traceparent) try: - cursor = connection.cursor() + cursor = new_internal_cursor(connection) try: with suppress_instrumentation(): result = cursor.execute(statement) @@ -297,10 +504,12 @@ def _update_application_name(connection, traceparent: str | None) -> None: logger.exception("Failed to set PostgreSQL application_name") -async def _update_application_name_async(connection, traceparent: str | None) -> None: +async def _update_application_name_async( + connection, traceparent: str | None, *, new_internal_cursor: Callable[[Any], Any] +) -> None: statement = _application_name_statement(traceparent) try: - cursor = connection.cursor() + cursor = new_internal_cursor(connection) try: with suppress_instrumentation(): result = cursor.execute(statement) diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py index 662d1155..4d7b9c31 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/smoke/postgresql_trace_context.py @@ -4,6 +4,7 @@ from contextlib import suppress from importlib import import_module from importlib.metadata import version +from unittest.mock import patch os.environ.setdefault("OTEL_LOGS_EXPORTER", "none") @@ -164,15 +165,16 @@ async def execute(self, query, *args, **kwargs): query = "SET nonexistent_correlation_parameter = 'failure'" return await super().execute(query, *args, **kwargs) - connection = await psycopg.AsyncConnection.connect(cursor_factory=RejectingSetCursor) + connection = await psycopg.AsyncConnection.connect() try: - async with connection.cursor() as cursor: - failed = False - try: - await cursor.execute("SELECT 1") - except psycopg.errors.InFailedSqlTransaction: - failed = True - assert failed + with patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)): + async with connection.cursor() as cursor: + failed = False + try: + await cursor.execute("SELECT 1") + except psycopg.errors.InFailedSqlTransaction: + failed = True + assert failed assert connection.info.transaction_status == psycopg.pq.TransactionStatus.INERROR finally: await connection.close() @@ -227,9 +229,12 @@ def execute(self, query, *args, **kwargs): query = "SET nonexistent_correlation_parameter = 'failure'" return super().execute(query, *args, **kwargs) - connection = driver.connect(cursor_factory=RejectingSetCursor) + connection = driver.connect() try: - with connection.cursor() as cursor: + with ( + patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)), + connection.cursor() as cursor, + ): try: cursor.execute("SELECT 1") except transaction_error: @@ -244,10 +249,13 @@ def execute(self, query, *args, **kwargs): finally: connection.close() - connection = driver.connect(cursor_factory=RejectingSetCursor) + connection = driver.connect() try: connection.autocommit = True - with connection.cursor() as cursor: + with ( + patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)), + connection.cursor() as cursor, + ): cursor.execute("SELECT 1") assert cursor.fetchone()[0] == 1 finally: @@ -426,12 +434,15 @@ def execute(self, query, *args, **kwargs): query = "SET nonexistent_correlation_parameter = 'failure'" return super().execute(query, *args, **kwargs) - connection = driver.connect(cursor_factory=RejectingCleanupCursor) + connection = driver.connect() try: with suppress_instrumentation(), connection.cursor() as cursor: cursor.execute("CREATE TEMP TABLE cleanup_failure (value int)") connection.commit() - with connection.cursor() as cursor: + with ( + patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingCleanupCursor)), + connection.cursor() as cursor, + ): cursor.execute("INSERT INTO cleanup_failure VALUES (1) RETURNING value") assert cursor.fetchone()[0] == 1 status = ( @@ -446,6 +457,221 @@ def execute(self, query, *args, **kwargs): connection.close() +def _parenthesized_queries(driver, exporter): + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, connection.cursor() as cursor: + for query in ( + "(SELECT current_setting('application_name'))", + "((SELECT current_setting('application_name'))) UNION (SELECT current_setting('application_name'))", + ): + exporter.clear() + cursor.execute(query) + _assert_database_span(exporter, cursor.fetchone()[0]) + if driver.__name__ == "psycopg2": + with suppress_instrumentation(): + cursor.execute( + "CREATE FUNCTION begin_job() RETURNS text LANGUAGE sql " + "AS $$ SELECT current_setting('application_name') $$" + ) + exporter.clear() + cursor.callproc("begin_job") + _assert_database_span(exporter, cursor.fetchone()[0]) + with suppress_instrumentation(): + cursor.execute("DROP FUNCTION begin_job()") + + +def _skipped_middle_query(driver, exporter): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + inner_values = [] + + class CustomerCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query == _QUERY: + with self.connection.cursor() as middle: + middle.execute("SELECT ';'") + assert middle.fetchone()[0] == ";" + elif query == "SELECT ';'": + with self.connection.cursor() as inner: + inner.execute("SELECT current_setting('application_name')") + inner_values.append(inner.fetchone()[0]) + return super().execute(query, *args, **kwargs) + + with ( + driver.connect(cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME) as connection, + connection.cursor() as cursor, + ): + exporter.clear() + cursor.execute(_QUERY) + outer = _assert_database_span(exporter, cursor.fetchone()[0]) + inner = _assert_database_span(exporter, inner_values[0]) + assert inner.context.span_id != outer.context.span_id + + +def _custom_commit_before_query(driver, exporter): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + + class CustomerCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query == _QUERY: + with self.connection.cursor() as auxiliary: + auxiliary.execute("COMMIT") + status = ( + self.connection.info.transaction_status + if driver.__name__ == "psycopg" + else self.connection.get_transaction_status() + ) + assert status == 0 + return super().execute(query, *args, **kwargs) + + for autocommit in (False, True): + with driver.connect(cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.autocommit = autocommit + exporter.clear() + with connection.cursor() as cursor: + cursor.execute(_QUERY) + _assert_database_span(exporter, cursor.fetchone()[0]) + connection.rollback() + with suppress_instrumentation(): + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +async def _delayed_sync_callback(driver, exporter): + cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor + loop = asyncio.get_running_loop() + completed = loop.create_future() + + def _execute_delayed_query(connection): + try: + with connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name')") + completed.set_result(cursor.fetchone()[0]) + except (AssertionError, driver.Error) as error: + completed.set_exception(error) + + class CustomerCursor(cursor_class): + def execute(self, query, *args, **kwargs): + if query == _QUERY: + loop.call_soon(_execute_delayed_query, self.connection) + return super().execute(query, *args, **kwargs) + + with driver.connect(cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME) as connection: + exporter.clear() + with connection.cursor() as cursor: + cursor.execute(_QUERY) + outer = _assert_database_span(exporter, cursor.fetchone()[0]) + delayed = _assert_database_span(exporter, await asyncio.wait_for(completed, timeout=5)) + assert delayed.context.span_id != outer.context.span_id + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +async def _async_driver_regressions(exporter): + import psycopg + + inner_values = [] + + class CustomerCursor(psycopg.AsyncCursor): + async def execute(self, query, *args, **kwargs): + if query == _QUERY: + async with self.connection.cursor() as middle: + await middle.execute("SELECT ';'") + assert (await middle.fetchone())[0] == ";" + elif query == "SELECT ';'": + async with self.connection.cursor() as inner: + await inner.execute("SELECT current_setting('application_name')") + inner_values.append((await inner.fetchone())[0]) + elif query == "SELECT current_setting('application_name') AS after_commit": + async with self.connection.cursor() as auxiliary: + await auxiliary.execute("COMMIT") + assert self.connection.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + return await super().execute(query, *args, **kwargs) + + for autocommit in (False, True): + async with ( + await psycopg.AsyncConnection.connect( + autocommit=autocommit, cursor_factory=CustomerCursor, application_name=_DEFAULT_APPLICATION_NAME + ) as connection, + connection.cursor() as cursor, + ): + for query in ( + "(SELECT current_setting('application_name'))", + "((SELECT current_setting('application_name'))) UNION (SELECT current_setting('application_name'))", + _QUERY, + "SELECT current_setting('application_name') AS after_commit", + ): + exporter.clear() + inner_values.clear() + await cursor.execute(query) + outer = _assert_database_span(exporter, (await cursor.fetchone())[0]) + if inner_values: + inner = _assert_database_span(exporter, inner_values[0]) + assert inner.context.span_id != outer.context.span_id + await connection.rollback() + with suppress_instrumentation(): + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + await connection.rollback() + + +def _manual_custom_cursor(exporter): + import psycopg + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + + hook_queries = [] + + class CustomerCursor(psycopg.Cursor): + def execute(self, query, *args, **kwargs): + hook_queries.append(query) + return super().execute(query, *args, **kwargs) + + with psycopg.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.cursor_factory = CustomerCursor + PsycopgInstrumentor().instrument_connection(connection) + instrumented_factory = connection.cursor_factory + PsycopgInstrumentor().instrument_connection(connection) + assert connection.cursor_factory is instrumented_factory + exporter.clear() + with connection.cursor() as cursor: + cursor.execute(_QUERY) + _assert_database_span(exporter, cursor.fetchone()[0]) + assert hook_queries == [_QUERY] + PsycopgInstrumentor().uninstrument_connection(connection) + assert connection.cursor_factory is CustomerCursor + with connection.cursor() as cursor: + cursor.execute("SELECT current_setting('application_name')") + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + + +async def _async_manual_custom_cursor(exporter): + import psycopg + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + + hook_queries = [] + + class CustomerCursor(psycopg.AsyncCursor): + async def execute(self, query, *args, **kwargs): + hook_queries.append(query) + return await super().execute(query, *args, **kwargs) + + async with await psycopg.AsyncConnection.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.cursor_factory = CustomerCursor + PsycopgInstrumentor().instrument_connection(connection) + instrumented_factory = connection.cursor_factory + PsycopgInstrumentor().instrument_connection(connection) + assert connection.cursor_factory is instrumented_factory + exporter.clear() + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + _assert_database_span(exporter, (await cursor.fetchone())[0]) + assert hook_queries == [_QUERY] + PsycopgInstrumentor().uninstrument_connection(connection) + assert connection.cursor_factory is CustomerCursor + async with connection.cursor() as cursor: + await cursor.execute("SELECT current_setting('application_name')") + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + + def _psycopg_server_cursor_boundary(exporter): import psycopg @@ -567,10 +793,17 @@ def main(): _failed_cleanup_command(driver) _nested_cursor_queries(driver, exporter) _transaction_sql(driver, exporter) + _parenthesized_queries(driver, exporter) + _skipped_middle_query(driver, exporter) + _custom_commit_before_query(driver, exporter) + asyncio.run(_delayed_sync_callback(driver, exporter)) asyncio.run(_async_queries(exporter, tracer)) asyncio.run(_async_nested_cursor_queries(exporter)) asyncio.run(_async_transaction_sql(exporter)) + asyncio.run(_async_driver_regressions(exporter)) _custom_cursor(exporter) + _manual_custom_cursor(exporter) + asyncio.run(_async_manual_custom_cursor(exporter)) _psycopg_server_cursor_boundary(exporter) _psycopg2_commit_boundary(exporter) asyncio.run(_async_query_cancellation()) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 04f73367..345228a2 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import inspect import threading from contextlib import nullcontext from types import SimpleNamespace @@ -24,6 +25,10 @@ from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags from splunk_otel.postgresql import ( + _AsyncDriverCursor, + _CursorAdapter, + _SyncDriverCursor, + _cursor_factory_wrapper, _execute_with_postgresql_trace_context, _execute_with_postgresql_trace_context_async, ) @@ -48,8 +53,15 @@ is_remote=False, trace_flags=TraceFlags(TraceFlags.SAMPLED), ) +INNER_DATABASE_CONTEXT = SpanContext( + trace_id=PARENT_CONTEXT.trace_id, + span_id=0x4444444444444444, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), +) TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01" SECOND_TRACEPARENT = "00-11111111111111111111111111111111-3333333333333333-01" +INNER_TRACEPARENT = "00-11111111111111111111111111111111-4444444444444444-01" class FakeConnection: @@ -77,7 +89,7 @@ def cursor(self): return cursor -class FakePsycopgCursor: +class _FakeDriverCursor: __module__ = "psycopg" def __init__(self, connection): @@ -98,8 +110,11 @@ def execute(self, query): def close(self): self.closed = True + def callproc(self, name): + return _FakeDriverCursor.execute(self, name) + -class FakeAsyncPsycopgCursor: +class _FakeAsyncDriverCursor: __module__ = "psycopg.fake" def __init__(self, connection): @@ -110,7 +125,9 @@ def __init__(self, connection): async def execute(self, query): self.calls.append(query) self.connection.calls.append(query) - self.connection.before_execute(query) + result = self.connection.before_execute(query) + if inspect.isawaitable(result): + await result if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name elif isinstance(query, str) and query.startswith("SET application_name = '"): @@ -121,6 +138,34 @@ async def close(self): self.closed = True +def _fake_internal_cursor_factory(connection): + return connection.cursor() + + +class FakePsycopgCursor(_SyncDriverCursor, _FakeDriverCursor): + _new_internal_cursor = staticmethod(_fake_internal_cursor_factory) + + +class FakeAsyncPsycopgCursor(_AsyncDriverCursor, _FakeAsyncDriverCursor): + _new_internal_cursor = staticmethod(_fake_internal_cursor_factory) + + +@pytest.fixture(autouse=True) +def fake_cursor_adapters(monkeypatch): + def _fake_cursor_factory(db_api=None, base_factory=None): # noqa: ARG001 + return base_factory + + module = SimpleNamespace(_new_cursor_factory=_fake_cursor_factory) + adapter = _CursorAdapter( + module=module, + factory_name="_new_cursor_factory", + base_cursor=_FakeDriverCursor, + new_internal_cursor=_fake_internal_cursor_factory, + ) + monkeypatch.setattr(postgresql, "_load_cursor_adapters", lambda: [adapter]) + return module + + def _fake_traced_execution(_cursor, query_method, *args, **kwargs): with trace.use_span(NonRecordingSpan(DATABASE_CONTEXT)): return query_method(*args, **kwargs) @@ -141,12 +186,17 @@ async def _fake_nested_traced_execution_async(_cursor, query_method, *args, **kw return await query_method(*args, **kwargs) +def _fake_inner_traced_execution(_cursor, query_method, *args, **kwargs): + with trace.use_span(NonRecordingSpan(INNER_DATABASE_CONTEXT)): + return query_method(*args, **kwargs) + + def _fake_untraced_execution(_cursor, query_method, *args, **kwargs): return query_method(*args, **kwargs) @pytest.mark.parametrize("nested_fails", [False, True]) -def test_nested_customer_query_restores_outer_context_without_deadlock(nested_fails): +def test_custom_cursor_queries_run_before_outer_driver_execution(nested_fails): connection = FakeConnection() error = RuntimeError("nested query failed") results = [] @@ -167,7 +217,7 @@ def execute(self, query): assert result == SECOND_TRACEPARENT if nested_fails: assert raised.value is error - assert connection.application_name == TRACEPARENT + assert connection.application_name == "startup-app" return super().execute(query) cursor = CustomerCursor(connection) @@ -185,9 +235,9 @@ def _execute_customer_query(): assert not thread.is_alive(), "Nested customer query deadlocked" assert results == [TRACEPARENT] assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", f"SET application_name = '{SECOND_TRACEPARENT}'", "SELECT nested", + "SET application_name = DEFAULT", f"SET application_name = '{TRACEPARENT}'", "SELECT outer", "SET application_name = DEFAULT", @@ -195,7 +245,7 @@ def _execute_customer_query(): @pytest.mark.parametrize("nested_fails", [False, True]) -def test_async_nested_customer_query_restores_outer_context(nested_fails): +def test_async_custom_cursor_queries_run_before_outer_driver_execution(nested_fails): async def run(): connection = FakeAsyncConnection() error = RuntimeError("nested query failed") @@ -216,7 +266,7 @@ async def execute(self, query): assert result == SECOND_TRACEPARENT if nested_fails: assert raised.value is error - assert connection.application_name == TRACEPARENT + assert connection.application_name == "startup-app" return await super().execute(query) cursor = CustomerCursor(connection) @@ -228,9 +278,9 @@ async def execute(self, query): ) assert result == TRACEPARENT assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", f"SET application_name = '{SECOND_TRACEPARENT}'", "SELECT nested", + "SET application_name = DEFAULT", f"SET application_name = '{TRACEPARENT}'", "SELECT outer", "SET application_name = DEFAULT", @@ -239,7 +289,7 @@ async def execute(self, query): asyncio.run(run()) -def test_async_child_task_waits_for_outer_connection_lock(): +def test_async_custom_cursor_can_await_auxiliary_task_before_outer_query(): async def run(): connection = FakeAsyncConnection() nested_started = asyncio.Event() @@ -256,8 +306,8 @@ class CustomerCursor(FakeAsyncPsycopgCursor): async def execute(self, query): children.append(asyncio.create_task(_execute_nested_customer_query())) await nested_started.wait() - assert not children[0].done() - assert connection.application_name == TRACEPARENT + assert await children[0] == SECOND_TRACEPARENT + assert connection.application_name == "startup-app" return await super().execute(query) cursor = CustomerCursor(connection) @@ -265,14 +315,13 @@ async def execute(self, query): _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT outer"), {} ) assert result == TRACEPARENT - assert await children[0] == SECOND_TRACEPARENT assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", - "SELECT outer", - "SET application_name = DEFAULT", f"SET application_name = '{SECOND_TRACEPARENT}'", "SELECT child", "SET application_name = DEFAULT", + f"SET application_name = '{TRACEPARENT}'", + "SELECT outer", + "SET application_name = DEFAULT", ] asyncio.run(asyncio.wait_for(run(), timeout=2)) @@ -339,6 +388,104 @@ def test_data_statement_after_comments_propagates_context(statement): ) +@pytest.mark.parametrize("statement", ["(SELECT 1)", "((SELECT 1)) UNION (SELECT 2)", "/* comment */ (SELECT 1)"]) +def test_parenthesized_queries_propagate_context(statement): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, statement), {} + ) + assert result == TRACEPARENT + assert connection.application_name == "startup-app" + + +def test_callproc_function_name_is_not_a_transaction_command(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.callproc, "begin_job"), {} + ) + assert result == TRACEPARENT + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "begin_job", + "SET application_name = DEFAULT", + ] + + +def test_delayed_callback_acquires_lock_and_does_not_restore_finished_context(monkeypatch): + async def run(): + connection = FakeConnection() + acquisitions = [] + results = [] + loop = asyncio.get_running_loop() + completed = loop.create_future() + + class ConnectionLock: + def __enter__(self): + acquisitions.append("acquired") + + def __exit__(self, *_args): + pass + + monkeypatch.setattr(postgresql, "_get_sync_connection_lock", lambda _connection: ConnectionLock()) + + def _execute_delayed_query(): + cursor = FakePsycopgCursor(connection) + results.append( + _execute_with_postgresql_trace_context( + _fake_nested_traced_execution, None, (cursor, cursor.execute, "SELECT delayed"), {} + ) + ) + completed.set_result(None) + + def _fake_schedule_callback(query): + if query == "SELECT outer": + loop.call_soon(_execute_delayed_query) + + connection.before_execute = _fake_schedule_callback + cursor = FakePsycopgCursor(connection) + assert ( + _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT outer"), {} + ) + == TRACEPARENT + ) + await asyncio.wait_for(completed, timeout=1) + assert results == [SECOND_TRACEPARENT] + assert acquisitions == ["acquired", "acquired"] + assert connection.application_name == "startup-app" + assert connection.calls[-1] == "SET application_name = DEFAULT" + + asyncio.run(run()) + + +def test_skipped_driver_query_keeps_inherited_outer_setting(): + connection = FakeConnection() + + def _fake_nested_driver_calls(query): + if query == "SELECT outer": + middle = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_nested_traced_execution, None, (middle, middle.execute, "SELECT ';'"), {} + ) + assert result == TRACEPARENT + elif query == "SELECT ';'": + inner = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_inner_traced_execution, None, (inner, inner.execute, "SELECT inner"), {} + ) + assert result == INNER_TRACEPARENT + + connection.before_execute = _fake_nested_driver_calls + cursor = FakePsycopgCursor(connection) + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT outer"), {} + ) + assert result == TRACEPARENT + assert connection.application_name == "startup-app" + + def test_query_sets_database_span_traceparent(): connection = FakeConnection() cursor = FakePsycopgCursor(connection) @@ -650,6 +797,91 @@ async def traced_execution_async(self, *, cursor, query_method): assert not postgresql.enable_postgresql_trace_context_propagation() +@pytest.mark.parametrize("missing_factory", [False, True]) +def test_public_enable_checks_all_factories_before_wrapping(monkeypatch, fake_cursor_adapters, missing_factory): + class CursorTracer(FakeCursorTracer): + pass + + original_execution = CursorTracer.traced_execution + factory_name = "_new_cursor_factory" + original_factory = getattr(fake_cursor_adapters, factory_name) + incompatible_module = SimpleNamespace() + if not missing_factory: + setattr(incompatible_module, factory_name, lambda wrong_parameter: wrong_parameter) + adapters = [ + _CursorAdapter(fake_cursor_adapters, "_new_cursor_factory", _FakeDriverCursor, _fake_internal_cursor_factory), + _CursorAdapter(incompatible_module, "_new_cursor_factory", _FakeDriverCursor, _fake_internal_cursor_factory), + ] + monkeypatch.setattr(postgresql, "_enabled", False) + monkeypatch.setattr(postgresql, "import_module", lambda _name: SimpleNamespace(CursorTracer=CursorTracer)) + monkeypatch.setattr(postgresql, "_load_cursor_adapters", lambda: adapters) + + assert not postgresql.enable_postgresql_trace_context_propagation() + assert CursorTracer.traced_execution is original_execution + assert getattr(fake_cursor_adapters, factory_name) is original_factory + + +@pytest.mark.parametrize("positional_factory", [False, True]) +def test_cursor_factory_places_driver_adapter_after_customer_hooks(positional_factory): + connection = FakeConnection() + calls = [] + + class CustomerCursor(_FakeDriverCursor): + def execute(self, query): + calls.append(self.connection.application_name) + return super().execute(query) + + adapter = _CursorAdapter(None, "_new_cursor_factory", _FakeDriverCursor, _fake_internal_cursor_factory) + + def _fake_instrumentation_cursor_factory(db_api, base_factory): # noqa: ARG001 + return base_factory + + args = (None, CustomerCursor) if positional_factory else () + kwargs = {} if positional_factory else {"db_api": None, "base_factory": CustomerCursor} + factory = _cursor_factory_wrapper(adapter)(_fake_instrumentation_cursor_factory, None, args, kwargs) + cursor = factory(connection) + result = _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + + assert result == TRACEPARENT + assert calls == ["customer-app"] + assert isinstance(cursor, CustomerCursor) + assert connection.calls[-1] == "SET application_name = DEFAULT" + + +def test_delayed_native_call_cannot_reuse_finished_span_lease(): + async def run(): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + loop = asyncio.get_running_loop() + completed = loop.create_future() + + def _execute_uninstrumented_callback(): + completed.set_result(cursor.execute("SELECT delayed")) + + def _fake_schedule_native_callback(query): + if query == "SELECT outer": + loop.call_soon(_execute_uninstrumented_callback) + + connection.before_execute = _fake_schedule_native_callback + assert ( + _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT outer"), {} + ) + == TRACEPARENT + ) + assert await asyncio.wait_for(completed, timeout=1) == "startup-app" + assert connection.calls == [ + f"SET application_name = '{TRACEPARENT}'", + "SELECT outer", + "SET application_name = DEFAULT", + "SELECT delayed", + ] + + asyncio.run(run()) + + @pytest.mark.parametrize("flags", [0, 1, 2, 3]) def test_query_propagates_valid_context_flags(flags): context = SpanContext( @@ -769,19 +1001,19 @@ async def run(): second_query_started = asyncio.Event() release_first_query = asyncio.Event() - class BlockingQueryCursor(FakeAsyncPsycopgCursor): - async def execute(self, query): - result = await super().execute(query) + async def _fake_blocking_driver_execute(query): + if query == "SELECT first": first_query_started.set() await release_first_query.wait() - return result + + connection.before_execute = _fake_blocking_driver_execute async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): with trace.use_span(NonRecordingSpan(SECOND_DATABASE_CONTEXT)): second_query_started.set() return await query_method(*args, **kwargs) - first_cursor = BlockingQueryCursor(connection) + first_cursor = FakeAsyncPsycopgCursor(connection) second_cursor = FakeAsyncPsycopgCursor(connection) first = asyncio.create_task( _execute_with_postgresql_trace_context_async( From bedc84b8fb9a10b095356b420d89aef245e7b696 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Thu, 15 Oct 2026 09:49:23 -0400 Subject: [PATCH 06/10] Align database trace propagation flag with the Node distro Rename the opt-in setting to SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED across runtime configuration, documentation, existing tests, and the PostgreSQL CI smoke matrix. Clarify that Python currently applies the setting only to PostgreSQL. Assisted-by: Codex (GPT-6) --- .github/workflows/ci-main.yml | 2 +- README.md | 2 +- docs/postgresql.md | 4 +++- src/splunk_otel/env.py | 2 +- src/splunk_otel/runtime.py | 4 ++-- tests/smoke/postgresql_trace_context.py | 2 +- tests/test_configurator.py | 2 +- tests/test_runtime.py | 2 +- 8 files changed, 11 insertions(+), 9 deletions(-) diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 76c18e3c..23476012 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -60,7 +60,7 @@ jobs: python-version: ['3.10', '3.14'] env: POSTGRESQL_SMOKE_INSTRUMENTATION: ${{ matrix.mode }} - SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED: ${{ matrix.mode != 'disabled' }} + SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED: ${{ matrix.mode != 'disabled' }} services: postgres: image: postgres:16-alpine diff --git a/README.md b/README.md index 2f5ea863..5425368e 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ metrics. ### Experimental PostgreSQL trace correlation -Set `SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true` to link PostgreSQL +Set `SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED=true` to link PostgreSQL query samples to APM traces and database client spans. This feature is disabled by default. See [PostgreSQL trace correlation](docs/postgresql.md) for setup, supported drivers, and limitations. diff --git a/docs/postgresql.md b/docs/postgresql.md index e525fb87..2469b37d 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -42,10 +42,12 @@ For PostgreSQL, it installs: Start the application with: ```sh -SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED=true \ +SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED=true \ opentelemetry-instrument python app.py ``` +Python currently applies this setting only to PostgreSQL. + The environment switch applies when `OTEL_CONFIG_FILE` is unset. Activation lasts for the process lifetime. To turn it off, unset the switch and restart. diff --git a/src/splunk_otel/env.py b/src/splunk_otel/env.py index b9dc0fda..2e8ff9f0 100644 --- a/src/splunk_otel/env.py +++ b/src/splunk_otel/env.py @@ -47,7 +47,7 @@ SPLUNK_OPAMP_ENABLED = "SPLUNK_OPAMP_ENABLED" SPLUNK_OPAMP_ENDPOINT = "SPLUNK_OPAMP_ENDPOINT" SPLUNK_OPAMP_POLLING_INTERVAL = "SPLUNK_OPAMP_POLLING_INTERVAL" -SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED = "SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED" +SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED = "SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED" SPLUNK_PROFILER_ENABLED = "SPLUNK_PROFILER_ENABLED" SPLUNK_PROFILER_CALL_STACK_INTERVAL = "SPLUNK_PROFILER_CALL_STACK_INTERVAL" SPLUNK_PROFILER_LOGS_ENDPOINT = "SPLUNK_PROFILER_LOGS_ENDPOINT" diff --git a/src/splunk_otel/runtime.py b/src/splunk_otel/runtime.py index 41963f2a..81a09cf6 100644 --- a/src/splunk_otel/runtime.py +++ b/src/splunk_otel/runtime.py @@ -17,7 +17,7 @@ from opentelemetry.instrumentation.environment_variables import OTEL_PYTHON_DISABLED_INSTRUMENTATIONS from opentelemetry.instrumentation.logging import LoggingInstrumentor -from splunk_otel.env import SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED, Env +from splunk_otel.env import SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED, Env from splunk_otel.postgresql import enable_postgresql_trace_context_propagation _DISABLED_INSTRUMENTATIONS_WILDCARD = "*" @@ -26,7 +26,7 @@ def configure_postgresql_trace_context_propagation(env: Env) -> None: - if env.is_true(SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED, "false"): + if env.is_true(SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED, "false"): enable_postgresql_trace_context_propagation() diff --git a/tests/smoke/postgresql_trace_context.py b/tests/smoke/postgresql_trace_context.py index 4d7b9c31..b533aefa 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/smoke/postgresql_trace_context.py @@ -14,7 +14,7 @@ os.environ.setdefault("OTEL_TRACES_EXPORTER", "none") os.environ.setdefault("OTEL_TRACES_SAMPLER", "parentbased_always_on") os.environ.setdefault("SPLUNK_OPAMP_ENABLED", "false") -os.environ.setdefault("SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") +os.environ.setdefault("SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") os.environ.setdefault("PGHOST", "127.0.0.1") os.environ.setdefault("PGPORT", "5432") os.environ.setdefault("PGDATABASE", "otel") diff --git a/tests/test_configurator.py b/tests/test_configurator.py index 089f1ca9..e45c3073 100644 --- a/tests/test_configurator.py +++ b/tests/test_configurator.py @@ -93,7 +93,7 @@ def test_declarative_server_timing_reads_splunk_http_configuration(): @pytest.mark.parametrize("use_config_file", [False, True]) def test_configurator_activates_postgresql_only_without_config_file(monkeypatch, use_config_file): enable_calls = [] - monkeypatch.setenv("SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") + monkeypatch.setenv("SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") if use_config_file: monkeypatch.setenv("OTEL_CONFIG_FILE", "config.yaml") else: diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 867c735e..d89bdd16 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -52,7 +52,7 @@ def test_postgresql_trace_context_propagation_requires_opt_in(monkeypatch, value "splunk_otel.runtime.enable_postgresql_trace_context_propagation", lambda: enable_calls.append(True), ) - env_store = {} if value is None else {"SPLUNK_POSTGRESQL_TRACE_CONTEXT_PROPAGATION_ENABLED": value} + env_store = {} if value is None else {"SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED": value} configure_postgresql_trace_context_propagation(Env(env_store)) From 9f3bebe9d403b7db2e4a0ef89e4f9c9a9263f3d1 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Thu, 15 Oct 2026 10:48:25 -0400 Subject: [PATCH 07/10] Run PostgreSQL correlation integration checks with oteltest Exercise automatic CLI startup, manual instrumentation, and disabled propagation with real PostgreSQL and the oteltest OTLP receiver. Match query application_name values to received database client span IDs and retain the existing database regressions in a shared helper. Update the PostgreSQL CI matrix, add callback failure checks, and document local execution. Make failed applications and missing or mismatched exports fail the runner explicitly. Assisted-by: Codex (GPT-6) --- .github/workflows/ci-main.yml | 18 +- docs/postgresql.md | 20 ++ .../postgresql/regressions.py} | 107 +++-------- tests/integration/postgresql/trace_context.py | 180 ++++++++++++++++++ tests/test_postgresql_oteltest.py | 119 ++++++++++++ 5 files changed, 353 insertions(+), 91 deletions(-) rename tests/{smoke/postgresql_trace_context.py => integration/postgresql/regressions.py} (90%) create mode 100644 tests/integration/postgresql/trace_context.py create mode 100644 tests/test_postgresql_oteltest.py diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 23476012..501d8002 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -50,8 +50,8 @@ jobs: - name: Run oteltests run: oteltest tests/integration/*.py - postgresql-smoke: - name: PostgreSQL trace-context smoke test (Python ${{ matrix.python-version }}, ${{ matrix.mode }}) + postgresql-integration: + name: PostgreSQL oteltest (Python ${{ matrix.python-version }}, ${{ matrix.mode }}) runs-on: ubuntu-latest strategy: fail-fast: false @@ -59,8 +59,7 @@ jobs: mode: [auto, manual, disabled] python-version: ['3.10', '3.14'] env: - POSTGRESQL_SMOKE_INSTRUMENTATION: ${{ matrix.mode }} - SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED: ${{ matrix.mode != 'disabled' }} + POSTGRESQL_TEST_INSTRUMENTATION: ${{ matrix.mode }} services: postgres: image: postgres:16-alpine @@ -84,12 +83,9 @@ jobs: with: python-version: ${{ matrix.python-version }} - - name: Install smoke-test dependencies - run: >- - pip install . "psycopg[binary]" psycopg2-binary - opentelemetry-instrumentation-psycopg - opentelemetry-instrumentation-psycopg2 + - name: Install oteltest + run: pip install 'oteltest>=0.49.0' - - name: Run PostgreSQL smoke test + - name: Run PostgreSQL oteltest timeout-minutes: 5 - run: python tests/smoke/postgresql_trace_context.py + run: oteltest tests/integration/postgresql/trace_context.py diff --git a/docs/postgresql.md b/docs/postgresql.md index 2469b37d..6dc9ed32 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -131,3 +131,23 @@ and declines incompatible versions before installing wrappers. Regression tests must also check behavior when instrumentation dependencies change. The current validation covers instrumentation 0.66b1, Psycopg 3.3.6, Psycopg2 2.9.13, and PostgreSQL 16; it does not establish a broader supported version range. + +## Integration tests + +The PostgreSQL CI job uses `oteltest` with PostgreSQL 16 on Python 3.10 and 3.14. +It covers automatic instrumentation, manual instrumentation, and disabled +propagation. Query results must match the database client spans received over +OTLP. The same job runs the transaction, custom cursor, and concurrency regressions. + +To run locally, start a test database, install `oteltest>=0.49.0`, and run: + +```sh +PGHOST=127.0.0.1 PGPORT=5432 PGDATABASE=otel PGUSER=otel PGPASSWORD=otel \ +POSTGRESQL_TEST_INSTRUMENTATION=auto \ +oteltest tests/integration/postgresql/trace_context.py +``` + +Repeat with `POSTGRESQL_TEST_INSTRUMENTATION=manual` and `disabled`. +The test installs this checkout and both driver instrumentations in an isolated +environment. It exports traces to the local `oteltest` receiver on port 4317. +Collector ingestion and Database Monitoring UI checks are separate manual tests. diff --git a/tests/smoke/postgresql_trace_context.py b/tests/integration/postgresql/regressions.py similarity index 90% rename from tests/smoke/postgresql_trace_context.py rename to tests/integration/postgresql/regressions.py index b533aefa..2105a05e 100644 --- a/tests/smoke/postgresql_trace_context.py +++ b/tests/integration/postgresql/regressions.py @@ -1,26 +1,10 @@ import asyncio -import os from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from importlib import import_module -from importlib.metadata import version from unittest.mock import patch -os.environ.setdefault("OTEL_LOGS_EXPORTER", "none") -os.environ.setdefault("OTEL_METRICS_EXPORTER", "none") -os.environ.setdefault("OTEL_PYTHON_DISABLED_INSTRUMENTATIONS", "system_metrics") -os.environ.setdefault("OTEL_SERVICE_NAME", "postgresql-trace-context-smoke-test") -os.environ.setdefault("OTEL_TRACES_EXPORTER", "none") -os.environ.setdefault("OTEL_TRACES_SAMPLER", "parentbased_always_on") -os.environ.setdefault("SPLUNK_OPAMP_ENABLED", "false") -os.environ.setdefault("SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED", "true") -os.environ.setdefault("PGHOST", "127.0.0.1") -os.environ.setdefault("PGPORT", "5432") -os.environ.setdefault("PGDATABASE", "otel") -os.environ.setdefault("PGUSER", "otel") -os.environ.setdefault("PGPASSWORD", "otel") - from opentelemetry import trace from opentelemetry.instrumentation.utils import suppress_instrumentation from opentelemetry.trace import NonRecordingSpan, SpanContext, SpanKind, TraceFlags @@ -749,70 +733,33 @@ async def _async_query_cancellation(): assert (await cursor.fetchone())[0] == 1 -def main(): - from opentelemetry.instrumentation.auto_instrumentation import initialize - from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor - from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter - from splunk_otel import init_splunk_otel - from splunk_otel.postgresql import enable_postgresql_trace_context_propagation - - mode = os.environ.get("POSTGRESQL_SMOKE_INSTRUMENTATION", "auto") - if mode == "manual": - init_splunk_otel() - PsycopgInstrumentor().instrument() - Psycopg2Instrumentor().instrument() - else: - initialize(swallow_exceptions=False) - - exporter = InMemorySpanExporter() - trace.get_tracer_provider().add_span_processor(SimpleSpanProcessor(exporter)) - tracer = trace.get_tracer(__name__) - +def run_regressions(exporter, tracer): import psycopg import psycopg2 + from splunk_otel.postgresql import enable_postgresql_trace_context_propagation - if mode == "disabled": - for driver in (psycopg, psycopg2): - with ( - driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, - connection.cursor() as cursor, - ): - cursor.execute(_QUERY) - assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME - else: - assert enable_postgresql_trace_context_propagation() - assert enable_postgresql_trace_context_propagation() - for driver in (psycopg, psycopg2): - _sync_queries(driver, exporter, tracer) - _sync_concurrency(driver, exporter, tracer) - _unsampled_queries(driver, exporter) - _reset_to_default(driver) - _failed_internal_command(driver) - _failed_cleanup_command(driver) - _nested_cursor_queries(driver, exporter) - _transaction_sql(driver, exporter) - _parenthesized_queries(driver, exporter) - _skipped_middle_query(driver, exporter) - _custom_commit_before_query(driver, exporter) - asyncio.run(_delayed_sync_callback(driver, exporter)) - asyncio.run(_async_queries(exporter, tracer)) - asyncio.run(_async_nested_cursor_queries(exporter)) - asyncio.run(_async_transaction_sql(exporter)) - asyncio.run(_async_driver_regressions(exporter)) - _custom_cursor(exporter) - _manual_custom_cursor(exporter) - asyncio.run(_async_manual_custom_cursor(exporter)) - _psycopg_server_cursor_boundary(exporter) - _psycopg2_commit_boundary(exporter) - asyncio.run(_async_query_cancellation()) - - print( # noqa: T201 - f"POSTGRESQL_CORRELATION_OK mode={mode} psycopg={version('psycopg')} " - f"psycopg2={version('psycopg2-binary')} instrumentation={version('opentelemetry-instrumentation')}" - ) - - -if __name__ == "__main__": - main() + assert enable_postgresql_trace_context_propagation() + assert enable_postgresql_trace_context_propagation() + for driver in (psycopg, psycopg2): + _sync_queries(driver, exporter, tracer) + _sync_concurrency(driver, exporter, tracer) + _unsampled_queries(driver, exporter) + _reset_to_default(driver) + _failed_internal_command(driver) + _failed_cleanup_command(driver) + _nested_cursor_queries(driver, exporter) + _transaction_sql(driver, exporter) + _parenthesized_queries(driver, exporter) + _skipped_middle_query(driver, exporter) + _custom_commit_before_query(driver, exporter) + asyncio.run(_delayed_sync_callback(driver, exporter)) + asyncio.run(_async_queries(exporter, tracer)) + asyncio.run(_async_nested_cursor_queries(exporter)) + asyncio.run(_async_transaction_sql(exporter)) + asyncio.run(_async_driver_regressions(exporter)) + _custom_cursor(exporter) + _manual_custom_cursor(exporter) + asyncio.run(_async_manual_custom_cursor(exporter)) + _psycopg_server_cursor_boundary(exporter) + _psycopg2_commit_boundary(exporter) + asyncio.run(_async_query_cancellation()) diff --git a/tests/integration/postgresql/trace_context.py b/tests/integration/postgresql/trace_context.py new file mode 100644 index 00000000..1da0cca0 --- /dev/null +++ b/tests/integration/postgresql/trace_context.py @@ -0,0 +1,180 @@ +import asyncio +import json +import os +import re +from importlib.metadata import version +from pathlib import Path + + +_QUERY = "SELECT current_setting('application_name')" +_DEFAULT_APPLICATION_NAME = "startup-app" +_RESULT_PREFIX = "POSTGRESQL_QUERY_RESULT " +_DRIVERS = {"psycopg", "psycopg2", "psycopg-async"} + + +def _print_query_result(driver, application_name, parent): + context = parent.get_span_context() + print( # noqa: T201 + _RESULT_PREFIX + + json.dumps( + { + "driver": driver, + "application_name": application_name, + "parent_trace_id": f"{context.trace_id:032x}", + "parent_span_id": f"{context.span_id:016x}", + } + ), + flush=True, + ) + + +def _sync_query(driver, tracer): + with ( + driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, + tracer.start_as_current_span(f"{driver.__name__} caller") as parent, + connection.cursor() as cursor, + ): + cursor.execute(_QUERY) + application_name = cursor.fetchone()[0] + _print_query_result(driver.__name__, application_name, parent) + + +async def _async_query(tracer): + import psycopg + + async with await psycopg.AsyncConnection.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + with tracer.start_as_current_span("psycopg async caller") as parent: + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + application_name = (await cursor.fetchone())[0] + _print_query_result("psycopg-async", application_name, parent) + + +def main(): + from opentelemetry import trace + + mode = os.environ["POSTGRESQL_TEST_INSTRUMENTATION"] + if mode == "manual": + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + from splunk_otel import init_splunk_otel + + init_splunk_otel() + PsycopgInstrumentor().instrument() + Psycopg2Instrumentor().instrument() + + import psycopg + import psycopg2 + + tracer = trace.get_tracer(__name__) + for driver in (psycopg, psycopg2): + _sync_query(driver, tracer) + asyncio.run(_async_query(tracer)) + + if mode != "disabled": + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from regressions import run_regressions + + # Keep per-query regression checks alongside normal OTLP export. + exporter = InMemorySpanExporter() + trace.get_tracer_provider().add_span_processor(SimpleSpanProcessor(exporter)) + run_regressions(exporter, tracer) + + if not trace.get_tracer_provider().force_flush(): + message = "PostgreSQL test traces did not finish exporting" + raise RuntimeError(message) + print( # noqa: T201 + f"POSTGRESQL_CORRELATION_OK mode={mode} psycopg={version('psycopg')} " + f"psycopg2={version('psycopg2-binary')} instrumentation={version('opentelemetry-instrumentation')}" + ) + + +class PostgreSQLTraceContextOtelTest: + def __init__(self): + self.mode = os.environ.get("POSTGRESQL_TEST_INSTRUMENTATION", "auto") + if self.mode not in {"auto", "manual", "disabled"}: + message = f"Unknown PostgreSQL test instrumentation mode: {self.mode}" + raise ValueError(message) + + def requirements(self): + return ( + str(Path(__file__).resolve().parents[3]), + "psycopg[binary]", + "psycopg2-binary", + "opentelemetry-instrumentation-psycopg", + "opentelemetry-instrumentation-psycopg2", + ) + + def environment_variables(self): + return { + "POSTGRESQL_TEST_INSTRUMENTATION": self.mode, + "SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED": str(self.mode != "disabled").lower(), + "SPLUNK_OPAMP_ENABLED": "false", + "OTEL_SERVICE_NAME": "postgresql-trace-context-test", + "OTEL_PYTHON_DISABLED_INSTRUMENTATIONS": "system_metrics", + "OTEL_TRACES_EXPORTER": "otlp", + "OTEL_TRACES_SAMPLER": "parentbased_always_on", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://127.0.0.1:4317", + "OTEL_EXPORTER_OTLP_PROTOCOL": "grpc", + "OTEL_METRICS_EXPORTER": "none", + "OTEL_LOGS_EXPORTER": "none", + "PGHOST": os.environ.get("PGHOST", "127.0.0.1"), + "PGPORT": os.environ.get("PGPORT", "5432"), + "PGDATABASE": os.environ.get("PGDATABASE", "otel"), + "PGUSER": os.environ.get("PGUSER", "otel"), + "PGPASSWORD": os.environ.get("PGPASSWORD", "otel"), + } + + def wrapper_command(self): + return "" if self.mode == "manual" else "opentelemetry-instrument" + + def is_http(self): + return False + + def on_start(self): + return 120 + + def on_stop(self, telemetry, stdout: str, stderr: str, returncode: int) -> None: + from opentelemetry.proto.trace.v1.trace_pb2 import Span + from oteltest.telemetry import get_spans + + # oteltest 0.49 logs AssertionError without failing its command. + if returncode != 0: + message = f"PostgreSQL test exited with {returncode}.\nstdout: {stdout}\nstderr: {stderr}" + raise RuntimeError(message) + results = [ + json.loads(line.removeprefix(_RESULT_PREFIX)) + for line in stdout.splitlines() + if line.startswith(_RESULT_PREFIX) + ] + if len(results) != len(_DRIVERS) or {result["driver"] for result in results} != _DRIVERS: + message = "Missing PostgreSQL query results" + raise RuntimeError(message) + + spans = get_spans(telemetry) + for result in results: + matches = [ + span + for span in spans + if span.kind == Span.SPAN_KIND_CLIENT + and span.trace_id.hex() == result["parent_trace_id"] + and span.parent_span_id.hex() == result["parent_span_id"] + ] + if len(matches) != 1: + message = f"Expected one exported database CLIENT span for {result['driver']}, got {len(matches)}" + raise RuntimeError(message) + span = matches[0] + expected = ( + _DEFAULT_APPLICATION_NAME + if self.mode == "disabled" + else rf"00-{span.trace_id.hex()}-{span.span_id.hex()}-[0-9a-f]{{2}}" + ) + # Some OTLP exporters omit traceparent flag bits; compare the IDs. + if re.fullmatch(expected, result["application_name"]) is None: + message = f"PostgreSQL application_name did not match the exported span for {result['driver']}" + raise RuntimeError(message) + + +if __name__ == "__main__": + main() diff --git a/tests/test_postgresql_oteltest.py b/tests/test_postgresql_oteltest.py new file mode 100644 index 00000000..f2997a5f --- /dev/null +++ b/tests/test_postgresql_oteltest.py @@ -0,0 +1,119 @@ +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from opentelemetry.proto.trace.v1.trace_pb2 import Span + + +_SCRIPT_PATH = Path(__file__).parent / "integration" / "postgresql" / "trace_context.py" +_SPEC = importlib.util.spec_from_file_location("postgresql_trace_context", _SCRIPT_PATH) +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) + + +@pytest.fixture(autouse=True) +def instrumentation_mode(monkeypatch): + monkeypatch.setenv("POSTGRESQL_TEST_INSTRUMENTATION", "auto") + + +def _fake_received_telemetry(spans): + scope = SimpleNamespace(spans=spans) + resource = SimpleNamespace(scope_spans=[scope]) + request = SimpleNamespace(pbreq=SimpleNamespace(resource_spans=[resource])) + return SimpleNamespace(trace_requests=[request]) + + +def _fake_query_results(*, propagation_enabled=True): + spans = [] + results = [] + for index, driver in enumerate(("psycopg", "psycopg2", "psycopg-async"), 1): + span = Span( + trace_id=index.to_bytes(16, "big"), + span_id=(index + 10).to_bytes(8, "big"), + parent_span_id=(index + 20).to_bytes(8, "big"), + kind=Span.SPAN_KIND_CLIENT, + flags=0x100, + ) + results.append( + { + "driver": driver, + "parent_trace_id": span.trace_id.hex(), + "parent_span_id": span.parent_span_id.hex(), + "application_name": ( + f"00-{span.trace_id.hex()}-{span.span_id.hex()}-03" if propagation_enabled else "startup-app" + ), + } + ) + spans.append(span) + return spans, results + + +def _fake_application_stdout(results): + return "\n".join("POSTGRESQL_QUERY_RESULT " + json.dumps(result) for result in results) + + +@pytest.mark.parametrize("mode", ["auto", "manual", "disabled"]) +def test_postgresql_oteltest_configures_startup_and_accepts_received_spans(monkeypatch, mode): + monkeypatch.setenv("POSTGRESQL_TEST_INSTRUMENTATION", mode) + monkeypatch.setenv("PGPORT", "55439") + test = _MODULE.PostgreSQLTraceContextOtelTest() + spans, results = _fake_query_results(propagation_enabled=mode != "disabled") + + test.on_stop(_fake_received_telemetry(spans), _fake_application_stdout(results), "", 0) + + environment = test.environment_variables() + assert environment["SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED"] == str(mode != "disabled").lower() + assert environment["POSTGRESQL_TEST_INSTRUMENTATION"] == mode + assert environment["PGPORT"] == "55439" + assert environment["OTEL_TRACES_EXPORTER"] == "otlp" + assert test.wrapper_command() == ("" if mode == "manual" else "opentelemetry-instrument") + assert test.requirements()[0] == str(Path(__file__).resolve().parent.parent) + + +def test_postgresql_oteltest_fails_for_application_errors(): + test = _MODULE.PostgreSQLTraceContextOtelTest() + + with pytest.raises(RuntimeError, match="exited with 1"): + test.on_stop(_fake_received_telemetry([]), "", "query failed", 1) + + +def test_postgresql_oteltest_fails_for_missing_query_results(): + test = _MODULE.PostgreSQLTraceContextOtelTest() + + with pytest.raises(RuntimeError, match="Missing PostgreSQL query results"): + test.on_stop(_fake_received_telemetry([]), "", "", 0) + + +@pytest.mark.parametrize("invalid_span", ["missing", "duplicate", "kind", "parent", "trace", "span", "flags"]) +def test_postgresql_oteltest_fails_for_missing_or_mismatched_received_spans(invalid_span): + test = _MODULE.PostgreSQLTraceContextOtelTest() + spans, results = _fake_query_results() + if invalid_span == "missing": + spans.pop() + elif invalid_span == "duplicate": + spans.append(spans[0]) + elif invalid_span == "kind": + spans[0].kind = Span.SPAN_KIND_SERVER + elif invalid_span == "parent": + spans[0].parent_span_id = b"mismatch" + elif invalid_span == "trace": + spans[0].trace_id = b"mismatch" + elif invalid_span == "span": + spans[0].span_id = b"mismatch" + else: + results[0]["application_name"] = results[0]["application_name"][:-2] + "zz" + + with pytest.raises(RuntimeError, match=r"CLIENT span|did not match"): + test.on_stop(_fake_received_telemetry(spans), _fake_application_stdout(results), "", 0) + + +def test_postgresql_oteltest_rejects_trace_context_when_disabled(monkeypatch): + monkeypatch.setenv("POSTGRESQL_TEST_INSTRUMENTATION", "disabled") + test = _MODULE.PostgreSQLTraceContextOtelTest() + spans, results = _fake_query_results() + + with pytest.raises(RuntimeError, match="did not match"): + test.on_stop(_fake_received_telemetry(spans), _fake_application_stdout(results), "", 0) From 5516c6eb31426e0da68844daf583791c57470f6c Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Mon, 19 Oct 2026 14:54:57 -0400 Subject: [PATCH 08/10] Refactor PostgreSQL trace correlation setup and execution Split activation and driver adapter construction into focused helpers. Inline signature forwarding helpers and clarify Psycopg names. Extract active query context management and application_name execution cleanup for both sync and async paths. Assisted-by: Codex (GPT-6) --- src/splunk_otel/postgresql.py | 277 +++++++++++++++++++++------------- 1 file changed, 171 insertions(+), 106 deletions(-) diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index ff2d3cd3..5cbbdb5e 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -20,7 +20,7 @@ import re import threading import weakref -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from contextvars import ContextVar from dataclasses import dataclass, field from functools import partial @@ -32,7 +32,7 @@ from opentelemetry.instrumentation.utils import is_instrumentation_enabled, suppress_instrumentation if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Iterator from opentelemetry.trace import SpanContext @@ -101,52 +101,74 @@ def enable_postgresql_trace_context_propagation() -> bool: if _enabled: return True - try: - dbapi = import_module("opentelemetry.instrumentation.dbapi") - except ImportError: - logger.warning("PostgreSQL trace context propagation requires DBAPI instrumentation") + cursor_tracer = _load_compatible_cursor_tracer() + if cursor_tracer is None: return False - cursor_tracer = getattr(dbapi, "CursorTracer", None) - for name in ("traced_execution", "traced_execution_async"): - method = getattr(cursor_tracer, name, None) - if not _has_execution_signature(method): - logger.warning("PostgreSQL trace context propagation requires compatible DBAPI execution methods") - return False - adapters = _load_cursor_adapters() - if not adapters: - logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") + psycopg3_instrumentation = next( + (adapter.module for adapter in adapters if hasattr(adapter.module, "psycopg")), None + ) + if not _validate_driver_instrumentation(adapters, psycopg3_instrumentation): return False - for adapter in adapters: - factory = getattr(adapter.module, adapter.factory_name, None) - if not _has_factory_signature(factory): - logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") - return False - - psycopg_adapter = next((adapter for adapter in adapters if hasattr(adapter.module, "psycopg")), None) - if psycopg_adapter is not None and not _has_parameter_signature( - getattr(psycopg_adapter.module.PsycopgInstrumentor, "instrument_connection", None), ("connection",) - ): - logger.warning("PostgreSQL trace context propagation requires compatible manual instrumentation") + + _install_postgresql_wrappers(cursor_tracer, adapters, psycopg3_instrumentation) + _enabled = True + return True + + +def _load_compatible_cursor_tracer() -> type | None: + try: + dbapi = import_module("opentelemetry.instrumentation.dbapi") + except ImportError: + logger.warning("PostgreSQL trace context propagation requires DBAPI instrumentation") + return None + + cursor_tracer = getattr(dbapi, "CursorTracer", None) + for name in ("traced_execution", "traced_execution_async"): + method = getattr(cursor_tracer, name, None) + if not _has_parameter_signature(method, ("self", "cursor", "query_method")): + logger.warning("PostgreSQL trace context propagation requires compatible DBAPI execution methods") + return None + return cursor_tracer + + +def _validate_driver_instrumentation(adapters: list[_CursorAdapter], psycopg3_instrumentation: Any) -> bool: + if not adapters: + logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") + return False + for adapter in adapters: + factory = getattr(adapter.module, adapter.factory_name, None) + if not _has_parameter_signature(factory, ("db_api", "base_factory")): + logger.warning("PostgreSQL trace context propagation requires compatible driver cursor factories") return False - for adapter in adapters: - wrapt.wrap_function_wrapper(adapter.module, adapter.factory_name, _cursor_factory_wrapper(adapter)) - if psycopg_adapter is not None: - wrapt.wrap_function_wrapper( - psycopg_adapter.module.PsycopgInstrumentor, - "instrument_connection", - _manual_connection_wrapper(psycopg_adapter.module), - ) - wrapt.wrap_function_wrapper(cursor_tracer, "traced_execution", _execute_with_postgresql_trace_context) + if psycopg3_instrumentation is not None and not _has_parameter_signature( + getattr(psycopg3_instrumentation.PsycopgInstrumentor, "instrument_connection", None), ("connection",) + ): + logger.warning("PostgreSQL trace context propagation requires compatible manual instrumentation") + return False + return True + + +def _install_postgresql_wrappers( + cursor_tracer: type, adapters: list[_CursorAdapter], psycopg3_instrumentation: Any +) -> None: + for adapter in adapters: + wrapt.wrap_function_wrapper(adapter.module, adapter.factory_name, _cursor_factory_wrapper(adapter)) + if psycopg3_instrumentation is not None: wrapt.wrap_function_wrapper( - cursor_tracer, - "traced_execution_async", - _execute_with_postgresql_trace_context_async, + psycopg3_instrumentation.PsycopgInstrumentor, + "instrument_connection", + _manual_connection_wrapper(psycopg3_instrumentation), ) - _enabled = True - return True + + wrapt.wrap_function_wrapper(cursor_tracer, "traced_execution", _execute_with_postgresql_trace_context) + wrapt.wrap_function_wrapper( + cursor_tracer, + "traced_execution_async", + _execute_with_postgresql_trace_context_async, + ) def _load_cursor_adapters() -> list[_CursorAdapter]: @@ -157,34 +179,36 @@ def _load_cursor_adapters() -> list[_CursorAdapter]: except ImportError: continue if name == "psycopg": - driver = module.psycopg - adapters.extend( - _CursorAdapter( - module=module, - factory_name=factory_name, - base_cursor=base_cursor, - new_internal_cursor=partial(base_cursor, row_factory=driver.rows.tuple_row), - asynchronous=asynchronous, - ) - for factory_name, base_cursor, asynchronous in ( - ("_new_cursor_factory", driver.Cursor, False), - ("_new_cursor_async_factory", driver.AsyncCursor, True), - ) - ) + adapters.extend(_create_psycopg_cursor_adapters(module)) else: - adapters.append( - _CursorAdapter( - module=module, - factory_name="_new_cursor_factory", - base_cursor=module.psycopg2.extensions.cursor, - new_internal_cursor=module.psycopg2.extensions.cursor, - ) - ) + adapters.append(_create_psycopg2_cursor_adapter(module)) return adapters -def _has_factory_signature(factory) -> bool: - return _has_parameter_signature(factory, ("db_api", "base_factory")) +def _create_psycopg_cursor_adapters(module) -> list[_CursorAdapter]: + driver = module.psycopg + return [ + _CursorAdapter( + module=module, + factory_name=factory_name, + base_cursor=base_cursor, + new_internal_cursor=partial(base_cursor, row_factory=driver.rows.tuple_row), + asynchronous=asynchronous, + ) + for factory_name, base_cursor, asynchronous in ( + ("_new_cursor_factory", driver.Cursor, False), + ("_new_cursor_async_factory", driver.AsyncCursor, True), + ) + ] + + +def _create_psycopg2_cursor_adapter(module) -> _CursorAdapter: + return _CursorAdapter( + module=module, + factory_name="_new_cursor_factory", + base_cursor=module.psycopg2.extensions.cursor, + new_internal_cursor=module.psycopg2.extensions.cursor, + ) def _cursor_factory_wrapper(adapter: _CursorAdapter): @@ -263,10 +287,6 @@ async def callproc(self, *args, **kwargs): ) -def _has_execution_signature(method: Any) -> bool: - return _has_parameter_signature(method, ("self", "cursor", "query_method")) - - def _has_parameter_signature(method: Any, names: tuple[str, ...]) -> bool: if not callable(method): return False @@ -324,6 +344,17 @@ async def correlated_query_method(*method_args, **method_kwargs): return await traced_execution_async(cursor, correlated_query_method, *query_args, **kwargs) +@contextmanager +def _active_query_context(connection, owner: object, traceparent: str | None) -> Iterator[None]: + execution = _QueryContext(connection, owner, traceparent) + token = _active_queries.set((*_active_queries.get(), execution)) + try: + yield + finally: + execution.lifetime.finish() + _active_queries.reset(token) + + def _execute_driver( cursor, query_method: Callable[..., Any], @@ -341,23 +372,17 @@ def _execute_driver( lock = nullcontext() if parent is not None else _get_sync_connection_lock(cursor.connection) with lock: traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None - execution = _QueryContext(cursor.connection, owner, traceparent) - token = _active_queries.set((*_active_queries.get(), execution)) - try: - if propagate: - _update_application_name(cursor.connection, traceparent, new_internal_cursor=new_internal_cursor) - return query_method(*args, **kwargs) - finally: - try: - if propagate and (traceparent is not None or parent is not None): - _update_application_name( - cursor.connection, - parent.traceparent if parent is not None else None, - new_internal_cursor=new_internal_cursor, - ) - finally: - execution.lifetime.finish() - _active_queries.reset(token) + with _active_query_context(cursor.connection, owner, traceparent): + return _execute_with_application_name( + cursor.connection, + query_method, + args, + kwargs, + traceparent=traceparent, + parent=parent, + propagate=propagate, + new_internal_cursor=new_internal_cursor, + ) async def _execute_driver_async( @@ -377,25 +402,65 @@ async def _execute_driver_async( lock = nullcontext() if parent is not None else _get_async_connection_lock(cursor.connection) async with lock: traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None - execution = _QueryContext(cursor.connection, owner, traceparent) - token = _active_queries.set((*_active_queries.get(), execution)) - try: - if propagate: - await _update_application_name_async( - cursor.connection, traceparent, new_internal_cursor=new_internal_cursor - ) - return await query_method(*args, **kwargs) - finally: - try: - if propagate and (traceparent is not None or parent is not None): - await _update_application_name_async( - cursor.connection, - parent.traceparent if parent is not None else None, - new_internal_cursor=new_internal_cursor, - ) - finally: - execution.lifetime.finish() - _active_queries.reset(token) + with _active_query_context(cursor.connection, owner, traceparent): + return await _execute_with_application_name_async( + cursor.connection, + query_method, + args, + kwargs, + traceparent=traceparent, + parent=parent, + propagate=propagate, + new_internal_cursor=new_internal_cursor, + ) + + +def _execute_with_application_name( + connection, + query_method: Callable[..., Any], + args, + kwargs, + *, + traceparent: str | None, + parent: _QueryContext | None, + propagate: bool, + new_internal_cursor: Callable[[Any], Any], +): + try: + if propagate: + _update_application_name(connection, traceparent, new_internal_cursor=new_internal_cursor) + return query_method(*args, **kwargs) + finally: + if propagate and (traceparent is not None or parent is not None): + _update_application_name( + connection, + parent.traceparent if parent is not None else None, + new_internal_cursor=new_internal_cursor, + ) + + +async def _execute_with_application_name_async( + connection, + query_method: Callable[..., Any], + args, + kwargs, + *, + traceparent: str | None, + parent: _QueryContext | None, + propagate: bool, + new_internal_cursor: Callable[[Any], Any], +): + try: + if propagate: + await _update_application_name_async(connection, traceparent, new_internal_cursor=new_internal_cursor) + return await query_method(*args, **kwargs) + finally: + if propagate and (traceparent is not None or parent is not None): + await _update_application_name_async( + connection, + parent.traceparent if parent is not None else None, + new_internal_cursor=new_internal_cursor, + ) def _live_span_lease(cursor, owner: object) -> _SpanLease | None: @@ -533,8 +598,8 @@ def _is_postgresql_client_cursor(cursor) -> bool: # Named cursors execute SQL during later fetches, outside this query's span. if getattr(cursor, "name", None): return False - return any(_is_postgresql_module(cursor_class.__module__) for cursor_class in type(cursor).__mro__) + return any(_is_psycopg_module(cursor_class.__module__) for cursor_class in type(cursor).__mro__) -def _is_postgresql_module(module_name: str) -> bool: +def _is_psycopg_module(module_name: str) -> bool: return any(module_name == name or module_name.startswith(f"{name}.") for name in _DRIVER_MODULE_NAMES) From a5422e0cf421b41beefe41385dc65c66065e62b9 Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Wed, 21 Oct 2026 15:23:13 -0400 Subject: [PATCH 09/10] Fix PostgreSQL correlation statistics, pipelines, and templates Normalize correlation commands in query statistics, skip propagation in pipelines, and handle Python 3.14 templates without converting bound values to text. Add regressions and simplify the setup documentation. --- .github/workflows/ci-main.yml | 14 ++ docs/postgresql.md | 190 ++++++----------- src/splunk_otel/postgresql.py | 46 ++++- tests/integration/postgresql/regressions.py | 214 +++++++++++++++++++- tests/test_postgresql.py | 120 ++++++++--- 5 files changed, 404 insertions(+), 180 deletions(-) diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 501d8002..e0e3db9f 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -86,6 +86,20 @@ jobs: - name: Install oteltest run: pip install 'oteltest>=0.49.0' + - name: Enable PostgreSQL query statistics + env: + POSTGRES_CONTAINER: ${{ job.services.postgres.id }} + run: | + docker exec "$POSTGRES_CONTAINER" psql -U otel -d otel -c "ALTER SYSTEM SET shared_preload_libraries = 'pg_stat_statements'" + docker restart "$POSTGRES_CONTAINER" + for attempt in {1..30}; do + if docker exec "$POSTGRES_CONTAINER" pg_isready -U otel -d otel; then + exit 0 + fi + sleep 1 + done + exit 1 + - name: Run PostgreSQL oteltest timeout-minutes: 5 run: oteltest tests/integration/postgresql/trace_context.py diff --git a/docs/postgresql.md b/docs/postgresql.md index 6dc9ed32..e5b453b7 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -1,145 +1,76 @@ # PostgreSQL trace correlation -This experimental feature connects sampled PostgreSQL queries in Splunk Database -Monitoring to the APM trace and database client span that executed them. It places -the database client span's W3C `traceparent` in PostgreSQL's `application_name`. -It sends only `traceparent`, without baggage or `tracestate`. Valid unsampled -span contexts also work. +This experimental feature lets Splunk Database Monitoring link PostgreSQL query +samples to APM traces. It puts the database span's trace and span IDs in +PostgreSQL's `application_name` while a query runs. -## Database Monitoring prerequisites +## Setup -Configure a Collector with the PostgreSQL receiver, database monitoring -permissions, and both `db.server.query_sample` and `db.server.top_query` enabled. -Export database metrics and query events to Splunk Database Monitoring. Query -events use `/v3/event` with the `X-splunk-instrumentation-library: dbmon` header. -See the [PostgreSQL receiver setup instructions](https://help.splunk.com/en/splunk-observability-cloud/monitor-databases/get-data-in/configure-receivers/postgresql-receiver). +Configure the [PostgreSQL receiver](https://help.splunk.com/en/splunk-observability-cloud/monitor-databases/get-data-in/configure-receivers/postgresql-receiver) +with `db.server.query_sample` and `db.server.top_query` enabled. The receiver must +read trace IDs from `application_name`; Splunk Collector 0.162.0 was tested. +Send database metrics, query events, and application traces to the same Splunk +organization. A query sample can link to a trace only if APM has that trace. -The receiver must support extracting trace and span IDs from `application_name`. -End-to-end correlation was validated with Splunk Collector 0.162.0. The minimum -compatible Collector version has not yet been established for this feature. -The application must also export traces to APM in the same organization; those -traces can use a separate export path. A query sample can link to a trace only -when that trace is available in APM. - -## Enablement - -The feature is disabled by default. In the same Python environment as your -application, install `splunk-opentelemetry` and your application's PostgreSQL -driver first. Then install instrumentation for the installed libraries: +Install `splunk-opentelemetry` and your PostgreSQL driver in the application's +Python environment. Then run: ```sh opentelemetry-bootstrap -a install +SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED=true \ +opentelemetry-instrument python app.py ``` -This command selects instrumentation versions that match the installed distro. -For PostgreSQL, it installs: +The feature is off by default. The setting works when `OTEL_CONFIG_FILE` is unset. +To turn it off, unset the setting and restart the application. -| Application driver | Instrumentation package | -|---|---| -| `psycopg` (Psycopg 3) | [opentelemetry-instrumentation-psycopg](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg/psycopg.html) | -| `psycopg2` or `psycopg2-binary` | [opentelemetry-instrumentation-psycopg2](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html) | +For manual instrumentation, call `init_splunk_otel()` before using +`PsycopgInstrumentor` or `Psycopg2Instrumentor` and before creating connections. +You can also enable it directly with +`splunk_otel.postgresql.enable_postgresql_trace_context_propagation()`. -Start the application with: +## Supported queries -```sh -SPLUNK_DB_TRACE_CONTEXT_PROPAGATION_ENABLED=true \ -opentelemetry-instrument python app.py -``` +The feature supports Psycopg 3 sync and async connections and Psycopg2 sync +connections. Python 3.14 template queries require Psycopg 3.3 or later. +Custom cursor methods must call the driver method through `super()`. + +The feature skips: + +- Psycopg2 async connections and Psycopg 3 queries inside explicit pipelines. +- Named/server cursors, fetch operations, and COPY. +- SQL transaction commands, such as BEGIN, COMMIT, ROLLBACK, and SAVEPOINT. +- SQL CALL, DO, and query batches. Any semicolon except one at the end causes + the query to be skipped, even inside a string or comment. + +Standalone `executemany` is supported. Queries in a pipeline still get normal +tracing, but no correlation commands. Correlation resumes after the pipeline exits. + +## Connection changes and cost + +Each correlated query adds two database commands: one to set `application_name` +and one to reset it. The set command uses parameters, so each new trace ID does +not create a separate `pg_stat_statements` entry. + +After the query, the feature resets `application_name` to the connection's +configured default. It does not restore a value the application set through SQL. + +Queries that use this feature share a lock on the connection. Raw cursors and +queries with instrumentation suppressed can bypass the lock and see another +query's trace IDs. Span duration includes lock waiting and the extra commands. + +The feature logs failures in its own commands. A failed command can abort the +application's transaction or leave trace IDs in `application_name`. The +application must handle recovery; the feature does not roll back transactions. +Cleanup errors do not replace the application's query result or error. + +## Tests + +Tested with PostgreSQL 16, Psycopg 3.3.6, Psycopg2 2.9.13, and OpenTelemetry +instrumentation 0.66b1. Other versions have not been validated. -Python currently applies this setting only to PostgreSQL. - -The environment switch applies when `OTEL_CONFIG_FILE` is unset. Activation lasts -for the process lifetime. To turn it off, unset the switch and restart. - -Manual instrumentation must initialize the Splunk distro first with -`init_splunk_otel()`, or call -`splunk_otel.postgresql.enable_postgresql_trace_context_propagation()` before using -`PsycopgInstrumentor` or `Psycopg2Instrumentor`. Enable the feature before creating -instrumented connections and cursors. - -## Driver coverage - -The feature targets Psycopg 3 synchronous and asynchronous connections and -Psycopg2 synchronous connections. Psycopg2 asynchronous connections pass through -without trace propagation. - -Custom cursor methods must delegate through `super()` to the driver method. -The feature sets context after those custom hooks, immediately before the driver -executes the query. Direct calls to native driver methods can bypass this adapter. - -Named/server cursors, cursor fetch operations, and COPY operations are outside -this feature's automatic coverage. Named cursors can execute SQL later during -fetch, after the execute span has ended, so this feature skips propagation for -them in both drivers. - -SQL transaction-control statements (including BEGIN, COMMIT, ROLLBACK, and -savepoint commands), SQL CALL, and DO run without propagation commands. Those -statements can change transaction boundaries, making SET and cleanup unsafe. -SQL batches also run without propagation. The adapter conservatively skips any -statement with a semicolon other than a single trailing semicolon, including -semicolons inside strings or comments. These calls still participate in the -connection lock. - -## Connection behavior - -The feature uses a native driver cursor on the same connection for internal -commands. It suppresses instrumentation and bypasses customer cursor hooks for -those commands. A per-connection lock covers SET, driver execution, and cleanup. -Customer hooks run outside this lock. Auxiliary queries from those hooks finish -their own propagation before the outer query starts driver execution. - -Nested driver executions in the same thread or async task can reuse the lock. -Cleanup restores the active outer driver's setting, including through skipped -intermediate statements. Other threads and tasks must acquire the lock. Context -copied into a delayed callback cannot grant ownership after its query has finished. -Each correlated query sends two additional commands: SET before driver execution -and context restoration afterward. - -This opt-in owns the connection's `application_name` during execution. Cleanup -of the outermost query resets the setting to the connection's configured default -before releasing the instrumentation lock, including when the query raises an -exception. If instrumentation creates no valid new database span for an eligible -statement, the feature sends SET DEFAULT before execution. It does not restore -a value that the application previously set through SQL. - -The lock only coordinates wrapped queries. Raw cursors and suppressed queries -can bypass it and observe another query's context during execution. Failed -cleanup can also leave trace context set until the application recovers the -connection or rolls back its transaction. - -## Failure behavior and query duration - -The feature logs internal command failures and attempts the customer query. A -failed setup command can leave a transaction aborted, so the customer query can -also fail. Cleanup failures do not replace the customer query's result or error; -they can still leave an otherwise successful transaction aborted. The feature -does not roll back application transactions. - -The extra commands add latency. Database span duration and database operation -duration metrics include lock waiting and the additional commands. - -## Instrumentation compatibility - -The Splunk adapter uses the installed OpenTelemetry driver instrumentors to -create spans. It wraps shared DBAPI execution callbacks to capture the database -span context, then adapts the Psycopg cursor factories to apply that context at -driver execution. It also adapts Psycopg 3 manual connection instrumentation to -include the application's custom cursor factory. - -These hooks use private instrumentation APIs. Activation checks their signatures -and declines incompatible versions before installing wrappers. Regression tests -must also check behavior when instrumentation dependencies change. The current -validation covers instrumentation 0.66b1, Psycopg 3.3.6, Psycopg2 2.9.13, and -PostgreSQL 16; it does not establish a broader supported version range. - -## Integration tests - -The PostgreSQL CI job uses `oteltest` with PostgreSQL 16 on Python 3.10 and 3.14. -It covers automatic instrumentation, manual instrumentation, and disabled -propagation. Query results must match the database client spans received over -OTLP. The same job runs the transaction, custom cursor, and concurrency regressions. - -To run locally, start a test database, install `oteltest>=0.49.0`, and run: +To run the integration tests, start PostgreSQL with `pg_stat_statements` preloaded +and utility tracking enabled. Install `oteltest>=0.49.0`, then run: ```sh PGHOST=127.0.0.1 PGPORT=5432 PGDATABASE=otel PGUSER=otel PGPASSWORD=otel \ @@ -148,6 +79,3 @@ oteltest tests/integration/postgresql/trace_context.py ``` Repeat with `POSTGRESQL_TEST_INSTRUMENTATION=manual` and `disabled`. -The test installs this checkout and both driver instrumentations in an isolated -environment. It exports traces to the local `oteltest` receiver on port 4317. -Collector ingestion and Database Monitoring UI checks are separate manual tests. diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index 5cbbdb5e..0062ce04 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -18,6 +18,7 @@ import inspect import logging import re +import sys import threading import weakref from contextlib import contextmanager, nullcontext @@ -31,6 +32,9 @@ from opentelemetry import trace from opentelemetry.instrumentation.utils import is_instrumentation_enabled, suppress_instrumentation +if sys.version_info >= (3, 14): + from string.templatelib import Interpolation, Template + if TYPE_CHECKING: from collections.abc import Callable, Iterator @@ -39,6 +43,7 @@ logger = logging.getLogger(__name__) _SET_APPLICATION_NAME_DEFAULT = "SET application_name = DEFAULT" +_SET_APPLICATION_NAME = "SELECT pg_catalog.set_config('application_name', %s, false)" _DRIVER_MODULE_NAMES = ("psycopg", "psycopg2") _CURSOR_AND_QUERY_METHOD_ARGUMENTS = 2 _TRANSACTION_OPERATIONS = frozenset( @@ -487,6 +492,9 @@ def _parent_query(connection, owner: object) -> _QueryContext | None: def _can_propagate_statement(cursor, query_method, args, kwargs) -> bool: + # Internal commands can consume errors from earlier queued customer queries. + if getattr(getattr(cursor.connection, "info", None), "pipeline_status", 0) != 0: + return False if query_method.__name__ == "callproc": return True statement = args[0] if args else kwargs.get("query") @@ -494,6 +502,8 @@ def _can_propagate_statement(cursor, query_method, args, kwargs) -> bool: statement = statement.decode("ascii", errors="replace") elif any(base.__module__ in ("psycopg.sql", "psycopg2.sql") for base in type(statement).__mro__): statement = statement.as_string(cursor.connection) + elif sys.version_info >= (3, 14) and isinstance(statement, Template): + statement = import_module("psycopg.sql").as_string(_template_for_statement_check(statement), cursor) if not isinstance(statement, str): return False @@ -509,6 +519,28 @@ def _can_propagate_statement(cursor, query_method, args, kwargs) -> bool: return operation is not None and operation[0].upper() not in _TRANSACTION_OPERATIONS +def _template_for_statement_check(template: Template) -> Template: + # Inspect SQL without adapting parameter values or requiring text dumpers. + parts = [] + for part in template: + if isinstance(part, str): + parts.append(part) + elif part.format_spec in ("", "s", "t", "b", "l"): + parts.append("NULL") + elif part.format_spec == "q" and isinstance(part.value, Template): + parts.append( + Interpolation( + _template_for_statement_check(part.value), + part.expression, + conversion=part.conversion, + format_spec=part.format_spec, + ) + ) + else: + parts.append(part) + return Template(*parts) + + def _strip_leading_comments(statement: str) -> str: statement = statement.lstrip() while statement.startswith(("--", "/*")): @@ -553,12 +585,12 @@ def _current_database_traceparent(span_context_before: SpanContext) -> str | Non def _update_application_name(connection, traceparent: str | None, *, new_internal_cursor: Callable[[Any], Any]) -> None: - statement = _application_name_statement(traceparent) + statement, parameters = _application_name_command(traceparent) try: cursor = new_internal_cursor(connection) try: with suppress_instrumentation(): - result = cursor.execute(statement) + result = cursor.execute(statement, parameters) if inspect.iscoroutine(result): result.close() if inspect.isawaitable(result): @@ -572,12 +604,12 @@ def _update_application_name(connection, traceparent: str | None, *, new_interna async def _update_application_name_async( connection, traceparent: str | None, *, new_internal_cursor: Callable[[Any], Any] ) -> None: - statement = _application_name_statement(traceparent) + statement, parameters = _application_name_command(traceparent) try: cursor = new_internal_cursor(connection) try: with suppress_instrumentation(): - result = cursor.execute(statement) + result = cursor.execute(statement, parameters) if inspect.isawaitable(result): await result finally: @@ -588,10 +620,10 @@ async def _update_application_name_async( logger.exception("Failed to set PostgreSQL application_name") -def _application_name_statement(traceparent: str | None) -> str: +def _application_name_command(traceparent: str | None) -> tuple[str, tuple[str] | None]: if traceparent is None: - return _SET_APPLICATION_NAME_DEFAULT - return f"SET application_name = '{traceparent}'" + return _SET_APPLICATION_NAME_DEFAULT, None + return _SET_APPLICATION_NAME, (traceparent,) def _is_postgresql_client_cursor(cursor) -> bool: diff --git a/tests/integration/postgresql/regressions.py b/tests/integration/postgresql/regressions.py index 2105a05e..51fe78fa 100644 --- a/tests/integration/postgresql/regressions.py +++ b/tests/integration/postgresql/regressions.py @@ -1,4 +1,5 @@ import asyncio +import sys from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from importlib import import_module @@ -19,6 +20,189 @@ def _traceparent(span) -> str: return f"00-{context.trace_id:032x}-{context.span_id:016x}-{int(context.trace_flags) & 0xFF:02x}" +def _bounded_query_statistics(driver, exporter): + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection: + connection.autocommit = True + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements") + cursor.execute("SHOW pg_stat_statements.track_utility") + assert cursor.fetchone()[0] == "on" + cursor.execute("SELECT pg_stat_statements_reset()") + + exporter.clear() + contexts = set() + with connection.cursor() as cursor: + for _ in range(64): + cursor.execute(_QUERY) + contexts.add(cursor.fetchone()[0]) + assert len(contexts) == 64 + for context in contexts: + _assert_database_span(exporter, context) + + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute( + "SELECT query, calls FROM pg_stat_statements " + "WHERE query LIKE 'SELECT pg_catalog.set_config(%' OR query LIKE 'SET application_name%'" + ) + entries = cursor.fetchall() + assert len(entries) == 2, entries + setup_calls = [calls for query, calls in entries if query.startswith("SELECT pg_catalog.set_config(")] + assert setup_calls == [64], entries + cursor.execute("SELECT calls FROM pg_stat_statements WHERE query LIKE 'SELECT current_setting(%'") + assert cursor.fetchall() == [(64,)] + cursor.execute("SELECT dealloc FROM pg_stat_statements_info") + assert cursor.fetchone()[0] == 0 + + +def _pipeline_errors(exporter): + import psycopg + + for autocommit in (False, True): + with psycopg.connect(autocommit=autocommit, application_name=_DEFAULT_APPLICATION_NAME) as connection: + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("CREATE TEMP TABLE pipeline_errors (id integer PRIMARY KEY)") + connection.commit() + exporter.clear() + constraint_name = None + try: + with connection.pipeline() as pipeline, connection.cursor() as cursor: + cursor.execute(_QUERY) + assert cursor.fetchone()[0] == _DEFAULT_APPLICATION_NAME + cursor.execute("INSERT INTO pipeline_errors VALUES (1)") + cursor.execute("INSERT INTO pipeline_errors VALUES (1)") + pipeline.sync() + except psycopg.errors.UniqueViolation as error: + constraint_name = error.diag.constraint_name + assert constraint_name == "pipeline_errors_pkey" + assert exporter.get_finished_spans() + connection.rollback() + exporter.clear() + with connection.cursor() as cursor: + cursor.execute(_QUERY) + _assert_database_span(exporter, cursor.fetchone()[0]) + + +async def _async_pipeline_errors(exporter): + import psycopg + + for autocommit in (False, True): + async with await psycopg.AsyncConnection.connect( + autocommit=autocommit, application_name=_DEFAULT_APPLICATION_NAME + ) as connection: + with suppress_instrumentation(): + async with connection.cursor() as cursor: + await cursor.execute("CREATE TEMP TABLE pipeline_errors (id integer PRIMARY KEY)") + await connection.commit() + exporter.clear() + constraint_name = None + try: + async with connection.pipeline() as pipeline, connection.cursor() as cursor: + await cursor.execute(_QUERY) + assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME + await cursor.execute("INSERT INTO pipeline_errors VALUES (1)") + await cursor.execute("INSERT INTO pipeline_errors VALUES (1)") + await pipeline.sync() + except psycopg.errors.UniqueViolation as error: + constraint_name = error.diag.constraint_name + assert constraint_name == "pipeline_errors_pkey" + assert exporter.get_finished_spans() + await connection.rollback() + exporter.clear() + async with connection.cursor() as cursor: + await cursor.execute(_QUERY) + _assert_database_span(exporter, (await cursor.fetchone())[0]) + + +def _template_queries(): + from string.templatelib import Interpolation, Template + + from psycopg import sql + + return ( + Template("SELECT current_setting('application_name')"), + Template("SELECT current_setting(", Interpolation("application_name", "setting"), ")"), + Template("SELECT current_setting('application_name') AS ", Interpolation("name", "column", format_spec="i")), + Template("", Interpolation(Template("SELECT current_setting('application_name')"), "query", format_spec="q")), + Template("SELECT current_setting('application_name'), ", Interpolation("O'Reilly", "value"), "::text"), + ), Template("", Interpolation(sql.SQL("COMMIT"), "command", format_spec="q")) + + +def _binary_template_queries(cursor): + import struct + from string.templatelib import Interpolation, Template + + from psycopg.adapt import Dumper + from psycopg.pq import Format + + class BinaryValue: + def __init__(self, value): + self.value = value + + class BinaryValueDumper(Dumper): + format = Format.BINARY + oid = 23 + + def dump(self, value): + return struct.pack("!i", value.value) + + cursor.adapters.register_dumper(BinaryValue, BinaryValueDumper) + value = BinaryValue(42) + queries = tuple( + Template("SELECT current_setting('application_name'), ", Interpolation(value, "value", format_spec=format_spec)) + for format_spec in ("", "b") + ) + return (*queries, Template("", Interpolation(queries[0], "query", format_spec="q"))) + + +def _sync_template_queries(exporter): + import psycopg + + queries, commit = _template_queries() + with psycopg.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, connection.cursor() as cursor: + for query in queries: + exporter.clear() + cursor.execute(query) + row = cursor.fetchone() + _assert_database_span(exporter, row[0]) + if len(row) > 1: + assert row[1] == "O'Reilly" + for query in _binary_template_queries(cursor): + exporter.clear() + cursor.execute(query) + application_name, value = cursor.fetchone() + assert value == 42 + _assert_database_span(exporter, application_name) + cursor.execute(commit) + assert connection.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + assert connection.info.parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + + +async def _async_template_queries(exporter): + import psycopg + + queries, commit = _template_queries() + async with ( + await psycopg.AsyncConnection.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, + connection.cursor() as cursor, + ): + for query in queries: + exporter.clear() + await cursor.execute(query) + row = await cursor.fetchone() + _assert_database_span(exporter, row[0]) + if len(row) > 1: + assert row[1] == "O'Reilly" + for query in _binary_template_queries(cursor): + exporter.clear() + await cursor.execute(query) + application_name, value = await cursor.fetchone() + assert value == 42 + _assert_database_span(exporter, application_name) + await cursor.execute(commit) + assert connection.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + assert connection.info.parameter_status("application_name") == _DEFAULT_APPLICATION_NAME + + def _assert_database_span(exporter, application_name, parent_context=None): matches = [span for span in exporter.get_finished_spans() if _traceparent(span) == application_name] assert len(matches) == 1, f"PostgreSQL value {application_name!r} did not identify one exported span" @@ -143,15 +327,17 @@ async def query(): assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 - class RejectingSetCursor(psycopg.AsyncCursor): + class RejectingApplicationNameCursor(psycopg.AsyncCursor): async def execute(self, query, *args, **kwargs): - if query.startswith("SET application_name"): - query = "SET nonexistent_correlation_parameter = 'failure'" + if query.startswith(("SET application_name", "SELECT pg_catalog.set_config(")): + return await super().execute("SET nonexistent_correlation_parameter = 'failure'") return await super().execute(query, *args, **kwargs) connection = await psycopg.AsyncConnection.connect() try: - with patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)): + with patch.object( + connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingApplicationNameCursor) + ): async with connection.cursor() as cursor: failed = False try: @@ -207,16 +393,18 @@ def _failed_internal_command(driver): cursor_class = driver.Cursor if driver.__name__ == "psycopg" else driver.extensions.cursor transaction_error = driver.errors.InFailedSqlTransaction - class RejectingSetCursor(cursor_class): + class RejectingApplicationNameCursor(cursor_class): def execute(self, query, *args, **kwargs): - if query.startswith("SET application_name"): - query = "SET nonexistent_correlation_parameter = 'failure'" + if query.startswith(("SET application_name", "SELECT pg_catalog.set_config(")): + return super().execute("SET nonexistent_correlation_parameter = 'failure'") return super().execute(query, *args, **kwargs) connection = driver.connect() try: with ( - patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)), + patch.object( + connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingApplicationNameCursor) + ), connection.cursor() as cursor, ): try: @@ -237,7 +425,9 @@ def execute(self, query, *args, **kwargs): try: connection.autocommit = True with ( - patch.object(connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingSetCursor)), + patch.object( + connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingApplicationNameCursor) + ), connection.cursor() as cursor, ): cursor.execute("SELECT 1") @@ -741,6 +931,7 @@ def run_regressions(exporter, tracer): assert enable_postgresql_trace_context_propagation() assert enable_postgresql_trace_context_propagation() for driver in (psycopg, psycopg2): + _bounded_query_statistics(driver, exporter) _sync_queries(driver, exporter, tracer) _sync_concurrency(driver, exporter, tracer) _unsampled_queries(driver, exporter) @@ -763,3 +954,8 @@ def run_regressions(exporter, tracer): _psycopg_server_cursor_boundary(exporter) _psycopg2_commit_boundary(exporter) asyncio.run(_async_query_cancellation()) + _pipeline_errors(exporter) + asyncio.run(_async_pipeline_errors(exporter)) + if sys.version_info >= (3, 14): + _sync_template_queries(exporter) + asyncio.run(_async_template_queries(exporter)) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 345228a2..4c489d84 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -97,14 +97,15 @@ def __init__(self, connection): self.calls = [] self.closed = False - def execute(self, query): - self.calls.append(query) - self.connection.calls.append(query) + def execute(self, query, parameters=None): + call = (query, parameters) if parameters is not None else query + self.calls.append(call) + self.connection.calls.append(call) self.connection.before_execute(query) if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name - elif isinstance(query, str) and query.startswith("SET application_name = '"): - self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") + elif query == "SELECT pg_catalog.set_config('application_name', %s, false)": + self.connection.application_name = parameters[0] return self.connection.application_name def close(self): @@ -122,16 +123,17 @@ def __init__(self, connection): self.calls = [] self.closed = False - async def execute(self, query): - self.calls.append(query) - self.connection.calls.append(query) + async def execute(self, query, parameters=None): + call = (query, parameters) if parameters is not None else query + self.calls.append(call) + self.connection.calls.append(call) result = self.connection.before_execute(query) if inspect.isawaitable(result): await result if query == "SET application_name = DEFAULT": self.connection.application_name = self.connection.default_application_name - elif isinstance(query, str) and query.startswith("SET application_name = '"): - self.connection.application_name = query.removeprefix("SET application_name = '").removesuffix("'") + elif query == "SELECT pg_catalog.set_config('application_name', %s, false)": + self.connection.application_name = parameters[0] return self.connection.application_name async def close(self): @@ -195,6 +197,43 @@ def _fake_untraced_execution(_cursor, query_method, *args, **kwargs): return query_method(*args, **kwargs) +@pytest.mark.parametrize("pipeline_status", [1, 2]) +@pytest.mark.parametrize("query_fails", [False, True]) +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_pipeline_queries_send_no_internal_commands(pipeline_status, query_fails, asynchronous, caplog): + connection = FakeAsyncConnection() if asynchronous else FakeConnection() + connection.info = SimpleNamespace(pipeline_status=pipeline_status) + cursor = connection.cursor() + error = RuntimeError("customer pipeline query failed") + + def _fake_customer_query_failure(_query): + if query_fails: + raise error + + connection.before_execute = _fake_customer_query_failure + + def run(): + if asynchronous: + return asyncio.run( + _execute_with_postgresql_trace_context_async( + _fake_traced_execution_async, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + ) + return _execute_with_postgresql_trace_context( + _fake_traced_execution, None, (cursor, cursor.execute, "SELECT 1"), {} + ) + + if query_fails: + with pytest.raises(RuntimeError) as raised: + run() + assert raised.value is error + else: + assert run() == "customer-app" + assert connection.calls == ["SELECT 1"] + assert connection.cursors == [cursor] + assert not caplog.records + + @pytest.mark.parametrize("nested_fails", [False, True]) def test_custom_cursor_queries_run_before_outer_driver_execution(nested_fails): connection = FakeConnection() @@ -235,10 +274,10 @@ def _execute_customer_query(): assert not thread.is_alive(), "Nested customer query deadlocked" assert results == [TRACEPARENT] assert connection.calls == [ - f"SET application_name = '{SECOND_TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), "SELECT nested", "SET application_name = DEFAULT", - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT outer", "SET application_name = DEFAULT", ] @@ -278,10 +317,10 @@ async def execute(self, query): ) assert result == TRACEPARENT assert connection.calls == [ - f"SET application_name = '{SECOND_TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), "SELECT nested", "SET application_name = DEFAULT", - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT outer", "SET application_name = DEFAULT", ] @@ -316,10 +355,10 @@ async def execute(self, query): ) assert result == TRACEPARENT assert connection.calls == [ - f"SET application_name = '{SECOND_TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), "SELECT child", "SET application_name = DEFAULT", - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT outer", "SET application_name = DEFAULT", ] @@ -407,7 +446,7 @@ def test_callproc_function_name_is_not_a_transaction_command(): ) assert result == TRACEPARENT assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "begin_job", "SET application_name = DEFAULT", ] @@ -500,11 +539,13 @@ def test_query_sets_database_span_traceparent(): assert result == TRACEPARENT assert cursor.calls == ["SELECT application_name"] - assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] + assert connection.cursors[0].calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)) + ] assert all(internal_cursor.closed for internal_cursor in connection.cursors) assert connection.application_name == "startup-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT application_name", "SET application_name = DEFAULT", ] @@ -624,11 +665,13 @@ async def run(): assert result == TRACEPARENT assert cursor.calls == ["SELECT application_name"] - assert connection.cursors[0].calls == [f"SET application_name = '{TRACEPARENT}'"] + assert connection.cursors[0].calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)) + ] assert all(internal_cursor.closed for internal_cursor in connection.cursors) assert connection.application_name == "startup-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT application_name", "SET application_name = DEFAULT", ] @@ -681,7 +724,7 @@ def wrapped(_cursor, query_method, *args, **kwargs): second_thread.start() assert second_query_waiting.wait(timeout=1) assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT first", ] finally: @@ -692,12 +735,16 @@ def wrapped(_cursor, query_method, *args, **kwargs): assert not first_thread.is_alive() assert not second_thread.is_alive() second_calls = ( - [f"SET application_name = '{SECOND_TRACEPARENT}'", second_query, "SET application_name = DEFAULT"] + [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), + second_query, + "SET application_name = DEFAULT", + ] if second_query == "SELECT second" else [second_query] ) assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT first", "SET application_name = DEFAULT", *second_calls, @@ -727,7 +774,7 @@ class CursorTracer(FakeCursorTracer): assert result == TRACEPARENT assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT 1", "SET application_name = DEFAULT", ] @@ -758,7 +805,7 @@ class CursorTracer(FakeCursorTracer): assert connection.commit() == "startup-app" assert cursor.execute("SELECT untraced") == "startup-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT 1", "SET application_name = DEFAULT", "COMMIT", @@ -873,7 +920,7 @@ def _fake_schedule_native_callback(query): ) assert await asyncio.wait_for(completed, timeout=1) == "startup-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT outer", "SET application_name = DEFAULT", "SELECT delayed", @@ -912,7 +959,7 @@ def test_failed_internal_command_attempts_customer_query_and_closes_cursor(caplo cursor = FakePsycopgCursor(connection) def _fake_internal_command_failure(query): - if query.startswith("SET "): + if query.startswith(("SET ", "SELECT pg_catalog.set_config(")): message = "SET rejected" raise RuntimeError(message) @@ -927,7 +974,7 @@ def _fake_internal_command_failure(query): assert result == "customer-app" assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT 1", "SET application_name = DEFAULT", ] @@ -1035,17 +1082,24 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): ) ) await asyncio.wait_for(second_query_started.wait(), timeout=1) - assert connection.calls == [f"SET application_name = '{TRACEPARENT}'", "SELECT first"] + assert connection.calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT first", + ] finally: release_first_query.set() await asyncio.wait_for(asyncio.gather(*([first, second] if second is not None else [first])), timeout=1) second_calls = ( - [f"SET application_name = '{SECOND_TRACEPARENT}'", second_query, "SET application_name = DEFAULT"] + [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), + second_query, + "SET application_name = DEFAULT", + ] if second_query == "SELECT second" else [second_query] ) assert connection.calls == [ - f"SET application_name = '{TRACEPARENT}'", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), "SELECT first", "SET application_name = DEFAULT", *second_calls, @@ -1161,7 +1215,7 @@ async def run(): internal_command_started = asyncio.Event() class BlockingInternalCursor(FakeAsyncPsycopgCursor): - async def execute(self, _query): + async def execute(self, _query, _parameters=None): internal_command_started.set() await asyncio.Event().wait() From 277f2ba3a3e817e6d81a041621033c12231892be Mon Sep 17 00:00:00 2001 From: Pablo Collins Date: Fri, 23 Oct 2026 16:15:18 -0400 Subject: [PATCH 10/10] Fix PostgreSQL transaction setup and concurrent correlation --- docs/postgresql.md | 3 + src/splunk_otel/postgresql.py | 25 +++++- tests/integration/postgresql/regressions.py | 85 +++++++++++++++++++++ tests/test_postgresql.py | 77 +++++++++++++++---- 4 files changed, 173 insertions(+), 17 deletions(-) diff --git a/docs/postgresql.md b/docs/postgresql.md index e5b453b7..0cc4a8e5 100644 --- a/docs/postgresql.md +++ b/docs/postgresql.md @@ -40,6 +40,9 @@ The feature skips: - Psycopg2 async connections and Psycopg 3 queries inside explicit pipelines. - Named/server cursors, fetch operations, and COPY. - SQL transaction commands, such as BEGIN, COMMIT, ROLLBACK, and SAVEPOINT. +- Transaction settings, such as SET TRANSACTION (including SNAPSHOT), + SET SESSION CHARACTERISTICS, and SET [SESSION | LOCAL] transaction_isolation, + transaction_read_only, or transaction_deferrable. - SQL CALL, DO, and query batches. Any semicolon except one at the end causes the query to be skipped, even inside a string or comment. diff --git a/src/splunk_otel/postgresql.py b/src/splunk_otel/postgresql.py index 0062ce04..19c02922 100644 --- a/src/splunk_otel/postgresql.py +++ b/src/splunk_otel/postgresql.py @@ -49,6 +49,9 @@ _TRANSACTION_OPERATIONS = frozenset( {"BEGIN", "START", "COMMIT", "END", "ROLLBACK", "ABORT", "SAVEPOINT", "RELEASE", "PREPARE", "CALL", "DO"} ) +_TRANSACTION_SETTINGS = frozenset( + {"TRANSACTION", "CHARACTERISTICS", "TRANSACTION_ISOLATION", "TRANSACTION_READ_ONLY", "TRANSACTION_DEFERRABLE"} +) class _Lifetime: @@ -372,10 +375,10 @@ def _execute_driver( lease = _live_span_lease(cursor, owner) if not is_instrumentation_enabled() or lease is None: return query_method(*args, **kwargs) - propagate = _can_propagate_statement(cursor, query_method, args, kwargs) parent = _parent_query(cursor.connection, owner) lock = nullcontext() if parent is not None else _get_sync_connection_lock(cursor.connection) with lock: + propagate = _can_propagate_statement(cursor, query_method, args, kwargs) traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None with _active_query_context(cursor.connection, owner, traceparent): return _execute_with_application_name( @@ -402,10 +405,10 @@ async def _execute_driver_async( lease = _live_span_lease(cursor, owner) if not is_instrumentation_enabled() or lease is None: return await query_method(*args, **kwargs) - propagate = _can_propagate_statement(cursor, query_method, args, kwargs) parent = _parent_query(cursor.connection, owner) lock = nullcontext() if parent is not None else _get_async_connection_lock(cursor.connection) async with lock: + propagate = _can_propagate_statement(cursor, query_method, args, kwargs) traceparent = lease.traceparent if propagate else parent.traceparent if parent is not None else None with _active_query_context(cursor.connection, owner, traceparent): return await _execute_with_application_name_async( @@ -516,7 +519,23 @@ def _can_propagate_statement(cursor, query_method, args, kwargs) -> bool: if statement.startswith("("): return True operation = re.match(r"[a-z_][a-z_0-9$]*", statement, re.IGNORECASE) - return operation is not None and operation[0].upper() not in _TRANSACTION_OPERATIONS + if operation is None or operation[0].upper() in _TRANSACTION_OPERATIONS: + return False + # Transaction settings and snapshot imports must precede the first SELECT. + return operation[0].upper() != "SET" or not _sets_transaction_characteristics(statement[operation.end() :]) + + +def _sets_transaction_characteristics(statement: str) -> bool: + for _ in range(2): + statement = _strip_leading_comments(statement) + setting = re.match(r'[a-z_][a-z_0-9$]*|"[^"]*"', statement, re.IGNORECASE) + if setting is None: + return False + name = setting[0].strip('"').upper() + if name not in {"SESSION", "LOCAL"}: + return name in _TRANSACTION_SETTINGS + statement = statement[setting.end() :] + return False def _template_for_statement_check(template: Template) -> Template: diff --git a/tests/integration/postgresql/regressions.py b/tests/integration/postgresql/regressions.py index 51fe78fa..57b88eb2 100644 --- a/tests/integration/postgresql/regressions.py +++ b/tests/integration/postgresql/regressions.py @@ -507,6 +507,89 @@ async def execute(self, query, *args, **kwargs): assert (await cursor.fetchone())[0] == _DEFAULT_APPLICATION_NAME +def _transaction_configuration_statements(driver): + sql = import_module(f"{driver.__name__}.sql") + return ( + ("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE", "transaction_isolation", "serializable"), + ("SET TRANSACTION READ ONLY", "transaction_read_only", "on"), + ("SET TRANSACTION DEFERRABLE", "transaction_deferrable", "on"), + ("SET transaction_isolation = 'serializable'", "transaction_isolation", "serializable"), + (b"SET SESSION transaction_isolation TO 'repeatable read'", "transaction_isolation", "repeatable read"), + ( + sql.SQL("set /* outer /* nested */ comment */ local \"transaction_isolation\" to 'serializable'"), + "transaction_isolation", + "serializable", + ), + ("SET LOCAL transaction_read_only = on", "transaction_read_only", "on"), + ("SET SESSION transaction_deferrable = on", "transaction_deferrable", "on"), + ) + + +def _transaction_configuration(driver, exporter): + sql = import_module(f"{driver.__name__}.sql") + for autocommit in (False, True): + with driver.connect(application_name=_DEFAULT_APPLICATION_NAME) as connection, connection.cursor() as cursor: + connection.autocommit = autocommit + for statement, setting, expected in _transaction_configuration_statements(driver): + cursor.execute("BEGIN") + cursor.execute(statement) + exporter.clear() + cursor.execute("SELECT current_setting(%s), current_setting('application_name')", (setting,)) + value, application_name = cursor.fetchone() + assert value == expected + _assert_database_span(exporter, application_name) + connection.rollback() + + with driver.connect() as source, source.cursor() as source_cursor: + source_cursor.execute("BEGIN ISOLATION LEVEL REPEATABLE READ") + source_cursor.execute("SELECT pg_export_snapshot(), pg_current_snapshot()::text") + snapshot_id, snapshot = source_cursor.fetchone() + cursor.execute("BEGIN ISOLATION LEVEL REPEATABLE READ") + cursor.execute(sql.SQL("SET TRANSACTION SNAPSHOT {}").format(sql.Literal(snapshot_id))) + exporter.clear() + cursor.execute("SELECT pg_current_snapshot()::text, current_setting('application_name')") + imported_snapshot, application_name = cursor.fetchone() + assert imported_snapshot == snapshot + _assert_database_span(exporter, application_name) + connection.rollback() + + +async def _async_transaction_configuration(exporter): + import psycopg + + for autocommit in (False, True): + async with ( + await psycopg.AsyncConnection.connect( + autocommit=autocommit, application_name=_DEFAULT_APPLICATION_NAME + ) as connection, + connection.cursor() as cursor, + ): + for statement, setting, expected in _transaction_configuration_statements(psycopg): + await cursor.execute("BEGIN") + await cursor.execute(statement) + exporter.clear() + await cursor.execute("SELECT current_setting(%s), current_setting('application_name')", (setting,)) + value, application_name = await cursor.fetchone() + assert value == expected + _assert_database_span(exporter, application_name) + await connection.rollback() + + async with await psycopg.AsyncConnection.connect() as source, source.cursor() as source_cursor: + await source_cursor.execute("BEGIN ISOLATION LEVEL REPEATABLE READ") + await source_cursor.execute("SELECT pg_export_snapshot(), pg_current_snapshot()::text") + snapshot_id, snapshot = await source_cursor.fetchone() + await cursor.execute("BEGIN ISOLATION LEVEL REPEATABLE READ") + await cursor.execute( + psycopg.sql.SQL("SET TRANSACTION SNAPSHOT {}").format(psycopg.sql.Literal(snapshot_id)) + ) + exporter.clear() + await cursor.execute("SELECT pg_current_snapshot()::text, current_setting('application_name')") + imported_snapshot, application_name = await cursor.fetchone() + assert imported_snapshot == snapshot + _assert_database_span(exporter, application_name) + await connection.rollback() + + def _transaction_sql(driver, exporter): sql = import_module(f"{driver.__name__}.sql") statements = ( @@ -940,6 +1023,7 @@ def run_regressions(exporter, tracer): _failed_cleanup_command(driver) _nested_cursor_queries(driver, exporter) _transaction_sql(driver, exporter) + _transaction_configuration(driver, exporter) _parenthesized_queries(driver, exporter) _skipped_middle_query(driver, exporter) _custom_commit_before_query(driver, exporter) @@ -947,6 +1031,7 @@ def run_regressions(exporter, tracer): asyncio.run(_async_queries(exporter, tracer)) asyncio.run(_async_nested_cursor_queries(exporter)) asyncio.run(_async_transaction_sql(exporter)) + asyncio.run(_async_transaction_configuration(exporter)) asyncio.run(_async_driver_regressions(exporter)) _custom_cursor(exporter) _manual_custom_cursor(exporter) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 4c489d84..13228ed4 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -70,6 +70,7 @@ def __init__(self): self.application_name = "customer-app" self.calls = [] self.cursors = [] + self.info = SimpleNamespace(pipeline_status=0) self.before_execute = lambda _query: None def cursor(self): @@ -111,6 +112,14 @@ def execute(self, query, parameters=None): def close(self): self.closed = True + def executemany(self, query, parameters_list): + self.connection.info.pipeline_status = 1 + try: + for parameters in parameters_list: + _FakeDriverCursor.execute(self, query, parameters) + finally: + self.connection.info.pipeline_status = 0 + def callproc(self, name): return _FakeDriverCursor.execute(self, name) @@ -139,6 +148,14 @@ async def execute(self, query, parameters=None): async def close(self): self.closed = True + async def executemany(self, query, parameters_list): + self.connection.info.pipeline_status = 1 + try: + for parameters in parameters_list: + await _FakeAsyncDriverCursor.execute(self, query, parameters) + finally: + self.connection.info.pipeline_status = 0 + def _fake_internal_cursor_factory(connection): return connection.cursor() @@ -381,6 +398,19 @@ async def execute(self, query): "ROLLBACK TO SAVEPOINT customer_savepoint", "PREPARE TRANSACTION 'customer_transaction'", "COMMIT PREPARED 'customer_transaction'", + "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "SET TRANSACTION READ ONLY", + "SET TRANSACTION DEFERRABLE", + "SET TRANSACTION SNAPSHOT '00000003-0000001B-1'", + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "SET transaction_isolation = 'serializable'", + "SET SESSION transaction_isolation TO 'repeatable read'", + "SET LOCAL \"transaction_isolation\" = 'serializable'", + "SET transaction_read_only = on", + "SET LOCAL transaction_deferrable = on", + "-- customer comment\n SET /* outer /* nested */ comment */ TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "set /* setting */ local /* name */ \"transaction_isolation\" to 'serializable';", + b"SET TRANSACTION ISOLATION LEVEL SERIALIZABLE", "CALL customer_procedure()", "DO $$ BEGIN COMMIT; END $$", "SELECT 1; COMMIT", @@ -416,9 +446,18 @@ async def run(): @pytest.mark.parametrize( - "statement", ["SELECT 1;", "-- comment\nSELECT 1", "/* outer /* nested */ comment */ SELECT 1"] + "statement", + [ + "SELECT 1;", + "-- comment\nSELECT 1", + "/* outer /* nested */ comment */ SELECT 1", + "SET statement_timeout = '1s'", + "SET LOCAL statement_timeout = '1s'", + "SET SESSION \"statement_timeout\" = '1s'", + "SET \"myapp.transaction_isolation\" = 'serializable'", + ], ) -def test_data_statement_after_comments_propagates_context(statement): +def test_statement_without_transaction_configuration_propagates_context(statement): connection = FakeConnection() cursor = FakePsycopgCursor(connection) assert ( @@ -677,8 +716,9 @@ async def run(): ] +@pytest.mark.parametrize("first_method", ["execute", "executemany"]) @pytest.mark.parametrize("second_query", ["SELECT second", "COMMIT"]) -def test_shared_connection_keeps_set_and_query_together(monkeypatch, second_query): +def test_shared_connection_keeps_set_and_query_together(monkeypatch, first_method, second_query): connection = FakeConnection() first_cursor = FakePsycopgCursor(connection) second_cursor = FakePsycopgCursor(connection) @@ -708,15 +748,19 @@ def before_execute(query): connection.before_execute = before_execute - def run(cursor, query, span_context): - def wrapped(_cursor, query_method, *args, **kwargs): + def run(cursor, method, query, span_context): + def _fake_traced_execution(_cursor, query_method, *args, **kwargs): with trace.use_span(NonRecordingSpan(span_context)): return query_method(*args, **kwargs) - _execute_with_postgresql_trace_context(wrapped, None, (cursor, cursor.execute, query), {}) + args = (cursor, getattr(cursor, method), query) + if method == "executemany": + args += ([()],) + _execute_with_postgresql_trace_context(_fake_traced_execution, None, args, {}) - first_thread = threading.Thread(target=run, args=(first_cursor, "SELECT first", DATABASE_CONTEXT)) - second_thread = threading.Thread(target=run, args=(second_cursor, second_query, SECOND_DATABASE_CONTEXT)) + first_thread = threading.Thread(target=run, args=(first_cursor, first_method, "SELECT first", DATABASE_CONTEXT)) + second_thread = threading.Thread(target=run, args=(second_cursor, "execute", second_query, SECOND_DATABASE_CONTEXT)) + first_call = ("SELECT first", ()) if first_method == "executemany" else "SELECT first" try: first_thread.start() @@ -725,7 +769,7 @@ def wrapped(_cursor, query_method, *args, **kwargs): assert second_query_waiting.wait(timeout=1) assert connection.calls == [ ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), - "SELECT first", + first_call, ] finally: release_first_query.set() @@ -745,7 +789,7 @@ def wrapped(_cursor, query_method, *args, **kwargs): ) assert connection.calls == [ ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), - "SELECT first", + first_call, "SET application_name = DEFAULT", *second_calls, ] @@ -1040,8 +1084,9 @@ def _fake_query_or_cleanup_failure(query): assert "Failed to set PostgreSQL application_name" in caplog.text +@pytest.mark.parametrize("first_method", ["execute", "executemany"]) @pytest.mark.parametrize("second_query", ["SELECT second", "COMMIT"]) -def test_async_shared_connection_keeps_set_and_query_together(second_query): +def test_async_shared_connection_keeps_set_and_query_together(first_method, second_query): async def run(): connection = FakeAsyncConnection() first_query_started = asyncio.Event() @@ -1062,11 +1107,15 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): first_cursor = FakeAsyncPsycopgCursor(connection) second_cursor = FakeAsyncPsycopgCursor(connection) + first_args = (first_cursor, getattr(first_cursor, first_method), "SELECT first") + if first_method == "executemany": + first_args += ([()],) + first_call = ("SELECT first", ()) if first_method == "executemany" else "SELECT first" first = asyncio.create_task( _execute_with_postgresql_trace_context_async( _fake_traced_execution_async, None, - (first_cursor, first_cursor.execute, "SELECT first"), + first_args, {}, ) ) @@ -1084,7 +1133,7 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): await asyncio.wait_for(second_query_started.wait(), timeout=1) assert connection.calls == [ ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), - "SELECT first", + first_call, ] finally: release_first_query.set() @@ -1100,7 +1149,7 @@ async def _fake_second_traced_execution(_cursor, query_method, *args, **kwargs): ) assert connection.calls == [ ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), - "SELECT first", + first_call, "SET application_name = DEFAULT", *second_calls, ]