diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml index 4997936e..e0e3db9f 100644 --- a/.github/workflows/ci-main.yml +++ b/.github/workflows/ci-main.yml @@ -49,3 +49,57 @@ jobs: - name: Run oteltests run: oteltest tests/integration/*.py + + postgresql-integration: + name: PostgreSQL oteltest (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_TEST_INSTRUMENTATION: ${{ matrix.mode }} + 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 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/README.md b/README.md index 9e46c0d7..5425368e 100644 --- a/README.md +++ b/README.md @@ -58,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_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. + ## 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..0cc4a8e5 --- /dev/null +++ b/docs/postgresql.md @@ -0,0 +1,84 @@ +# PostgreSQL trace correlation + +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. + +## Setup + +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. + +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 +``` + +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. + +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()`. + +## Supported queries + +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. +- 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. + +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. + +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 \ +POSTGRESQL_TEST_INSTRUMENTATION=auto \ +oteltest tests/integration/postgresql/trace_context.py +``` + +Repeat with `POSTGRESQL_TEST_INSTRUMENTATION=manual` and `disabled`. 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..2e8ff9f0 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_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/postgresql.py b/src/splunk_otel/postgresql.py new file mode 100644 index 00000000..19c02922 --- /dev/null +++ b/src/splunk_otel/postgresql.py @@ -0,0 +1,656 @@ +# 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 re +import sys +import threading +import weakref +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar +from dataclasses import dataclass, field +from functools import partial +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 sys.version_info >= (3, 14): + from string.templatelib import Interpolation, Template + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from opentelemetry.trace import SpanContext + +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( + {"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: + 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 +_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 at driver execution.""" + global _enabled # noqa: PLW0603 + + with _enable_lock: + if _enabled: + return True + + cursor_tracer = _load_compatible_cursor_tracer() + if cursor_tracer is None: + return False + + adapters = _load_cursor_adapters() + 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 + + _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 + + 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( + psycopg3_instrumentation.PsycopgInstrumentor, + "instrument_connection", + _manual_connection_wrapper(psycopg3_instrumentation), + ) + + 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]: + adapters = [] + for name in _DRIVER_MODULE_NAMES: + try: + module = import_module(f"opentelemetry.instrumentation.{name}") + except ImportError: + continue + if name == "psycopg": + adapters.extend(_create_psycopg_cursor_adapters(module)) + else: + adapters.append(_create_psycopg2_cursor_adapter(module)) + return adapters + + +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): + 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_parameter_signature(method: Any, names: tuple[str, ...]) -> bool: + if not callable(method): + return False + try: + parameters = tuple(inspect.signature(method).parameters.values())[: len(names)] + except (TypeError, ValueError): + return False + 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 + ) + + +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_client_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): + 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) + + +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_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() + + async def correlated_query_method(*method_args, **method_kwargs): + 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) + + +@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], + args, + kwargs, + *, + 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) + 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( + cursor.connection, + query_method, + args, + kwargs, + traceparent=traceparent, + parent=parent, + propagate=propagate, + new_internal_cursor=new_internal_cursor, + ) + + +async def _execute_driver_async( + cursor, + query_method: Callable[..., Any], + args, + kwargs, + *, + 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) + 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( + 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: + 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 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 and query.lifetime.is_active() + ), + 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") + 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) + 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 + + # 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) + if statement.startswith("("): + return True + operation = re.match(r"[a-z_][a-z_0-9$]*", statement, re.IGNORECASE) + 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: + # 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(("--", "/*")): + 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: + 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, *, new_internal_cursor: Callable[[Any], Any]) -> None: + statement, parameters = _application_name_command(traceparent) + try: + cursor = new_internal_cursor(connection) + try: + with suppress_instrumentation(): + result = cursor.execute(statement, parameters) + 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, *, new_internal_cursor: Callable[[Any], Any] +) -> None: + statement, parameters = _application_name_command(traceparent) + try: + cursor = new_internal_cursor(connection) + try: + with suppress_instrumentation(): + result = cursor.execute(statement, parameters) + 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_command(traceparent: str | None) -> tuple[str, tuple[str] | None]: + if traceparent is None: + return _SET_APPLICATION_NAME_DEFAULT, None + return _SET_APPLICATION_NAME, (traceparent,) + + +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_psycopg_module(cursor_class.__module__) for cursor_class in type(cursor).__mro__) + + +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) diff --git a/src/splunk_otel/runtime.py b/src/splunk_otel/runtime.py index fe19a54a..81a09cf6 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_DB_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_DB_TRACE_CONTEXT_PROPAGATION_ENABLED, "false"): + enable_postgresql_trace_context_propagation() + + class _LoggingInstrumentor(Protocol): def instrument(self) -> None: ... diff --git a/tests/integration/postgresql/regressions.py b/tests/integration/postgresql/regressions.py new file mode 100644 index 00000000..57b88eb2 --- /dev/null +++ b/tests/integration/postgresql/regressions.py @@ -0,0 +1,1046 @@ +import asyncio +import sys +from concurrent.futures import ThreadPoolExecutor +from contextlib import suppress +from importlib import import_module +from unittest.mock import patch + + +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 _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" + 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 application_name == _DEFAULT_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] == _DEFAULT_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(application_name=_DEFAULT_APPLICATION_NAME) 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, application_name=_DEFAULT_APPLICATION_NAME + ) 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] == _DEFAULT_APPLICATION_NAME + assert len([span for span in exporter.get_finished_spans() if span.kind == SpanKind.CLIENT]) == 12 + + class RejectingApplicationNameCursor(psycopg.AsyncCursor): + async def execute(self, query, *args, **kwargs): + 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(RejectingApplicationNameCursor) + ): + 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 RejectingApplicationNameCursor(cursor_class): + def execute(self, query, *args, **kwargs): + 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(RejectingApplicationNameCursor) + ), + 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() + try: + connection.autocommit = True + with ( + patch.object( + connection.cursor_factory, "_new_internal_cursor", staticmethod(RejectingApplicationNameCursor) + ), + 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 _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_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 = ( + ("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 + + 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() + try: + with suppress_instrumentation(), connection.cursor() as cursor: + cursor.execute("CREATE TEMP TABLE cleanup_failure (value int)") + connection.commit() + 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 = ( + 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 _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 + + 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] == _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) + assert state == "active", "Customer query did not start before cancellation" + assert statement == "SELECT pg_sleep(10)" + finally: + query.cancel() + with suppress(asyncio.CancelledError): + await query + + 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 run_regressions(exporter, tracer): + import psycopg + import psycopg2 + from splunk_otel.postgresql import enable_postgresql_trace_context_propagation + + 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) + _reset_to_default(driver) + _failed_internal_command(driver) + _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) + 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_transaction_configuration(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()) + _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/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_configurator.py b/tests/test_configurator.py index cf11db99..e45c3073 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_DB_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..13228ed4 --- /dev/null +++ b/tests/test_postgresql.py @@ -0,0 +1,1307 @@ +# 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 inspect +import threading +from contextlib import nullcontext +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 ( + _AsyncDriverCursor, + _CursorAdapter, + _SyncDriverCursor, + _cursor_factory_wrapper, + _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), +) +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: + def __init__(self): + self.default_application_name = "startup-app" + self.application_name = "customer-app" + self.calls = [] + self.cursors = [] + self.info = SimpleNamespace(pipeline_status=0) + self.before_execute = lambda _query: None + + def cursor(self): + cursor = FakePsycopgCursor(self) + self.cursors.append(cursor) + return cursor + + def commit(self): + self.calls.append("COMMIT") + return self.application_name + + +class FakeAsyncConnection(FakeConnection): + def cursor(self): + cursor = FakeAsyncPsycopgCursor(self) + self.cursors.append(cursor) + return cursor + + +class _FakeDriverCursor: + __module__ = "psycopg" + + def __init__(self, connection): + self.connection = connection + self.calls = [] + self.closed = False + + 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 query == "SELECT pg_catalog.set_config('application_name', %s, false)": + self.connection.application_name = parameters[0] + return self.connection.application_name + + 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) + + +class _FakeAsyncDriverCursor: + __module__ = "psycopg.fake" + + def __init__(self, connection): + self.connection = connection + self.calls = [] + self.closed = False + + 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 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): + 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() + + +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) + + +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_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_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("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() + 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 == "startup-app" + 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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), + "SELECT nested", + "SET application_name = DEFAULT", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT outer", + "SET application_name = DEFAULT", + ] + + +@pytest.mark.parametrize("nested_fails", [False, True]) +def test_async_custom_cursor_queries_run_before_outer_driver_execution(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 == "startup-app" + 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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), + "SELECT nested", + "SET application_name = DEFAULT", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT outer", + "SET application_name = DEFAULT", + ] + + asyncio.run(run()) + + +def test_async_custom_cursor_can_await_auxiliary_task_before_outer_query(): + 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 await children[0] == SECOND_TRACEPARENT + assert connection.application_name == "startup-app" + 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 connection.calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (SECOND_TRACEPARENT,)), + "SELECT child", + "SET application_name = DEFAULT", + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT outer", + "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'", + "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", + "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", + "SET statement_timeout = '1s'", + "SET LOCAL statement_timeout = '1s'", + "SET SESSION \"statement_timeout\" = '1s'", + "SET \"myapp.transaction_isolation\" = 'serializable'", + ], +) +def test_statement_without_transaction_configuration_propagates_context(statement): + connection = FakeConnection() + cursor = FakePsycopgCursor(connection) + assert ( + _execute_with_postgresql_trace_context(_fake_traced_execution, None, (cursor, cursor.execute, statement), {}) + == TRACEPARENT + ) + + +@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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (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) + + 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 == [ + ("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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (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) + result = _execute_with_postgresql_trace_context( + _fake_untraced_execution, + None, + (cursor, cursor.execute, "SELECT 2"), + {}, + ) + + assert result == "startup-app" + assert connection.calls == [ + "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_sends_only_customer_query(): + class FakePsycopg2NamedCursor(FakePsycopgCursor): + __module__ = "psycopg2.extensions" + + def execute(self, query): + if self.calls: + raise RuntimeError + return super().execute(query) + + connection = FakeConnection() + cursor = FakePsycopg2NamedCursor(connection) + cursor.name = "server_cursor" + + with trace.use_span(NonRecordingSpan(PARENT_CONTEXT)): + 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_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 == [ + ("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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT application_name", + "SET application_name = DEFAULT", + ] + + +@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, first_method, second_query): + 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, 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) + + 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, 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() + assert first_query_started.wait(timeout=1) + second_thread.start() + assert second_query_waiting.wait(timeout=1) + assert connection.calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + first_call, + ] + 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() + second_calls = ( + [ + ("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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + first_call, + "SET application_name = DEFAULT", + *second_calls, + ] + + +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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + "SELECT 1", + "SET application_name = DEFAULT", + ] + + +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_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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (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): + 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("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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (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( + 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 ", "SELECT pg_catalog.set_config(")): + 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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (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 + + +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.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") + cleanup_error = RuntimeError("cleanup failed") + + def _fake_query_or_cleanup_failure(query): + if query == "SET application_name = DEFAULT": + raise cleanup_error + 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 + + +@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(first_method, second_query): + async def run(): + connection = FakeAsyncConnection() + first_query_started = asyncio.Event() + second_query_started = asyncio.Event() + release_first_query = asyncio.Event() + + async def _fake_blocking_driver_execute(query): + if query == "SELECT first": + first_query_started.set() + await release_first_query.wait() + + 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 = 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_args, + {}, + ) + ) + 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, second_query), + {}, + ) + ) + await asyncio.wait_for(second_query_started.wait(), timeout=1) + assert connection.calls == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + first_call, + ] + finally: + release_first_query.set() + await asyncio.wait_for(asyncio.gather(*([first, second] if second is not None else [first])), timeout=1) + second_calls = ( + [ + ("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 == [ + ("SELECT pg_catalog.set_config('application_name', %s, false)", (TRACEPARENT,)), + first_call, + "SET application_name = DEFAULT", + *second_calls, + ] + + 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_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) + 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()) + + +@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") + cleanup_error = RuntimeError("cleanup failed") + + def _fake_query_or_cleanup_failure(query): + if query == "SET application_name = DEFAULT": + raise cleanup_error + 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() + + class BlockingInternalCursor(FakeAsyncPsycopgCursor): + async def execute(self, _query, _parameters=None): + 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_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) diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 234d6662..d89bdd16 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_DB_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)