sqlspec 0.16.0__cp313-cp313-macosx_14_0_arm64.whl

sqlspec/extensions/litestar/_utils.py

@@ -0,0 +1,52 @@
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from litestar.types import Scope
+
+__all__ = ("delete_sqlspec_scope_state", "get_sqlspec_scope_state", "set_sqlspec_scope_state")
+
+_SCOPE_NAMESPACE = "_sqlspec"
+
+
+def get_sqlspec_scope_state(scope: "Scope", key: str, default: Any = None, pop: bool = False) -> Any:
+    """Get an internal value from connection scope state.
+
+    Note:
+        If called with a default value, this method behaves like to `dict.set_default()`, both setting the key in the
+        namespace to the default value, and returning it.
+
+        If called without a default value, the method behaves like `dict.get()`, returning ``None`` if the key does not
+        exist.
+
+    Args:
+        scope: The connection scope.
+        key: Key to get from internal namespace in scope state.
+        default: Default value to return.
+        pop: Boolean flag dictating whether the value should be deleted from the state.
+
+    Returns:
+        Value mapped to ``key`` in internal connection scope namespace.
+    """
+    namespace = scope.setdefault(_SCOPE_NAMESPACE, {})  # type: ignore[misc]
+    return namespace.pop(key, default) if pop else namespace.get(key, default)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
+
+
+def set_sqlspec_scope_state(scope: "Scope", key: str, value: Any) -> None:
+    """Set an internal value in connection scope state.
+
+    Args:
+        scope: The connection scope.
+        key: Key to set under internal namespace in scope state.
+        value: Value for key.
+    """
+    scope.setdefault(_SCOPE_NAMESPACE, {})[key] = value  # type: ignore[misc]
+
+
+def delete_sqlspec_scope_state(scope: "Scope", key: str) -> None:
+    """Remove an internal value from connection scope state.
+
+    Args:
+        scope: The connection scope.
+        key: Key to set under internal namespace in scope state.
+    """
+    del scope.setdefault(_SCOPE_NAMESPACE, {})[key]  # type: ignore[misc]

sqlspec/extensions/litestar/cli.py

@@ -0,0 +1,48 @@
+"""Litestar CLI integration for SQLSpec migrations."""
+
+from contextlib import suppress
+from typing import TYPE_CHECKING
+
+from litestar.cli._utils import LitestarGroup
+
+from sqlspec.cli import add_migration_commands
+
+try:
+    import rich_click as click
+except ImportError:
+    import click  # type: ignore[no-redef]
+
+if TYPE_CHECKING:
+    from litestar import Litestar
+
+    from sqlspec.extensions.litestar.plugin import SQLSpec
+
+
+def get_database_migration_plugin(app: "Litestar") -> "SQLSpec":
+    """Retrieve the SQLSpec plugin from the Litestar application's plugins.
+
+    Args:
+        app: The Litestar application
+
+    Returns:
+        The SQLSpec plugin
+
+    Raises:
+        ImproperConfigurationError: If the SQLSpec plugin is not found
+    """
+    from sqlspec.exceptions import ImproperConfigurationError
+    from sqlspec.extensions.litestar.plugin import SQLSpec
+
+    with suppress(KeyError):
+        return app.plugins.get(SQLSpec)
+    msg = "Failed to initialize database migrations. The required SQLSpec plugin is missing."
+    raise ImproperConfigurationError(msg)
+
+
+@click.group(cls=LitestarGroup, name="db")
+def database_group(ctx: "click.Context") -> None:
+    """Manage SQLSpec database components."""
+    ctx.obj = {"app": ctx.obj, "configs": get_database_migration_plugin(ctx.obj.app).config}
+
+
+add_migration_commands(database_group)

