import logging
from collections.abc import Iterable
from contextlib import suppress
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, NoReturn, TypeAlias, cast, overload
from litestar.di import Provide
from litestar.exceptions import NotFoundException
from litestar.middleware import DefineMiddleware
from litestar.plugins import CLIPlugin, InitPluginProtocol, OpenAPISchemaPlugin
from sqlspec.base import SQLSpec
from sqlspec.config import (
AsyncConfigT,
AsyncDatabaseConfig,
DatabaseConfigProtocol,
DriverT,
NoPoolAsyncConfig,
NoPoolSyncConfig,
SyncConfigT,
SyncDatabaseConfig,
)
from sqlspec.core import CorrelationExtractor, OffsetPagination
from sqlspec.core.sqlcommenter import SQLCommenterContext
from sqlspec.exceptions import ImproperConfigurationError, NotFoundError
from sqlspec.extensions.litestar._utils import (
delete_sqlspec_scope_state,
get_sqlspec_scope_state,
set_sqlspec_scope_state,
)
from sqlspec.extensions.litestar.cli import database_group
from sqlspec.extensions.litestar.handlers import (
autocommit_handler_maker,
connection_provider_maker,
lifespan_handler_maker,
manual_handler_maker,
pool_provider_maker,
session_provider_maker,
)
from sqlspec.typing import NUMPY_INSTALLED, ConnectionT, PoolT, SchemaT, import_optional_attr
from sqlspec.utils.correlation import CorrelationContext
from sqlspec.utils.logging import get_logger, log_with_context
from sqlspec.utils.serializers import DEFAULT_TYPE_ENCODERS, numpy_array_dec_hook
if TYPE_CHECKING:
from collections.abc import AsyncGenerator, Callable
from contextlib import AbstractAsyncContextManager
from litestar import Litestar, Request
from litestar._openapi.schema_generation.schema import SchemaCreator
from litestar.config.app import AppConfig
from litestar.datastructures.state import State
from litestar.openapi.spec import Schema
from litestar.types import ASGIApp, BeforeMessageSendHookHandler, Receive, Scope, Send
from litestar.typing import FieldDefinition
from rich_click import Group
from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
from sqlspec.loader import SQLFileLoader
AnyDatabaseConfig: TypeAlias = (
SyncDatabaseConfig[Any, Any, Any]
| NoPoolSyncConfig[Any, Any]
| AsyncDatabaseConfig[Any, Any, Any]
| NoPoolAsyncConfig[Any, Any]
)
__all__ = (
"CORRELATION_STATE_KEY",
"DEFAULT_COMMIT_MODE",
"DEFAULT_CONNECTION_KEY",
"DEFAULT_CORRELATION_HEADER",
"DEFAULT_POOL_KEY",
"DEFAULT_SESSION_KEY",
"TRACE_CONTEXT_FALLBACK_HEADERS",
"CommitMode",
"CorrelationMiddleware",
"PluginConfigState",
"SQLSpecPlugin",
"_OffsetPaginationSchemaPlugin",
"not_found_error_handler",
)
logger = get_logger("sqlspec.extensions.litestar")
CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
DEFAULT_COMMIT_MODE: CommitMode = "manual"
DEFAULT_CONNECTION_KEY = "db_connection"
DEFAULT_POOL_KEY = "db_pool"
DEFAULT_SESSION_KEY = "db_session"
DEFAULT_CORRELATION_HEADER = "x-request-id"
TRACE_CONTEXT_FALLBACK_HEADERS: tuple[str, ...] = (
DEFAULT_CORRELATION_HEADER,
"x-correlation-id",
"traceparent",
"x-cloud-trace-context",
"grpc-trace-bin",
"x-amzn-trace-id",
"x-b3-traceid",
"x-client-trace-id",
)
CORRELATION_STATE_KEY = "sqlspec_correlation_id"
_LITESTAR_NUMPY_ARRAY_TYPE: type[Any] | None = None
def not_found_error_handler(_request: "Request[Any, Any, Any]", exc: NotFoundError) -> NoReturn:
"""Translate :class:`sqlspec.exceptions.NotFoundError` into Litestar's HTTP 404.
Re-raised as :class:`litestar.exceptions.NotFoundException` so the standard
Litestar exception-handler chain renders it (including any RFC 7807 handler
the user has registered) and the OpenAPI 404 schema stays consistent.
"""
detail = str(exc) or "Not Found"
raise NotFoundException(detail=detail) from exc
class CorrelationMiddleware:
__slots__ = ("_app", "_extractor", "_headers")
def __init__(self, app: "ASGIApp", *, headers: tuple[str, ...]) -> None:
self._app = app
self._headers = headers
self._extractor = (
CorrelationExtractor(
primary_header=headers[0],
additional_headers=headers[1:] if len(headers) > 1 else None,
auto_trace_headers=False,
)
if headers
else None
)
async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
scope_type = scope.get("type")
if str(scope_type) != "http" or self._extractor is None:
await self._app(scope, receive, send)
return
raw_headers = scope.get("headers") or []
header_dict = {name.decode().lower(): value.decode() for name, value in raw_headers}
header_value = self._extractor.extract(lambda header: header_dict.get(header))
previous_correlation_id = CorrelationContext.get()
CorrelationContext.set(header_value)
set_sqlspec_scope_state(scope, CORRELATION_STATE_KEY, header_value)
try:
await self._app(scope, receive, send)
finally:
with suppress(KeyError):
delete_sqlspec_scope_state(scope, CORRELATION_STATE_KEY)
CorrelationContext.set(previous_correlation_id)
@dataclass
class PluginConfigState:
"""Internal state for each database configuration."""
config: "DatabaseConfigProtocol[Any, Any, Any]"
connection_key: str
pool_key: str
session_key: str
commit_mode: CommitMode
extra_commit_statuses: "set[int] | None"
extra_rollback_statuses: "set[int] | None"
enable_correlation_middleware: bool
correlation_header: str
enable_sqlcommenter_middleware: bool
correlation_headers: tuple[str, ...] = field(init=False)
disable_di: bool
connection_provider: "Callable[[State, Scope], AsyncGenerator[Any, None]] | None" = field(default=None, init=False)
pool_provider: "Callable[[State, Scope], Any] | None" = field(default=None, init=False)
session_provider: "Callable[..., AsyncGenerator[Any, None]] | None" = field(default=None, init=False)
before_send_handler: "BeforeMessageSendHookHandler | None" = field(default=None, init=False)
lifespan_handler: "Callable[[Litestar], AbstractAsyncContextManager[None]] | None" = field(default=None, init=False)
annotation: "type[DatabaseConfigProtocol[Any, Any, Any]] | None" = field(default=None, init=False)
class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
"""Litestar plugin for SQLSpec database integration.
Automatically configures NumPy array serialization when NumPy is installed,
enabling seamless bidirectional conversion between NumPy arrays and JSON
for vector embedding workflows.
Session Table Migrations:
The Litestar extension includes migrations for creating session storage tables.
To include these migrations in your database migration workflow, add 'litestar'
to the include_extensions list in your migration configuration.
"""
__slots__ = ("_correlation_headers", "_enable_sqlcommenter_middleware", "_plugin_configs", "_sqlspec")
def __init__(self, sqlspec: SQLSpec, *, loader: "SQLFileLoader | None" = None) -> None:
"""Initialize SQLSpec plugin.
Args:
sqlspec: Pre-configured SQLSpec instance with registered database configs.
loader: Optional SQL file loader instance (SQLSpec may already have one).
"""
self._sqlspec = sqlspec
self._plugin_configs: list[PluginConfigState] = []
for cfg in self._sqlspec.configs.values():
config_union = cast("AnyDatabaseConfig", cfg)
settings = self._extract_extension_settings(config_union)
state = self._config_state(config_union, settings)
self._plugin_configs.append(state)
correlation_headers: list[str] = []
enable_sqlcommenter = False
for state in self._plugin_configs:
if state.enable_sqlcommenter_middleware and state.config.statement_config.enable_sqlcommenter:
enable_sqlcommenter = True
if not state.enable_correlation_middleware:
continue
for header in state.correlation_headers:
if header not in correlation_headers:
correlation_headers.append(header)
self._correlation_headers = tuple(correlation_headers)
self._enable_sqlcommenter_middleware = enable_sqlcommenter
log_with_context(
logger,
logging.DEBUG,
"extension.init",
framework="litestar",
stage="init",
config_count=len(self._plugin_configs),
correlation_headers=len(self._correlation_headers),
)