sqlspec 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

sqlspec/exceptions.py

@@ -1,3 +1,5 @@
+from collections.abc import Generator
+from contextlib import contextmanager
 from typing import Any, Optional
 
 __all__ = (
@@ -6,7 +8,9 @@ __all__ = (
     "MissingDependencyError",
     "MultipleResultsFoundError",
     "NotFoundError",
+    "ParameterStyleMismatchError",
     "RepositoryError",
+    "SQLParsingError",
     "SQLSpecError",
     "SerializationError",
 )
@@ -74,6 +78,20 @@ class SQLParsingError(SQLSpecError):
         super().__init__(message)
 
 
+class ParameterStyleMismatchError(SQLSpecError):
+    """Error when parameter style doesn't match SQL placeholder style.
+
+    This exception is raised when there's a mismatch between the parameter type
+    (dictionary, tuple, etc.) and the placeholder style in the SQL query
+    (named, positional, etc.).
+    """
+
+    def __init__(self, message: Optional[str] = None) -> None:
+        if message is None:
+            message = "Parameter style mismatch: dictionary parameters provided but no named placeholders found in SQL."
+        super().__init__(message)
+
+
 class ImproperConfigurationError(SQLSpecError):
     """Improper Configuration error.
 
@@ -99,3 +117,15 @@ class NotFoundError(RepositoryError):
 
 class MultipleResultsFoundError(RepositoryError):
     """A single database result was required but more than one were found."""
+
+
+@contextmanager
+def wrap_exceptions(wrap_exceptions: bool = True) -> Generator[None, None, None]:
+    try:
+        yield
+
+    except Exception as exc:
+        if wrap_exceptions is False:
+            raise
+        msg = "An error occurred during the operation."
+        raise RepositoryError(detail=msg) from exc

sqlspec/extensions/litestar/__init__.py

@@ -0,0 +1,19 @@
+from sqlspec.extensions.litestar.config import DatabaseConfig
+from sqlspec.extensions.litestar.handlers import (
+    autocommit_handler_maker,
+    connection_provider_maker,
+    lifespan_handler_maker,
+    manual_handler_maker,
+    pool_provider_maker,
+)
+from sqlspec.extensions.litestar.plugin import SQLSpec
+
+__all__ = (
+    "DatabaseConfig",
+    "SQLSpec",
+    "autocommit_handler_maker",
+    "connection_provider_maker",
+    "lifespan_handler_maker",
+    "manual_handler_maker",
+    "pool_provider_maker",
+)

sqlspec/extensions/litestar/_utils.py

@@ -0,0 +1,56 @@
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from litestar.types import Scope
+
+__all__ = (
+    "delete_sqlspec_scope_state",
+    "get_sqlspec_scope_state",
+    "set_sqlspec_scope_state",
+)
+
+_SCOPE_NAMESPACE = "_sqlspec"
+
+
+def get_sqlspec_scope_state(scope: "Scope", key: str, default: Any = None, pop: bool = False) -> Any:
+    """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
+        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
+        exist.
+
+    Args:
+        scope: The connection scope.
+        key: Key to get from internal namespace in scope state.
+        default: Default value to return.
+        pop: Boolean flag dictating whether the value should be deleted from the state.
+
+    Returns:
+        Value mapped to ``key`` in internal connection scope namespace.
+    """
+    namespace = scope.setdefault(_SCOPE_NAMESPACE, {})  # type: ignore[misc]
+    return namespace.pop(key, default) if pop else namespace.get(key, default)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
+
+
+def set_sqlspec_scope_state(scope: "Scope", key: str, value: Any) -> None:
+    """Set an internal value in connection scope state.
+
+    Args:
+        scope: The connection scope.
+        key: Key to set under internal namespace in scope state.
+        value: Value for key.
+    """
+    scope.setdefault(_SCOPE_NAMESPACE, {})[key] = value  # type: ignore[misc]
+
+
+def delete_sqlspec_scope_state(scope: "Scope", key: str) -> None:
+    """Remove an internal value from connection scope state.
+
+    Args:
+        scope: The connection scope.
+        key: Key to set under internal namespace in scope state.
+    """
+    del scope.setdefault(_SCOPE_NAMESPACE, {})[key]  # type: ignore[misc]

sqlspec/extensions/litestar/config.py

@@ -0,0 +1,87 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+
+from sqlspec.base import (
+    ConnectionT,
+    PoolT,
+)
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar.handlers import (
+    autocommit_handler_maker,
+    connection_provider_maker,
+    lifespan_handler_maker,
+    manual_handler_maker,
+    pool_provider_maker,
+    session_provider_maker,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable
+    from contextlib import AbstractAsyncContextManager
+
+    from litestar import Litestar
+    from litestar.datastructures.state import State
+    from litestar.types import BeforeMessageSendHookHandler, Scope
+
+    from sqlspec.base import (
+        AsyncConfigT,
+        ConnectionT,
+        DriverT,
+        PoolT,
+        SyncConfigT,
+    )
+
+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"
+
+
+@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)
+    connection_provider: "Callable[[State,Scope], Awaitable[ConnectionT]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
+    pool_provider: "Callable[[State,Scope], Awaitable[PoolT]]" = field(init=False, repr=False, hash=False)  # pyright: ignore[reportGeneralTypeIssues]
+    session_provider: "Callable[[State,Scope], Awaitable[DriverT]]" = 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.support_connection_pooling and self.pool_key == DEFAULT_POOL_KEY:  # type: ignore[union-attr,unused-ignore]
+            """If the database configuration does not support connection pooling, the pool key must be unique.  We just automatically generate a unique identify so it won't conflict with other configs that may get added"""
+            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}"  # type: ignore[unreachable]
+            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, config=self.config)
+        self.pool_provider = pool_provider_maker(pool_key=self.pool_key, config=self.config)
+        self.session_provider = session_provider_maker(session_key=self.session_key, config=self.config)

