sqlspec 0.16.1__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

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)