Add PostgreSQL query-to-trace correlation

.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

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

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`.

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()
 

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"