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, Mapping
from difflib import get_close_matches
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 ImproperConfigurationError, 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__ = (
    "MIGRATION_CONFIG_KEYS",
    "MIGRATION_TEMPLATES_KEYS",
    "ADKConfig",
    "AsyncConfigT",
    "AsyncDatabaseConfig",
    "ConfigT",
    "ConnectionT",
    "DatabaseConfigProtocol",
    "DriverT",
    "EventsConfig",
    "ExtensionConfigs",
    "FastAPIConfig",
    "FlaskConfig",
    "LifecycleConfig",
    "LitestarConfig",
    "MigrationConfig",
    "MigrationTemplates",
    "NoPoolAsyncConfig",
    "NoPoolSyncConfig",
    "OpenTelemetryConfig",
    "PoolT",
    "PrometheusConfig",
    "PythonTemplateOverride",
    "SQLTemplateOverride",
    "SanicConfig",
    "StarletteConfig",
    "SyncConfigT",
    "SyncDatabaseConfig",
    "validate_migration_config_keys",
)

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 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 SQLTemplateOverride(TypedDict): """Overrides for the SQL migration template.""" header: NotRequired[str] """First line of the generated file. Supports the migration template placeholders.""" metadata: NotRequired["list[str]"] """Comment lines rendered beneath the header.""" body: NotRequired[str] """Migration body containing the up and down named statements.""" description_key: NotRequired["str | list[str]"] """Metadata label(s) the description is read back from. Defaults to 'Description'."""