sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.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 `dict.setdefault()`, 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/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.
+    """