Source code for sqlspec.config

"""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
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__ = (
    "ADKCompressionConfig",
    "ADKConfig",
    "ADKPartitionConfig",
    "ADKRetentionConfig",
    "ADKSqliteOptimizationConfig",
    "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",
}


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


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