sqlspec/extensions/litestar/config.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar.handlers import (
+    autocommit_handler_maker,
+    connection_provider_maker,
+    lifespan_handler_maker,
+    manual_handler_maker,
+    pool_provider_maker,
+    session_provider_maker,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Awaitable
+    from contextlib import AbstractAsyncContextManager
+
+    from litestar import Litestar
+    from litestar.datastructures.state import State
+    from litestar.types import BeforeMessageSendHookHandler, Scope
+
+    from sqlspec.config import AsyncConfigT, DriverT, SyncConfigT
+    from sqlspec.typing import ConnectionT, PoolT
+
+
+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"
+
+__all__ = (
+    "DEFAULT_COMMIT_MODE",
+    "DEFAULT_CONNECTION_KEY",
+    "DEFAULT_POOL_KEY",
+    "DEFAULT_SESSION_KEY",
+    "CommitMode",
+    "DatabaseConfig",
+)
+
+
+@dataclass
+class DatabaseConfig:
+    config: "Union[SyncConfigT, AsyncConfigT]" = field()  # type: ignore[valid-type]   # pyright: ignore[reportGeneralTypeIssues]
+    connection_key: str = field(default=DEFAULT_CONNECTION_KEY)
+    pool_key: str = field(default=DEFAULT_POOL_KEY)
+    session_key: str = field(default=DEFAULT_SESSION_KEY)
+    commit_mode: "CommitMode" = field(default=DEFAULT_COMMIT_MODE)
+    extra_commit_statuses: "Optional[set[int]]" = field(default=None)
+    extra_rollback_statuses: "Optional[set[int]]" = field(default=None)
+    enable_correlation_middleware: bool = field(default=True)
+    connection_provider: "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]" = field(  # pyright: ignore[reportGeneralTypeIssues]
+        init=False, repr=False, hash=False
+    )
+    pool_provider: "Callable[[State,Scope], Awaitable[PoolT]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
+    session_provider: "Callable[[Any], AsyncGenerator[DriverT, None]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
+    before_send_handler: "BeforeMessageSendHookHandler" = field(init=False, repr=False, hash=False)
+    lifespan_handler: "Callable[[Litestar], AbstractAsyncContextManager[None]]" = field(
+        init=False, repr=False, hash=False
+    )
+    annotation: "type[Union[SyncConfigT, AsyncConfigT]]" = field(init=False, repr=False, hash=False)  # type: ignore[valid-type]   # pyright: ignore[reportGeneralTypeIssues]
+
+    def __post_init__(self) -> None:
+        if not self.config.supports_connection_pooling and self.pool_key == DEFAULT_POOL_KEY:  # type: ignore[union-attr,unused-ignore]
+            self.pool_key = f"_{self.pool_key}_{id(self.config)}"
+        if self.commit_mode == "manual":
+            self.before_send_handler = manual_handler_maker(connection_scope_key=self.connection_key)
+        elif self.commit_mode == "autocommit":
+            self.before_send_handler = autocommit_handler_maker(
+                commit_on_redirect=False,
+                extra_commit_statuses=self.extra_commit_statuses,
+                extra_rollback_statuses=self.extra_rollback_statuses,
+                connection_scope_key=self.connection_key,
+            )
+        elif self.commit_mode == "autocommit_include_redirect":
+            self.before_send_handler = autocommit_handler_maker(
+                commit_on_redirect=True,
+                extra_commit_statuses=self.extra_commit_statuses,
+                extra_rollback_statuses=self.extra_rollback_statuses,
+                connection_scope_key=self.connection_key,
+            )
+        else:
+            msg = f"Invalid commit mode: {self.commit_mode}"
+            raise ImproperConfigurationError(detail=msg)
+        self.lifespan_handler = lifespan_handler_maker(config=self.config, pool_key=self.pool_key)
+        self.connection_provider = connection_provider_maker(
+            connection_key=self.connection_key, pool_key=self.pool_key, config=self.config
+        )
+        self.pool_provider = pool_provider_maker(config=self.config, pool_key=self.pool_key)
+        self.session_provider = session_provider_maker(
+            config=self.config, connection_dependency_key=self.connection_key
+        )

sqlspec/extensions/litestar/handlers.py

@@ -0,0 +1,260 @@
+import contextlib
+import inspect
+from collections.abc import AsyncGenerator
+from contextlib import AbstractAsyncContextManager
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast
+
+from litestar.constants import HTTP_DISCONNECT, HTTP_RESPONSE_START, WEBSOCKET_CLOSE, WEBSOCKET_DISCONNECT
+
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar._utils import (
+    delete_sqlspec_scope_state,
+    get_sqlspec_scope_state,
+    set_sqlspec_scope_state,
+)
+from sqlspec.utils.sync_tools import ensure_async_
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Coroutine
+
+    from litestar import Litestar
+    from litestar.datastructures.state import State
+    from litestar.types import Message, Scope
+
+    from sqlspec.config import DatabaseConfigProtocol, DriverT
+    from sqlspec.typing import ConnectionT, PoolT
+
+SESSION_TERMINUS_ASGI_EVENTS = {HTTP_RESPONSE_START, HTTP_DISCONNECT, WEBSOCKET_DISCONNECT, WEBSOCKET_CLOSE}
+
+__all__ = (
+    "SESSION_TERMINUS_ASGI_EVENTS",
+    "autocommit_handler_maker",
+    "connection_provider_maker",
+    "lifespan_handler_maker",
+    "manual_handler_maker",
+    "pool_provider_maker",
+    "session_provider_maker",
+)
+
+
+def manual_handler_maker(connection_scope_key: str) -> "Callable[[Message, Scope], Coroutine[Any, Any, None]]":
+    """Create handler for manual connection management.
+
+    Args:
+        connection_scope_key: The key used to store the connection in the ASGI scope.
+
+    Returns:
+        The handler callable.
+    """
+
+    async def handler(message: "Message", scope: "Scope") -> None:
+        """Handle closing and cleaning up connections before sending the response.
+
+        Args:
+            message: ASGI Message.
+            scope: ASGI Scope.
+        """
+        connection = get_sqlspec_scope_state(scope, connection_scope_key)
+        if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
+            if hasattr(connection, "close") and callable(connection.close):
+                await ensure_async_(connection.close)()
+            delete_sqlspec_scope_state(scope, connection_scope_key)
+
+    return handler
+
+
+def autocommit_handler_maker(
+    connection_scope_key: str,
+    commit_on_redirect: bool = False,
+    extra_commit_statuses: "Optional[set[int]]" = None,
+    extra_rollback_statuses: "Optional[set[int]]" = None,
+) -> "Callable[[Message, Scope], Coroutine[Any, Any, None]]":
+    """Create handler for automatic transaction commit/rollback based on response status.
+
+    Args:
+        connection_scope_key: The key used to store the connection in the ASGI scope.
+        commit_on_redirect: Issue a commit when the response status is a redirect (3XX).
+        extra_commit_statuses: A set of additional status codes that trigger a commit.
+        extra_rollback_statuses: A set of additional status codes that trigger a rollback.
+
+    Raises:
+        ImproperConfigurationError: If extra_commit_statuses and extra_rollback_statuses share status codes.
+
+    Returns:
+        The handler callable.
+    """
+    if extra_commit_statuses is None:
+        extra_commit_statuses = set()
+
+    if extra_rollback_statuses is None:
+        extra_rollback_statuses = set()
+
+    if len(extra_commit_statuses & extra_rollback_statuses) > 0:
+        msg = "Extra rollback statuses and commit statuses must not share any status codes"
+        raise ImproperConfigurationError(msg)
+
+    commit_range = range(200, 400 if commit_on_redirect else 300)
+
+    async def handler(message: "Message", scope: "Scope") -> None:
+        """Handle commit/rollback, closing and cleaning up connections before sending.
+
+        Args:
+            message: ASGI Message.
+            scope: ASGI Scope.
+        """
+        connection = get_sqlspec_scope_state(scope, connection_scope_key)
+        try:
+            if connection is not None and message["type"] == HTTP_RESPONSE_START:
+                if (message["status"] in commit_range or message["status"] in extra_commit_statuses) and message[
+                    "status"
+                ] not in extra_rollback_statuses:
+                    if hasattr(connection, "commit") and callable(connection.commit):
+                        await ensure_async_(connection.commit)()
+                elif hasattr(connection, "rollback") and callable(connection.rollback):
+                    await ensure_async_(connection.rollback)()
+        finally:
+            if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
+                if hasattr(connection, "close") and callable(connection.close):
+                    await ensure_async_(connection.close)()
+                delete_sqlspec_scope_state(scope, connection_scope_key)
+
+    return handler
+
+
+def lifespan_handler_maker(
+    config: "DatabaseConfigProtocol[Any, Any, Any]", pool_key: str
+) -> "Callable[[Litestar], AbstractAsyncContextManager[None]]":
+    """Create lifespan handler for managing database connection pool lifecycle.
+
+    Args:
+        config: The database configuration object.
+        pool_key: The key under which the connection pool will be stored in `app.state`.
+
+    Returns:
+        The lifespan handler function.
+    """
+
+    @contextlib.asynccontextmanager
+    async def lifespan_handler(app: "Litestar") -> "AsyncGenerator[None, None]":
+        """Manage database pool lifecycle for the application.
+
+        Args:
+            app: The Litestar application instance.
+
+        Yields:
+            Control to application during pool lifetime.
+        """
+        db_pool = await ensure_async_(config.create_pool)()
+        app.state.update({pool_key: db_pool})
+        try:
+            yield
+        finally:
+            app.state.pop(pool_key, None)
+            try:
+                await ensure_async_(config.close_pool)()
+            except Exception as e:
+                if app.logger:
+                    app.logger.warning("Error closing database pool for %s. Error: %s", pool_key, e)
+
+    return lifespan_handler
+
+
+def pool_provider_maker(
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", pool_key: str
+) -> "Callable[[State, Scope], Awaitable[PoolT]]":
+    """Create provider for injecting the application-level database pool.
+
+    Args:
+        config: The database configuration object.
+        pool_key: The key used to store the connection pool in `app.state`.
+
+    Returns:
+        The pool provider function.
+    """
+
+    async def provide_pool(state: "State", scope: "Scope") -> "PoolT":
+        """Provide the database pool from application state.
+
+        Args:
+            state: The Litestar application State object.
+            scope: The ASGI scope (unused for app-level pool).
+
+        Returns:
+            The database connection pool.
+
+        Raises:
+            ImproperConfigurationError: If the pool is not found in `app.state`.
+        """
+        if (db_pool := state.get(pool_key)) is None:
+            msg = (
+                f"Database pool with key '{pool_key}' not found in application state. "
+                "Ensure the SQLSpec lifespan handler is correctly configured and has run."
+            )
+            raise ImproperConfigurationError(msg)
+        return cast("PoolT", db_pool)
+
+    return provide_pool
+
+
+def connection_provider_maker(
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", pool_key: str, connection_key: str
+) -> "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]":
+    async def provide_connection(state: "State", scope: "Scope") -> "AsyncGenerator[ConnectionT, None]":
+        if (db_pool := state.get(pool_key)) is None:
+            msg = f"Database pool with key '{pool_key}' not found. Cannot create a connection."
+            raise ImproperConfigurationError(msg)
+
+        connection_cm = config.provide_connection(db_pool)
+
+        if not isinstance(connection_cm, AbstractAsyncContextManager):
+            conn_instance: ConnectionT
+            if hasattr(connection_cm, "__await__"):
+                conn_instance = await cast("Awaitable[ConnectionT]", connection_cm)
+            else:
+                conn_instance = cast("ConnectionT", connection_cm)
+            set_sqlspec_scope_state(scope, connection_key, conn_instance)
+            yield conn_instance
+            return
+
+        entered_connection = await connection_cm.__aenter__()
+        try:
+            set_sqlspec_scope_state(scope, connection_key, entered_connection)
+            yield entered_connection
+        finally:
+            await connection_cm.__aexit__(None, None, None)
+            delete_sqlspec_scope_state(scope, connection_key)
+
+    return provide_connection
+
+
+def session_provider_maker(
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", connection_dependency_key: str
+) -> "Callable[[Any], AsyncGenerator[DriverT, None]]":
+    async def provide_session(*args: Any, **kwargs: Any) -> "AsyncGenerator[DriverT, None]":
+        yield cast("DriverT", config.driver_type(connection=args[0] if args else kwargs.get(connection_dependency_key)))  # pyright: ignore
+
+    conn_type_annotation = config.connection_type
+
+    from litestar.params import Dependency
+
+    db_conn_param = inspect.Parameter(
+        name=connection_dependency_key,
+        kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        annotation=conn_type_annotation,
+        default=Dependency(skip_validation=True),
+    )
+
+    provider_signature = inspect.Signature(
+        parameters=[db_conn_param],
+        return_annotation=AsyncGenerator[config.driver_type, None],  # type: ignore[name-defined]
+    )
+
+    provide_session.__signature__ = provider_signature  # type: ignore[attr-defined]
+
+    if not hasattr(provide_session, "__annotations__") or provide_session.__annotations__ is None:
+        provide_session.__annotations__ = {}
+
+    provide_session.__annotations__[connection_dependency_key] = conn_type_annotation
+    provide_session.__annotations__["return"] = AsyncGenerator[config.driver_type, None]  # type: ignore[name-defined]
+
+    return provide_session

sqlspec/extensions/litestar/plugin.py

@@ -0,0 +1,145 @@
+from typing import TYPE_CHECKING, Any, Union
+
+from litestar.di import Provide
+from litestar.plugins import CLIPlugin, InitPluginProtocol
+
+from sqlspec.base import SQLSpec as SQLSpecBase
+from sqlspec.config import AsyncConfigT, DatabaseConfigProtocol, DriverT, SyncConfigT
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar.config import DatabaseConfig
+from sqlspec.typing import ConnectionT, PoolT
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from click import Group
+    from litestar.config.app import AppConfig
+
+logger = get_logger("extensions.litestar")
+
+
+class SQLSpec(InitPluginProtocol, CLIPlugin, SQLSpecBase):
+    """Litestar plugin for SQLSpec database integration."""
+
+    __slots__ = ("_config", "_plugin_configs")
+
+    def __init__(self, config: Union["SyncConfigT", "AsyncConfigT", "DatabaseConfig", list["DatabaseConfig"]]) -> None:
+        """Initialize SQLSpec plugin.
+
+        Args:
+            config: Database configuration for SQLSpec plugin.
+        """
+        self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
+        if isinstance(config, DatabaseConfigProtocol):
+            self._plugin_configs: list[DatabaseConfig] = [DatabaseConfig(config=config)]
+        elif isinstance(config, DatabaseConfig):
+            self._plugin_configs = [config]
+        else:
+            self._plugin_configs = config
+
+    @property
+    def config(self) -> "list[DatabaseConfig]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Return the plugin configuration.
+
+        Returns:
+            List of database configurations.
+        """
+        return self._plugin_configs
+
+    def on_cli_init(self, cli: "Group") -> None:
+        """Configure CLI commands for SQLSpec database operations.
+
+        Args:
+            cli: The Click command group to add commands to.
+        """
+        from sqlspec.extensions.litestar.cli import database_group
+
+        cli.add_command(database_group)
+
+    def on_app_init(self, app_config: "AppConfig") -> "AppConfig":
+        """Configure Litestar application with SQLSpec database integration.
+
+        Args:
+            app_config: The Litestar application configuration instance.
+
+        Returns:
+            The updated application configuration instance.
+        """
+
+        self._validate_dependency_keys()
+
+        def store_sqlspec_in_state() -> None:
+            app_config.state.sqlspec = self
+
+        app_config.on_startup.append(store_sqlspec_in_state)
+        app_config.signature_types.extend(
+            [SQLSpec, ConnectionT, PoolT, DriverT, DatabaseConfig, DatabaseConfigProtocol, SyncConfigT, AsyncConfigT]
+        )
+
+        signature_namespace = {}
+
+        for c in self._plugin_configs:
+            c.annotation = self.add_config(c.config)
+            app_config.signature_types.append(c.annotation)
+            app_config.signature_types.append(c.config.connection_type)  # type: ignore[union-attr]
+            app_config.signature_types.append(c.config.driver_type)  # type: ignore[union-attr]
+
+            if hasattr(c.config, "get_signature_namespace"):
+                signature_namespace.update(c.config.get_signature_namespace())  # type: ignore[attr-defined]
+
+            app_config.before_send.append(c.before_send_handler)
+            app_config.lifespan.append(c.lifespan_handler)  # pyright: ignore[reportUnknownMemberType]
+            app_config.dependencies.update(
+                {
+                    c.connection_key: Provide(c.connection_provider),
+                    c.pool_key: Provide(c.pool_provider),
+                    c.session_key: Provide(c.session_provider),
+                }
+            )
+
+        if signature_namespace:
+            app_config.signature_namespace.update(signature_namespace)
+
+        return app_config
+
+    def get_annotations(self) -> "list[type[Union[SyncConfigT, AsyncConfigT]]]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Return the list of annotations.
+
+        Returns:
+            List of annotations.
+        """
+        return [c.annotation for c in self.config]
+
+    def get_annotation(
+        self, key: "Union[str, SyncConfigT, AsyncConfigT, type[Union[SyncConfigT, AsyncConfigT]]]"
+    ) -> "type[Union[SyncConfigT, AsyncConfigT]]":
+        """Return the annotation for the given configuration.
+
+        Args:
+            key: The configuration instance or key to lookup
+
+        Raises:
+            KeyError: If no configuration is found for the given key.
+
+        Returns:
+            The annotation for the configuration.
+        """
+        for c in self.config:
+            if key == c.config or key in {c.annotation, c.connection_key, c.pool_key}:
+                return c.annotation
+        msg = f"No configuration found for {key}"
+        raise KeyError(msg)
+
+    def _validate_dependency_keys(self) -> None:
+        """Validate that connection and pool keys are unique across configurations.
+
+        Raises:
+            ImproperConfigurationError: If connection keys or pool keys are not unique.
+        """
+        connection_keys = [c.connection_key for c in self.config]
+        pool_keys = [c.pool_key for c in self.config]
+        if len(set(connection_keys)) != len(connection_keys):
+            msg = "When using multiple database configuration, each configuration must have a unique `connection_key`."
+            raise ImproperConfigurationError(detail=msg)
+        if len(set(pool_keys)) != len(pool_keys):
+            msg = "When using multiple database configuration, each configuration must have a unique `pool_key`."
+            raise ImproperConfigurationError(detail=msg)