sqlspec 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

sqlspec/driver/mixins/_sql_translator.py

@@ -1,3 +1,5 @@
+"""SQL translation mixin for cross-database compatibility."""
+
 from typing import Final, NoReturn, Optional
 
 from mypy_extensions import trait
@@ -33,8 +35,7 @@ class SQLTranslatorMixin:
         Returns:
             SQL string in target dialect
 
-        Raises:
-            SQLConversionError: If parsing or conversion fails
+
         """
 
         parsed_expression: Optional[exp.Expression] = None
@@ -53,7 +54,15 @@ class SQLTranslatorMixin:
         return self._generate_sql_safely(parsed_expression, target_dialect, pretty)
 
     def _parse_statement_safely(self, statement: "Statement") -> "exp.Expression":
-        """Parse statement with copy=False optimization and proper error handling."""
+        """Parse statement with error handling.
+
+        Args:
+            statement: SQL statement to parse
+
+        Returns:
+            Parsed expression
+
+        """
         try:
             sql_string = str(statement)
 
@@ -62,23 +71,52 @@ class SQLTranslatorMixin:
             self._raise_parse_error(e)
 
     def _generate_sql_safely(self, expression: "exp.Expression", dialect: DialectType, pretty: bool) -> str:
-        """Generate SQL with proper error handling."""
+        """Generate SQL with error handling.
+
+        Args:
+            expression: Parsed expression to convert
+            dialect: Target SQL dialect
+            pretty: Whether to format the output SQL
+
+        Returns:
+            Generated SQL string
+
+        """
         try:
             return expression.sql(dialect=dialect, pretty=pretty)
         except Exception as e:
             self._raise_conversion_error(dialect, e)
 
     def _raise_statement_parse_error(self) -> NoReturn:
-        """Raise error for unparsable statements."""
+        """Raise error for unparsable statements.
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         msg = "Statement could not be parsed"
         raise SQLConversionError(msg)
 
     def _raise_parse_error(self, e: Exception) -> NoReturn:
-        """Raise error for parsing failures."""
+        """Raise error for parsing failures.
+
+        Args:
+            e: Original exception that caused the failure
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         error_msg = f"Failed to parse SQL statement: {e!s}"
         raise SQLConversionError(error_msg) from e
 
     def _raise_conversion_error(self, dialect: DialectType, e: Exception) -> NoReturn:
-        """Raise error for conversion failures."""
+        """Raise error for conversion failures.
+
+        Args:
+            dialect: Target dialect that caused the failure
+            e: Original exception that caused the failure
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         error_msg = f"Failed to convert SQL expression to {dialect}: {e!s}"
         raise SQLConversionError(error_msg) from e

sqlspec/extensions/aiosql/adapter.py

@@ -2,7 +2,7 @@
 
 This module provides adapter classes that implement the aiosql adapter protocols
 while using SQLSpec drivers under the hood. This enables users to load SQL queries
-from files using aiosql while leveraging all of SQLSpec's advanced features.
+from files using aiosql while using SQLSpec's features for execution and type mapping.
 """
 
 import logging

sqlspec/extensions/litestar/_utils.py

@@ -12,7 +12,7 @@ def get_sqlspec_scope_state(scope: "Scope", key: str, default: Any = None, pop:
     """Get an internal value from connection scope state.
 
     Note:
-        If called with a default value, this method behaves like to `dict.set_default()`, both setting the key in the
+        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

sqlspec/extensions/litestar/config.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, cast
 
 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,
