sqlspec 0.26.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/extensions/litestar/config.py

@@ -1,276 +1,69 @@
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, cast
+"""Configuration types for Litestar session store extension."""
 
-from sqlspec.exceptions import ImproperConfigurationError
-from sqlspec.extensions.litestar._utils import get_sqlspec_scope_state, set_sqlspec_scope_state
-from sqlspec.extensions.litestar.handlers import (
-    autocommit_handler_maker,
-    connection_provider_maker,
-    lifespan_handler_maker,
-    manual_handler_maker,
-    pool_provider_maker,
-    session_provider_maker,
-)
+from typing_extensions import NotRequired, TypedDict
 
-if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Awaitable
-    from contextlib import AbstractAsyncContextManager, AbstractContextManager
+__all__ = ("LitestarConfig",)
 
-    from litestar import Litestar
-    from litestar.datastructures.state import State
-    from litestar.types import BeforeMessageSendHookHandler, Scope
 
-    from sqlspec.config import AsyncConfigT, DriverT, SyncConfigT
-    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
-    from sqlspec.typing import ConnectionT, PoolT
+class LitestarConfig(TypedDict):
+    """Configuration options for Litestar session store extension.
 
+    All fields are optional with sensible defaults. Use in extension_config["litestar"]:
 
-CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
-DEFAULT_COMMIT_MODE: CommitMode = "manual"
-DEFAULT_CONNECTION_KEY = "db_connection"
-DEFAULT_POOL_KEY = "db_pool"
-DEFAULT_SESSION_KEY = "db_session"
+    Example:
+        from sqlspec.adapters.oracledb import OracleAsyncConfig
 
-__all__ = (
-    "DEFAULT_COMMIT_MODE",
-    "DEFAULT_CONNECTION_KEY",
-    "DEFAULT_POOL_KEY",
-    "DEFAULT_SESSION_KEY",
-    "AsyncDatabaseConfig",
-    "CommitMode",
-    "DatabaseConfig",
-    "SyncDatabaseConfig",
-)
-
-
-@dataclass
-class DatabaseConfig:
-    config: "Union[SyncConfigT, AsyncConfigT]" = field()  # type: ignore[valid-type]   # pyright: ignore[reportGeneralTypeIssues]
-    connection_key: str = field(default=DEFAULT_CONNECTION_KEY)
-    pool_key: str = field(default=DEFAULT_POOL_KEY)
-    session_key: str = field(default=DEFAULT_SESSION_KEY)
-    commit_mode: "CommitMode" = field(default=DEFAULT_COMMIT_MODE)
-    extra_commit_statuses: "Optional[set[int]]" = field(default=None)
-    extra_rollback_statuses: "Optional[set[int]]" = field(default=None)
-    enable_correlation_middleware: bool = field(default=True)
-    connection_provider: "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]" = field(  # pyright: ignore[reportGeneralTypeIssues]
-        init=False, repr=False, hash=False
-    )
-    pool_provider: "Callable[[State,Scope], Awaitable[PoolT]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
-    session_provider: "Callable[[Any], AsyncGenerator[DriverT, None]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
-    before_send_handler: "BeforeMessageSendHookHandler" = field(init=False, repr=False, hash=False)
-    lifespan_handler: "Callable[[Litestar], AbstractAsyncContextManager[None]]" = field(
-        init=False, repr=False, hash=False
-    )
-    annotation: "type[Union[SyncConfigT, AsyncConfigT]]" = field(init=False, repr=False, hash=False)  # type: ignore[valid-type]   # pyright: ignore[reportGeneralTypeIssues]
-
-    def __post_init__(self) -> None:
-        if not self.config.supports_connection_pooling and self.pool_key == DEFAULT_POOL_KEY:  # type: ignore[union-attr,unused-ignore]
-            self.pool_key = f"_{self.pool_key}_{id(self.config)}"
-        if self.commit_mode == "manual":
-            self.before_send_handler = manual_handler_maker(connection_scope_key=self.connection_key)
-        elif self.commit_mode == "autocommit":
-            self.before_send_handler = autocommit_handler_maker(
-                commit_on_redirect=False,
-                extra_commit_statuses=self.extra_commit_statuses,
-                extra_rollback_statuses=self.extra_rollback_statuses,
-                connection_scope_key=self.connection_key,
-            )
-        elif self.commit_mode == "autocommit_include_redirect":
-            self.before_send_handler = autocommit_handler_maker(
-                commit_on_redirect=True,
-                extra_commit_statuses=self.extra_commit_statuses,
-                extra_rollback_statuses=self.extra_rollback_statuses,
-                connection_scope_key=self.connection_key,
-            )
-        else:
-            msg = f"Invalid commit mode: {self.commit_mode}"
-            raise ImproperConfigurationError(detail=msg)
-        self.lifespan_handler = lifespan_handler_maker(config=self.config, pool_key=self.pool_key)
-        self.connection_provider = connection_provider_maker(
-            connection_key=self.connection_key, pool_key=self.pool_key, config=self.config
-        )
-        self.pool_provider = pool_provider_maker(config=self.config, pool_key=self.pool_key)
-        self.session_provider = session_provider_maker(
-            config=self.config, connection_dependency_key=self.connection_key
+        config = OracleAsyncConfig(
+            pool_config={"dsn": "oracle://localhost/XEPDB1"},
+            extension_config={
+                "litestar": {
+                    "session_table": "my_sessions",
+                    "in_memory": True
+                }
+            }
         )
 
-    def get_request_session(
-        self, state: "State", scope: "Scope"
-    ) -> "Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]":
-        """Get a session instance from the current request.
-
-        This method provides access to the database session that has been added to the request
-        scope, similar to Advanced Alchemy's provide_session method. It first looks for an
-        existing session in the request scope state, and if not found, creates a new one using
-        the connection from the scope.
-
-        Args:
-            state: The Litestar application State object.
-            scope: The ASGI scope containing the request context.
-
-        Returns:
-            A driver session instance.
-
-        Raises:
-            ImproperConfigurationError: If no connection is available in the scope.
-        """
-        # Create a unique scope key for sessions to avoid conflicts
-        session_scope_key = f"{self.session_key}_instance"
-
-        # Try to get existing session from scope
-        session = get_sqlspec_scope_state(scope, session_scope_key)
-        if session is not None:
-            return cast("Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]", session)
-
-        # Get connection from scope state
-        connection = get_sqlspec_scope_state(scope, self.connection_key)
-        if connection is None:
-            msg = f"No database connection found in scope for key '{self.connection_key}'. "
-            msg += "Ensure the connection dependency is properly configured and available."
-            raise ImproperConfigurationError(detail=msg)
-
-        # Create new session using the connection
-        # Access driver_type which is available on all config types
-        session = self.config.driver_type(connection=connection)  # type: ignore[union-attr]
-
-        # Store session in scope for future use
-        set_sqlspec_scope_state(scope, session_scope_key, session)
-
-        return cast("Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]", session)
-
-    def get_request_connection(self, state: "State", scope: "Scope") -> "Any":
-        """Get a connection instance from the current request.
-
-        This method provides access to the database connection that has been added to the request
-        scope. This is useful in guards, middleware, or other contexts where you need direct
-        access to the connection that's been established for the current request.
-
-        Args:
-            state: The Litestar application State object.
-            scope: The ASGI scope containing the request context.
-
-        Returns:
-            A database connection instance.
-
-        Raises:
-            ImproperConfigurationError: If no connection is available in the scope.
-        """
-        connection = get_sqlspec_scope_state(scope, self.connection_key)
-        if connection is None:
-            msg = f"No database connection found in scope for key '{self.connection_key}'. "
-            msg += "Ensure the connection dependency is properly configured and available."
-            raise ImproperConfigurationError(detail=msg)
-
-        return cast("Any", connection)
-
-
-# Add passthrough methods to both specialized classes for convenience
-class SyncDatabaseConfig(DatabaseConfig):
-    """Sync-specific DatabaseConfig with better typing for get_request_session."""
-
-    def get_request_session(self, state: "State", scope: "Scope") -> "SyncDriverAdapterBase":
-        """Get a sync session instance from the current request.
-
-        This method provides access to the database session that has been added to the request
-        scope, similar to Advanced Alchemy's provide_session method. It first looks for an
-        existing session in the request scope state, and if not found, creates a new one using
-        the connection from the scope.
-
-        Args:
-            state: The Litestar application State object.
-            scope: The ASGI scope containing the request context.
-
-        Returns:
-            A sync driver session instance.
-        """
-        session = super().get_request_session(state, scope)
-        return cast("SyncDriverAdapterBase", session)
-
-    def provide_session(self) -> "AbstractContextManager[SyncDriverAdapterBase]":
-        """Provide a database session context manager.
-
-        This is a passthrough to the underlying config's provide_session method
-        for convenient access to database sessions.
-
-        Returns:
-            Context manager that yields a sync driver session.
-        """
-        return self.config.provide_session()  # type: ignore[union-attr,no-any-return]
-
-    def provide_connection(self) -> "AbstractContextManager[Any]":
-        """Provide a database connection context manager.
-
-        This is a passthrough to the underlying config's provide_connection method
-        for convenient access to database connections.
-
-        Returns:
-            Context manager that yields a sync database connection.
-        """
-        return self.config.provide_connection()  # type: ignore[union-attr,no-any-return]
-
-    def create_connection(self) -> "Any":
-        """Create and return a new database connection.
-
-        This is a passthrough to the underlying config's create_connection method
-        for direct connection creation without context management.
-
-        Returns:
-            A new sync database connection.
-        """
-        return self.config.create_connection()  # type: ignore[union-attr]
-
-
-class AsyncDatabaseConfig(DatabaseConfig):
-    """Async-specific DatabaseConfig with better typing for get_request_session."""
-
-    def get_request_session(self, state: "State", scope: "Scope") -> "AsyncDriverAdapterBase":
-        """Get an async session instance from the current request.
-
-        This method provides access to the database session that has been added to the request
-        scope, similar to Advanced Alchemy's provide_session method. It first looks for an
-        existing session in the request scope state, and if not found, creates a new one using
-        the connection from the scope.
-
-        Args:
-            state: The Litestar application State object.
-            scope: The ASGI scope containing the request context.
-
-        Returns:
-            An async driver session instance.
-        """
-        session = super().get_request_session(state, scope)
-        return cast("AsyncDriverAdapterBase", session)
-
-    def provide_session(self) -> "AbstractAsyncContextManager[AsyncDriverAdapterBase]":
-        """Provide a database session context manager.
-
-        This is a passthrough to the underlying config's provide_session method
-        for convenient access to database sessions.
-
-        Returns:
-            Context manager that yields an async driver session.
-        """
-        return self.config.provide_session()  # type: ignore[union-attr,no-any-return]
-
-    def provide_connection(self) -> "AbstractAsyncContextManager[Any]":
-        """Provide a database connection context manager.
-
-        This is a passthrough to the underlying config's provide_connection method
-        for convenient access to database connections.
-
-        Returns:
-            Context manager that yields an async database connection.
-        """
-        return self.config.provide_connection()  # type: ignore[union-attr,no-any-return]
-
-    async def create_connection(self) -> "Any":
-        """Create and return a new database connection.
-
-        This is a passthrough to the underlying config's create_connection method
-        for direct connection creation without context management.
+    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 INMEMORY clause for Oracle Database,
+    which stores table data in columnar format in memory for faster query performance.
+
+    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(
+                pool_config={"dsn": "oracle://..."},
+                extension_config={
+                    "litestar": {
+                        "in_memory": True
+                    }
+                }
+            )
 
