feat(observability): enable telemetry integrations (#237)

Introduce runtime helpers for lifecycle and span orchestration, improve storage telemetry diagnostics, and add comprehensive unit and integration tests. Implement redaction for SQL statements and parameters, and provide telemetry snapshot functions for monitoring recent storage events.

Author: cofin

docs/extensions/litestar/api.rst

@@ -34,7 +34,8 @@ Configure the plugin via ``extension_config`` in database configuration:
                "commit_mode": "autocommit",
                "extra_commit_statuses": {201, 204},
                "extra_rollback_statuses": {409},
-               "enable_correlation_middleware": True
+               "enable_correlation_middleware": True,
+               "correlation_header": "x-correlation-id",
            }
        }
    )
@@ -74,10 +75,22 @@ Configuration Options
      - ``set[int]``
      - ``None``
      - Additional HTTP status codes that trigger rollbacks
-   * - ``enable_correlation_middleware``
+  * - ``enable_correlation_middleware``
      - ``bool``
      - ``True``
      - Enable request correlation tracking
+  * - ``correlation_header``
+     - ``str``
+     - ``"X-Request-ID"``
+     - HTTP header to read when populating the correlation ID middleware
+   * - ``correlation_headers``
+     - ``list[str]``
+     - ``[]``
+     - Additional headers to consider (auto-detected headers are appended unless disabled)
+   * - ``auto_trace_headers``
+     - ``bool``
+     - ``True``
+     - Toggle automatic detection of standard tracing headers (`Traceparent`, `X-Cloud-Trace-Context`, etc.)
 
 Session Stores
 ==============

docs/guides/architecture/architecture.md

