Source code for sqlspec.extensions.litestar.store

"""Base session store classes for Litestar integration."""

import re
from abc import abstractmethod
from datetime import datetime, timedelta, timezone
from typing import Any, ClassVar, Final, Generic, TypeVar, cast

from litestar.stores.base import Store

from sqlspec.exceptions import ImproperConfigurationError
from sqlspec.migrations.schema import SchemaTarget, ensure_schema_async, ensure_schema_sync
from sqlspec.observability import resolve_db_system
from sqlspec.utils.logging import get_logger
from sqlspec.utils.sync_tools import async_
from sqlspec.utils.type_guards import has_extension_config

__all__ = ("BaseSQLSpecStore",)


ConfigT = TypeVar("ConfigT")


logger = get_logger("sqlspec.extensions.litestar.store")


VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
MAX_TABLE_NAME_LENGTH: Final = 63



[docs]
class BaseSQLSpecStore(Store, Generic[ConfigT]):
    """Base class for SQLSpec-backed Litestar session stores.

    Inherits Litestar's ``Store`` abstract base class for server-side session
    storage backed by SQLSpec database adapters.

    This abstract base class provides common functionality for all database-specific
    store implementations including:
    - Connection management via SQLSpec configs
    - Session expiration calculation
    - Table creation utilities

    Subclasses must implement dialect-specific SQL queries.

    Args:
        config: SQLSpec database configuration with extension_config["litestar"] settings.
    """

    __slots__ = ("_config", "_table_name")

    extension_config_options: ClassVar[frozenset[str]] = frozenset({
        "auto_trace_headers",
        "commit_mode",
        "connection_key",
        "correlation_header",
        "correlation_headers",
        "create_schema",
        "disable_di",
        "enable_correlation_middleware",
        "enable_sqlcommenter_middleware",
        "extra_commit_statuses",
        "extra_rollback_statuses",
        "manage_schema",
        "migrations_path",
        "pool_key",
        "run_migrations",
        "session_key",
        "session_table",
    })


[docs]
    def __init__(self, config: ConfigT) -> None:
        """Initialize the session store.

        Args:
            config: SQLSpec database configuration.
        """
        self._config = config
        self._validate_extension_config()
        self._table_name = self._table_name_from_config()
        self._ensure_table_name(self._table_name)

@property def config(self) -> ConfigT: """Return the database configuration.""" return self._config @property def table_name(self) -> str: """Return the session table name.""" return self._table_name @property def create_schema_enabled(self) -> bool: """Return whether adapter-level creation should run.""" manage_schema, create_schema = self._schema_management_flags() return manage_schema and create_schema

[docs]

@abstractmethod async def delete_expired(self) -> int: """Delete all expired sessions. Returns: Number of sessions deleted. """ raise NotImplementedError