sqlspec 0.10.1__py3-none-any.whl → 0.11.1__py3-none-any.whl

sqlspec/extensions/litestar/handlers.py

@@ -1,4 +1,8 @@
+# ruff: noqa: PLC2801
 import contextlib
+import inspect
+from collections.abc import AsyncGenerator
+from contextlib import AbstractAsyncContextManager
 from typing import TYPE_CHECKING, Any, Callable, Optional, cast
 
 from litestar.constants import HTTP_DISCONNECT, HTTP_RESPONSE_START, WEBSOCKET_CLOSE, WEBSOCKET_DISCONNECT
@@ -9,11 +13,10 @@ from sqlspec.extensions.litestar._utils import (
     get_sqlspec_scope_state,
     set_sqlspec_scope_state,
 )
-from sqlspec.utils.sync_tools import maybe_async_
+from sqlspec.utils.sync_tools import ensure_async_
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Awaitable, Coroutine
-    from contextlib import AbstractAsyncContextManager
+    from collections.abc import Awaitable, Coroutine
 
     from litestar import Litestar
     from litestar.datastructures.state import State
@@ -26,28 +29,38 @@ if TYPE_CHECKING:
 SESSION_TERMINUS_ASGI_EVENTS = {HTTP_RESPONSE_START, HTTP_DISCONNECT, WEBSOCKET_DISCONNECT, WEBSOCKET_CLOSE}
 """ASGI events that terminate a session scope."""
 
+__all__ = (
+    "SESSION_TERMINUS_ASGI_EVENTS",
+    "autocommit_handler_maker",
+    "connection_provider_maker",
+    "lifespan_handler_maker",
+    "manual_handler_maker",
+    "pool_provider_maker",
+    "session_provider_maker",
+)
+
 
 def manual_handler_maker(connection_scope_key: str) -> "Callable[[Message, Scope], Coroutine[Any, Any, None]]":
-    """Set up the handler to issue a transaction commit or rollback based on specified status codes
+    """Set up the handler to close the connection.
+
     Args:
-        connection_scope_key: The key to use within the application state
+        connection_scope_key: The key used to store the connection in the ASGI scope.
 
     Returns:
-        The handler callable
+        The handler callable.
     """
 
     async def handler(message: "Message", scope: "Scope") -> None:
-        """Handle commit/rollback, closing and cleaning up sessions before sending.
+        """Handle closing and cleaning up connections before sending the response.
 
         Args:
-            message: ASGI-``Message``
-            scope: An ASGI-``Scope``
-
+            message: ASGI Message.
+            scope: ASGI Scope.
         """
         connection = get_sqlspec_scope_state(scope, connection_scope_key)
         if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
-            with contextlib.suppress(Exception):
-                await maybe_async_(connection.close)()
+            if hasattr(connection, "close") and callable(connection.close):
+                await ensure_async_(connection.close)()
             delete_sqlspec_scope_state(scope, connection_scope_key)
 
     return handler
@@ -59,18 +72,19 @@ def autocommit_handler_maker(
     extra_commit_statuses: "Optional[set[int]]" = None,
     extra_rollback_statuses: "Optional[set[int]]" = None,
 ) -> "Callable[[Message, Scope], Coroutine[Any, Any, None]]":
-    """Set up the handler to issue a transaction commit or rollback based on specified status codes
+    """Set up the handler to issue a transaction commit or rollback based on response status codes.
+
     Args:
-        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
-        connection_scope_key: The key to use within the application state
+        connection_scope_key: The key used to store the connection in the ASGI scope.
+        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.
 
     Raises:
-        ImproperConfigurationError: If extra_commit_statuses and extra_rollback_statuses share any status codes
+        ImproperConfigurationError: If extra_commit_statuses and extra_rollback_statuses share status codes.
 
     Returns:
-        The handler callable
+        The handler callable.
     """
     if extra_commit_statuses is None:
         extra_commit_statuses = set()
@@ -85,12 +99,11 @@ def autocommit_handler_maker(
     commit_range = range(200, 400 if commit_on_redirect else 300)
 
     async def handler(message: "Message", scope: "Scope") -> None:
-        """Handle commit/rollback, closing and cleaning up sessions before sending.
+        """Handle commit/rollback, closing and cleaning up connections before sending.
 
         Args:
-            message: ASGI-``litestar.types.Message``
-            scope: An ASGI-``litestar.types.Scope``
-
+            message: ASGI Message.
+            scope: ASGI Scope.
         """
         connection = get_sqlspec_scope_state(scope, connection_scope_key)
         try:
@@ -98,13 +111,14 @@ 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:
-                    await maybe_async_(connection.commit)()
-                else:
-                    await maybe_async_(connection.rollback)()
+                    if hasattr(connection, "commit") and callable(connection.commit):
+                        await ensure_async_(connection.commit)()
+                elif hasattr(connection, "rollback") and callable(connection.rollback):
+                    await ensure_async_(connection.rollback)()
         finally:
             if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
-                with contextlib.suppress(Exception):
-                    await maybe_async_(connection.close)()
+                if hasattr(connection, "close") and callable(connection.close):
+                    await ensure_async_(connection.close)()
                 delete_sqlspec_scope_state(scope, connection_scope_key)
 
     return handler
@@ -114,101 +128,144 @@ def lifespan_handler_maker(
     config: "DatabaseConfigProtocol[Any, Any, Any]",
     pool_key: str,
 ) -> "Callable[[Litestar], AbstractAsyncContextManager[None]]":
-    """Build the lifespan handler for the database configuration.
+    """Build the lifespan handler for managing the database connection pool.
+
+    The pool is created on application startup and closed on shutdown.
 
     Args:
-        config: The database configuration.
-        pool_key: The key to use for the connection pool within Litestar.
+        config: The database configuration object.
+        pool_key: The key under which the connection pool will be stored in `app.state`.
 
     Returns:
-        The generated lifespan handler for the connection.
+        The generated lifespan handler.
     """
 
     @contextlib.asynccontextmanager
     async def lifespan_handler(app: "Litestar") -> "AsyncGenerator[None, None]":
-        db_pool = await maybe_async_(config.create_pool)()
+        """Manages the database pool lifecycle.
+
+        Args:
+            app: The Litestar application instance.
+
+        Yields:
+            The generated lifespan handler.
+        """
+        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 maybe_async_(config.close_pool)()
+                await ensure_async_(config.close_pool)()
             except Exception as e:  # noqa: BLE001
-                if app.logger:
+                if app.logger:  # pragma: no cover
                     app.logger.warning("Error closing database pool for %s. Error: %s", pool_key, e)
 
     return lifespan_handler
 
 
-def connection_provider_maker(
-    connection_key: str,
-    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]",
-) -> "Callable[[State,Scope], Awaitable[ConnectionT]]":
-    """Build the connection provider for the database configuration.
+def pool_provider_maker(
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", pool_key: str
+) -> "Callable[[State, Scope], Awaitable[PoolT]]":
+    """Build the pool provider to inject the application-level database pool.
 
     Args:
-        connection_key: The dependency key to use for the session within Litestar.
-        config: The database configuration.
+        config: The database configuration object.
+        pool_key: The key used to store the connection pool in `app.state`.
 
     Returns:
-        The generated connection provider for the connection.
+        The generated pool provider.
     """
 
-    async def provide_connection(state: "State", scope: "Scope") -> "ConnectionT":
-        connection = get_sqlspec_scope_state(scope, connection_key)
-        if connection is None:
-            connection = await maybe_async_(config.create_connection)()
-            set_sqlspec_scope_state(scope, connection_key, connection)
-        return cast("ConnectionT", connection)
+    async def provide_pool(state: "State", scope: "Scope") -> "PoolT":
+        """Provides the database pool from `app.state`.
 
-    return provide_connection
+        Args:
+            state: The Litestar application State object.
+            scope: The ASGI scope (unused for app-level pool).
 
 
-def pool_provider_maker(
-    pool_key: str,
-    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]",
-) -> "Callable[[State,Scope], Awaitable[PoolT]]":
-    """Build the pool provider for the database configuration.
+        Returns:
+            The database connection pool.
 
-    Args:
-        pool_key: The dependency key to use for the pool within Litestar.
-        config: The database configuration.
+        Raises:
+            ImproperConfigurationError: If the pool is not found in `app.state`.
+        """
+        # The pool is stored in app.state by the lifespan handler.
+        # state.get(key) accesses app.state[key]
+        db_pool = state.get(pool_key)
+        if db_pool is None:
+            # This case should ideally not happen if the lifespan handler ran correctly.
+            msg = (
+                f"Database pool with key '{pool_key}' not found in application state. "
+                "Ensure the SQLSpec lifespan handler is correctly configured and has run."
+            )
+            raise ImproperConfigurationError(msg)
+        return cast("PoolT", db_pool)
 
-    Returns:
-        The generated connection pool for the database.
-    """
+    return provide_pool
 
