sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/extensions/litestar/channels.py

@@ -0,0 +1,165 @@
+"""Litestar channels backend backed by SQLSpec's EventChannel."""
+
+import asyncio
+import base64
+import hashlib
+import re
+from typing import TYPE_CHECKING, Any
+
+from litestar.channels.backends.base import ChannelsBackend
+
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+
+    from sqlspec.extensions.events import AsyncEventChannel
+
+logger = get_logger("sqlspec.extensions.litestar.channels")
+
+_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+class SQLSpecChannelsBackend(ChannelsBackend):
+    """A Litestar Channels backend implemented on top of SQLSpec's EventChannel.
+
+    This backend allows Litestar's ChannelsPlugin to use a SQLSpec database as the
+    broker. Under the hood it relies on SQLSpec's events extension, which can be
+    configured to use a durable table queue or native adapter backends.
+
+    Notes:
+        Litestar channels may use arbitrary string names. SQLSpec event channel
+        names must be valid identifiers. This backend maps Litestar channel names
+        to deterministic database channel identifiers via hashing.
+    """
+
+    def __init__(
+        self, event_channel: "AsyncEventChannel", *, channel_prefix: str = "litestar", poll_interval: float = 0.2
+    ) -> None:
+        if not _IDENTIFIER_PATTERN.match(channel_prefix):
+            msg = f"channel_prefix must be a valid identifier, got: {channel_prefix!r}"
+            raise ValueError(msg)
+        if poll_interval <= 0:
+            msg = "poll_interval must be greater than zero"
+            raise ValueError(msg)
+        self._event_channel = event_channel
+        self._channel_prefix = channel_prefix
+        self._poll_interval = poll_interval
+        self._output_queue: asyncio.Queue[tuple[str, bytes]] | None = None
+        self._shutdown = asyncio.Event()
+        self._tasks: dict[str, asyncio.Task[None]] = {}
+        self._to_db_channel: dict[str, str] = {}
+        self._to_litestar_channel: dict[str, str] = {}
+
+    async def on_startup(self) -> None:
+        self._shutdown.clear()
+        if self._output_queue is None:
+            self._output_queue = asyncio.Queue()
+
+    async def on_shutdown(self) -> None:
+        self._shutdown.set()
+        tasks = list(self._tasks.values())
+        self._tasks.clear()
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        self._to_db_channel.clear()
+        self._to_litestar_channel.clear()
+        await self._event_channel.shutdown()
+
+    async def publish(self, data: bytes, channels: "Iterable[str]") -> None:
+        payload = {"data_b64": base64.b64encode(data).decode("ascii")}
+        for channel in channels:
+            db_channel = self._db_channel_name(channel)
+            await self._event_channel.publish(db_channel, payload)
+
+    async def subscribe(self, channels: "Iterable[str]") -> None:
+        for channel in channels:
+            if channel in self._tasks:
+                continue
+            db_channel = self._db_channel_name(channel)
+            task = asyncio.create_task(self._stream_channel(channel, db_channel))
+            self._tasks[channel] = task
+
+    async def unsubscribe(self, channels: "Iterable[str]") -> None:
+        cancelled: list[asyncio.Task[None]] = []
+        for channel in channels:
+            task = self._tasks.pop(channel, None)
+            if task is None:
+                continue
+            task.cancel()
+            cancelled.append(task)
+        if cancelled:
+            await asyncio.gather(*cancelled, return_exceptions=True)
+        self._cleanup_channel_mappings()
+
+    def stream_events(self) -> "AsyncGenerator[tuple[str, bytes], None]":
+        return self._event_generator()
+
+    async def get_history(self, channel: str, limit: int | None = None) -> list[bytes]:
+        """Return history entries for a channel.
+
+        SQLSpec's event queue is primarily designed for durable delivery, not
+        for history replay. For now, this backend does not expose history.
+        """
+
+        return []
+
+    def _cleanup_channel_mappings(self) -> None:
+        active = set(self._tasks)
+        removed = [name for name in self._to_db_channel if name not in active]
+        for name in removed:
+            db_name = self._to_db_channel.pop(name, None)
+            if db_name:
+                self._to_litestar_channel.pop(db_name, None)
+
+    async def _event_generator(self) -> "AsyncGenerator[tuple[str, bytes], None]":
+        if self._output_queue is None:
+            msg = "SQLSpecChannelsBackend not started - call ChannelsPlugin.on_startup() first"
+            raise RuntimeError(msg)
+        queue = self._output_queue
+        while True:
+            item = await queue.get()
+            yield item
+
+    def _db_channel_name(self, channel: str) -> str:
+        existing = self._to_db_channel.get(channel)
+        if existing:
+            return existing
+        digest = hashlib.sha256(channel.encode("utf-8")).hexdigest()[:24]
+        db_channel = f"{self._channel_prefix}_{digest}"
+        self._to_db_channel[channel] = db_channel
+        self._to_litestar_channel[db_channel] = channel
+        return db_channel
+
+    async def _stream_channel(self, channel: str, db_channel: str) -> None:
+        try:
+            async for message in self._event_channel.iter_events(db_channel, poll_interval=self._poll_interval):
+                if self._shutdown.is_set():
+                    return
+                payload = message.payload
+                decoded = self._decode_payload(payload)
+                if decoded is None:
+                    logger.warning("litestar channel %s dropped malformed payload: %r", channel, payload)
+                    await self._event_channel.ack(message.event_id)
+                    continue
+                assert self._output_queue is not None
+                await self._output_queue.put((channel, decoded))
+                await self._event_channel.ack(message.event_id)
+        except asyncio.CancelledError:
+            raise
+        except Exception as error:  # pragma: no cover - defensive
+            logger.warning("litestar channel %s stream worker error: %s", channel, error)
+
+    @staticmethod
+    def _decode_payload(payload: Any) -> bytes | None:
+        if not isinstance(payload, dict):
+            return None
+        encoded = payload.get("data_b64")
+        if not isinstance(encoded, str) or not encoded:
+            return None
+        try:
+            return base64.b64decode(encoded.encode("ascii"))
+        except (ValueError, UnicodeEncodeError):
+            return None

