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

sqlspec/extensions/aiosql/adapter.py

@@ -0,0 +1,461 @@
+"""AioSQL adapter implementation for SQLSpec.
+
+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.
+"""
+
+import logging
+from collections.abc import AsyncGenerator, Generator
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypeVar, Union, cast
+
+from sqlspec.core.result import SQLResult
+from sqlspec.core.statement import SQL, StatementConfig
+from sqlspec.exceptions import MissingDependencyError
+from sqlspec.typing import AIOSQL_INSTALLED
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
+
+logger = logging.getLogger("sqlspec.extensions.aiosql")
+
+__all__ = ("AiosqlAsyncAdapter", "AiosqlSyncAdapter")
+
+T = TypeVar("T")
+
+
+class AsyncCursorLike:
+    def __init__(self, result: Any) -> None:
+        self.result = result
+
+    async def fetchall(self) -> list[Any]:
+        if isinstance(self.result, SQLResult) and self.result.data is not None:
+            return list(self.result.data)
+        return []
+
+    async def fetchone(self) -> Optional[Any]:
+        rows = await self.fetchall()
+        return rows[0] if rows else None
+
+
+class CursorLike:
+    def __init__(self, result: Any) -> None:
+        self.result = result
+
+    def fetchall(self) -> list[Any]:
+        if isinstance(self.result, SQLResult) and self.result.data is not None:
+            return list(self.result.data)
+        return []
+
+    def fetchone(self) -> Optional[Any]:
+        rows = self.fetchall()
+        return rows[0] if rows else None
+
+
+def _check_aiosql_available() -> None:
+    if not AIOSQL_INSTALLED:
+        msg = "aiosql"
+        raise MissingDependencyError(msg, "aiosql")
+
+
+def _normalize_dialect(dialect: "Union[str, Any, None]") -> str:
+    """Normalize dialect name for SQLGlot compatibility.
+
+    Args:
+        dialect: Original dialect name (can be str, Dialect, type[Dialect], or None)
+
+    Returns:
+        Converted dialect name compatible with SQLGlot
+    """
+    if dialect is None:
+        return "sql"
+
+    if hasattr(dialect, "__name__"):
+        dialect_str = str(dialect.__name__).lower()  # pyright: ignore
+    elif hasattr(dialect, "name"):
+        dialect_str = str(dialect.name).lower()  # pyright: ignore
+    elif isinstance(dialect, str):
+        dialect_str = dialect.lower()
+    else:
+        dialect_str = str(dialect).lower()
+
+    dialect_mapping = {
+        "postgresql": "postgres",
+        "psycopg": "postgres",
+        "asyncpg": "postgres",
+        "psqlpy": "postgres",
+        "sqlite3": "sqlite",
+        "aiosqlite": "sqlite",
+    }
+    return dialect_mapping.get(dialect_str, dialect_str)
+
+
+class _AiosqlAdapterBase:
+    """Base adapter class providing common functionality for aiosql integration."""
+
+    def __init__(self, driver: "Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]") -> None:
+        """Initialize the base adapter.
+
+        Args:
+            driver: SQLSpec driver to use for execution.
+        """
+        _check_aiosql_available()
+        self.driver = driver
+
+    def process_sql(self, query_name: str, op_type: "Any", sql: str) -> str:
+        """Process SQL string for aiosql compatibility.
+
+        Args:
+            query_name: Name of the query
+            op_type: Operation type
+            sql: SQL string to process
+
+        Returns:
+            Processed SQL string
+        """
+        return sql
+
+    def _create_sql_object(self, sql: str, parameters: "Any" = None) -> SQL:
+        """Create SQL object with proper configuration.
+
+        Args:
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            Configured SQL object
+        """
+        return SQL(
+            sql,
+            parameters,
+            config=StatementConfig(enable_validation=False),
+            dialect=_normalize_dialect(getattr(self.driver, "dialect", "sqlite")),
+        )
+
+
+class AiosqlSyncAdapter(_AiosqlAdapterBase):
+    """Synchronous adapter that implements aiosql protocol using SQLSpec drivers.
+
+    This adapter bridges aiosql's synchronous driver protocol with SQLSpec's sync drivers,
+    enabling queries loaded by aiosql to be executed with SQLSpec drivers.
+    """
+
+    is_aio_driver: ClassVar[bool] = False
+
+    def __init__(self, driver: "SyncDriverAdapterBase") -> None:
+        """Initialize the sync adapter.
+
+        Args:
+            driver: SQLSpec sync driver to use for execution
+        """
+        super().__init__(driver)
+
+    def select(
+        self, conn: Any, query_name: str, sql: str, parameters: "Any", record_class: Optional[Any] = None
+    ) -> Generator[Any, None, None]:
+        """Execute a SELECT query and return results as generator.
+
+        Args:
+            conn: Database connection (passed through to SQLSpec driver)
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+            record_class: Deprecated - use schema_type in driver.execute instead
+
+        Yields:
+            Query result rows
+
+        Note:
+            The record_class parameter is ignored for compatibility. Use schema_type
+            in driver.execute or _sqlspec_schema_type in parameters for type mapping.
+        """
+        if record_class is not None:
+            logger.warning(
+                "record_class parameter is deprecated and ignored. "
+                "Use schema_type in driver.execute or _sqlspec_schema_type in parameters."
+            )
+
+        result = self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)
+
+        if isinstance(result, SQLResult) and result.data is not None:
+            yield from result.data
+
+    def select_one(
+        self, conn: Any, query_name: str, sql: str, parameters: "Any", record_class: Optional[Any] = None
+    ) -> Optional[dict[str, Any]]:
+        """Execute a SELECT query and return first result.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+            record_class: Deprecated - use schema_type in driver.execute instead
+
+        Returns:
+            First result row or None
+
+        Note:
+            The record_class parameter is ignored for compatibility. Use schema_type
+            in driver.execute or _sqlspec_schema_type in parameters for type mapping.
+        """
+        if record_class is not None:
+            logger.warning(
+                "record_class parameter is deprecated and ignored. "
+                "Use schema_type in driver.execute or _sqlspec_schema_type in parameters."
+            )
+
+        result = cast("SQLResult", self.driver.execute(self._create_sql_object(sql, parameters), connection=conn))
+
+        if hasattr(result, "data") and result.data and isinstance(result, SQLResult):
+            return cast("Optional[dict[str, Any]]", result.data[0])
+        return None
+
+    def select_value(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> Optional[Any]:
+        """Execute a SELECT query and return first value of first row.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            First value of first row or None
+        """
+        row = self.select_one(conn, query_name, sql, parameters)
+        if row is None:
+            return None
+
+        if isinstance(row, dict):
+            return next(iter(row.values())) if row else None
+        if hasattr(row, "__getitem__"):
+            return row[0] if len(row) > 0 else None
+        return row
+
+    @contextmanager
+    def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> Generator[Any, None, None]:
+        """Execute a SELECT query and return cursor context manager.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Yields:
+            Cursor-like object with results
+        """
+        result = self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)
+
+        yield CursorLike(result)
+
+    def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> int:
+        """Execute INSERT/UPDATE/DELETE and return affected rows.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            Number of affected rows
+        """
+        result = cast("SQLResult", self.driver.execute(self._create_sql_object(sql, parameters), connection=conn))
+
+        return result.rows_affected if hasattr(result, "rows_affected") else 0
+
+    def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> int:
+        """Execute INSERT/UPDATE/DELETE with many parameter sets.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Sequence of parameter sets
+
+        Returns:
+            Number of affected rows
+        """
+        result = cast(
+            "SQLResult", self.driver.execute_many(self._create_sql_object(sql), parameters=parameters, connection=conn)
+        )
+
+        return result.rows_affected if hasattr(result, "rows_affected") else 0
+
+    def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> Optional[Any]:
+        """Execute INSERT with RETURNING and return result.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            Returned value or None
+        """
+        return self.select_one(conn, query_name, sql, parameters)
+
+
+class AiosqlAsyncAdapter(_AiosqlAdapterBase):
+    """Asynchronous adapter that implements aiosql protocol using SQLSpec drivers.
+
+    This adapter bridges aiosql's async driver protocol with SQLSpec's async drivers,
+    enabling queries loaded by aiosql to be executed with SQLSpec async drivers.
+    """
+
+    is_aio_driver: ClassVar[bool] = True
+
+    def __init__(self, driver: "AsyncDriverAdapterBase") -> None:
+        """Initialize the async adapter.
+
+        Args:
+            driver: SQLSpec async driver to use for execution
+        """
+        super().__init__(driver)
+
+    async def select(
+        self, conn: Any, query_name: str, sql: str, parameters: "Any", record_class: Optional[Any] = None
+    ) -> list[Any]:
+        """Execute a SELECT query and return results as list.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+            record_class: Deprecated - use schema_type in driver.execute instead
+
+        Returns:
+            List of query result rows
+
+        Note:
+            The record_class parameter is ignored for compatibility. Use schema_type
+            in driver.execute or _sqlspec_schema_type in parameters for type mapping.
+        """
+        if record_class is not None:
+            logger.warning(
+                "record_class parameter is deprecated and ignored. "
+                "Use schema_type in driver.execute or _sqlspec_schema_type in parameters."
+            )
+
+        result = await self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)  # type: ignore[misc]
+
+        if hasattr(result, "data") and result.data is not None and isinstance(result, SQLResult):
+            return list(result.data)
+        return []
+
+    async def select_one(
+        self, conn: Any, query_name: str, sql: str, parameters: "Any", record_class: Optional[Any] = None
+    ) -> Optional[Any]:
+        """Execute a SELECT query and return first result.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+            record_class: Deprecated - use schema_type in driver.execute instead
+
+        Returns:
+            First result row or None
+
+        Note:
+            The record_class parameter is ignored for compatibility. Use schema_type
+            in driver.execute or _sqlspec_schema_type in parameters for type mapping.
+        """
+        if record_class is not None:
+            logger.warning(
+                "record_class parameter is deprecated and ignored. "
+                "Use schema_type in driver.execute or _sqlspec_schema_type in parameters."
+            )
+
+        result = await self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)  # type: ignore[misc]
+
+        if hasattr(result, "data") and result.data and isinstance(result, SQLResult):
+            return result.data[0]
+        return None
+
+    async def select_value(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> Optional[Any]:
+        """Execute a SELECT query and return first value of first row.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            First value of first row or None
+        """
+        row = await self.select_one(conn, query_name, sql, parameters)
+        if row is None:
+            return None
+
+        if isinstance(row, dict):
+            return next(iter(row.values())) if row else None
+        if hasattr(row, "__getitem__"):
+            return row[0] if len(row) > 0 else None
+        return row
+
+    @asynccontextmanager
+    async def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> AsyncGenerator[Any, None]:
+        """Execute a SELECT query and return cursor context manager.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Yields:
+            Cursor-like object with results
+        """
+        result = await self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)  # type: ignore[misc]
+
+        yield AsyncCursorLike(result)
+
+    async def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> None:
+        """Execute INSERT/UPDATE/DELETE.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Note:
+            Returns None per aiosql async protocol
+        """
+        await self.driver.execute(self._create_sql_object(sql, parameters), connection=conn)  # type: ignore[misc]
+
+    async def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> None:
+        """Execute INSERT/UPDATE/DELETE with many parameter sets.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Sequence of parameter sets
+
+        Note:
+            Returns None per aiosql async protocol
+        """
+        await self.driver.execute_many(self._create_sql_object(sql), parameters=parameters, connection=conn)  # type: ignore[misc]
+
+    async def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: "Any") -> Optional[Any]:
+        """Execute INSERT with RETURNING and return result.
+
+        Args:
+            conn: Database connection
+            query_name: Name of the query
+            sql: SQL string
+            parameters: Query parameters
+
+        Returns:
+            Returned value or None
+        """
+        return await self.select_one(conn, query_name, sql, parameters)

sqlspec/extensions/litestar/__init__.py

@@ -0,0 +1,6 @@
+from sqlspec.extensions.litestar import handlers, providers
+from sqlspec.extensions.litestar.cli import database_group
+from sqlspec.extensions.litestar.config import DatabaseConfig
+from sqlspec.extensions.litestar.plugin import SQLSpec
+
+__all__ = ("DatabaseConfig", "SQLSpec", "database_group", "handlers", "providers")

sqlspec/extensions/litestar/_utils.py

@@ -0,0 +1,52 @@
+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/cli.py

@@ -0,0 +1,48 @@
+"""Litestar CLI integration for SQLSpec migrations."""
+
+from contextlib import suppress
+from typing import TYPE_CHECKING
+
+from litestar.cli._utils import LitestarGroup
+
+from sqlspec.cli import add_migration_commands
+
+try:
+    import rich_click as click
+except ImportError:
+    import click  # type: ignore[no-redef]
+
+if TYPE_CHECKING:
+    from litestar import Litestar
+
+    from sqlspec.extensions.litestar.plugin import SQLSpec
+
+
+def get_database_migration_plugin(app: "Litestar") -> "SQLSpec":
+    """Retrieve the SQLSpec plugin from the Litestar application's plugins.
+
+    Args:
+        app: The Litestar application
+
+    Returns:
+        The SQLSpec plugin
+
+    Raises:
+        ImproperConfigurationError: If the SQLSpec plugin is not found
+    """
+    from sqlspec.exceptions import ImproperConfigurationError
+    from sqlspec.extensions.litestar.plugin import SQLSpec
+
+    with suppress(KeyError):
+        return app.plugins.get(SQLSpec)
+    msg = "Failed to initialize database migrations. The required SQLSpec plugin is missing."
+    raise ImproperConfigurationError(msg)
+
+
+@click.group(cls=LitestarGroup, name="db")
+def database_group(ctx: "click.Context") -> None:
+    """Manage SQLSpec database components."""
+    ctx.obj = {"app": ctx.obj, "configs": get_database_migration_plugin(ctx.obj.app).config}
+
+
+add_migration_commands(database_group)

sqlspec/extensions/litestar/config.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+
+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 AsyncGenerator, Awaitable
+    from contextlib import AbstractAsyncContextManager
+
+    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.typing import ConnectionT, PoolT
+
+
+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"
+
+__all__ = (
+    "DEFAULT_COMMIT_MODE",
+    "DEFAULT_CONNECTION_KEY",
+    "DEFAULT_POOL_KEY",
+    "DEFAULT_SESSION_KEY",
+    "CommitMode",
+    "DatabaseConfig",
+)
+
+
+@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
+        )