-    async def provide_pool(state: "State", scope: "Scope") -> "PoolT":
-        pool = get_sqlspec_scope_state(scope, pool_key)
-        if pool is None:
-            pool = await maybe_async_(config.create_pool)()
-            set_sqlspec_scope_state(scope, pool_key, pool)
-        return cast("PoolT", pool)
 
-    return provide_pool
+def connection_provider_maker(
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]",
+    pool_key: str,
+    connection_key: str,
+) -> "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]":
+    async def provide_connection(state: "State", scope: "Scope") -> "AsyncGenerator[ConnectionT, None]":
+        db_pool = state.get(pool_key)
+        if db_pool is None:
+            msg = f"Database pool with key '{pool_key}' not found. Cannot create a connection."
+            raise ImproperConfigurationError(msg)
+
+        connection_cm = config.provide_connection(db_pool)
+
+        if not isinstance(connection_cm, AbstractAsyncContextManager):
+            conn_instance: ConnectionT
+            if hasattr(connection_cm, "__await__"):
+                conn_instance = await cast("Awaitable[ConnectionT]", connection_cm)
+            else:
+                conn_instance = cast("ConnectionT", connection_cm)
+            set_sqlspec_scope_state(scope, connection_key, conn_instance)
+            yield conn_instance
+            return
+
+        entered_connection: Optional[ConnectionT] = None
+        try:
+            entered_connection = await connection_cm.__aenter__()
+            set_sqlspec_scope_state(scope, connection_key, entered_connection)
+            yield entered_connection
+        finally:
+            if entered_connection is not None:
+                await connection_cm.__aexit__(None, None, None)
+            delete_sqlspec_scope_state(scope, connection_key)  # Optional: clear from scope
+
+    return provide_connection
 
 
 def session_provider_maker(
-    session_key: str,
-    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]",
-) -> "Callable[[State,Scope], Awaitable[DriverT]]":
-    """Build the session provider for the database configuration.
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", connection_dependency_key: str
+) -> "Callable[[Any], AsyncGenerator[DriverT, None]]":
+    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
 
-    Args:
-        session_key: The dependency key to use for the session within Litestar.
-        config: The database configuration.
+    conn_type_annotation = config.connection_type
 
-    Returns:
-        The generated session provider for the database.
-    """
+    db_conn_param = inspect.Parameter(
+        name=connection_dependency_key, kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=conn_type_annotation
+    )
+
+    provider_signature = inspect.Signature(
+        parameters=[db_conn_param],
+        return_annotation=AsyncGenerator[config.driver_type, None],  # type: ignore[name-defined]
+    )
+
+    provide_session.__signature__ = provider_signature  # type: ignore[attr-defined]
+
+    if not hasattr(provide_session, "__annotations__") or provide_session.__annotations__ is None:
+        provide_session.__annotations__ = {}
 
-    async def provide_session(state: "State", scope: "Scope") -> "DriverT":
-        session = get_sqlspec_scope_state(scope, session_key)
-        if session is None:
-            connection = await maybe_async_(config.create_connection)()
-            session = config.driver_type(connection=connection)  # pyright: ignore[reportCallIssue]
-            set_sqlspec_scope_state(scope, session_key, session)
-        return cast("DriverT", session)
+    provide_session.__annotations__[connection_dependency_key] = conn_type_annotation
+    provide_session.__annotations__["return"] = config.driver_type
 
     return provide_session

sqlspec/extensions/litestar/plugin.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Literal, Union
+from typing import TYPE_CHECKING, Any, Union
 
 from litestar.di import Provide
 from litestar.plugins import InitPluginProtocol
@@ -19,13 +19,6 @@ if TYPE_CHECKING:
     from litestar.config.app import AppConfig
 
 
-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"
-
-
 class SQLSpec(InitPluginProtocol, SQLSpecBase):
     """SQLSpec plugin."""
 
@@ -70,6 +63,13 @@ class SQLSpec(InitPluginProtocol, SQLSpecBase):
             The updated :class:`AppConfig <.config.app.AppConfig>` instance.
         """
         self._validate_dependency_keys()
+
+        def store_sqlspec_in_state() -> None:
+            app_config.state.sqlspec = self
+
+        app_config.on_startup.append(store_sqlspec_in_state)
+
+        # Register types for injection
         app_config.signature_types.extend(
             [
                 SQLSpec,
@@ -82,6 +82,7 @@ class SQLSpec(InitPluginProtocol, SQLSpecBase):
                 AsyncConfigT,
             ]
         )
+
         for c in self._plugin_configs:
             c.annotation = self.add_config(c.config)
             app_config.signature_types.append(c.annotation)