sqlspec/extensions/litestar/handlers.py

@@ -0,0 +1,213 @@
+import contextlib
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast
+
+from litestar.constants import HTTP_DISCONNECT, HTTP_RESPONSE_START, WEBSOCKET_CLOSE, WEBSOCKET_DISCONNECT
+
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar._utils import (
+    delete_sqlspec_scope_state,
+    get_sqlspec_scope_state,
+    set_sqlspec_scope_state,
+)
+from sqlspec.utils.sync_tools import maybe_async_
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Awaitable, Coroutine
+    from contextlib import AbstractAsyncContextManager
+
+    from litestar import Litestar
+    from litestar.datastructures.state import State
+    from litestar.types import Message, Scope
+
+    from sqlspec.base import ConnectionT, DatabaseConfigProtocol, DriverT, PoolT
+
+
+SESSION_TERMINUS_ASGI_EVENTS = {HTTP_RESPONSE_START, HTTP_DISCONNECT, WEBSOCKET_DISCONNECT, WEBSOCKET_CLOSE}
+"""ASGI events that terminate a session scope."""
+
+
+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
+    Args:
+        connection_scope_key: The key to use within the application state
+
+    Returns:
+        The handler callable
+    """
+
+    async def handler(message: "Message", scope: "Scope") -> None:
+        """Handle commit/rollback, closing and cleaning up sessions before sending.
+
+        Args:
+            message: ASGI-``Message``
+            scope: An 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)()
+            delete_sqlspec_scope_state(scope, connection_scope_key)
+
+    return handler
+
+
+def autocommit_handler_maker(
+    connection_scope_key: str,
+    commit_on_redirect: bool = False,
+    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
+    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
+
+    Raises:
+        ImproperConfigurationError: If extra_commit_statuses and extra_rollback_statuses share any status codes
+
+    Returns:
+        The handler callable
+    """
+    if extra_commit_statuses is None:
+        extra_commit_statuses = set()
+
+    if extra_rollback_statuses is None:
+        extra_rollback_statuses = set()
+
+    if len(extra_commit_statuses & extra_rollback_statuses) > 0:
+        msg = "Extra rollback statuses and commit statuses must not share any status codes"
+        raise ImproperConfigurationError(msg)
+
+    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.
+
+        Args:
+            message: ASGI-``litestar.types.Message``
+            scope: An ASGI-``litestar.types.Scope``
+
+        """
+        connection = get_sqlspec_scope_state(scope, connection_scope_key)
+        try:
+            if connection is not None and message["type"] == HTTP_RESPONSE_START:
+                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)()
+        finally:
+            if connection and message["type"] in SESSION_TERMINUS_ASGI_EVENTS:
+                with contextlib.suppress(Exception):
+                    await maybe_async_(connection.close)()
+                delete_sqlspec_scope_state(scope, connection_scope_key)
+
+    return handler
+
+
+def lifespan_handler_maker(
+    config: "DatabaseConfigProtocol[Any, Any, Any]",
+    pool_key: str,
+) -> "Callable[[Litestar], AbstractAsyncContextManager[None]]":
+    """Build the lifespan handler for the database configuration.
+
+    Args:
+        config: The database configuration.
+        pool_key: The key to use for the connection pool within Litestar.
+
+    Returns:
+        The generated lifespan handler for the connection.
+    """
+
+    @contextlib.asynccontextmanager
+    async def lifespan_handler(app: "Litestar") -> "AsyncGenerator[None, None]":
+        db_pool = await maybe_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)()
+            except Exception as e:  # noqa: BLE001
+                if app.logger:
+                    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.
+
+    Args:
+        connection_key: The dependency key to use for the session within Litestar.
+        config: The database configuration.
+
+    Returns:
+        The generated connection provider for the connection.
+    """
+
+    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)
+
+    return provide_connection
+
+
+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.
+
+    Args:
+        pool_key: The dependency key to use for the pool within Litestar.
+        config: The database configuration.
+
+    Returns:
+        The generated connection pool for the database.
+    """
+
+    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 session_provider_maker(
+    session_key: str,
+    config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]",
+) -> "Callable[[State,Scope], Awaitable[DriverT]]":
+    """Build the session provider for the database configuration.
+
+    Args:
+        session_key: The dependency key to use for the session within Litestar.
+        config: The database configuration.
+
+    Returns:
+        The generated session provider for the database.
+    """
+
+    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)
+
+    return provide_session

sqlspec/extensions/litestar/plugin.py

@@ -1,34 +1,64 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Literal, Union
 
+from litestar.di import Provide
 from litestar.plugins import InitPluginProtocol
 
+from sqlspec.base import (
+    AsyncConfigT,
+    ConnectionT,
+    DatabaseConfigProtocol,
+    PoolT,
+    SyncConfigT,
+)
+from sqlspec.base import SQLSpec as SQLSpecBase
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.litestar.config import DatabaseConfig
+
 if TYPE_CHECKING:
+    from click import Group
     from litestar.config.app import AppConfig
 
-    from sqlspec.base import ConfigManager
+
+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 SQLSpecPlugin(InitPluginProtocol):
+class SQLSpec(InitPluginProtocol, SQLSpecBase):
     """SQLSpec plugin."""
 
-    __slots__ = ("_config",)
+    __slots__ = ("_config", "_plugin_configs")
 
-    def __init__(self, config: "ConfigManager") -> None:
+    def __init__(
+        self,
+        config: Union["SyncConfigT", "AsyncConfigT", "DatabaseConfig", list["DatabaseConfig"]],
+    ) -> None:
         """Initialize ``SQLSpecPlugin``.
 
         Args:
             config: configure SQLSpec plugin for use with Litestar.
         """
-        self._config = config
+        self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
+        if isinstance(config, DatabaseConfigProtocol):
+            self._plugin_configs: list[DatabaseConfig] = [DatabaseConfig(config=config)]
+        elif isinstance(config, DatabaseConfig):
+            self._plugin_configs = [config]
+        else:
+            self._plugin_configs = config
 
     @property
-    def config(self) -> "ConfigManager":
+    def config(self) -> "list[DatabaseConfig]":  # pyright: ignore[reportInvalidTypeVarUse]
         """Return the plugin config.
 
         Returns:
             ConfigManager.
         """
-        return self._config
+        return self._plugin_configs
+
+    def on_cli_init(self, cli: "Group") -> None:
+        """Configure the CLI for use with SQLSpec."""
 
     def on_app_init(self, app_config: "AppConfig") -> "AppConfig":
         """Configure application for use with SQLSpec.
@@ -39,8 +69,72 @@ class SQLSpecPlugin(InitPluginProtocol):
         Returns:
             The updated :class:`AppConfig <.config.app.AppConfig>` instance.
         """
+        self._validate_dependency_keys()
+        app_config.signature_types.extend(
+            [
+                SQLSpec,
+                ConnectionT,
+                PoolT,
+                DatabaseConfig,
+                DatabaseConfigProtocol,
+                SyncConfigT,
+                AsyncConfigT,
+            ]
+        )
+        for c in self._plugin_configs:
+            c.annotation = self.add_config(c.config)
+            app_config.before_send.append(c.before_send_handler)
+            app_config.lifespan.append(c.lifespan_handler)  # pyright: ignore[reportUnknownMemberType]
+            app_config.dependencies.update(
+                {
+                    c.connection_key: Provide(c.connection_provider),
+                    c.pool_key: Provide(c.pool_provider),
+                    c.session_key: Provide(c.session_provider),
+                },
+            )
 
-        from sqlspec.base import ConfigManager
-
-        app_config.signature_types.append(ConfigManager)
         return app_config
+
+    def get_annotations(self) -> "list[type[Union[SyncConfigT, AsyncConfigT]]]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Return the list of annotations.
+
+        Returns:
+            List of annotations.
+        """
+        return [c.annotation for c in self.config]
+
+    def get_annotation(
+        self,
+        key: "Union[str, SyncConfigT, AsyncConfigT, type[Union[SyncConfigT, AsyncConfigT]]]",
+    ) -> "type[Union[SyncConfigT, AsyncConfigT]]":
+        """Return the annotation for the given configuration.
+
+        Args:
+            key: The configuration instance or key to lookup
+
+        Raises:
+            KeyError: If no configuration is found for the given key.
+
+        Returns:
+            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}:
+                return c.annotation
+        msg = f"No configuration found for {key}"
+        raise KeyError(msg)
+
+    def _validate_dependency_keys(self) -> None:
+        """Verify uniqueness of ``connection_key`` and ``pool_key``.
+
+        Raises:
+            ImproperConfigurationError: If session keys or pool keys are not unique.
+        """
+        connection_keys = [c.connection_key for c in self.config]
+        pool_keys = [c.pool_key for c in self.config]
+        if len(set(connection_keys)) != len(connection_keys):
+            msg = "When using multiple database configuration, each configuration must have a unique `connection_key`."
+            raise ImproperConfigurationError(detail=msg)
+        if len(set(pool_keys)) != len(pool_keys):
+            msg = "When using multiple database configuration, each configuration must have a unique `pool_key`."
+            raise ImproperConfigurationError(detail=msg)