sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/extensions/litestar/store.py

@@ -0,0 +1,296 @@
+"""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, Final, Generic, TypeVar, cast
+
+from sqlspec.observability import resolve_db_system
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.type_guards import has_extension_config
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+
+ConfigT = TypeVar("ConfigT")
+
+
+logger = get_logger("sqlspec.extensions.litestar.store")
+
+__all__ = ("BaseSQLSpecStore",)
+
+VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+MAX_TABLE_NAME_LENGTH: Final = 63
+
+
+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.
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore
+
+        config = AsyncpgConfig(
+            connection_config={"dsn": "postgresql://..."},
+            extension_config={"litestar": {"session_table": "my_sessions"}}
+        )
+        store = AsyncpgStore(config)
+        await store.create_table()
+
+    Notes:
+        Configuration is read from config.extension_config["litestar"]:
+        - session_table: Table name (default: "litestar_session")
+    """
+
+    __slots__ = ("_config", "_table_name")
+
+    def __init__(self, config: ConfigT) -> None:
+        """Initialize the session store.
+
+        Args:
+            config: SQLSpec database configuration.
+
+        Notes:
+            Reads table_name from config.extension_config["litestar"]["session_table"].
+            Defaults to "litestar_session" if not specified.
+        """
+        self._config = config
+        self._table_name = self._get_table_name_from_config()
+        self._validate_table_name(self._table_name)
+
+    def _get_table_name_from_config(self) -> str:
+        """Extract table name from config.extension_config.
+
+        Returns:
+            Table name for the session store.
+
+        Notes:
+            Accepts ``session_table: True`` for default name or a string for custom 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
+
+    @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
+
+    @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
+
+    @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
+
+    @abstractmethod
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        raise NotImplementedError
+
+    @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
+
+    @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
+
+    @abstractmethod
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_create_table_sql(self) -> str:
+        """Get the CREATE TABLE SQL for this database dialect.
+
+        Returns:
+            SQL statement to create the sessions table.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_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.
+
+        Notes:
+            Should use IF EXISTS or dialect-specific error handling
+            to allow idempotent migrations.
+        """
+        raise NotImplementedError
+
+    async def __aenter__(self) -> "BaseSQLSpecStore":
+        """Enter context manager."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> None:
+        """Exit context manager."""
+        return
+
+    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 _validate_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.
+
+        Notes:
+            - Must start with letter or underscore
+            - Can only contain letters, numbers, and underscores
+            - Maximum length is 63 characters (PostgreSQL limit)
+            - Prevents SQL injection in table names
+        """
+        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)

sqlspec/extensions/otel/__init__.py

@@ -0,0 +1,58 @@
+"""Optional helpers for enabling OpenTelemetry spans via ObservabilityConfig."""
+
+from collections.abc import Callable
+from typing import Any
+
+from sqlspec.observability import ObservabilityConfig, TelemetryConfig
+from sqlspec.typing import trace
+from sqlspec.utils.module_loader import ensure_opentelemetry
+
+__all__ = ("enable_tracing",)
+
+
+def _wrap_provider(provider: Any | None) -> Callable[[], Any] | None:
+    if provider is None:
+        return None
+
+    def _factory() -> Any:
+        return provider
+
+    return _factory
+
+
+def enable_tracing(
+    *,
+    base_config: ObservabilityConfig | None = None,
+    tracer_provider: Any | None = None,
+    tracer_provider_factory: Callable[[], Any] | None = None,
+    resource_attributes: dict[str, Any] | None = None,
+    enable_spans: bool = True,
+) -> ObservabilityConfig:
+    """Return an ObservabilityConfig with OpenTelemetry spans enabled.
+
+    Args:
+        base_config: Existing observability config to extend. When omitted a new instance is created.
+        tracer_provider: Optional provider instance to reuse. Mutually exclusive with tracer_provider_factory.
+        tracer_provider_factory: Callable that returns a tracer provider when spans are first used.
+        resource_attributes: Additional attributes to attach to every span.
+        enable_spans: Allow disabling spans while keeping the rest of the config.
+
+    Returns:
+        ObservabilityConfig with telemetry options configured for OpenTelemetry.
+    """
+
+    ensure_opentelemetry()
+
+    if tracer_provider is not None and tracer_provider_factory is not None:
+        msg = "Provide either tracer_provider or tracer_provider_factory, not both"
+        raise ValueError(msg)
+
+    telemetry = TelemetryConfig(
+        enable_spans=enable_spans,
+        provider_factory=tracer_provider_factory or _wrap_provider(tracer_provider) or trace.get_tracer_provider,
+        resource_attributes=resource_attributes,
+    )
+
+    config = base_config.copy() if base_config else ObservabilityConfig()
+    config.telemetry = telemetry
+    return config

sqlspec/extensions/prometheus/__init__.py

@@ -0,0 +1,113 @@
+"""Prometheus metrics helpers that integrate with the observability statement observers."""
+
+from collections.abc import Iterable
+from typing import Any
+
+from sqlspec.observability import ObservabilityConfig, StatementEvent, StatementObserver, resolve_db_system
+from sqlspec.typing import Counter, Histogram
+from sqlspec.utils.module_loader import ensure_prometheus
+
+__all__ = ("PrometheusStatementObserver", "enable_metrics")
+
+
+class PrometheusStatementObserver:
+    """Statement observer that records Prometheus metrics."""
+
+    __slots__ = ("_counters", "_duration", "_label_names", "_rows")
+
+    def __init__(
+        self,
+        *,
+        namespace: str = "sqlspec",
+        subsystem: str = "driver",
+        registry: Any | None = None,
+        label_names: Iterable[str] = ("db_system", "operation"),
+        duration_buckets: tuple[float, ...] | None = None,
+    ) -> None:
+        self._label_names = tuple(label_names)
+        self._counters = Counter(
+            "query_total",
+            "Total SQL statements executed",
+            labelnames=self._label_names,
+            namespace=namespace,
+            subsystem=subsystem,
+            registry=registry,
+        )
+        histogram_kwargs: dict[str, Any] = {}
+        if duration_buckets is not None:
+            histogram_kwargs["buckets"] = duration_buckets
+
+        self._duration = Histogram(
+            "query_duration_seconds",
+            "SQL execution time in seconds",
+            labelnames=self._label_names,
+            namespace=namespace,
+            subsystem=subsystem,
+            registry=registry,
+            **histogram_kwargs,
+        )
+        self._rows = Histogram(
+            "query_rows",
+            "Rows affected per statement",
+            labelnames=self._label_names,
+            namespace=namespace,
+            subsystem=subsystem,
+            registry=registry,
+        )
+
+    def __call__(self, event: StatementEvent) -> None:
+        label_values = self._label_values(event)
+        self._counters.labels(*label_values).inc()
+        self._duration.labels(*label_values).observe(max(event.duration_s, 0.0))
+        if event.rows_affected is not None:
+            self._rows.labels(*label_values).observe(float(event.rows_affected))
+
+    def _label_values(self, event: StatementEvent) -> tuple[str, ...]:
+        values: list[str] = []
+        payload = event.as_dict()
+        for name in self._label_names:
+            if name == "db_system":
+                db_system = event.db_system
+                if db_system is None:
+                    db_system = resolve_db_system(event.adapter)
+                values.append(db_system)
+            elif name == "driver":
+                values.append(event.driver)
+            elif name == "operation":
+                values.append(event.operation or "EXECUTE")
+            elif name == "adapter":
+                values.append(event.adapter)
+            elif name == "bind_key":
+                values.append(event.bind_key or "default")
+            else:
+                value = payload.get(name)
+                values.append("" if value is None else str(value))
+        return tuple(values)
+
+
+def enable_metrics(
+    *,
+    base_config: ObservabilityConfig | None = None,
+    namespace: str = "sqlspec",
+    subsystem: str = "driver",
+    registry: Any | None = None,
+    label_names: Iterable[str] = ("db_system", "operation"),
+    duration_buckets: tuple[float, ...] | None = None,
+) -> ObservabilityConfig:
+    """Attach a Prometheus-backed statement observer to the provided config."""
+
+    ensure_prometheus()
+
+    observer = PrometheusStatementObserver(
+        namespace=namespace,
+        subsystem=subsystem,
+        registry=registry,
+        label_names=label_names,
+        duration_buckets=duration_buckets,
+    )
+
+    config = base_config.copy() if base_config else ObservabilityConfig()
+    existing: list[StatementObserver] = list(config.statement_observers or ())
+    existing.append(observer)
+    config.statement_observers = tuple(existing)
+    return config

sqlspec/extensions/starlette/__init__.py

@@ -0,0 +1,19 @@
+"""Starlette extension for SQLSpec.
+
+Provides middleware-based session management, automatic transaction handling,
+and connection pooling lifecycle management for Starlette applications.
+"""
+
+from sqlspec.extensions.starlette._state import SQLSpecConfigState
+from sqlspec.extensions.starlette._utils import get_connection_from_request, get_or_create_session
+from sqlspec.extensions.starlette.extension import SQLSpecPlugin
+from sqlspec.extensions.starlette.middleware import SQLSpecAutocommitMiddleware, SQLSpecManualMiddleware
+
+__all__ = (
+    "SQLSpecAutocommitMiddleware",
+    "SQLSpecConfigState",
+    "SQLSpecManualMiddleware",
+    "SQLSpecPlugin",
+    "get_connection_from_request",
+    "get_or_create_session",
+)

sqlspec/extensions/starlette/_state.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from sqlspec.config import DatabaseConfigProtocol
+
+__all__ = ("CommitMode", "SQLSpecConfigState")
+
+CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
+
+
+@dataclass
+class SQLSpecConfigState:
+    """Internal state for each database configuration.
+
+    Tracks all configuration parameters needed for middleware and session management.
+    """
+
+    config: "DatabaseConfigProtocol[Any, Any, Any]"
+    connection_key: str
+    pool_key: str
+    session_key: str
+    commit_mode: CommitMode
+    extra_commit_statuses: "set[int] | None"
+    extra_rollback_statuses: "set[int] | None"
+    disable_di: bool
+    enable_correlation_middleware: bool = False
+    correlation_header: str = "x-request-id"
+    correlation_headers: "tuple[str, ...] | None" = None
+    auto_trace_headers: bool = True

sqlspec/extensions/starlette/_utils.py

@@ -0,0 +1,96 @@
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+
+    from sqlspec.extensions.starlette._state import SQLSpecConfigState
+
+__all__ = (
+    "get_connection_from_request",
+    "get_or_create_session",
+    "get_state_value",
+    "has_state_value",
+    "pop_state_value",
+    "set_state_value",
+)
+
+_MISSING = object()
+
+
+def _get_state_dict(state: Any) -> dict[str, Any]:
+    """Return the underlying state dictionary."""
+    try:
+        return cast("dict[str, Any]", object.__getattribute__(state, "_state"))
+    except AttributeError:
+        return cast("dict[str, Any]", state.__dict__)
+
+
+def get_state_value(state: Any, key: str, default: Any = _MISSING) -> Any:
+    """Get a value from a Starlette state object."""
+    data = _get_state_dict(state)
+    if default is _MISSING:
+        try:
+            return data[key]
+        except KeyError as exc:
+            msg = f"'{state.__class__.__name__}' object has no attribute '{key}'"
+            raise AttributeError(msg) from exc
+    return data.get(key, default)
+
+
+def set_state_value(state: Any, key: str, value: Any) -> None:
+    """Set a value on a Starlette state object."""
+    _get_state_dict(state)[key] = value
+
+
+def pop_state_value(state: Any, key: str) -> Any | None:
+    """Remove a value from a Starlette state object."""
+    return _get_state_dict(state).pop(key, None)
+
+
+def has_state_value(state: Any, key: str) -> bool:
+    """Check if a Starlette state object has a stored value."""
+    return key in _get_state_dict(state)
+
+
+def get_connection_from_request(request: "Request", config_state: "SQLSpecConfigState") -> Any:
+    """Get database connection from request state.
+
+    Args:
+        request: Starlette request instance.
+        config_state: Configuration state for the database.
+
+    Returns:
+        Database connection object.
+    """
+    return get_state_value(request.state, config_state.connection_key)
+
+
+def get_or_create_session(request: "Request", config_state: "SQLSpecConfigState") -> Any:
+    """Get or create database session for request.
+
+    Sessions are cached per request to ensure the same session instance
+    is returned for multiple calls within the same request.
+
+    Args:
+        request: Starlette request instance.
+        config_state: Configuration state for the database.
+
+    Returns:
+        Database session (driver instance).
+    """
+    session_instance_key = f"{config_state.session_key}_instance"
+
+    existing_session = get_state_value(request.state, session_instance_key, None)
+    if existing_session is not None:
+        return existing_session
+
+    connection = get_connection_from_request(request, config_state)
+
+    session = config_state.config.driver_type(
+        connection=connection,
+        statement_config=config_state.config.statement_config,
+        driver_features=config_state.config.driver_features,
+    )
+
+    set_state_value(request.state, session_instance_key, session)
+    return session