-        Returns:
-            A new async database connection.
-        """
-        return await self.config.create_connection()  # type: ignore[union-attr]
+    Notes:
+        - Improves query performance for session lookups (10-100x faster)
+        - Tables created with INMEMORY clause
+        - Requires Oracle Database In-Memory option license
+        - Ignored by non-Oracle adapters
+    """

sqlspec/extensions/litestar/handlers.py

@@ -1,10 +1,11 @@
 import contextlib
 import inspect
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from contextlib import AbstractAsyncContextManager
-from typing import TYPE_CHECKING, Any, Callable, Optional, cast
+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 (
@@ -37,11 +38,14 @@ __all__ = (
 )
 
 
-def manual_handler_maker(connection_scope_key: str) -> "Callable[[Message, Scope], Coroutine[Any, Any, None]]":
+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.
@@ -56,7 +60,9 @@ def manual_handler_maker(connection_scope_key: str) -> "Callable[[Message, Scope
         """
         connection = get_sqlspec_scope_state(scope, connection_scope_key)
         if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
-            if hasattr(connection, "close") and callable(connection.close):
+            if is_async:
+                await connection.close()
+            else:
                 await ensure_async_(connection.close)()
             delete_sqlspec_scope_state(scope, connection_scope_key)
 
@@ -65,14 +71,16 @@ def manual_handler_maker(connection_scope_key: str) -> "Callable[[Message, Scope
 
 def autocommit_handler_maker(
     connection_scope_key: str,
+    is_async: bool = False,
     commit_on_redirect: bool = False,
-    extra_commit_statuses: "Optional[set[int]]" = None,
-    extra_rollback_statuses: "Optional[set[int]]" = None,
+    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.
@@ -108,13 +116,19 @@ def autocommit_handler_maker(
                 if (message["status"] in commit_range or message["status"] in extra_commit_statuses) and message[
                     "status"
                 ] not in extra_rollback_statuses:
-                    if hasattr(connection, "commit") and callable(connection.commit):
+                    if is_async:
+                        await connection.commit()
+                    else:
                         await ensure_async_(connection.commit)()
-                elif hasattr(connection, "rollback") and callable(connection.rollback):
+                elif is_async:
+                    await connection.rollback()
+                else:
                     await ensure_async_(connection.rollback)()
         finally:
             if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
-                if hasattr(connection, "close") and callable(connection.close):
+                if is_async:
+                    await connection.close()
+                else:
                     await ensure_async_(connection.close)()
                 delete_sqlspec_scope_state(scope, connection_scope_key)
 
@@ -144,14 +158,23 @@ def lifespan_handler_maker(
         Yields:
             Control to application during pool lifetime.
         """
-        db_pool = await ensure_async_(config.create_pool)()
+        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:
-                await ensure_async_(config.close_pool)()
+                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)
@@ -215,11 +238,11 @@ def connection_provider_maker(
             msg = f"Database pool with key '{pool_key}' not found. Cannot create a connection."
             raise ImproperConfigurationError(msg)
 
-        connection_cm = config.provide_connection(db_pool)
+        connection_cm: Any = config.provide_connection(db_pool)
 
         if not isinstance(connection_cm, AbstractAsyncContextManager):
             conn_instance: ConnectionT
-            if hasattr(connection_cm, "__await__"):
+            if inspect.isawaitable(connection_cm):
                 conn_instance = await cast("Awaitable[ConnectionT]", connection_cm)
             else:
                 conn_instance = cast("ConnectionT", connection_cm)
@@ -252,12 +275,18 @@ def session_provider_maker(
     """
 
     async def provide_session(*args: Any, **kwargs: Any) -> "AsyncGenerator[DriverT, None]":
-        yield cast("DriverT", config.driver_type(connection=args[0] if args else kwargs.get(connection_dependency_key)))  # pyright: ignore
+        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
 
-    from litestar.params import Dependency
-
     db_conn_param = inspect.Parameter(
         name=connection_dependency_key,
         kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -272,7 +301,7 @@ def session_provider_maker(
 
     provide_session.__signature__ = provider_signature  # type: ignore[attr-defined]
 
-    if not hasattr(provide_session, "__annotations__") or provide_session.__annotations__ is None:
+    if provide_session.__annotations__ is None:
         provide_session.__annotations__ = {}
 
     provide_session.__annotations__[connection_dependency_key] = conn_type_annotation

sqlspec/extensions/litestar/migrations/0001_create_session_table.py

@@ -0,0 +1,137 @@
+"""Create Litestar session table migration using store DDL definitions."""
+
+from typing import TYPE_CHECKING, NoReturn
+
+from sqlspec.exceptions import SQLSpecError
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.module_loader import import_string
+
+if TYPE_CHECKING:
+    from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+    from sqlspec.migrations.context import MigrationContext
+
+logger = get_logger("migrations.litestar.session")
+
+__all__ = ("down", "up")
+
+
+def _get_store_class(context: "MigrationContext | None") -> "type[BaseSQLSpecStore]":
+    """Get the appropriate store class based on the config's module path.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        Store class matching the config's adapter.
+
+    Notes:
+        Dynamically imports the store class from the config's module path.
+        For example, AsyncpgConfig at 'sqlspec.adapters.asyncpg.config'
+        maps to AsyncpgStore at 'sqlspec.adapters.asyncpg.litestar.store.AsyncpgStore'.
+    """
+    if not context or not context.config:
+        _raise_missing_config()
+
+    config_class = type(context.config)
+    config_module = config_class.__module__
+    config_name = config_class.__name__
+
+    if not config_module.startswith("sqlspec.adapters."):
+        _raise_unsupported_config(f"{config_module}.{config_name}")
+
+    adapter_name = config_module.split(".")[2]
+    store_class_name = config_name.replace("Config", "Store")
+
+    store_path = f"sqlspec.adapters.{adapter_name}.litestar.store.{store_class_name}"
+
+    try:
+        store_class: type[BaseSQLSpecStore] = import_string(store_path)
+    except ImportError as e:
+        _raise_store_import_failed(store_path, e)
+
+    return store_class
+
+
+def _raise_missing_config() -> NoReturn:
+    """Raise error when migration context has no config.
+
+    Raises:
+        SQLSpecError: Always raised.
+    """
+    msg = "Migration context must have a config to determine store class"
+    raise SQLSpecError(msg)
+
+
+def _raise_unsupported_config(config_type: str) -> NoReturn:
+    """Raise error for unsupported config type.
+
+    Args:
+        config_type: The unsupported config type name.
+
+    Raises:
+        SQLSpecError: Always raised with config type info.
+    """
+    msg = f"Unsupported config type for Litestar session migration: {config_type}"
+    raise SQLSpecError(msg)
+
+
+def _raise_store_import_failed(store_path: str, error: ImportError) -> NoReturn:
+    """Raise error when store class import fails.
+
+    Args:
+        store_path: The import path that failed.
+        error: The original import error.
+
+    Raises:
+        SQLSpecError: Always raised with import details.
+    """
+    msg = f"Failed to import Litestar store class from {store_path}: {error}"
+    raise SQLSpecError(msg) from error
+
+
+async def up(context: "MigrationContext | None" = None) -> "list[str]":
+    """Create the litestar session table using store DDL definitions.
+
+    This migration delegates to the appropriate store class to generate
+    dialect-specific DDL. The store classes contain the single source of
+    truth for session table schemas.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        List of SQL statements to execute for upgrade.
+
+    Notes:
+        Table configuration is read from context.config.extension_config["litestar"].
+    """
+    store_class = _get_store_class(context)
+    if context is None or context.config is None:
+        _raise_missing_config()
+    store = store_class(config=context.config)
+
+    return [store._get_create_table_sql()]  # pyright: ignore[reportPrivateUsage]
+
+
+async def down(context: "MigrationContext | None" = None) -> "list[str]":
+    """Drop the litestar session table using store DDL definitions.
+
+    This migration delegates to the appropriate store class to generate
+    dialect-specific DROP statements. The store classes contain the single
+    source of truth for session table schemas.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        List of SQL statements to execute for downgrade.
+
+    Notes:
+        Table configuration is read from context.config.extension_config["litestar"].
+    """
+    store_class = _get_store_class(context)
+    if context is None or context.config is None:
+        _raise_missing_config()
+    store = store_class(config=context.config)
+
+    return store._get_drop_table_sql()  # pyright: ignore[reportPrivateUsage]

sqlspec/extensions/litestar/migrations/__init__.py

@@ -0,0 +1,3 @@
+"""Litestar extension migrations for session table creation."""
+
+__all__ = ()