sqlspec/extensions/litestar/cli.py

@@ -0,0 +1,102 @@
+"""Litestar CLI integration for SQLSpec migrations."""
+
+from collections.abc import Callable
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, cast
+
+import anyio
+import rich_click as click
+from litestar.cli._utils import LitestarGroup
+
+from sqlspec.cli import add_migration_commands
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+
+if TYPE_CHECKING:
+    from litestar import Litestar
+
+    from sqlspec.extensions.litestar.plugin import SQLSpecPlugin
+
+
+def _safe_group(*, aliases: "list[str] | None" = None, **kwargs: Any) -> Callable[[Callable[..., Any]], Any]:
+    if aliases is None:
+        return cast("Callable[[Callable[..., Any]], Any]", click.group(**kwargs))
+    try:
+        return cast("Callable[[Callable[..., Any]], Any]", click.group(aliases=aliases, **kwargs))
+    except TypeError:
+        return cast("Callable[[Callable[..., Any]], Any]", click.group(**kwargs))
+
+
+def get_database_migration_plugin(app: "Litestar") -> "SQLSpecPlugin":
+    """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.extensions.litestar.plugin import SQLSpecPlugin
+
+    with suppress(KeyError):
+        return app.plugins.get(SQLSpecPlugin)
+    msg = "Failed to initialize database migrations. The required SQLSpec plugin is missing."
+    raise ImproperConfigurationError(msg)
+
+
+@_safe_group(cls=LitestarGroup, name="db", aliases=["database"])
+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)
+
+
+def add_sessions_delete_expired_command() -> None:
+    """Add delete-expired command to Litestar's sessions CLI group."""
+    try:
+        from litestar.cli._utils import console
+        from litestar.cli.commands.sessions import get_session_backend, sessions_group
+    except ImportError:
+        return
+
+    @sessions_group.command("delete-expired")  # type: ignore[untyped-decorator]
+    @click.option(
+        "--verbose", is_flag=True, default=False, help="Show detailed information about the cleanup operation"
+    )
+    def delete_expired_sessions_command(app: "Litestar", verbose: bool) -> None:
+        """Delete expired sessions from the session store.
+
+        This command removes all sessions that have passed their expiration time.
+        It can be scheduled via cron or systemd timers for automatic maintenance.
+
+        Examples:
+            litestar sessions delete-expired
+            litestar sessions delete-expired --verbose
+        """
+        backend = get_session_backend(app)
+        store = backend.config.get_store_from_app(app)
+
+        if not isinstance(store, BaseSQLSpecStore):
+            console.print(f"[red]{type(store).__name__} does not support deleting expired sessions")
+            return
+
+        async def _delete_expired() -> int:
+            return await store.delete_expired()
+
+        count = anyio.run(_delete_expired)
+
+        if count > 0:
+            if verbose:
+                console.print(f"[green]Successfully deleted {count} expired session(s)")
+            else:
+                console.print(f"[green]Deleted {count} expired session(s)")
+        else:
+            console.print("[yellow]No expired sessions found")
+
+
+add_sessions_delete_expired_command()

sqlspec/extensions/litestar/config.py

@@ -0,0 +1,90 @@
+"""Configuration types for Litestar session store extension."""
+
+from typing_extensions import NotRequired, TypedDict
+
+__all__ = ("LitestarConfig",)
+
+
+class LitestarConfig(TypedDict):
+    """Configuration options for Litestar session store extension.
+
+    All fields are optional with sensible defaults. Use in extension_config["litestar"]:
+
+    Example:
+        from sqlspec.adapters.oracledb import OracleAsyncConfig
+
+        config = OracleAsyncConfig(
+            connection_config={"dsn": "oracle://localhost/XEPDB1"},
+            extension_config={
+                "litestar": {
+                    "session_table": "my_sessions",
+                    "in_memory": True
+                }
+            }
+        )
+
+    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: 'litestar_session'
+
+    Examples:
+        "app_sessions"
+        "user_sessions"
+        "tenant_acme_sessions"
+    """
+
+    in_memory: NotRequired[bool]
+    """Enable in-memory table storage (Oracle-specific). Default: False.
+
+    When enabled, tables are created with the in-memory attribute for databases that support it.
+
+    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(
+                connection_config={"dsn": "oracle://..."},
+                extension_config={
+                    "litestar": {
+                        "in_memory": True
+                    }
+                }
+            )
+
+    Notes:
+        - Tables created with INMEMORY PRIORITY HIGH clause
+        - Ignored by unsupported adapters
+    """
+
+    shard_count: NotRequired[int]
+    """Optional hash shard count for session table primary key.
+
+    When set (>1), adapters that support computed shard columns (e.g., Spanner)
+    will create a generated shard_id using MOD(FARM_FINGERPRINT(session_id), shard_count)
+    and include it in the primary key to reduce hotspotting. Ignored by adapters
+    that do not support computed shards.
+    """
+
+    table_options: NotRequired[str]
+    """Optional raw OPTIONS/engine-specific table options string.
+
+    Passed verbatim when the adapter supports table-level OPTIONS/clauses
+    (e.g., Spanner columnar/tiered storage). Ignored by adapters that do not
+    support table options.
+    """
+
+    index_options: NotRequired[str]
+    """Optional raw OPTIONS/engine-specific options for the expires_at index.
+
+    Passed verbatim to the index definition for adapters that support index
+    OPTIONS/clauses. Ignored by adapters that do not support index options.
+    """

sqlspec/extensions/litestar/handlers.py

@@ -0,0 +1,316 @@
+import contextlib
+import inspect
+from collections.abc import AsyncGenerator, Callable
+from contextlib import AbstractAsyncContextManager, AbstractContextManager
+from typing import TYPE_CHECKING, Any, cast
+
+from litestar.constants import HTTP_DISCONNECT, HTTP_RESPONSE_START, WEBSOCKET_CLOSE, WEBSOCKET_DISCONNECT
+from litestar.params import Dependency
+
+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_, with_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, is_async: bool = False
+) -> "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.
+        is_async: Whether the database driver is async (uses direct await) or sync (uses ensure_async_).
+
+    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 is_async:
+                await connection.close()
+            else:
+                await ensure_async_(connection.close)()
+            delete_sqlspec_scope_state(scope, connection_scope_key)
+
+    return handler
+
+
+def autocommit_handler_maker(
+    connection_scope_key: str,
+    is_async: bool = False,
+    commit_on_redirect: bool = False,
+    extra_commit_statuses: "set[int] | None" = None,
+    extra_rollback_statuses: "set[int] | None" = 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.
+        is_async: Whether the database driver is async (uses direct await) or sync (uses ensure_async_).
+        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 is_async:
+                        await connection.commit()
+                    else:
+                        await ensure_async_(connection.commit)()
+                elif is_async:
+                    await connection.rollback()
+                else:
+                    await ensure_async_(connection.rollback)()
+        finally:
+            if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
+                if is_async:
+                    await connection.close()
+                else:
+                    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: Any
+        if config.is_async:
+            db_pool = await config.create_pool()
+        else:
+            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:
+                if config.is_async:
+                    close_result = config.close_pool()
+                    if close_result is not None:
+                        await close_result
+                else:
+                    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]]":
+    """Create provider for database connections with proper lifecycle management.
+
+    Args:
+        config: The database configuration object.
+        pool_key: The key used to retrieve the connection pool from `app.state`.
+        connection_key: The key used to store the connection in the ASGI scope.
+
+    Returns:
+        The connection provider function.
+    """
+
+    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: Any = config.provide_connection(db_pool)
+        context_manager: AbstractAsyncContextManager[ConnectionT] | None = None
+
+        if isinstance(connection_cm, AbstractAsyncContextManager):
+            context_manager = connection_cm
+        elif isinstance(connection_cm, AbstractContextManager):
+            context_manager = with_ensure_async_(connection_cm)
+
+        if context_manager is None:
+            conn_instance: ConnectionT
+            if inspect.isawaitable(connection_cm):
+                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 context_manager.__aenter__()
+        try:
+            set_sqlspec_scope_state(scope, connection_key, entered_connection)
+            yield entered_connection
+        finally:
+            await context_manager.__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]]":
+    """Create provider for database driver sessions.
+
+    Args:
+        config: The database configuration object.
+        connection_dependency_key: The key used for connection dependency injection.
+
+    Returns:
+        The session provider function.
+    """
+
+    async def provide_session(*args: Any, **kwargs: Any) -> "AsyncGenerator[DriverT, None]":
+        connection_obj = args[0] if args else kwargs.get(connection_dependency_key)
+        yield cast(
+            "DriverT",
+            config.driver_type(
+                connection=connection_obj,
+                statement_config=config.statement_config,
+                driver_features=config.driver_features,
+            ),
+        )  # pyright: ignore
+
+    conn_type_annotation = config.connection_type
+
+    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 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