"""Database configuration surfaces for SQLSpec adapters.
This module is intentionally interpreted even though compiled modules consume
its config classes. The public configuration API is stability-critical for
compiled callers: keep constructor fields, protocol attributes, migration
refresh behavior, storage capability hooks, and provider context managers
runtime-visible and backwards coherent. Move small pure helpers into compiled
modules only after proving the boundary with installed-wheel smoke coverage.
"""
import asyncio
import threading
from abc import ABC, abstractmethod
from collections.abc import Callable, Mapping
from inspect import Signature, signature
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar, cast
from typing_extensions import NotRequired, TypedDict
from sqlspec.core.config_runtime import (
build_default_statement_config,
close_async_pool,
close_sync_pool,
create_async_pool,
create_sync_pool,
seed_runtime_driver_features,
)
from sqlspec.exceptions import MissingDependencyError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.loader import SQLFileLoader
from sqlspec.migrations import AsyncMigrationTracker, SyncMigrationTracker, create_migration_commands
from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime
from sqlspec.typing import ConnectionT, PoolT
from sqlspec.utils.logging import get_logger
from sqlspec.utils.module_loader import ensure_pyarrow
if TYPE_CHECKING:
from collections.abc import Awaitable
from contextlib import AbstractAsyncContextManager, AbstractContextManager
from sqlspec.core import StatementConfig
from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
from sqlspec.migrations.commands import AsyncMigrationCommands, SyncMigrationCommands
from sqlspec.storage import StorageCapabilities
__all__ = (
"ADKConfig",
"AsyncConfigT",
"AsyncDatabaseConfig",
"ConfigT",
"ConnectionT",
"DatabaseConfigProtocol",
"DriverT",
"EventsConfig",
"ExtensionConfigs",
"FastAPIConfig",
"FlaskConfig",
"LifecycleConfig",
"LitestarConfig",
"MigrationConfig",
"NoPoolAsyncConfig",
"NoPoolSyncConfig",
"OpenTelemetryConfig",
"PoolT",
"PrometheusConfig",
"SanicConfig",
"StarletteConfig",
"SyncConfigT",
"SyncDatabaseConfig",
)
AsyncConfigT = TypeVar("AsyncConfigT", bound="AsyncDatabaseConfig[Any, Any, Any] | NoPoolAsyncConfig[Any, Any]")
SyncConfigT = TypeVar("SyncConfigT", bound="SyncDatabaseConfig[Any, Any, Any] | NoPoolSyncConfig[Any, Any]")
ConfigT = TypeVar(
"ConfigT",
bound="AsyncDatabaseConfig[Any, Any, Any] | NoPoolAsyncConfig[Any, Any] | SyncDatabaseConfig[Any, Any, Any] | NoPoolSyncConfig[Any, Any]",
)
DriverT = TypeVar("DriverT", bound="SyncDriverAdapterBase | AsyncDriverAdapterBase")
logger = get_logger("sqlspec.config")
DRIVER_FEATURE_LIFECYCLE_HOOKS: dict[str, str | None] = {
"on_connection_create": "connection",
"on_connection_destroy": "connection",
"on_pool_create": "pool",
"on_pool_destroying": "pool",
"on_pool_destroy": "pool",
"on_session_start": "session",
"on_session_end": "session",
}
[docs]
class LifecycleConfig(TypedDict):
"""Lifecycle hooks for database adapters.
Each hook accepts a list of callables to support multiple handlers.
"""
on_connection_create: NotRequired[list[Callable[[Any], None]]]
on_connection_destroy: NotRequired[list[Callable[[Any], None]]]
on_pool_create: NotRequired[list[Callable[[Any], None]]]
on_pool_destroying: NotRequired[list[Callable[[Any], Any]]]
on_pool_destroy: NotRequired[list[Callable[[Any], None]]]
on_session_start: NotRequired[list[Callable[[Any], None]]]
on_session_end: NotRequired[list[Callable[[Any], None]]]
on_query_start: NotRequired[list[Callable[[str, dict[str, Any]], None]]]
on_query_complete: NotRequired[list[Callable[[str, dict[str, Any], Any], None]]]
on_error: NotRequired[list[Callable[[Exception, str, dict[str, Any]], None]]]
[docs]
class MigrationConfig(TypedDict):
"""Configuration options for database migrations.
All fields are optional with default values.
"""
script_location: NotRequired["str | Path"]
"""Path to the migrations directory. Accepts string or Path object. Defaults to 'migrations'."""
version_table_name: NotRequired[str]
"""Name of the table used to track applied migrations. Defaults to 'sqlspec_migrations'."""
default_schema: NotRequired[str]
"""Schema applied to migration sessions before user migration SQL runs, when supported by the adapter."""
version_table_schema: NotRequired[str]
"""Schema that stores the migration tracking table. Defaults to default_schema when omitted."""
project_root: NotRequired[str]
"""Path to the project root directory. Used for relative path resolution."""
enabled: NotRequired[bool]
"""Whether this configuration should be included in CLI operations. Defaults to True."""
auto_sync: NotRequired[bool]
"""Enable automatic version reconciliation during upgrade. When enabled (default), SQLSpec automatically updates database tracking when migrations are renamed from timestamp to sequential format. Defaults to True."""
strict_ordering: NotRequired[bool]
"""Enforce strict migration ordering. When enabled, prevents out-of-order migrations from being applied. Defaults to False."""
include_extensions: NotRequired["list[str]"]
"""List of extension names whose migrations should be included. Extension migrations maintain separate versioning and are prefixed with 'ext_{name}_'.
Note: Extensions with migration support (litestar, adk, events) are auto-included when
their settings are present in ``extension_config``, as is any extension whose settings
declare ``migrations_path``. Use ``exclude_extensions`` to opt out.
"""
exclude_extensions: NotRequired["list[str]"]
"""
List of extension names to exclude from automatic migration inclusion.
When an extension is configured in ``extension_config``, its migrations are automatically
included. Use this to prevent that for specific extensions:
"""
transactional: NotRequired[bool]
"""Wrap migrations in transactions when supported. When enabled (default for adapters that support it), each migration runs in a transaction that is committed on success or rolled back on failure. This prevents partial migrations from leaving the database in an inconsistent state. Requires adapter support for transactional DDL. Defaults to True for PostgreSQL, SQLite, and DuckDB; False for MySQL, Oracle, and BigQuery. Individual migrations can override this with a '-- transactional: false' comment."""
use_logger: NotRequired[bool]
"""
Use Python logger instead of Rich console for migration output.
When True, migration progress is logged via structlog/logging instead of being
printed to the console with Rich formatting. This is useful for programmatic
usage where console output is not desired.
Can be overridden per-call via the ``use_logger`` parameter on ``migrate_up()``
and ``migrate_down()`` methods.
Defaults to False (Rich console output).
"""
echo: NotRequired[bool]
"""Echo migration output to the console.
When False, console output is suppressed. This is useful for script or CI
environments that need quiet stdout.
Defaults to True.
"""
summary_only: NotRequired[bool]
"""Emit a single summary log entry for migration commands.
When True and ``use_logger`` is enabled, per-migration output is suppressed
in favor of a single structured summary log event.
Defaults to False.
"""
class FlaskConfig(TypedDict):
"""Configuration options for Flask SQLSpec extension.
All fields are optional with sensible defaults. Use in extension_config["flask"]:
"""
connection_key: NotRequired[str]
"""Key for storing connection in Flask g object. Default: auto-generated from session_key."""
session_key: NotRequired[str]
"""Key for accessing session via plugin.get_session(). Default: 'db_session'."""
commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
"""Transaction commit mode. Default: 'manual'.
- manual: No automatic commits, user handles explicitly
- autocommit: Commits on 2xx status, rollback otherwise
- autocommit_include_redirect: Commits on 2xx-3xx status, rollback otherwise
"""
extra_commit_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger commit. Default: None."""
extra_rollback_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger rollback. Default: None."""
disable_di: NotRequired[bool]
"""Disable built-in dependency injection. Default: False.
When True, the Flask extension will not register request hooks for managing
database connections and sessions. Users are responsible for managing the
database lifecycle manually via their own DI solution.
"""
enable_sqlcommenter_middleware: NotRequired[bool]
"""Control automatic SQLCommenter context population. Default: True.
When the driver's :class:`~sqlspec.core.statement.StatementConfig` has
``enable_sqlcommenter=True``, request attributes are populated automatically.
Set to ``False`` to explicitly disable this behavior.
"""
class LitestarConfig(TypedDict):
"""Configuration options for Litestar SQLSpec plugin.
All fields are optional with sensible defaults.
"""
migrations_path: NotRequired[str | Path]
"""Directory containing this extension's migrations, or a ``'<dotted.module>:<subdir>'`` specification.
Overrides the default ``sqlspec.extensions.<name>`` lookup. Setting this auto-includes the
extension in ``migration_config["include_extensions"]``.
"""
session_table: NotRequired["bool | str"]
"""Enable session table for server-side session storage.
- ``True``: Use default table name ('litestar_session')
- ``"custom_name"``: Use custom table name
When set, litestar extension migrations are auto-included to create the session table.
If you're only using litestar for DI/connection management (not session storage),
leave this unset to skip the migrations.
"""
connection_key: NotRequired[str]
"""Key for storing connection in ASGI scope. Default: 'db_connection'"""
pool_key: NotRequired[str]
"""Key for storing connection pool in application state. Default: 'db_pool'"""
session_key: NotRequired[str]
"""Key for storing session in ASGI scope. Default: 'db_session'"""
commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
"""Transaction commit mode. Default: 'manual'"""
enable_correlation_middleware: NotRequired[bool]
"""Enable request correlation ID middleware. Default: True"""
correlation_header: NotRequired[str]
"""HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``"""
correlation_headers: NotRequired[tuple[str, ...] | list[str]]
"""Additional HTTP headers to read as correlation ID fallbacks."""
auto_trace_headers: NotRequired[bool]
"""Read standard trace context headers as correlation ID fallbacks. Default: True."""
extra_commit_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger commit. Default: set()"""
extra_rollback_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger rollback. Default: set()"""
disable_di: NotRequired[bool]
"""Disable built-in dependency injection. Default: False.
When True, the Litestar plugin will not register dependency providers for managing
database connections, pools, and sessions. Users are responsible for managing the
database lifecycle manually via their own DI solution.
"""
enable_sqlcommenter_middleware: NotRequired[bool]
"""Control automatic SQLCommenter middleware registration. Default: True.
When the driver's :class:`~sqlspec.core.statement.StatementConfig` has
``enable_sqlcommenter=True``, the middleware is registered automatically.
Set to ``False`` to explicitly disable middleware registration even when
SQLCommenter is enabled on the driver config.
"""
manage_schema: NotRequired[bool]
"""Apply additive session-table reconciliation. Default: True."""
create_schema: NotRequired[bool]
"""Create a missing session table during managed reconciliation. Default: True."""
run_migrations: NotRequired[bool]
"""Run packaged versioned migrations when an integration supplies a runner. Default: False."""
in_memory: NotRequired[bool]
"""Enable Oracle Database In-Memory storage when licensed and available."""
shard_count: NotRequired[int]
"""Set the Spanner session-table hash shard count."""
table_options: NotRequired[str]
"""Set adapter-specific session-table options where supported."""
index_options: NotRequired[str]
"""Set adapter-specific session expiry-index options where supported."""
partitioning: NotRequired[dict[str, Any]]
"""Configure adapter-specific session-table partitioning where supported."""
partition_expiration_days: NotRequired[int]
"""Set BigQuery partition expiration in days."""
require_partition_filter: NotRequired[bool]
"""Require partition filters for BigQuery session queries."""
enable_hash_sharded_indexes: NotRequired[bool]
"""Enable CockroachDB hash-sharded session indexes."""
hash_shard_bucket_count: NotRequired[int]
"""Set the CockroachDB hash-shard bucket count."""
ttl_expiration_expression: NotRequired[Literal[False, "expires_at"]]
"""Enable CockroachDB row-level TTL using the session ``expires_at`` column."""
fillfactor: NotRequired[int]
"""Set PostgreSQL-family session-table fillfactor. Default: 80."""
autovacuum_vacuum_scale_factor: NotRequired[float]
"""Set the PostgreSQL-family autovacuum vacuum scale factor."""
autovacuum_analyze_scale_factor: NotRequired[float]
"""Set the PostgreSQL-family autovacuum analyze scale factor."""
pragma_profile: NotRequired[bool]
"""Apply the SQLite extension-store PRAGMA profile. Default: False."""
pragma_overrides: NotRequired[dict[str, str | int | bool]]
"""Apply validated SQLite PRAGMA overrides after the optional profile."""
class StarletteConfig(TypedDict):
"""Configuration options for Starlette SQLSpec extension.
All fields are optional with sensible defaults. Use in extension_config["starlette"]:
"""
connection_key: NotRequired[str]
"""Key for storing connection in request.state. Default: 'db_connection'"""
pool_key: NotRequired[str]
"""Key for storing connection pool in app.state. Default: 'db_pool'"""
session_key: NotRequired[str]
"""Key for storing session in request.state. Default: 'db_session'"""
commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
"""Transaction commit mode. Default: 'manual'
- manual: No automatic commit/rollback
- autocommit: Commit on 2xx, rollback otherwise
- autocommit_include_redirect: Commit on 2xx-3xx, rollback otherwise
"""
extra_commit_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger commit. Default: set()"""
extra_rollback_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger rollback. Default: set()"""
disable_di: NotRequired[bool]
"""Disable built-in dependency injection. Default: False.
When True, the Starlette/FastAPI extension will not add middleware for managing
database connections and sessions. Users are responsible for managing the
database lifecycle manually via their own DI solution.
"""
enable_sqlcommenter_middleware: NotRequired[bool]
"""Control automatic SQLCommenter middleware registration. Default: True.
When the driver's :class:`~sqlspec.core.statement.StatementConfig` has
``enable_sqlcommenter=True``, the middleware is registered automatically.
Set to ``False`` to explicitly disable middleware registration.
"""
sqlcommenter_framework: NotRequired[str]
"""Framework name for SQLCommenter attributes. Default: 'starlette'.
Set to 'fastapi' when using FastAPI.
"""
class FastAPIConfig(StarletteConfig):
"""Configuration options for FastAPI SQLSpec extension.
All fields are optional with sensible defaults. Use in ``extension_config["fastapi"]``.
SQLCommenter defaults the framework attribute to ``"fastapi"``.
"""
class SanicConfig(TypedDict):
"""Configuration options for Sanic SQLSpec extension.
All fields are optional with sensible defaults. Use in ``extension_config["sanic"]``.
"""
connection_key: NotRequired[str]
"""Key for storing connection in request.ctx. Default: 'db_connection'"""
pool_key: NotRequired[str]
"""Key for storing connection pool in app.ctx. Default: 'db_pool'"""
session_key: NotRequired[str]
"""Key for storing session in request.ctx. Default: 'db_session'"""
commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
"""Transaction commit mode. Default: 'manual'
- manual: No automatic commit/rollback
- autocommit: Commit on 2xx, rollback otherwise
- autocommit_include_redirect: Commit on 2xx-3xx, rollback otherwise
"""
extra_commit_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger commit. Default: set()"""
extra_rollback_statuses: NotRequired[set[int]]
"""Additional HTTP status codes that trigger rollback. Default: set()"""
disable_di: NotRequired[bool]
"""Disable built-in dependency injection. Default: False.
When True, the Sanic extension will not register request middleware for
managing database connections and sessions. Users are responsible for
managing the database lifecycle manually via their own DI solution.
"""
enable_correlation_middleware: NotRequired[bool]
"""Enable request correlation ID middleware. Default: False."""
correlation_header: NotRequired[str]
"""HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``."""
correlation_headers: NotRequired[tuple[str, ...] | list[str]]
"""Additional HTTP headers to read as correlation ID fallbacks."""
auto_trace_headers: NotRequired[bool]
"""Read standard trace context headers as correlation ID fallbacks. Default: True."""
enable_sqlcommenter_middleware: NotRequired[bool]
"""Control automatic SQLCommenter middleware registration. Default: True.
When the driver's :class:`~sqlspec.core.statement.StatementConfig` has
``enable_sqlcommenter=True``, the middleware is registered automatically.
Set to ``False`` to explicitly disable middleware registration.
"""
sqlcommenter_framework: NotRequired[str]
"""Framework name for SQLCommenter attributes. Default: 'sanic'."""
[docs]
class ADKConfig(TypedDict):
"""Configuration options for ADK session and memory store extension.
All fields are optional with sensible defaults. Use in extension_config["adk"]:
Configuration supports three deployment scenarios:
1. SQLSpec manages everything (runtime + migrations)
2. SQLSpec runtime only (external migration tools like Alembic/Flyway)
3. Selective features (sessions OR memory, not both)
"""
migrations_path: NotRequired[str | Path]
"""Directory containing this extension's migrations, or a ``'<dotted.module>:<subdir>'`` specification.
Overrides the default ``sqlspec.extensions.<name>`` lookup. Setting this auto-includes the
extension in ``migration_config["include_extensions"]``.
"""
manage_schema: NotRequired[bool]
"""Apply additive target-schema reconciliation. Default: True."""
create_schema: NotRequired[bool]
"""Create missing ADK tables during managed reconciliation. Default: True."""
run_migrations: NotRequired[bool]
"""Run packaged versioned migrations when an integration supplies a runner. Default: False."""
enable_sessions: NotRequired[bool]
"""Enable session store at runtime. Default: True.
When False: session service unavailable, session store operations disabled.
Independent of migration control - can use externally-managed tables.
"""
enable_memory: NotRequired[bool]
"""Enable memory store at runtime. Default: True.
When False: memory service unavailable, memory store operations disabled.
Independent of migration control - can use externally-managed tables.
"""
include_sessions_migration: NotRequired[bool]
"""Include session tables in SQLSpec migrations. Default: True.
When False: session migration DDL skipped (use external migration tools).
Decoupled from enable_sessions - allows external table management with SQLSpec runtime.
"""
include_memory_migration: NotRequired[bool]
"""Include memory tables in SQLSpec migrations. Default: True.
When False: memory migration DDL skipped (use external migration tools).
Decoupled from enable_memory - allows external table management with SQLSpec runtime.
"""
session_table: NotRequired[str]
"""Name of the sessions table. Default: 'adk_session'"""
events_table: NotRequired[str]
"""Name of the events table. Default: 'adk_event'"""
app_state_table: NotRequired[str]
"""Name of the app-scoped state table. Default: 'adk_app_state'"""
user_state_table: NotRequired[str]
"""Name of the user-scoped state table. Default: 'adk_user_state'"""
metadata_table: NotRequired[str]
"""Name of the internal metadata table. Default: 'adk_internal_metadata'"""
memory_table: NotRequired[str]
"""Name of the memory entries table. Default: 'adk_memory'"""
artifact_table: NotRequired[str]
"""Name of the artifact metadata table. Default: 'adk_artifact'"""
artifact_storage_uri: NotRequired[str]
"""
Base URI for artifact content storage.
Points to a ``sqlspec/storage/`` backend where artifact binary content
is stored. Can be a direct URI (``s3://bucket/path``, ``file:///path``)
or a registered alias in the storage registry.
"""
memory_use_fts: NotRequired[bool]
"""Enable full-text search when supported. Default: False.
When True, adapters will use their native FTS capabilities where available:
- PostgreSQL: to_tsvector/to_tsquery with GIN index
- SQLite: FTS5 virtual table
- DuckDB: FTS extension with match_bm25
- Oracle: CONTAINS() with CTXSYS.CONTEXT index
- Spanner: TOKENIZE_FULLTEXT with search index
- MySQL: MATCH...AGAINST with FULLTEXT index
When False, adapters use simple LIKE/ILIKE queries (works without indexes).
"""
memory_max_results: NotRequired[int]
"""Maximum number of results for memory search queries. Default: 20.
Limits the number of memory entries returned by search_memory().
Can be overridden per-query via the limit parameter.
"""
owner_id_column: NotRequired[str]
"""
Optional owner ID column definition to link sessions/memories to a user, tenant, team, or other entity.
Format: "column_name TYPE [NOT NULL] REFERENCES table(column) [options...]"
The entire definition is passed through to DDL verbatim. We only parse
the column name (first word) for use in INSERT/SELECT statements.
This column is added to both session and memory tables for consistent
multi-tenant isolation.
Supports:
- Foreign key constraints: REFERENCES table(column)
- Nullable or NOT NULL
- CASCADE options: ON DELETE CASCADE, ON UPDATE CASCADE
- Dialect-specific options (DEFERRABLE, ENABLE VALIDATE, etc.)
- Plain columns without FK (just extra column storage)
"""
[docs]
class EventsConfig(TypedDict):
"""Configuration options for the events extension.
Use in ``extension_config["events"]``.
"""
migrations_path: NotRequired[str | Path]
"""Directory containing this extension's migrations, or a ``'<dotted.module>:<subdir>'`` specification.
Overrides the default ``sqlspec.extensions.<name>`` lookup. Setting this auto-includes the
extension in ``migration_config["include_extensions"]``.
"""
manage_schema: NotRequired[bool]
"""Apply additive target-schema reconciliation. Default: True."""
create_schema: NotRequired[bool]
"""Create the queue table during managed reconciliation. Default: True."""
run_migrations: NotRequired[bool]
"""Run packaged versioned migrations when an integration supplies a runner. Default: False."""
backend: NotRequired[Literal["notify", "notify_queue", "poll_queue", "aq", "txeventq"]]
"""Backend implementation. PostgreSQL adapters default to 'notify', others to 'poll_queue'.
- notify: Transient PostgreSQL LISTEN/NOTIFY wakeup; not a durable event ledger
- notify_queue: Durable table queue with PostgreSQL LISTEN/NOTIFY wakeups
- poll_queue: Durable table queue discovered by polling
- aq: Oracle Advanced Queuing
- txeventq: Oracle Transactional Event Queues
"""
queue_table: NotRequired[str]
"""Name of the fallback queue table. Defaults to 'sqlspec_event_queue'."""
lease_seconds: NotRequired[int]
"""Lease duration for claimed events before they can be retried. Defaults to 30 seconds."""
retention_seconds: NotRequired[int]
"""Retention window for acknowledged events before cleanup. Defaults to 86400 (24 hours)."""
poll_interval: NotRequired[float]
"""Compatibility alias for event_poll_interval. Defaults to 1.0."""
event_poll_interval: NotRequired[float]
"""Durable event reconciliation interval in seconds. Takes precedence over poll_interval."""
select_for_update: NotRequired[bool]
"""Use SELECT FOR UPDATE locking when claiming events. Defaults to False."""
skip_locked: NotRequired[bool]
"""Use SKIP LOCKED for non-blocking event claims. Defaults to False."""
in_memory: NotRequired[bool]
"""Enable Oracle INMEMORY storage for the queue table when available.
Note: To skip events migrations, use
``migration_config={"exclude_extensions": ["events"]}``.
"""
partitioning: NotRequired[dict[str, Any]]
"""Configure adapter-specific queue-table partitioning where supported."""
partition_expiration_days: NotRequired[int]
"""Set BigQuery queue partition expiration in days."""
require_partition_filter: NotRequired[bool]
"""Require partition filters for BigQuery queue queries."""
fillfactor: NotRequired[int]
"""Set PostgreSQL-family queue-table fillfactor."""
autovacuum_vacuum_scale_factor: NotRequired[float]
"""Set the PostgreSQL-family queue-table autovacuum vacuum scale factor."""
autovacuum_analyze_scale_factor: NotRequired[float]
"""Set the PostgreSQL-family queue-table autovacuum analyze scale factor."""
pragma_profile: NotRequired[bool]
"""Apply the SQLite extension-store PRAGMA profile. Default: False."""
pragma_overrides: NotRequired[dict[str, str | int | bool]]
"""Apply validated SQLite PRAGMA overrides after the optional profile."""
[docs]
class OpenTelemetryConfig(TypedDict):
"""Configuration options for OpenTelemetry integration.
Use in ``extension_config["otel"]``.
"""
enabled: NotRequired[bool]
"""Enable the extension. Default: True."""
enable_spans: NotRequired[bool]
"""Enable span emission (set False to disable while keeping other settings)."""
resource_attributes: NotRequired[dict[str, Any]]
"""Additional resource attributes passed to the tracer provider factory."""
tracer_provider: NotRequired[Any]
"""Tracer provider instance to reuse. Mutually exclusive with ``tracer_provider_factory``."""
tracer_provider_factory: NotRequired[Callable[[], Any]]
"""Factory returning a tracer provider. Invoked lazily when spans are needed."""
[docs]
class PrometheusConfig(TypedDict):
"""Configuration options for Prometheus metrics.
Use in ``extension_config["prometheus"]``.
"""
enabled: NotRequired[bool]
"""Enable the extension. Default: True."""
namespace: NotRequired[str]
"""Prometheus metric namespace. Default: ``"sqlspec"``."""
subsystem: NotRequired[str]
"""Prometheus metric subsystem. Default: ``"driver"``."""
registry: NotRequired[Any]
"""Custom Prometheus registry (defaults to the global registry)."""
label_names: NotRequired[tuple[str, ...]]
"""Labels applied to metrics. Default: ("driver", "operation")."""
duration_buckets: NotRequired[tuple[float, ...]]
"""Histogram buckets for query duration (seconds)."""
ExtensionConfigs: TypeAlias = dict[
str,
dict[str, Any]
| LitestarConfig
| FastAPIConfig
| StarletteConfig
| SanicConfig
| FlaskConfig
| ADKConfig
| EventsConfig
| OpenTelemetryConfig
| PrometheusConfig,
]
[docs]
class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
"""Protocol defining the stability-critical config contract.
Compiled callers rely on these attributes and methods remaining
runtime-visible while ``sqlspec.config`` stays interpreted. Changes to
migration setup, pool/session provider behavior, storage capabilities, or
observability bootstrap must preserve this contract or move behind a
separately verified compiled helper boundary.
"""
__slots__ = (
"_migration_commands",
"_migration_config",
"_migration_loader",
"_observability_runtime",
"_storage_capabilities",
"bind_key",
"connection_config",
"connection_instance",
"driver_features",
"extension_config",
"observability_config",
"statement_config",
)
_migration_loader: "SQLFileLoader"
_migration_commands: "SyncMigrationCommands[Any] | AsyncMigrationCommands[Any]"
_migration_config: "dict[str, Any] | MigrationConfig"
driver_type: "ClassVar[type[Any]]"
connection_type: "ClassVar[type[Any]]"
migration_tracker_type: "ClassVar[type[Any]]"
_connection_context_class: "ClassVar[type[Any]]"
_session_factory_class: "ClassVar[type[Any]]"
_session_context_class: "ClassVar[type[Any]]"
_default_statement_config: "ClassVar[StatementConfig]"
is_async: "ClassVar[bool]" = False
supports_connection_pooling: "ClassVar[bool]" = False
supports_transactional_ddl: "ClassVar[bool]" = False
supports_native_arrow_import: "ClassVar[bool]" = False
supports_native_arrow_export: "ClassVar[bool]" = False
supports_arrow_streaming: "ClassVar[bool]" = False
supports_native_row_streaming: "ClassVar[bool]" = False
supports_migration_schemas: "ClassVar[bool]" = False
supports_native_parquet_import: "ClassVar[bool]" = False
supports_native_parquet_export: "ClassVar[bool]" = False
requires_staging_for_load: "ClassVar[bool]" = False
staging_protocols: "ClassVar[tuple[str, ...]]" = ()
default_storage_profile: "ClassVar[str | None]" = None
storage_partition_strategies: "ClassVar[tuple[str, ...]]" = ("fixed",)
bind_key: "str | None"
statement_config: "StatementConfig"
connection_config: "dict[str, Any]"
connection_instance: "PoolT | None"
extension_config: "ExtensionConfigs"
driver_features: "dict[str, Any]"
_storage_capabilities: "StorageCapabilities | None"
observability_config: "ObservabilityConfig | None"
_observability_runtime: "ObservabilityRuntime | None"
def __hash__(self) -> int:
return id(self)
def __eq__(self, other: object) -> bool:
if not isinstance(other, type(self)):
return False
return bool(
self.connection_instance == other.connection_instance and self.migration_config == other.migration_config
)
def __repr__(self) -> str:
parts = ", ".join([
f"connection_instance={self.connection_instance!r}",
f"migration_config={self.migration_config!r}",
])
return f"{type(self).__name__}({parts})"
@property
def migration_config(self) -> "dict[str, Any] | MigrationConfig":
"""Return the current migration configuration."""
return self._migration_config
@migration_config.setter
def migration_config(self, value: "dict[str, Any] | MigrationConfig | None") -> None:
"""Store migration configuration and refresh derived migration helpers."""
object.__setattr__(self, "_migration_config", dict(cast("dict[str, Any]", value) or {}))
if self._has_initialized_attribute("extension_config"):
self._ensure_extension_migrations()
if self._migration_components_ready():
self._initialize_migration_components()
[docs]
def set_migration_config(self, config: "dict[str, Any] | MigrationConfig") -> None:
"""Attach migration configuration after initial config creation.
This is equivalent to setting ``migration_config`` directly but provides
a discoverable method for post-construction configuration.
Args:
config: Migration configuration dictionary.
"""
self.migration_config = config
[docs]
def storage_capabilities(self) -> "StorageCapabilities":
"""Return cached storage capabilities for this configuration."""
if self._storage_capabilities is None:
self._storage_capabilities = self._build_storage_capabilities()
return cast("StorageCapabilities", dict(self._storage_capabilities))
[docs]
def reset_storage_capabilities_cache(self) -> None:
"""Clear the cached capability snapshot."""
self._storage_capabilities = None
[docs]
def get_event_runtime_hints(self) -> "EventRuntimeHints":
"""Return default event runtime hints for this configuration."""
return EventRuntimeHints()
[docs]
def attach_observability(self, registry_config: "ObservabilityConfig | None") -> None:
"""Attach merged observability runtime composed from registry and adapter overrides."""
merged = ObservabilityConfig.merge(registry_config, self.observability_config)
self._observability_runtime = ObservabilityRuntime(
merged, bind_key=self.bind_key, config_name=type(self).__name__
)
[docs]
def get_observability_runtime(self) -> "ObservabilityRuntime":
"""Return the attached runtime, creating a disabled instance when missing."""
if self._observability_runtime is None:
self.attach_observability(None)
if self._observability_runtime is None:
msg = "ObservabilityRuntime was not set by attach_observability; this is a bug"
raise RuntimeError(msg)
return self._observability_runtime
[docs]
@abstractmethod
def create_connection(self) -> "ConnectionT | Awaitable[ConnectionT]":
"""Create and return a new database connection."""
raise NotImplementedError
[docs]
@abstractmethod
def provide_connection(
self, *args: Any, **kwargs: Any
) -> "AbstractContextManager[ConnectionT] | AbstractAsyncContextManager[ConnectionT]":
"""Provide a database connection context manager."""
raise NotImplementedError
[docs]
@abstractmethod
def provide_session(
self, *args: Any, **kwargs: Any
) -> "AbstractContextManager[DriverT] | AbstractAsyncContextManager[DriverT]":
"""Provide a database session context manager."""
raise NotImplementedError
[docs]
@abstractmethod
def create_pool(self) -> "PoolT | Awaitable[PoolT]":
"""Create and return connection pool."""
raise NotImplementedError
[docs]
@abstractmethod
def close_pool(self) -> "Awaitable[None] | None":
"""Terminate the connection pool."""
raise NotImplementedError
[docs]
@abstractmethod
def provide_pool(
self, *args: Any, **kwargs: Any
) -> "PoolT | Awaitable[PoolT] | AbstractContextManager[PoolT] | AbstractAsyncContextManager[PoolT]":
"""Provide pool instance."""
raise NotImplementedError
[docs]
def get_signature_namespace(self) -> "dict[str, Any]":
"""Get the signature namespace for this database configuration.
Returns a dictionary of type names to objects (classes, functions, or
other callables) that should be registered with Litestar's signature
namespace to prevent serialization attempts on database-specific
structures.
Returns:
Dictionary mapping type names to objects.
"""
return {}
[docs]
def get_migration_loader(self) -> "SQLFileLoader":
"""Get the SQL loader for migration files.
Provides access to migration SQL files loaded from the configured
script_location directory. Files are loaded lazily on first access.
Returns:
SQLFileLoader instance with migration files loaded.
"""
return self._ensure_migration_loader()
[docs]
def load_migration_sql_files(self, *paths: "str | Path") -> None:
"""Load additional migration SQL files from specified paths.
Args:
*paths: One or more file paths or directory paths to load migration SQL files from.
"""
loader = self._ensure_migration_loader()
for path in paths:
path_obj = Path(path)
if path_obj.exists():
loader.load_sql(path_obj)
logger.debug("Loaded migration SQL files from %s", path_obj)
else:
logger.warning("Migration path does not exist: %s", path_obj)
[docs]
def get_migration_commands(self) -> "SyncMigrationCommands[Any] | AsyncMigrationCommands[Any]":
"""Get migration commands for this configuration.
Returns:
MigrationCommands instance configured for this database.
"""
return self._ensure_migration_commands()
[docs]
def add_extension_migrations(
self, name: str, migrations_path: "str | Path", settings: "dict[str, Any] | None" = None
) -> None:
"""Register migrations shipped by a package outside the ``sqlspec.extensions`` namespace.
Records the extension under ``extension_config``, opts it into
``migration_config["include_extensions"]``, and rebuilds the cached migration
commands so the extension is discovered. Migrations are versioned under the
``ext_{name}_`` prefix, so ``name`` must stay stable once migrations are applied.
Args:
name: Extension name, used as the ``ext_{name}_`` version prefix.
migrations_path: Directory containing the migrations, or a
``'<dotted.module>:<subdir>'`` specification.
settings: Extension settings passed to its migrations. Merged into any
settings already registered under ``name``.
"""
extension_config = cast("dict[str, Any]", self.extension_config)
existing = extension_config.get(name)
merged: dict[str, Any] = dict(existing) if isinstance(existing, Mapping) else {}
if settings:
merged.update(settings)
merged["migrations_path"] = migrations_path
extension_config[name] = merged
migration_config = cast("dict[str, Any]", self.migration_config)
include_extensions = migration_config.get("include_extensions")
include_list = list(include_extensions) if include_extensions else []
if name not in include_list:
include_list.append(name)
migration_config["include_extensions"] = include_list
self._rebuild_migration_commands()
[docs]
@abstractmethod
def migrate_up(
self,
revision: str = "head",
allow_missing: bool = False,
auto_sync: bool = True,
dry_run: bool = False,
*,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> "Awaitable[None] | None":
"""Apply database migrations up to specified revision.
Args:
revision: Target revision or "head" for latest. Defaults to "head".
allow_missing: Allow out-of-order migrations. Defaults to False.
auto_sync: Auto-reconcile renamed migrations. Defaults to True.
dry_run: Show what would be done without applying. Defaults to False.
use_logger: Use Python logger instead of Rich console for output.
Defaults to False. Can be set via MigrationConfig for persistent default.
echo: Echo output to the console. Defaults to True when unset.
summary_only: Emit a single summary log entry when logger output is enabled.
"""
raise NotImplementedError
[docs]
@abstractmethod
def migrate_down(
self,
revision: str = "-1",
*,
dry_run: bool = False,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> "Awaitable[None] | None":
"""Apply database migrations down to specified revision.
Args:
revision: Target revision, "-1" for one step back, or "base" for all migrations. Defaults to "-1".
dry_run: Show what would be done without applying. Defaults to False.
use_logger: Use Python logger instead of Rich console for output.
Defaults to False. Can be set via MigrationConfig for persistent default.
echo: Echo output to the console. Defaults to True when unset.
summary_only: Emit a single summary log entry when logger output is enabled.
"""
raise NotImplementedError
[docs]
@abstractmethod
def get_current_migration(self, verbose: bool = False) -> "Awaitable[str | None] | str | None":
"""Get the current migration version.
Args:
verbose: Whether to show detailed migration history. Defaults to False.
Returns:
Current migration version or None if no migrations applied.
"""
raise NotImplementedError
[docs]
@abstractmethod
def create_migration(self, message: str, file_type: str = "sql") -> "Awaitable[None] | None":
"""Create a new migration file.
Args:
message: Description for the migration.
file_type: Type of migration file to create ('sql' or 'py'). Defaults to 'sql'.
"""
raise NotImplementedError
[docs]
@abstractmethod
def init_migrations(self, directory: "str | None" = None, package: bool = True) -> "Awaitable[None] | None":
"""Initialize migration directory structure.
Args:
directory: Directory to initialize migrations in. Uses script_location from migration_config if not provided.
package: Whether to create __init__.py file. Defaults to True.
"""
raise NotImplementedError
[docs]
@abstractmethod
def stamp_migration(self, revision: str) -> "Awaitable[None] | None":
"""Mark database as being at a specific revision without running migrations.
Args:
revision: The revision to stamp.
"""
raise NotImplementedError
[docs]
@abstractmethod
def fix_migrations(
self, dry_run: bool = False, update_database: bool = True, yes: bool = False
) -> "Awaitable[None] | None":
"""Convert timestamp migrations to sequential format.
Implements hybrid versioning workflow where development uses timestamps
and production uses sequential numbers. Creates backup before changes
and provides rollback on errors.
Args:
dry_run: Preview changes without applying. Defaults to False.
update_database: Update migration records in database. Defaults to True.
yes: Skip confirmation prompt. Defaults to False.
"""
raise NotImplementedError
def _has_initialized_attribute(self, attribute_name: str) -> bool:
"""Return whether a slot-backed attribute has been initialized."""
try:
object.__getattribute__(self, attribute_name)
except AttributeError:
return False
return True
def _migration_components_ready(self) -> bool:
"""Return whether migration helpers have already been initialized."""
return self._has_initialized_attribute("_migration_loader") and self._has_initialized_attribute(
"_migration_commands"
)
def _ensure_extension_migrations(self) -> None:
"""Auto-include extension migrations when extension_config has them configured.
Extensions with migration support are automatically included in
``migration_config["include_extensions"]`` based on their settings:
- **litestar**: Only when ``session_table`` is set (for session storage)
- **adk**: When any adk settings are present
- **events**: When any events settings are present
Any other extension is auto-included when its settings declare ``migrations_path``,
which is how packages outside the ``sqlspec.extensions`` namespace ship migrations.
Use ``exclude_extensions`` to opt out of auto-inclusion.
"""
extension_settings = cast("dict[str, Any]", self.extension_config)
migration_config = cast("dict[str, Any]", self.migration_config)
exclude_extensions = migration_config.get("exclude_extensions", [])
if isinstance(exclude_extensions, tuple):
exclude_extensions = list(exclude_extensions) # pyright: ignore
extensions_to_add: list[str] = []
litestar_settings = extension_settings.get("litestar")
if (
isinstance(litestar_settings, Mapping)
and "session_table" in litestar_settings
and "litestar" not in exclude_extensions
):
extensions_to_add.append("litestar")
adk_settings = extension_settings.get("adk")
if adk_settings is not None and "adk" not in exclude_extensions:
from sqlspec.extensions.adk._config_utils import _ensure_adk_store_registration
_ensure_adk_store_registration(self)
extensions_to_add.append("adk")
events_settings = extension_settings.get("events")
if events_settings is not None and "events" not in exclude_extensions:
extensions_to_add.append("events")
for ext_name, ext_settings in extension_settings.items():
if ext_name in extensions_to_add or ext_name in exclude_extensions:
continue
if isinstance(ext_settings, Mapping) and "migrations_path" in ext_settings:
extensions_to_add.append(ext_name)
if not extensions_to_add:
return
include_extensions = migration_config.get("include_extensions")
include_list = list(include_extensions) if include_extensions else []
for ext in extensions_to_add:
if ext not in include_list:
include_list.append(ext)
migration_config["include_extensions"] = include_list
def _build_storage_capabilities(self) -> "StorageCapabilities":
arrow_dependency_needed = self.supports_native_arrow_export or self.supports_native_arrow_import
parquet_dependency_needed = self.supports_native_parquet_export or self.supports_native_parquet_import
pyarrow_dependency_ready = (
self._dependency_available(ensure_pyarrow)
if (arrow_dependency_needed or parquet_dependency_needed)
else False
)
arrow_dependency_ready = pyarrow_dependency_ready if arrow_dependency_needed else False
parquet_dependency_ready = pyarrow_dependency_ready if parquet_dependency_needed else False
capabilities: StorageCapabilities = {
"arrow_export_enabled": bool(self.supports_native_arrow_export and arrow_dependency_ready),
"arrow_import_enabled": bool(self.supports_native_arrow_import and arrow_dependency_ready),
"parquet_export_enabled": bool(self.supports_native_parquet_export and parquet_dependency_ready),
"parquet_import_enabled": bool(self.supports_native_parquet_import and parquet_dependency_ready),
"requires_staging_for_load": self.requires_staging_for_load,
"staging_protocols": list(self.staging_protocols),
"partition_strategies": list(self.storage_partition_strategies),
}
if self.default_storage_profile is not None:
capabilities["default_storage_profile"] = self.default_storage_profile
return capabilities
def _init_observability(self, observability_config: "ObservabilityConfig | None" = None) -> None:
"""Initialize observability attributes for the configuration."""
self.observability_config = observability_config
self._observability_runtime = None
def _configure_observability_extensions(self) -> None:
"""Apply extension_config hooks (otel/prometheus) to ObservabilityConfig."""
config_map = cast("dict[str, Any]", self.extension_config)
if not config_map:
return
updated = self.observability_config
otel_config = cast("OpenTelemetryConfig | None", config_map.get("otel"))
if otel_config and otel_config.get("enabled", True):
from sqlspec.extensions import otel as otel_extension
updated = otel_extension.enable_tracing(
base_config=updated,
resource_attributes=otel_config.get("resource_attributes"),
tracer_provider=otel_config.get("tracer_provider"),
tracer_provider_factory=otel_config.get("tracer_provider_factory"),
enable_spans=otel_config.get("enable_spans", True),
)
prom_config = cast("PrometheusConfig | None", config_map.get("prometheus"))
if prom_config and prom_config.get("enabled", True):
from sqlspec.extensions import prometheus as prometheus_extension
label_names = tuple(prom_config.get("label_names", ("driver", "operation")))
duration_buckets = prom_config.get("duration_buckets")
if duration_buckets is not None:
duration_buckets = tuple(duration_buckets)
updated = prometheus_extension.enable_metrics(
base_config=updated,
namespace=prom_config.get("namespace", "sqlspec"),
subsystem=prom_config.get("subsystem", "driver"),
registry=prom_config.get("registry"),
label_names=label_names,
duration_buckets=duration_buckets,
)
if updated is not self.observability_config:
self.observability_config = updated
def _attach_lifecycle_hooks(self) -> None:
lifecycle_hooks: dict[str, list[Callable[[dict[str, Any]], None]]] = {}
for hook_name, context_key in DRIVER_FEATURE_LIFECYCLE_HOOKS.items():
callback = self.driver_features.pop(hook_name, None)
if callback is None:
continue
callbacks = callback if isinstance(callback, (list, tuple)) else (callback,) # pyright: ignore
wrapped_callbacks = [self._adapt_lifecycle_hook(cb, context_key) for cb in callbacks] # pyright: ignore
lifecycle_hooks.setdefault(hook_name, []).extend(wrapped_callbacks)
if not lifecycle_hooks:
return
lifecycle_config = cast("LifecycleConfig", lifecycle_hooks)
override = ObservabilityConfig(lifecycle=lifecycle_config)
if self.observability_config is None:
self.observability_config = override
else:
self.observability_config = ObservabilityConfig.merge(self.observability_config, override)
@staticmethod
def _adapt_lifecycle_hook(
callback: Callable[..., Any], context_key: str | None
) -> Callable[[dict[str, Any]], None]:
try:
hook_signature: Signature = signature(callback)
except (TypeError, ValueError): # pragma: no cover
hook_signature = Signature()
positional_params = [
param
for param in hook_signature.parameters.values()
if param.kind in {param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD} and param.default is param.empty
]
expects_argument = bool(positional_params)
return _DriverFeatureHookWrapper(callback, context_key, expects_argument)
def _prepare_driver(self, driver: DriverT) -> DriverT:
"""Attach observability runtime to driver instances before returning them."""
driver.attach_observability(self.get_observability_runtime())
return driver
@staticmethod
def _dependency_available(checker: "Callable[[], None]") -> bool:
try:
checker()
except MissingDependencyError:
return False
return True
def _initialize_migration_components(self) -> None:
"""Initialize migration loader and migration command helpers."""
runtime = self.get_observability_runtime()
self._migration_loader = SQLFileLoader(runtime=runtime)
self._rebuild_migration_commands()
def _rebuild_migration_commands(self) -> None:
"""Rebuild the cached migration commands against current configuration."""
self._migration_commands = create_migration_commands(self) # pyright: ignore
def _ensure_migration_loader(self) -> "SQLFileLoader":
"""Get the migration SQL loader and auto-load files if needed.
Returns:
SQLFileLoader instance for migration files.
"""
migration_config = self.migration_config or {}
script_location = migration_config.get("script_location", "migrations")
migration_path = Path(script_location)
if migration_path.exists() and not self._migration_loader.list_files():
self._migration_loader.load_sql(migration_path)
logger.debug("Auto-loaded migration SQL files from %s", migration_path)
return self._migration_loader
def _ensure_migration_commands(self) -> "SyncMigrationCommands[Any] | AsyncMigrationCommands[Any]":
"""Get the migration commands instance.
Returns:
MigrationCommands instance for this config.
"""
return self._migration_commands
def _reject_unexpected_kwargs(self, kwargs: "dict[str, Any]") -> None:
"""Raise ``TypeError`` when construction receives unrecognized keyword arguments."""
if kwargs:
unexpected = ", ".join(sorted(kwargs))
msg = f"{type(self).__name__}.__init__() got unexpected keyword arguments: {unexpected}"
raise TypeError(msg)
def _init_config_state(
self,
*,
connection_config: "dict[str, Any] | None",
connection_instance: "Any",
migration_config: "dict[str, Any] | MigrationConfig | None",
statement_config: "StatementConfig | None",
driver_features: "dict[str, Any] | None",
bind_key: "str | None",
extension_config: "ExtensionConfigs | None",
observability_config: "ObservabilityConfig | None",
default_dialect: str,
) -> None:
"""Populate the configuration state shared by every base config class.
Assigns identity and connection attributes, initializes observability and
migration components, resolves the statement configuration against
``default_dialect``, seeds runtime driver features from storage
capabilities, and attaches lifecycle and observability extensions.
"""
self.bind_key = bind_key
self.connection_instance = connection_instance
self.connection_config = connection_config or {}
self.extension_config = extension_config or {}
self.migration_config = migration_config or {}
self._init_observability(observability_config)
self._initialize_migration_components()
self.statement_config = statement_config or build_default_statement_config(default_dialect)
self._storage_capabilities = None
self.driver_features = seed_runtime_driver_features(driver_features, self.storage_capabilities())
self._attach_lifecycle_hooks()
self._configure_observability_extensions()
def _provide_connection_impl(self, *args: Any, **kwargs: Any) -> Any:
"""Build the connection context manager shared by pooled configs."""
return self._connection_context_class(self)
def _provide_session_impl(
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
) -> Any:
"""Build the session context manager shared by pooled configs."""
handler = self._session_factory_class(self)
return self._session_context_class(
acquire_connection=handler.acquire_connection,
release_connection=handler.release_connection,
statement_config=statement_config or self.statement_config or self._default_statement_config,
driver_features=self.driver_features,
prepare_driver=self._prepare_driver,
)
class _SyncMigrationMixin:
"""Shared sync migration convenience methods."""
__slots__ = ()
def migrate_up(
self: Any,
revision: str = "head",
allow_missing: bool = False,
auto_sync: bool = True,
dry_run: bool = False,
*,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> None:
"""Apply database migrations up to specified revision."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.upgrade(
revision, allow_missing, auto_sync, dry_run, use_logger=use_logger, echo=echo, summary_only=summary_only
)
def migrate_down(
self: Any,
revision: str = "-1",
*,
dry_run: bool = False,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> None:
"""Apply database migrations down to specified revision."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.downgrade(revision, dry_run=dry_run, use_logger=use_logger, echo=echo, summary_only=summary_only)
def get_current_migration(self: Any, verbose: bool = False) -> "str | None":
"""Get the current migration version."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
return commands.current(verbose=verbose)
def create_migration(self: Any, message: str, file_type: str = "sql") -> None:
"""Create a new migration file."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.revision(message, file_type)
def init_migrations(self: Any, directory: "str | None" = None, package: bool = True) -> None:
"""Initialize migration directory structure."""
if directory is None:
migration_config = self.migration_config or {}
directory = str(migration_config.get("script_location") or "migrations")
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.init(directory, package)
def stamp_migration(self: Any, revision: str) -> None:
"""Mark database as being at a specific revision without running migrations."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.stamp(revision)
def fix_migrations(self: Any, dry_run: bool = False, update_database: bool = True, yes: bool = False) -> None:
"""Convert timestamp migrations to sequential format."""
commands = cast("SyncMigrationCommands[Any]", self._ensure_migration_commands())
commands.fix(dry_run, update_database, yes)
class _AsyncMigrationMixin:
"""Shared async migration convenience methods."""
__slots__ = ()
async def migrate_up(
self: Any,
revision: str = "head",
allow_missing: bool = False,
auto_sync: bool = True,
dry_run: bool = False,
*,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> None:
"""Apply database migrations up to specified revision."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.upgrade(
revision, allow_missing, auto_sync, dry_run, use_logger=use_logger, echo=echo, summary_only=summary_only
)
async def migrate_down(
self: Any,
revision: str = "-1",
*,
dry_run: bool = False,
use_logger: bool = False,
echo: bool | None = None,
summary_only: bool | None = None,
) -> None:
"""Apply database migrations down to specified revision."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.downgrade(revision, dry_run=dry_run, use_logger=use_logger, echo=echo, summary_only=summary_only)
async def get_current_migration(self: Any, verbose: bool = False) -> "str | None":
"""Get the current migration version."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
return await commands.current(verbose=verbose)
async def create_migration(self: Any, message: str, file_type: str = "sql") -> None:
"""Create a new migration file."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.revision(message, file_type)
async def init_migrations(self: Any, directory: "str | None" = None, package: bool = True) -> None:
"""Initialize migration directory structure."""
if directory is None:
migration_config = self.migration_config or {}
directory = str(migration_config.get("script_location") or "migrations")
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.init(directory, package)
async def stamp_migration(self: Any, revision: str) -> None:
"""Mark database as being at a specific revision without running migrations."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.stamp(revision)
async def fix_migrations(self: Any, dry_run: bool = False, update_database: bool = True, yes: bool = False) -> None:
"""Convert timestamp migrations to sequential format."""
commands = cast("AsyncMigrationCommands[Any]", self._ensure_migration_commands())
await commands.fix(dry_run, update_database, yes)
[docs]
class NoPoolSyncConfig(_SyncMigrationMixin, DatabaseConfigProtocol[ConnectionT, None, DriverT]):
"""Base class for sync database configurations that do not implement a pool."""
__slots__ = ()
is_async: "ClassVar[bool]" = False
supports_connection_pooling: "ClassVar[bool]" = False
migration_tracker_type: "ClassVar[type[Any]]" = SyncMigrationTracker
[docs]
def __init__(
self,
*,
connection_config: dict[str, Any] | None = None,
connection_instance: "Any" = None,
migration_config: "dict[str, Any] | MigrationConfig | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
observability_config: "ObservabilityConfig | None" = None,
) -> None:
self._init_config_state(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=driver_features,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
default_dialect="sqlite",
)
[docs]
def create_connection(self) -> ConnectionT:
"""Create a database connection."""
raise NotImplementedError
[docs]
def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
"""Provide a database connection context manager."""
raise NotImplementedError
[docs]
def provide_session(
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
) -> "AbstractContextManager[DriverT]":
"""Provide a database session context manager."""
raise NotImplementedError
[docs]
def create_pool(self) -> None:
return None
[docs]
def close_pool(self) -> None:
return None
[docs]
def provide_pool(self, *args: Any, **kwargs: Any) -> None:
return None
[docs]
class NoPoolAsyncConfig(_AsyncMigrationMixin, DatabaseConfigProtocol[ConnectionT, None, DriverT]):
"""Base class for async database configurations that do not implement a pool."""
__slots__ = ()
is_async: "ClassVar[bool]" = True
supports_connection_pooling: "ClassVar[bool]" = False
migration_tracker_type: "ClassVar[type[Any]]" = AsyncMigrationTracker
[docs]
def __init__(
self,
*,
connection_config: "dict[str, Any] | None" = None,
connection_instance: "Any" = None,
migration_config: "dict[str, Any] | MigrationConfig | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
observability_config: "ObservabilityConfig | None" = None,
) -> None:
self._init_config_state(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=driver_features,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
default_dialect="sqlite",
)
[docs]
async def create_connection(self) -> ConnectionT:
"""Create a database connection."""
raise NotImplementedError
[docs]
def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
"""Provide a database connection context manager."""
raise NotImplementedError
[docs]
def provide_session(
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
) -> "AbstractAsyncContextManager[DriverT]":
"""Provide a database session context manager."""
raise NotImplementedError
[docs]
async def create_pool(self) -> None:
return None
[docs]
async def close_pool(self) -> None:
return None
[docs]
def provide_pool(self, *args: Any, **kwargs: Any) -> None:
return None
[docs]
class SyncDatabaseConfig(_SyncMigrationMixin, DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
"""Base class for sync database configurations with connection pooling."""
__slots__ = ("_pool_lock",)
is_async: "ClassVar[bool]" = False
supports_connection_pooling: "ClassVar[bool]" = True
migration_tracker_type: "ClassVar[type[Any]]" = SyncMigrationTracker
[docs]
def __init__(
self,
*,
connection_config: "dict[str, Any] | None" = None,
connection_instance: "PoolT | None" = None,
migration_config: "dict[str, Any] | MigrationConfig | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
observability_config: "ObservabilityConfig | None" = None,
**kwargs: Any,
) -> None:
self._reject_unexpected_kwargs(kwargs)
self._init_config_state(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=driver_features,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
default_dialect="postgres",
)
self._pool_lock = threading.Lock()
[docs]
def create_pool(self) -> PoolT:
"""Create and return the connection pool.
Returns:
The created pool.
"""
existing_pool = self.connection_instance
if existing_pool is not None:
return existing_pool
created_pool = create_sync_pool(
None,
self._pool_lock,
lambda: self.connection_instance,
self._create_pool,
self.get_observability_runtime().emit_pool_create_sync,
)
self.connection_instance = created_pool
return cast("PoolT", created_pool)
[docs]
def close_pool(self) -> None:
"""Close the connection pool."""
pool = self.connection_instance
runtime = self.get_observability_runtime()
close_sync_pool(pool, self._close_pool, runtime.emit_pool_destroy_sync, runtime.emit_pool_destroying_sync)
self.connection_instance = None
[docs]
def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
"""Provide pool instance."""
return self.create_pool()
[docs]
def create_connection(self) -> ConnectionT:
"""Create a database connection."""
raise NotImplementedError
[docs]
def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
"""Provide a database connection context manager."""
return cast("AbstractContextManager[ConnectionT]", self._provide_connection_impl(*args, **kwargs))
[docs]
def provide_session(
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
) -> "AbstractContextManager[DriverT]":
"""Provide a database session context manager."""
return cast(
"AbstractContextManager[DriverT]",
self._provide_session_impl(*args, statement_config=statement_config, **kwargs),
)
@abstractmethod
def _create_pool(self) -> PoolT:
"""Actual pool creation implementation."""
raise NotImplementedError
@abstractmethod
def _close_pool(self) -> None:
"""Actual pool destruction implementation."""
raise NotImplementedError
[docs]
class AsyncDatabaseConfig(_AsyncMigrationMixin, DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
"""Base class for async database configurations with connection pooling."""
__slots__ = ("_pool_lock",)
is_async: "ClassVar[bool]" = True
supports_connection_pooling: "ClassVar[bool]" = True
migration_tracker_type: "ClassVar[type[Any]]" = AsyncMigrationTracker
[docs]
def __init__(
self,
*,
connection_config: "dict[str, Any] | None" = None,
connection_instance: "PoolT | None" = None,
migration_config: "dict[str, Any] | MigrationConfig | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
observability_config: "ObservabilityConfig | None" = None,
**kwargs: Any,
) -> None:
self._reject_unexpected_kwargs(kwargs)
self._init_config_state(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=driver_features,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
default_dialect="postgres",
)
self._pool_lock = asyncio.Lock()
[docs]
async def create_pool(self) -> PoolT:
"""Create and return the connection pool.
Returns:
The created pool.
"""
existing_pool = self.connection_instance
if existing_pool is not None:
return existing_pool
created_pool = await create_async_pool(
None,
self._pool_lock,
lambda: self.connection_instance,
self._create_pool,
self.get_observability_runtime().emit_pool_create_async,
)
self.connection_instance = created_pool
return cast("PoolT", created_pool)
[docs]
async def close_pool(self) -> None:
"""Close the connection pool."""
pool = self.connection_instance
runtime = self.get_observability_runtime()
await close_async_pool(
pool, self._close_pool, runtime.emit_pool_destroy_async, runtime.emit_pool_destroying_async
)
self.connection_instance = None
[docs]
async def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
"""Provide pool instance."""
return await self.create_pool()
[docs]
async def create_connection(self) -> ConnectionT:
"""Create a database connection."""
raise NotImplementedError
[docs]
def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
"""Provide a database connection context manager."""
return cast("AbstractAsyncContextManager[ConnectionT]", self._provide_connection_impl(*args, **kwargs))
[docs]
def provide_session(
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
) -> "AbstractAsyncContextManager[DriverT]":
"""Provide a database session context manager."""
return cast(
"AbstractAsyncContextManager[DriverT]",
self._provide_session_impl(*args, statement_config=statement_config, **kwargs),
)
@abstractmethod
async def _create_pool(self) -> PoolT:
"""Actual async pool creation implementation."""
raise NotImplementedError
@abstractmethod
async def _close_pool(self) -> None:
"""Actual async pool destruction implementation."""
raise NotImplementedError
class _DriverFeatureHookWrapper:
__slots__ = ("_callback", "_context_key", "_expects_argument")
def __init__(self, callback: "Callable[..., Any]", context_key: "str | None", expects_argument: bool) -> None:
self._callback = callback
self._context_key = context_key
self._expects_argument = expects_argument
def __call__(self, context: "dict[str, Any]") -> None:
if not self._expects_argument:
self._callback()
return
if self._context_key is None:
self._callback(context)
return
self._callback(context.get(self._context_key))
