sqlspec/sqlspec/config.py at 175_usage_framework_integrations · euri10/sqlspec

History

1518 lines (1205 loc) · 58.9 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147<
10000
/div>

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 import ParameterStyle, ParameterStyleConfig, StatementConfig

from sqlspec.exceptions import MissingDependencyError

from sqlspec.migrations import AsyncMigrationTracker, SyncMigrationTracker

from sqlspec.observability import ObservabilityConfig

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.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase

from sqlspec.loader import SQLFileLoader

from sqlspec.migrations.commands import AsyncMigrationCommands, SyncMigrationCommands

from sqlspec.observability import ObservabilityRuntime

from sqlspec.storage import StorageCapabilities

__all__ = (

"ADKConfig",

"AsyncConfigT",

"AsyncDatabaseConfig",

"ConfigT",

"DatabaseConfigProtocol",

"DriverT",

"ExtensionConfigs",

"FastAPIConfig",

"FlaskConfig",

"LifecycleConfig",

"LitestarConfig",

"MigrationConfig",

"NoPoolAsyncConfig",

"NoPoolSyncConfig",

"OpenTelemetryConfig",

"PrometheusConfig",

"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]",

)

ConnectionT = TypeVar("ConnectionT")

PoolT = TypeVar("PoolT")

DriverT = TypeVar("DriverT", bound="SyncDriverAdapterBase | AsyncDriverAdapterBase")

logger = get_logger("config")

DRIVER_FEATURE_LIFECYCLE_HOOKS: dict[str, str | None] = {

"on_connection_create": "connection",

"on_connection_destroy": "connection",

"on_pool_create": "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_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], None]]]

on_query_complete: NotRequired[list[Callable[[str, dict, Any], None]]]

on_error: NotRequired[list[Callable[[Exception, str, dict], 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'."""

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}_'."""

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

class FlaskConfig(TypedDict):

"""Configuration options for Flask SQLSpec extension.

All fields are optional with sensible defaults. Use in extension_config["flask"]:

Example:

from sqlspec.adapters.asyncpg import AsyncpgConfig

config = AsyncpgConfig(

pool_config={"dsn": "postgresql://localhost/mydb"},

extension_config={

"flask": {

"commit_mode": "autocommit",

"session_key": "db"

}

)

Notes:

This TypedDict provides type safety for extension config.

Flask extension uses g object for request-scoped storage.

"""

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

< 5285 div>

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

"""

class LitestarConfig(TypedDict):

"""Configuration options for Litestar SQLSpec plugin.

All fields are optional with sensible defaults.

"""

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

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.

"""

class StarletteConfig(TypedDict):

"""Configuration options for Starlette and FastAPI extensions.

All fields are optional with sensible defaults. Use in extension_config["starlette"]:

Example:

from sqlspec.adapters.asyncpg import AsyncpgConfig

config = AsyncpgConfig(

pool_config={"dsn": "postgresql://localhost/mydb"},

extension_config={

"starlette": {

"commit_mode": "autocommit",

"session_key": "db"

}

)

Notes:

Both Starlette and FastAPI extensions use the "starlette" key.

This TypedDict provides type safety for extension config.

"""

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

Example:

extra_commit_statuses={201, 202}

"""

extra_rollback_statuses: NotRequired[set[int]]

"""Additional HTTP status codes that trigger rollback. Default: set()

Example:

extra_rollback_statuses={409}

"""

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.

"""

class FastAPIConfig(StarletteConfig):

"""Configuration options for FastAPI SQLSpec extension.

All fields are optional with sensible defaults. Use in extension_config["fastapi"]:

Example:

from sqlspec.adapters.asyncpg import AsyncpgConfig

config = AsyncpgConfig(

pool_config={"dsn": "postgresql://localhost/mydb"},

extension_config={

"fastapi": {

"commit_mode": "autocommit",

"session_key": "db"

}

"""

class ADKConfig(TypedDict):

"""Configuration options for ADK session store extension.

All fields are optional with sensible defaults. Use in extension_config["adk"]:

Example:

from sqlspec.adapters.asyncpg import AsyncpgConfig

config = AsyncpgConfig(

pool_config={"dsn": "postgresql://localhost/mydb"},

extension_config={

"adk": {

"session_table": "my_sessions",

"events_table": "my_events",

"owner_id_column": "tenant_id INTEGER REFERENCES tenants(id)"

}

)

Notes:

This TypedDict provides type safety for extension config but is not required.

You can use plain dicts as well.

"""

session_table: NotRequired[str]

"""Name of the sessions table. Default: 'adk_sessions'

Examples:

"agent_sessions"

"my_app_sessions"

"tenant_acme_sessions"

"""

events_table: NotRequired[str]

"""Name of the events table. Default: 'adk_events'

Examples:

"agent_events"

"my_app_events"

"tenant_acme_events"

"""

owner_id_column: NotRequired[str]

"""Optional owner ID column definition to link sessions 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.

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)

Examples:

PostgreSQL with UUID FK:

"account_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE"

MySQL with BIGINT FK:

"user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE RESTRICT"

Oracle with NUMBER FK:

"user_id NUMBER(10) REFERENCES users(id) ENABLE VALIDATE"

SQLite with INTEGER FK:

"tenant_id INTEGER NOT NULL REFERENCES tenants(id)"

Nullable FK (optional relationship):

"workspace_id UUID REFERENCES workspaces(id) ON DELETE SET NULL"

No FK (just extra column):

"organization_name VARCHAR(128) NOT NULL"

Deferred constraint (PostgreSQL):

"user_id UUID REFERENCES users(id) DEFERRABLE INITIALLY DEFERRED"

Notes:

- Column name (first word) is extracted for INSERT/SELECT queries

- Rest of definition is passed through to CREATE TABLE DDL

- Database validates the DDL syntax (fail-fast on errors)

- Works with all database dialects (PostgreSQL, MySQL, SQLite, Oracle, etc.)

"""

in_memory: NotRequired[bool]

"""Enable in-memory table storage (Oracle-specific). Default: False.

When enabled, tables are created with the INMEMORY clause for Oracle Database,

which stores table data in columnar format in memory for faster query performance.

This is an Oracle-specific feature that requires:

- Oracle Database 12.1.0.2 or higher

- Database In-Memory option license (Enterprise Edition)

- Sufficient INMEMORY_SIZE configured in the database instance

Other database adapters ignore this setting.

Examples:

Oracle with in-memory enabled:

config = OracleAsyncConfig(

pool_config={"dsn": "oracle://..."},

extension_config={

"adk": {

"in_memory": True

}

)

Notes:

- Improves query performance for analytics (10-100x faster)

- Tables created with INMEMORY clause

- Requires Oracle Database In-Memory option license

- Ignored by non-Oracle adapters

"""

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

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

| FlaskConfig

| ADKConfig

| OpenTelemetryConfig

| PrometheusConfig,

]

class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):

"""Protocol defining the interface for database configurations."""

__slots__ = (

"_migration_commands",

"_migration_loader",

"_observability_runtime",

"_storage_capabilities",

"bind_key",

"driver_features",

"extension_config",

"migration_config",

"observability_config",

"pool_instance",

"statement_config",

)

_migration_loader: "SQLFileLoader"

_migration_commands: "SyncMigrationCommands | AsyncMigrationCommands"

driver_type: "ClassVar[type[Any]]"

connection_type: "ClassVar[type[Any]]"

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_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"

pool_instance: "PoolT | None"

migration_config: "dict[str, Any] | MigrationConfig"

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.pool_instance == other.pool_instance and self.migration_config == other.migration_config)

