Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions docs/extensions/litestar/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}
}
)
Expand Down Expand Up @@ -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
==============
Expand Down
10 changes: 10 additions & 0 deletions docs/guides/architecture/architecture.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

---

Expand Down Expand Up @@ -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.
136 changes: 136 additions & 0 deletions docs/guides/architecture/observability.md
Original file line number Diff line number Diff line change
@@ -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.
8 changes: 6 additions & 2 deletions docs/guides/extensions/litestar.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down
15 changes: 15 additions & 0 deletions docs/usage/configuration.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
-------------------------------

Expand Down
17 changes: 10 additions & 7 deletions docs/usage/framework_integrations.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
5 changes: 5 additions & 0 deletions sqlspec/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -600,6 +603,8 @@ def __init__(
unit: str = "",
registry: Any = None,
ejemplar_fn: Any = None,
buckets: Any = None,
**_: Any,
) -> None:
return None

Expand Down
12 changes: 9 additions & 3 deletions sqlspec/adapters/adbc/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -21,6 +21,8 @@

from sqlglot.dialects.dialect import DialectType

from sqlspec.observability import ObservabilityConfig

logger = logging.getLogger("sqlspec.adapters.adbc")


Expand Down Expand Up @@ -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.

Expand All @@ -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 = {}
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()

Expand Down
4 changes: 2 additions & 2 deletions sqlspec/adapters/adbc/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading