Merge branch 'main' into 173_usage_configuration

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

@@ -296,6 +296,18 @@ Litestar Plugin Configuration
    :lines: 2-26
    :dedent: 2
 
+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

@@ -7,7 +7,7 @@ maintainers = [{ name = "Litestar Developers", email = "hello@litestar.dev" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.10, <4.0"
-version = "0.28.1"
+version = "0.29.0"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"
@@ -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",
@@ -209,7 +217,7 @@ opt_level = "3"   # Maximum optimization (0-3)
 allow_dirty = true
 commit = false
 commit_args = "--no-verify"
-current_version = "0.28.1"
+current_version = "0.29.0"
 ignore_missing_files = false
 ignore_missing_version = false
 message = "chore(release): bump to v{new_version}"

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)