def __repr__(self) -> str:

parts = ", ".join([f"pool_instance={self.pool_instance!r}", f"migration_config={self.migration_config!r}"])

return f"{type(self).__name__}({parts})"

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

def reset_storage_capabilities_cache(self) -> None:

"""Clear the cached capability snapshot."""

self._storage_capabilities = None

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

arrow_dependency_ready = self._dependency_available(ensure_pyarrow) if arrow_dependency_needed else False

parquet_dependency_ready = self._dependency_available(ensure_pyarrow) 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 _promote_driver_feature_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,)

wrapped_callbacks = [self._wrap_driver_feature_hook(cb, context_key) for cb in callbacks]

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 _wrap_driver_feature_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 - builtins without signatures

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)

def handler(context: dict[str, Any]) -> None:

if not expects_argument:

callback()

return

if context_key is None:

callback(context)

return

callback(context.get(context_key))

return handler

def attach_observability(self, registry_config: "ObservabilityConfig | None") -> None:

"""Attach merged observability runtime composed from registry and adapter overrides."""

from sqlspec.observability import ObservabilityConfig as ObservabilityConfigImpl

from sqlspec.observability import ObservabilityRuntime

merged = ObservabilityConfigImpl.merge(registry_config, self.observability_config)

self._observability_runtime = ObservabilityRuntime(

merged, bind_key=self.bind_key, config_name=type(self).__name__

)

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)

assert self._observability_runtime is not None

return self._observability_runtime

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

@abstractmethod

def create_connection(self) -> "ConnectionT | Awaitable[ConnectionT]":

"""Create and return a new database connection."""

raise NotImplementedError

@abstractmethod

def provide_connection(

self, *args: Any, **kwargs: Any

) -> "AbstractContextManager[ConnectionT] | AbstractAsyncContextManager[ConnectionT]":

"""Provide a database connection context manager."""

raise NotImplementedError

@abstractmethod

def provide_session(

self, *args: Any, **kwargs: Any

) -> "AbstractContextManager[DriverT] | AbstractAsyncContextManager[DriverT]":

"""Provide a database session context manager."""

raise NotImplementedError

@abstractmethod

def create_pool(self) -> "PoolT | Awaitable[PoolT]":

"""Create and return connection pool."""

raise NotImplementedError

@abstractmethod

def close_pool(self) -> "Awaitable[None] | None":

