"""Base session store classes for Litestar integration."""
import re
from abc import ABC, abstractmethod
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, TypeVar, cast
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
if TYPE_CHECKING:
from types import TracebackType
from typing_extensions import Self
__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(ABC, Generic[ConfigT]):
"""Base class for SQLSpec-backed Litestar session stores.
Implements the litestar.stores.base.Store protocol for server-side session
storage using 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)
[docs]
async def __aenter__(self) -> "Self":
"""Enter context manager."""
return self
[docs]
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> None:
"""Exit context manager."""
return
@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 get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
"""Get a session value by key.
Args:
key: Session ID to retrieve.
renew_for: If given and the value had an initial expiry time set, renew the
expiry time for ``renew_for`` seconds. If the value has not been set
with an expiry time this is a no-op.
Returns:
Session data as bytes if found and not expired, None otherwise.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
"""Store a session value.
Args:
key: Session ID.
value: Session data (will be converted to bytes if string).
expires_in: Time in seconds or timedelta before expiration.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def delete(self, key: str) -> None:
"""Delete a session by key.
Args:
key: Session ID to delete.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def delete_all(self) -> None:
"""Delete all sessions from the store."""
raise NotImplementedError
[docs]
@abstractmethod
async def exists(self, key: str) -> bool:
"""Check if a session key exists and is not expired.
Args:
key: Session ID to check.
Returns:
True if the session exists and is not expired.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def expires_in(self, key: str) -> "int | None":
"""Get the time in seconds until the session expires.
Args:
key: Session ID to check.
Returns:
Seconds until expiration, or None if no expiry or key doesn't exist.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def delete_expired(self) -> int:
"""Delete all expired sessions.
Returns:
Number of sessions deleted.
"""
raise NotImplementedError
[docs]
@abstractmethod
async def create_table(self) -> None:
"""Create the session table if it doesn't exist."""
raise NotImplementedError
[docs]
def prepare_schema_sync(self, driver: Any) -> None:
"""Prepare adapter-specific schema decisions with a synchronous driver."""
[docs]
async def prepare_schema_async(self, driver: Any) -> None:
"""Prepare adapter-specific schema decisions with an asynchronous driver."""
[docs]
async def reconcile_schema(self, *, assume_existing: bool = False) -> None:
"""Apply additive session-table changes from the canonical store DDL.
Args:
assume_existing: Skip table discovery after adapter-level creation.
"""
manage_schema, create_schema = self._schema_management_flags()
if not manage_schema:
return
statement_config = getattr(self._config, "statement_config", None)
dialect = getattr(statement_config, "dialect", None)
target = SchemaTarget.from_ddl(self._table_name, self._table_ddl(), dialect=dialect)
session_context = cast("Any", self._config).provide_session()
if hasattr(session_context, "__aenter__"):
async with session_context as driver:
await ensure_schema_async(
driver, [target], manage_schema=True, create_schema=create_schema, assume_existing=assume_existing
)
return
await async_(_reconcile_schema_sync)(
self._config, target, create_schema=create_schema, assume_existing=assume_existing
)
def _table_name_from_config(self) -> str:
"""Extract the configured session table name."""
default_name = "litestar_session"
if has_extension_config(self._config):
extension_config = cast("dict[str, dict[str, Any]]", self._config.extension_config)
litestar_config: dict[str, Any] = extension_config.get("litestar", {})
session_table = litestar_config.get("session_table", default_name)
if session_table is True:
return default_name
return str(session_table)
return default_name
def _validate_extension_config(self) -> None:
"""Reject Litestar options that this adapter store cannot honor."""
if not has_extension_config(self._config):
return
extension_config = cast("dict[str, dict[str, Any]]", self._config.extension_config)
settings = extension_config.get("litestar", {})
unsupported = sorted(set(settings).difference(type(self).extension_config_options))
if unsupported:
adapter = type(self).__module__.split(".")[2]
keys = ", ".join(repr(key) for key in unsupported)
msg = f"Unsupported Litestar configuration key(s) for {adapter}: {keys}"
raise ImproperConfigurationError(msg)
def _schema_management_flags(self) -> "tuple[bool, bool]":
"""Return automatic-management and missing-table creation flags."""
if not has_extension_config(self._config):
return True, True
extension_config = cast("dict[str, dict[str, Any]]", self._config.extension_config)
settings = extension_config.get("litestar", {})
return bool(settings.get("manage_schema", True)), bool(settings.get("create_schema", True))
@abstractmethod
def _table_ddl(self) -> str:
"""Get the CREATE TABLE SQL for this database dialect.
Returns:
SQL statement to create the sessions table.
"""
raise NotImplementedError
@abstractmethod
def _drop_table_sql(self) -> "list[str]":
"""Get the DROP TABLE SQL statements for this database dialect.
Returns:
List of SQL statements to drop the table and all indexes.
Order matters: drop indexes before table.
"""
raise NotImplementedError
def _log_table_created(self) -> None:
logger.debug(
"Litestar session table ready",
extra={"db.system": resolve_db_system(type(self).__name__), "session_table": self._table_name},
)
def _log_delete_all(self) -> None:
logger.debug(
"Litestar sessions cleared",
extra={"db.system": resolve_db_system(type(self).__name__), "session_table": self._table_name},
)
def _log_delete_expired(self, count: int) -> None:
logger.debug(
"Litestar sessions expired cleanup",
extra={
"db.system": resolve_db_system(type(self).__name__),
"session_table": self._table_name,
"deleted_sessions": count,
},
)
def _calculate_expires_at(self, expires_in: "int | timedelta | None") -> "datetime | None":
"""Calculate expiration timestamp from expires_in.
Args:
expires_in: Seconds or timedelta until expiration.
Returns:
UTC datetime of expiration, or None if no expiration.
"""
if expires_in is None:
return None
expires_in_seconds = int(expires_in.total_seconds()) if isinstance(expires_in, timedelta) else expires_in
if expires_in_seconds <= 0:
return None
return datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)
def _value_to_bytes(self, value: "str | bytes") -> bytes:
"""Convert value to bytes if needed.
Args:
value: String or bytes value.
Returns:
Value as bytes.
"""
if isinstance(value, str):
return value.encode("utf-8")
return value
@staticmethod
def _ensure_table_name(table_name: str) -> None:
"""Validate table name for SQL safety.
Args:
table_name: Table name to validate.
Raises:
ValueError: If table name is invalid.
"""
if not table_name:
msg = "Table name cannot be empty"
raise ValueError(msg)
if len(table_name) > MAX_TABLE_NAME_LENGTH:
msg = f"Table name too long: {len(table_name)} chars (max {MAX_TABLE_NAME_LENGTH})"
raise ValueError(msg)
if not VALID_TABLE_NAME_PATTERN.match(table_name):
msg = f"Invalid table name: {table_name!r}. Must start with letter/underscore and contain only alphanumeric characters and underscores"
raise ValueError(msg)
def _reconcile_schema_sync(config: Any, target: SchemaTarget, *, create_schema: bool, assume_existing: bool) -> None:
"""Run additive reconciliation for a synchronous Litestar store."""
with config.provide_session() as driver:
ensure_schema_sync(
driver, [target], manage_schema=True, create_schema=create_schema, assume_existing=assume_existing
)