@@ -13,13 +14,14 @@ from sqlspec.extensions.litestar.handlers import (
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable
-    from contextlib import AbstractAsyncContextManager
+    from contextlib import AbstractAsyncContextManager, AbstractContextManager
 
     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
 
 
@@ -34,8 +36,10 @@ __all__ = (
     "DEFAULT_CONNECTION_KEY",
     "DEFAULT_POOL_KEY",
     "DEFAULT_SESSION_KEY",
+    "AsyncDatabaseConfig",
     "CommitMode",
     "DatabaseConfig",
+    "SyncDatabaseConfig",
 )
 
 
@@ -90,3 +94,183 @@ class DatabaseConfig:
         self.session_provider = session_provider_maker(
             config=self.config, connection_dependency_key=self.connection_key
         )
+
+    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.
+
+        Returns:
+            A new async database connection.
+        """
+        return await self.config.create_connection()  # type: ignore[union-attr]

sqlspec/extensions/litestar/handlers.py

@@ -199,6 +199,17 @@ def pool_provider_maker(
 def connection_provider_maker(
     config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", pool_key: str, connection_key: str
 ) -> "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]":
+    """Create provider for database connections with proper lifecycle management.
+
+    Args:
+        config: The database configuration object.
+        pool_key: The key used to retrieve the connection pool from `app.state`.
+        connection_key: The key used to store the connection in the ASGI scope.
+
+    Returns:
+        The connection provider function.
+    """
+
     async def provide_connection(state: "State", scope: "Scope") -> "AsyncGenerator[ConnectionT, None]":
         if (db_pool := state.get(pool_key)) is None:
             msg = f"Database pool with key '{pool_key}' not found. Cannot create a connection."
@@ -230,6 +241,16 @@ def connection_provider_maker(
 def session_provider_maker(
     config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", connection_dependency_key: str
 ) -> "Callable[[Any], AsyncGenerator[DriverT, None]]":
+    """Create provider for database driver sessions.
+
+    Args:
+        config: The database configuration object.
+        connection_dependency_key: The key used for connection dependency injection.
+
+    Returns:
+        The session provider function.
+    """
+
     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
 

sqlspec/extensions/litestar/plugin.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
 from litestar.di import Provide
 from litestar.plugins import CLIPlugin, InitPluginProtocol
@@ -6,14 +6,17 @@ from litestar.plugins import CLIPlugin, InitPluginProtocol
 from sqlspec.base import SQLSpec as SQLSpecBase
 from sqlspec.config import AsyncConfigT, DatabaseConfigProtocol, DriverT, SyncConfigT
 from sqlspec.exceptions import ImproperConfigurationError
-from sqlspec.extensions.litestar.config import DatabaseConfig
+from sqlspec.extensions.litestar.config import AsyncDatabaseConfig, DatabaseConfig, SyncDatabaseConfig
 from sqlspec.typing import ConnectionT, PoolT
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
     from click import Group
     from litestar.config.app import AppConfig
+    from litestar.datastructures.state import State
+    from litestar.types import Scope
 
+    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
     from sqlspec.loader import SQLFileLoader
 
 logger = get_logger("extensions.litestar")
@@ -131,11 +134,242 @@ class SQLSpec(SQLSpecBase, InitPluginProtocol, CLIPlugin):
             The annotation for the configuration.
         """
         for c in self.config:
-            if key == c.config or key in {c.annotation, c.connection_key, c.pool_key}:
+            # Check annotation only if it's been set (during on_app_init)
+            annotation_match = hasattr(c, "annotation") and key == c.annotation
+            if key == c.config or annotation_match or key in {c.connection_key, c.pool_key}:
+                if not hasattr(c, "annotation"):
+                    msg = (
+                        "Annotation not set for configuration. Ensure the plugin has been initialized with on_app_init."
+                    )
+                    raise AttributeError(msg)
                 return c.annotation
         msg = f"No configuration found for {key}"
         raise KeyError(msg)
 
+    @overload
+    def get_config(self, name: "type[SyncConfigT]") -> "SyncConfigT": ...
+
+    @overload
+    def get_config(self, name: "type[AsyncConfigT]") -> "AsyncConfigT": ...
+
+    @overload
+    def get_config(self, name: str) -> "DatabaseConfig": ...
+
+    @overload
+    def get_config(self, name: "type[SyncDatabaseConfig]") -> "SyncDatabaseConfig": ...
+
+    @overload
+    def get_config(self, name: "type[AsyncDatabaseConfig]") -> "AsyncDatabaseConfig": ...
+
+    def get_config(
+        self, name: "Union[type[DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]], str, Any]"
+    ) -> "Union[DatabaseConfigProtocol[ConnectionT, PoolT, DriverT], DatabaseConfig, SyncDatabaseConfig, AsyncDatabaseConfig]":
+        """Get a configuration instance by name, supporting both base behavior and Litestar extensions.
+
+        This method extends the base get_config to support Litestar-specific lookup patterns
+        while maintaining compatibility with the base class signature. It supports lookup by
+        connection key, pool key, session key, config instance, or annotation type.
+
+        Args:
+            name: The configuration identifier - can be:
+                - Type annotation (base class behavior)
+                - connection_key (e.g., "auth_db_connection")
+                - pool_key (e.g., "analytics_db_pool")
+                - session_key (e.g., "reporting_db_session")
+                - config instance
+                - annotation type
+
+        Raises:
+            KeyError: If no configuration is found for the given name.
+
+        Returns:
+            The configuration instance for the specified name.
+        """
+        # First try base class behavior for type-based lookup
+        # Only call super() if name matches the expected base class types
+        if not isinstance(name, str):
+            try:
+                return super().get_config(name)  # type: ignore[no-any-return]
+            except (KeyError, AttributeError):
+                # Fall back to Litestar-specific lookup patterns
+                pass
+
+        # Litestar-specific lookups by string keys
+        if isinstance(name, str):
+            for c in self.config:
+                if name in {c.connection_key, c.pool_key, c.session_key}:
+                    return c  # Return the DatabaseConfig wrapper for string lookups
+
+        # Lookup by config instance or annotation
+        for c in self.config:
+            annotation_match = hasattr(c, "annotation") and name == c.annotation
+            if name == c.config or annotation_match:
+                return c.config  # Return the underlying config for type-based lookups
+
+        msg = f"No database configuration found for name '{name}'. Available keys: {self._get_available_keys()}"
+        raise KeyError(msg)
+
+    def provide_request_session(
+        self,
+        key: "Union[str, SyncConfigT, AsyncConfigT, type[Union[SyncConfigT, AsyncConfigT]]]",
+        state: "State",
+        scope: "Scope",
+    ) -> "Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]":
+        """Provide a database session for the specified configuration key from request scope.
+
+        This is a convenience method that combines get_config and get_request_session
+        into a single call, similar to Advanced Alchemy's provide_session pattern.
+
+        Args:
+            key: The configuration identifier (same as get_config)
+            state: The Litestar application State object
+            scope: The ASGI scope containing the request context
+
+        Returns:
+            A driver session instance for the specified database configuration
+
+        Example:
+            >>> sqlspec_plugin = connection.app.state.sqlspec
+            >>> # Direct session access by key
+            >>> auth_session = sqlspec_plugin.provide_request_session(
+            ...     "auth_db", state, scope
+            ... )
+            >>> analytics_session = sqlspec_plugin.provide_request_session(
+            ...     "analytics_db", state, scope
+            ... )
+        """
+        # Get DatabaseConfig wrapper for Litestar methods
+        db_config = self._get_database_config(key)
+        return db_config.get_request_session(state, scope)
+
+    def provide_sync_request_session(
+        self, key: "Union[str, SyncConfigT, type[SyncConfigT]]", state: "State", scope: "Scope"
+    ) -> "SyncDriverAdapterBase":
+        """Provide a sync database session for the specified configuration key from request scope.
+
+        This method provides better type hints for sync database sessions, ensuring the returned
+        session is properly typed as SyncDriverAdapterBase for better IDE support and type safety.
+
+        Args:
+            key: The sync configuration identifier
+            state: The Litestar application State object
+            scope: The ASGI scope containing the request context
+
+        Returns:
+            A sync driver session instance for the specified database configuration
+
+        Example:
+            >>> sqlspec_plugin = connection.app.state.sqlspec
+            >>> auth_session = sqlspec_plugin.provide_sync_request_session(
+            ...     "auth_db", state, scope
+            ... )
+            >>> # auth_session is now correctly typed as SyncDriverAdapterBase
+        """
+        # Get DatabaseConfig wrapper for Litestar methods
+        db_config = self._get_database_config(key)
+        session = db_config.get_request_session(state, scope)
+        return cast("SyncDriverAdapterBase", session)
+
+    def provide_async_request_session(
+        self, key: "Union[str, AsyncConfigT, type[AsyncConfigT]]", state: "State", scope: "Scope"
+    ) -> "AsyncDriverAdapterBase":
+        """Provide an async database session for the specified configuration key from request scope.
+
+        This method provides better type hints for async database sessions, ensuring the returned
+        session is properly typed as AsyncDriverAdapterBase for better IDE support and type safety.
+
+        Args:
+            key: The async configuration identifier
+            state: The Litestar application State object
+            scope: The ASGI scope containing the request context
+
+        Returns:
+            An async driver session instance for the specified database configuration
+
+        Example:
+            >>> sqlspec_plugin = connection.app.state.sqlspec
+            >>> auth_session = sqlspec_plugin.provide_async_request_session(
+            ...     "auth_db", state, scope
+            ... )
+            >>> # auth_session is now correctly typed as AsyncDriverAdapterBase
+        """
+        # Get DatabaseConfig wrapper for Litestar methods
+        db_config = self._get_database_config(key)
+        session = db_config.get_request_session(state, scope)
+        return cast("AsyncDriverAdapterBase", session)
+
+    def provide_request_connection(
+        self,
+        key: "Union[str, SyncConfigT, AsyncConfigT, type[Union[SyncConfigT, AsyncConfigT]]]",
+        state: "State",
+        scope: "Scope",
+    ) -> Any:
+        """Provide a database connection for the specified configuration key from request scope.
+
+        This is a convenience method that combines get_config and get_request_connection
+        into a single call.
+
+        Args:
+            key: The configuration identifier (same as get_config)
+            state: The Litestar application State object
+            scope: The ASGI scope containing the request context
+
+        Returns:
+            A database connection instance for the specified database configuration
+
+        Example:
+            >>> sqlspec_plugin = connection.app.state.sqlspec
+            >>> # Direct connection access by key
+            >>> auth_conn = sqlspec_plugin.provide_request_connection(
+            ...     "auth_db", state, scope
+            ... )
+            >>> analytics_conn = sqlspec_plugin.provide_request_connection(
+            ...     "analytics_db", state, scope
+            ... )
+        """
+        # Get DatabaseConfig wrapper for Litestar methods
+        db_config = self._get_database_config(key)
+        return db_config.get_request_connection(state, scope)
+
+    def _get_database_config(
+        self, key: "Union[str, SyncConfigT, AsyncConfigT, type[Union[SyncConfigT, AsyncConfigT]]]"
+    ) -> DatabaseConfig:
+        """Get a DatabaseConfig wrapper instance by name.
+
+        This is used internally by provide_request_session and provide_request_connection
+        to get the DatabaseConfig wrapper that has the request session methods.
+
+        Args:
+            key: The configuration identifier
+
+        Returns:
+            The DatabaseConfig wrapper instance
+
+        Raises:
+            KeyError: If no configuration is found for the given key
+        """
+        # For string keys, lookup by connection/pool/session keys
+        if isinstance(key, str):
+            for c in self.config:
+                if key in {c.connection_key, c.pool_key, c.session_key}:
+                    return c
+
+        # For other keys, lookup by config instance or annotation
+        for c in self.config:
+            annotation_match = hasattr(c, "annotation") and key == c.annotation
+            if key == c.config or annotation_match:
+                return c
+
+        msg = f"No database configuration found for name '{key}'. Available keys: {self._get_available_keys()}"
+        raise KeyError(msg)
+
+    def _get_available_keys(self) -> "list[str]":
+        """Get a list of all available configuration keys for error messages."""
+        keys = []
+        for c in self.config:
+            keys.extend([c.connection_key, c.pool_key, c.session_key])
+        return keys
+
     def _validate_dependency_keys(self) -> None:
         """Validate that connection and pool keys are unique across configurations.