@@ -14,6 +14,7 @@ orphan: true
 4. [Driver Implementation](#driver-implementation)
 5. [Parameter Handling](#parameter-handling)
 6. [Testing & Development](#testing--development)
+7. [Observability Runtime](#observability-runtime)
 
 ---
 
@@ -308,3 +309,12 @@ make install     # Standard development installation
 2. Implement the `config.py` and `driver.py` files.
 3. Add integration tests for the new adapter.
 4. Document any special cases or configurations.
+
+## Observability Runtime
+
+The observability subsystem (lifecycle dispatcher, statement observers, span manager, diagnostics) now sits alongside the driver architecture. Refer to the dedicated [Observability Runtime guide](./observability.md) for:
+
+- configuration sources (`ObservabilityConfig`, adapter overrides, and `driver_features` compatibility),
+- the full list of lifecycle events emitted by SQLSpec,
+- guidance on statement observers, redaction, and OpenTelemetry spans,
+- the Phase 4/5 roadmap for spans + diagnostics.

docs/guides/architecture/observability.md

@@ -0,0 +1,136 @@
+# SQLSpec Observability Runtime
+
+This guide explains how the consolidated observability stack works after the Lifecycle Dispatcher + Statement Observer integration. Use it as the single source of truth when wiring new adapters, features, or docs.
+
+## Goals
+
+1. **Unified Hooks** – every pool, connection, session, and query event is emitted through one dispatcher with zero work when no listeners exist.
+2. **Structured Statement Events** – observers receive normalized payloads (`StatementEvent`) for printing, logging, or exporting to tracing systems.
+3. **Optional OpenTelemetry Spans** – span creation is lazy and never imports `opentelemetry` unless spans are enabled.
+4. **Diagnostics** – storage bridge + serializer metrics + lifecycle counters roll up under `SQLSpec.telemetry_snapshot()` (Phase 5).
+5. **Loader & Migration Telemetry** – SQL file loader, caching, and migration runners emit metrics/spans without additional plumbing (Phase 7).
+
+## Configuration Sources
+
+There are three ways to enable observability today:
+
+1. **Registry-Level** – pass `observability_config=ObservabilityConfig(...)` to `SQLSpec()`.
+2. **Adapter Override** – each config constructor accepts `observability_config=` for adapter-specific knobs.
+3. **`driver_features` Compatibility** – existing keys such as `"on_connection_create"`, `"on_pool_destroy"`, and `"on_session_start"` are automatically promoted into lifecycle observers, so user-facing APIs do **not** change.
+
+```python
+from sqlspec import SQLSpec
+from sqlspec.adapters.duckdb import DuckDBConfig
+
+def ensure_extensions(connection):
+    connection.execute("INSTALL http_client; LOAD http_client;")
+
+config = DuckDBConfig(
+    pool_config={"database": ":memory:"},
+    driver_features={
+        "extensions": [{"name": "http_client"}],
+        "on_connection_create": ensure_extensions,  # promoted to observability runtime
+    },
+)
+
+sql = SQLSpec(observability_config=ObservabilityConfig(print_sql=True))
+sql.add_config(config)
+```
+
+> **Implementation note:** During config initialization we inspect `driver_features` for known hook keys and wrap them into `ObservabilityConfig` callbacks. Hooks that accepted a raw resource (e.g., connection) continue to do so without additional adapter plumbing.
+
+## Lifecycle Events
+
+The dispatcher exposes the following events (all opt-in and guard-checked):
+
+| Event | Context contents |
+| --- | --- |
+| `on_pool_create` / `on_pool_destroy` | `pool`, `config`, `bind_key`, `correlation_id` |
+| `on_connection_create` / `on_connection_destroy` | `connection`, plus base context |
+| `on_session_start` / `on_session_end` | `session` / driver instance |
+| `on_query_start` / `on_query_complete` | SQL text, parameters, metadata |
+| `on_error` | `exception` plus last query context |
+
+`SQLSpec.provide_connection()` and `SQLSpec.provide_session()` now emit these events automatically, regardless of whether the caller uses registry helpers or adapter helpers directly.
+
+## Statement Observers & Print SQL
+
+Statement observers receive `StatementEvent` objects. Typical uses:
+
+* enable `print_sql=True` to attach the built-in logger.
+* add custom redaction rules via `RedactionConfig` (mask parameters, mask literals, allow-list names).
+* forward events to bespoke loggers or telemetry exporters.
+
+```python
+def log_statement(event: StatementEvent) -> None:
+    logger.info("%s (%s) -> %ss", event.operation, event.driver, event.duration_s)
+
+ObservabilityConfig(
+    print_sql=False,
+    statement_observers=(log_statement,),
+    redaction=RedactionConfig(mask_parameters=True, parameter_allow_list=("tenant_id",)),
+)
+```
+
+### Optional Exporters (OpenTelemetry & Prometheus)
+
+Two helper modules wire optional dependencies into the runtime without forcing unconditional imports:
+
+* `sqlspec.extensions.otel.enable_tracing()` ensures `opentelemetry-api` is installed, then returns an `ObservabilityConfig` whose `TelemetryConfig` enables spans and (optionally) injects a tracer provider factory.
+* `sqlspec.extensions.prometheus.enable_metrics()` ensures `prometheus-client` is installed and appends a `PrometheusStatementObserver` that emits counters and histograms for every `StatementEvent`.
+
+Both helpers rely on the conditional stubs defined in `sqlspec/typing.py`, so they remain safe to import even when the extras are absent.
+
+```python
+from sqlspec.extensions import otel, prometheus
+
+config = otel.enable_tracing(resource_attributes={"service.name": "orders-api"})
+config = prometheus.enable_metrics(base_config=config, label_names=("driver", "operation", "adapter"))
+sql = SQLSpec(observability_config=config)
+```
+
+You can also opt in per adapter by passing `extension_config["otel"]` or `extension_config["prometheus"]` when constructing a config; the helpers above are invoked automatically during initialization.
+
+## Loader & Migration Telemetry
+
+`SQLSpec` instantiates a dedicated `ObservabilityRuntime` for the SQL file loader and shares it with every migration command/runner. Instrumentation highlights:
+
+- Loader metrics such as `SQLFileLoader.loader.load.invocations`, `.cache.hit`, `.files.loaded`, `.statements.loaded`, and `.directories.scanned` fire automatically when queries are loaded or cache state is inspected.
+- Migration runners publish cache stats (`{Config}.migrations.listing.cache_hit`, `.cache_miss`, `.metadata.cache_hit`), command metrics (`{Config}.migrations.command.upgrade.invocations`, `.downgrade.errors`), and per-migration execution metrics (`{Config}.migrations.upgrade.duration_ms`, `.downgrade.applied`).
+- Command and migration spans (`sqlspec.migration.command.upgrade`, `sqlspec.migration.upgrade`) include version numbers, bind keys, and correlation IDs; they end with duration attributes even when exceptions occur.
+
+All metrics surface through `SQLSpec.telemetry_snapshot()` under the adapter key, so exporters observe a flat counter space regardless of which subsystem produced the events.
+
+## Span Manager & Diagnostics
+
+* **Span Manager:** Query spans ship today, lifecycle events emit `sqlspec.lifecycle.*` spans, storage bridge helpers wrap reads/writes with `sqlspec.storage.*` spans, and migration runners create `sqlspec.migration.*` spans for both commands and individual revisions. Mocked span tests live in `tests/unit/test_observability.py`.
+* **Diagnostics:** `TelemetryDiagnostics` aggregates lifecycle counters, loader/migration metrics, storage bridge telemetry, and serializer cache stats. Storage telemetry carries backend IDs, bind key, and correlation IDs so snapshots/spans inherit the same context, and `SQLSpec.telemetry_snapshot()` exposes that data via flat counters plus a `storage_bridge.recent_jobs` list detailing the last 25 operations.
+
+Example snapshot payload:
+
+```
+{
+  "storage_bridge.bytes_written": 2048,
+  "storage_bridge.recent_jobs": [
+    {
+      "destination": "alias://warehouse/users.parquet",
+      "backend": "s3",
+      "bytes_processed": 2048,
+      "rows_processed": 16,
+      "config": "AsyncpgConfig",
+      "bind_key": "analytics",
+      "correlation_id": "8f64c0f6",
+      "format": "parquet"
+    }
+  ],
+  "serializer.hits": 12,
+  "serializer.misses": 2,
+  "AsyncpgConfig.lifecycle.query_start": 4
+}
+```
+
+## Next Steps (2025 Q4)
+
+1. **Exporter Validation:** Exercise the OpenTelemetry/Prometheus helpers against the new loader + migration metrics and document recommended dashboards.
+2. **Adapter Audit:** Confirm every adapter’s migration tracker benefits from the instrumentation (especially Oracle/BigQuery fixtures) and extend coverage where needed.
+3. **Performance Budgets:** Add guard-path benchmarks/tests to ensure disabled observability remains near-zero overhead now that migration/loader events emit metrics by default.

docs/guides/extensions/litestar.md

@@ -13,7 +13,7 @@ Explains how to wire SQLSpec into Litestar using the official plugin, covering d
 - Commit strategies: `manual`, `autocommit`, and `autocommit_include_redirect`, configured via `extension_config["litestar"]["commit_mode"]`.
 - Session storage uses adapter-specific stores built on `BaseSQLSpecStore` (e.g., `AsyncpgStore`, `AiosqliteStore`).
 - CLI support registers `litestar db ...` commands by including `database_group` in the Litestar CLI app.
-- Correlation middleware emits request IDs in query logs (`enable_correlation_middleware=True` by default).
+- Correlation middleware emits request IDs in query logs (`enable_correlation_middleware=True` by default). It auto-detects standard tracing headers (`X-Request-ID`, `Traceparent`, `X-Cloud-Trace-Context`, `X-Amzn-Trace-Id`, etc.) unless you override the set via `correlation_header` / `correlation_headers`.
 
 ## Installation
 
@@ -95,6 +95,10 @@ config = AsyncpgConfig(
         }
     },
 )
+
+## Correlation IDs
+
+Enable request-level correlation tracking (on by default) to thread Litestar requests into SQLSpec's observability runtime. The plugin inspects `X-Request-ID`, `Traceparent`, `X-Cloud-Trace-Context`, `X-Amzn-Trace-Id`, `grpc-trace-bin`, and `X-Correlation-ID` automatically, then falls back to generating a UUID if none are present. Override the primary header with `correlation_header`, append more via `correlation_headers`, or set `auto_trace_headers=False` to opt out of the auto-detection list entirely. Observers (print SQL, custom hooks, OpenTelemetry spans) automatically attach the current `correlation_id` to their payloads. Disable the middleware with `enable_correlation_middleware=False` when another piece of infrastructure manages IDs.
 ```
 
 ## Transaction Management
@@ -162,7 +166,7 @@ Commands include `db migrate`, `db upgrade`, `db downgrade`, and `db status`. Th
 
 ## Middleware and Observability
 
-- Correlation middleware annotates query logs with request-scoped IDs. Disable by setting `enable_correlation_middleware=False`.
+- Correlation middleware annotates query logs with request-scoped IDs. Disable by setting `enable_correlation_middleware=False`, override the primary header via `correlation_header`, add more with `correlation_headers`, or disable auto-detection using `auto_trace_headers=False`.
 - The plugin enforces graceful shutdown by closing pools during Litestar’s lifespan events.
 - Combine with Litestar’s `TelemetryConfig` to emit tracing spans around database calls.
 

docs/usage/configuration.rst

@@ -497,10 +497,25 @@ Litestar Plugin Configuration
                "pool_key": "db_pool",
                "commit_mode": "autocommit",
                "enable_correlation_middleware": True,
+               "correlation_header": "x-correlation-id",
+               "correlation_headers": ["x-custom-trace"],
+               "auto_trace_headers": True,  # Detect Traceparent, X-Cloud-Trace-Context, etc.
            }
        }
    )
 
+Telemetry Snapshot
+~~~~~~~~~~~~~~~~~~
+
+Call ``SQLSpec.telemetry_snapshot()`` to inspect lifecycle counters, serializer metrics, and recent storage jobs:
+
+.. code-block:: python
+
+   snapshot = spec.telemetry_snapshot()
+   print(snapshot["storage_bridge.bytes_written"])
+   for job in snapshot.get("storage_bridge.recent_jobs", []):
+       print(job["destination"], job.get("correlation_id"))
+
 Environment-Based Configuration
 -------------------------------
 

docs/usage/framework_integrations.rst

@@ -339,13 +339,16 @@ Enable request correlation tracking via ``extension_config``:
            pool_config={"dsn": "postgresql://..."},
            extension_config={
                "litestar": {
-                   "enable_correlation_middleware": True  # Default: True
-               }
-           }
-       )
-   )
-
-   # Queries will include correlation IDs in logs
+                   "enable_correlation_middleware": True,  # Default: True
+                   "correlation_header": "x-request-id",
+                   "correlation_headers": ["x-client-trace"],
+                   "auto_trace_headers": True,
+                }
+            }
+        )
+    )
+
+   # Queries will include correlation IDs in logs (header or generated UUID)
    # Format: [correlation_id=abc123] SELECT * FROM users
 
 FastAPI Integration

pyproject.toml

@@ -179,6 +179,14 @@ include = [
   "sqlspec/utils/fixtures.py",            # File fixture loading
   "sqlspec/utils/data_transformation.py", # Data transformation utilities
 
+  # === OBSERVABILITY ===
+  "sqlspec/observability/_config.py",
+  "sqlspec/observability/_diagnostics.py",
+  "sqlspec/observability/_dispatcher.py",
+  "sqlspec/observability/_observer.py",
+  "sqlspec/observability/_runtime.py",
+  "sqlspec/observability/_spans.py",
+
   # === STORAGE LAYER ===
   "sqlspec/storage/_utils.py",
   "sqlspec/storage/registry.py",

sqlspec/_typing.py

@@ -568,6 +568,9 @@ def get_tracer(
         ) -> Tracer:
             return Tracer()  # type: ignore[abstract] # pragma: no cover
 
+        def get_tracer_provider(self) -> Any:  # pragma: no cover
+            return None
+
         TracerProvider = type(None)  # Shim for TracerProvider if needed elsewhere
         StatusCode = type(None)  # Shim for StatusCode
         Status = type(None)  # Shim for Status
@@ -600,6 +603,8 @@ def __init__(
             unit: str = "",
             registry: Any = None,
             ejemplar_fn: Any = None,
+            buckets: Any = None,
+            **_: Any,
         ) -> None:
             return None
 

sqlspec/adapters/adbc/config.py

@@ -9,7 +9,7 @@
 
 from sqlspec.adapters.adbc._types import AdbcConnection
 from sqlspec.adapters.adbc.driver import AdbcCursor, AdbcDriver, AdbcExceptionHandler, get_adbc_statement_config
-from sqlspec.config import ADKConfig, FastAPIConfig, FlaskConfig, LitestarConfig, NoPoolSyncConfig, StarletteConfig
+from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig
 from sqlspec.core import StatementConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.utils.module_loader import import_string
@@ -21,6 +21,8 @@
 
     from sqlglot.dialects.dialect import DialectType
 
+    from sqlspec.observability import ObservabilityConfig
+
 logger = logging.getLogger("sqlspec.adapters.adbc")
 
 
@@ -116,7 +118,8 @@ def __init__(
         statement_config: StatementConfig | None = None,
         driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None,
         bind_key: str | None = None,
-        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
+        extension_config: "ExtensionConfigs | None" = None,
+        observability_config: "ObservabilityConfig | None" = None,
     ) -> None:
         """Initialize configuration.
 
@@ -127,6 +130,7 @@ def __init__(
             driver_features: Driver feature configuration (AdbcDriverFeatures)
             bind_key: Optional unique identifier for this configuration
             extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
+            observability_config: Adapter-level observability overrides for lifecycle hooks and observers
         """
         if connection_config is None:
             connection_config = {}
@@ -168,6 +172,7 @@ def __init__(
             driver_features=processed_driver_features,
             bind_key=bind_key,
             extension_config=extension_config,
+            observability_config=observability_config,
         )
 
     def _resolve_driver_name(self) -> str:
@@ -366,9 +371,10 @@ def session_manager() -> "Generator[AdbcDriver, None, None]":
                     or self.statement_config
                     or get_adbc_statement_config(str(self._get_dialect() or "sqlite"))
                 )
-                yield self.driver_type(
+                driver = self.driver_type(
                     connection=connection, statement_config=final_statement_config, driver_features=self.driver_features
                 )
+                yield self._prepare_driver(driver)
 
         return session_manager()
 

sqlspec/adapters/adbc/driver.py

@@ -706,8 +706,8 @@ def select_to_storage(
         self._require_capability("arrow_export_enabled")
         arrow_result = self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs)
         sync_pipeline: SyncStoragePipeline = cast("SyncStoragePipeline", self._storage_pipeline())
-        telemetry_payload = arrow_result.write_to_storage_sync(
-            destination, format_hint=format_hint, pipeline=sync_pipeline
+        telemetry_payload = self._write_result_to_storage_sync(
+            arrow_result, destination, format_hint=format_hint, pipeline=sync_pipeline
         )
         self._attach_partition_telemetry(telemetry_payload, partitioner)
         return self._create_storage_job(telemetry_payload, telemetry)