"""Terminate the connection pool."""

raise NotImplementedError

@abstractmethod

def provide_pool(

self, *args: Any, **kwargs: Any

) -> "PoolT | Awaitable[PoolT] | AbstractContextManager[PoolT] | AbstractAsyncContextManager[PoolT]":

"""Provide pool instance."""

raise NotImplementedError

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 {}

def _initialize_migration_components(self) -> None:

"""Initialize migration loader and commands with necessary imports.

Handles the circular import between config and commands by importing

at runtime when needed.

"""

from sqlspec.loader import SQLFileLoader

from sqlspec.migrations import create_migration_commands

runtime = self.get_observability_runtime()

self._migration_loader = SQLFileLoader(runtime=runtime)

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:

< 85EA div id="LC747" class="react-file-line html-div" data-testid="code-cell" data-line-number="747" style="position:relative"> 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 | AsyncMigrationCommands":

"""Get the migration commands instance.

Returns:

MigrationCommands instance for this config.

"""

return self._migration_commands

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

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)

def get_migration_commands(self) -> "SyncMigrationCommands | AsyncMigrationCommands":

"""Get migration commands for this configuration.

Returns:

MigrationCommands instance configured for this database.

"""

return self._ensure_migration_commands()

@abstractmethod

def migrate_up(

self, revision: str = "head", allow_missing: bool = False, auto_sync: bool = True, dry_run: bool = False

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

"""

raise NotImplementedError

@abstractmethod

def migrate_down(self, revision: str = "-1", *, dry_run: bool = False) -> "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.

"""

raise NotImplementedError

@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

@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

@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

@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

@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 provid 3359 es 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

class NoPoolSyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):

"""Base class for sync database configurations that do not implement a pool."""

__slots__ = ("connection_config",)

is_async: "ClassVar[bool]" = False

supports_connection_pooling: "ClassVar[bool]" = False

migration_tracker_type: "ClassVar[type[Any]]" = SyncMigrationTracker

def __init__(

self,

connection_config: dict[str, Any] | 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,

) -> None:

self.bind_key = bind_key

self.pool_instance = None

self.connection_config = connection_config or {}

self.extension_config = extension_config or {}

self.migration_config: dict[str, Any] | MigrationConfig = migration_config or {}

self._init_observability(observability_config)

self._initialize_migration_components()

if statement_config is None:

default_parameter_config = ParameterStyleConfig(

default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}

)

self.statement_config = StatementConfig(dialect="sqlite", parameter_config=default_parameter_config)

else:

self.statement_config = statement_config

self.driver_features = driver_features or {}

self._storage_capabilities = None

self.driver_features.setdefault("storage_capabilities", self.storage_capabilities())

self._promote_driver_feature_hooks()

self._configure_observability_extensions()

def create_connection(self) -> ConnectionT:

"""Create a database connection."""

raise NotImplementedError

def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":

"""Provide a database connection context manager."""

raise NotImplementedError

def provide_session(

self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any

) -> "AbstractContextManager[DriverT]":

"""Provide a database session context manager."""

raise NotImplementedError

def create_pool(self) -> None:

return None

def close_pool(self) -> None:

return None

def provide_pool(self, *args: Any, **kwargs: Any) -> None:

return None

def migrate_up(

self, revision: str = "head", allow_missing: bool = False, auto_sync: bool = True, dry_run: bool = False

) -> None:

"""Apply database migrations up to specified revision.

Args:

revision: Target revision or "head" for latest.

allow_missing: Allow out-of-order migrations.

auto_sync: Auto-reconcile renamed migrations.

dry_run: Show what would be done without applying.

"""

commands = self._ensure_migration_commands()

commands.upgrade(revision, allow_missing, auto_sync, dry_run)

def migrate_down(self, revision: str = "-1", *, dry_run: bool = False) -> None:

"""Apply database migrations down to specified revision.

Args:

revision: Target revision, "-1" for one step back, or "base" for all migrations.

dry_run: Show what would be done without applying.

"""

commands = self._ensure_migration_commands()

commands.downgrade(revision, dry_run=dry_run)

def get_current_migration(self, verbose: bool = False) -> "str | None":

"""Get the current migration version.

Args:

verbose: Whether to show detailed migration history.

Returns:

Current migration version or None if no migrations applied.

"""

commands = cast("SyncMigrationCommands", self._ensure_migration_commands())

return commands.current(verbose=verbose)

def create_migration(self, message: str, file_type: str = "sql") -> None:

"""Create a new migration file.

Args:

message: Description for the migration.

file_type: Type of migration file to create ('sql' or 'py').

"""

commands = self._ensure_migration_commands()

commands.revision(message, file_type)

def init_migrations(self, directory: "str | None" = None, package: bool = True) -> None:

"""Initialize migration directory structure.

Args:

directory: Directory to initialize migrations in.

package: Whether to create __init__.py file.

"""

View remainder of file in raw view

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

config.py

Latest commit

History

config.py

File metadata and controls