sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/aiomysql/litestar/store.py

@@ -0,0 +1,314 @@
+"""aiomysql session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Final
+
+import pymysql.err
+
+from sqlspec.adapters.aiomysql._typing import AiomysqlCursor, AiomysqlRawCursor
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.aiomysql.config import AiomysqlConfig
+
+logger = get_logger("sqlspec.adapters.aiomysql.litestar.store")
+
+__all__ = ("AiomysqlStore",)
+
+MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146
+
+
+class AiomysqlStore(BaseSQLSpecStore["AiomysqlConfig"]):
+    """MySQL/MariaDB session store using aiomysql driver.
+
+    Implements server-side session storage for Litestar using MySQL/MariaDB
+    via the aiomysql driver. Provides efficient session management with:
+    - Native async MySQL operations
+    - UPSERT support using ON DUPLICATE KEY UPDATE
+    - Automatic expiration handling
+    - Efficient cleanup of expired sessions
+    - Timezone-aware expiration (stored as UTC in DATETIME)
+
+    Args:
+        config: AiomysqlConfig instance.
+
+    Example:
+        from sqlspec.adapters.aiomysql import AiomysqlConfig
+        from sqlspec.adapters.aiomysql.litestar.store import AiomysqlStore
+
+        config = AiomysqlConfig(connection_config={"host": "localhost", ...})
+        store = AiomysqlStore(config)
+        await store.create_table()
+
+    Notes:
+        MySQL DATETIME is timezone-naive, so UTC datetimes are stored without
+        timezone info and timezone conversion is handled in Python layer.
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "AiomysqlConfig") -> None:
+        """Initialize aiomysql session store.
+
+        Args:
+            config: AiomysqlConfig instance.
+
+        Notes:
+            Table name is read from config.extension_config["litestar"]["session_table"].
+        """
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get MySQL CREATE TABLE SQL with optimized schema.
+
+        Returns:
+            SQL statement to create the sessions table with proper indexes.
+
+        Notes:
+            - Uses DATETIME(6) for microsecond precision timestamps
+            - MySQL doesn't have TIMESTAMPTZ, so we store UTC as timezone-naive
+            - LONGBLOB for large session data support (up to 4GB)
+            - InnoDB engine for ACID compliance and proper transaction support
+            - UTF8MB4 for full Unicode support (including emoji)
+            - Index on expires_at for efficient cleanup queries
+            - Auto-update of updated_at timestamp on row modification
+            - Table name is internally controlled, not user input (S608 suppressed)
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id VARCHAR(255) PRIMARY KEY,
+            data LONGBLOB NOT NULL,
+            expires_at DATETIME(6),
+            created_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6),
+            updated_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+            INDEX idx_{self._table_name}_expires_at (expires_at)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        """Get MySQL/MariaDB DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop indexes and table.
+        """
+        return [
+            f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}",
+            f"DROP TABLE IF EXISTS {self._table_name}",
+        ]
+
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        sql = self._get_create_table_sql()
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(sql)
+        self._log_table_created()
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        """Get a session value by key.
+
+        Args:
+            key: Session ID to retrieve.
+            renew_for: If given, renew the expiry time for this duration.
+
+        Returns:
+            Session data as bytes if found and not expired, None otherwise.
+
+        Notes:
+            Uses UTC_TIMESTAMP(6) for microsecond precision current time in MySQL.
+            Compares expires_at as UTC datetime (timezone-naive in MySQL).
+        """
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            async with (
+                self._config.provide_connection() as conn,
+                AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+            ):
+                await cursor.execute(sql, (key,))
+                row = await cursor.fetchone()
+
+                if row is None:
+                    return None
+
+                data_value, expires_at = row
+
+                if renew_for is not None and expires_at is not None:
+                    new_expires_at = self._calculate_expires_at(renew_for)
+                    if new_expires_at is not None:
+                        naive_expires_at = new_expires_at.replace(tzinfo=None)
+                        update_sql = f"""
+                            UPDATE {self._table_name}
+                            SET expires_at = %s, updated_at = UTC_TIMESTAMP(6)
+                            WHERE session_id = %s
+                            """
+                        await cursor.execute(update_sql, (naive_expires_at, key))
+                        await conn.commit()
+
+                return bytes(data_value)
+        except pymysql.err.ProgrammingError as e:
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return None
+            raise
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        """Store a session value.
+
+        Args:
+            key: Session ID.
+            value: Session data.
+            expires_in: Time until expiration.
+
+        Notes:
+            Uses INSERT ... ON DUPLICATE KEY UPDATE for efficient UPSERT.
+            Stores UTC datetime as timezone-naive DATETIME in MySQL.
+            Uses alias syntax (AS new) instead of deprecated VALUES() function.
+        """
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+        naive_expires_at = expires_at.replace(tzinfo=None) if expires_at else None
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES (%s, %s, %s) AS new
+        ON DUPLICATE KEY UPDATE
+            data = new.data,
+            expires_at = new.expires_at,
+            updated_at = UTC_TIMESTAMP(6)
+        """
+
+        async with (
+            self._config.provide_connection() as conn,
+            AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+        ):
+            await cursor.execute(sql, (key, data, naive_expires_at))
+            await conn.commit()
+
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        async with (
+            self._config.provide_connection() as conn,
+            AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+        ):
+            await cursor.execute(sql, (key,))
+            await conn.commit()
+
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        sql = f"DELETE FROM {self._table_name}"
+
+        try:
+            async with (
+                self._config.provide_connection() as conn,
+                AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+            ):
+                await cursor.execute(sql)
+                await conn.commit()
+            self._log_delete_all()
+        except pymysql.err.ProgrammingError as e:
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                logger.debug("Table %s does not exist, skipping delete_all", self._table_name)
+                return
+            raise
+
+    async def exists(self, key: str) -> bool:
+        """Check if a session key exists and is not expired.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            True if the session exists and is not expired.
+
+        Notes:
+            Uses UTC_TIMESTAMP(6) for microsecond precision current time comparison.
+        """
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            async with (
+                self._config.provide_connection() as conn,
+                AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+            ):
+                await cursor.execute(sql, (key,))
+                result = await cursor.fetchone()
+                return result is not None
+        except pymysql.err.ProgrammingError as e:
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return False
+            raise
+
+    async def expires_in(self, key: str) -> "int | None":
+        """Get the time in seconds until the session expires.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            Seconds until expiration, or None if no expiry or key doesn't exist.
+
+        Notes:
+            MySQL DATETIME is timezone-naive, but we treat it as UTC.
+            Compare against UTC now in Python layer for accuracy.
+        """
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        async with (
+            self._config.provide_connection() as conn,
+            AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+        ):
+            await cursor.execute(sql, (key,))
+            row = await cursor.fetchone()
+
+            if row is None or row[0] is None:
+                return None
+
+            expires_at_naive = row[0]
+            expires_at_utc = expires_at_naive.replace(tzinfo=timezone.utc)
+            now = datetime.now(timezone.utc)
+
+            if expires_at_utc <= now:
+                return 0
+
+            delta = expires_at_utc - now
+            return int(delta.total_seconds())
+
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+
+        Notes:
+            Uses UTC_TIMESTAMP(6) for microsecond precision current time comparison.
+            ROW_COUNT() returns the number of affected rows.
+        """
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= UTC_TIMESTAMP(6)"
+
+        async with (
+            self._config.provide_connection() as conn,
+            AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor,
+        ):
+            await cursor.execute(sql)
+            await conn.commit()
+            count: int = cursor.rowcount
+            if count > 0:
+                self._log_delete_expired(count)
+            return count

sqlspec/adapters/aiosqlite/__init__.py

@@ -0,0 +1,26 @@
+from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection, AiosqliteCursor, AiosqliteRawCursor
+from sqlspec.adapters.aiosqlite.config import AiosqliteConfig, AiosqliteConnectionParams, AiosqlitePoolParams
+from sqlspec.adapters.aiosqlite.core import default_statement_config
+from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver, AiosqliteExceptionHandler
+from sqlspec.adapters.aiosqlite.pool import (
+    AiosqliteConnectionPool,
+    AiosqliteConnectTimeoutError,
+    AiosqlitePoolClosedError,
+    AiosqlitePoolConnection,
+)
+
+__all__ = (
+    "AiosqliteConfig",
+    "AiosqliteConnectTimeoutError",
+    "AiosqliteConnection",
+    "AiosqliteConnectionParams",
+    "AiosqliteConnectionPool",
+    "AiosqliteCursor",
+    "AiosqliteDriver",
+    "AiosqliteExceptionHandler",
+    "AiosqlitePoolClosedError",
+    "AiosqlitePoolConnection",
+    "AiosqlitePoolParams",
+    "AiosqliteRawCursor",
+    "default_statement_config",
+)

sqlspec/adapters/aiosqlite/_typing.py

@@ -0,0 +1,109 @@
+# pyright: reportCallIssue=false, reportAttributeAccessIssue=false, reportArgumentType=false
+"""AIOSQLite adapter type definitions.
+
+This module contains type aliases and classes that are excluded from mypyc
+compilation to avoid ABI boundary issues.
+"""
+
+import contextlib
+from typing import TYPE_CHECKING, Any
+
+import aiosqlite
+
+_AiosqliteConnection = aiosqlite.Connection
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+    from typing import TypeAlias
+
+    from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver
+    from sqlspec.core import StatementConfig
+
+    AiosqliteConnection: TypeAlias = _AiosqliteConnection
+    AiosqliteRawCursor: TypeAlias = aiosqlite.Cursor
+
+if not TYPE_CHECKING:
+    AiosqliteConnection = _AiosqliteConnection
+    AiosqliteRawCursor = aiosqlite.Cursor
+
+
+class AiosqliteCursor:
+    """Async context manager for AIOSQLite cursors."""
+
+    __slots__ = ("connection", "cursor")
+
+    def __init__(self, connection: "AiosqliteConnection") -> None:
+        self.connection = connection
+        self.cursor: AiosqliteRawCursor | None = None
+
+    async def __aenter__(self) -> "AiosqliteRawCursor":
+        self.cursor = await self.connection.cursor()
+        return self.cursor
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> None:
+        if exc_type is not None:
+            return
+        if self.cursor is not None:
+            with contextlib.suppress(Exception):
+                await self.cursor.close()
+
+
+class AiosqliteSessionContext:
+    """Async context manager for AIOSQLite sessions.
+
+    This class is intentionally excluded from mypyc compilation to avoid ABI
+    boundary issues. It receives callables from uncompiled config classes and
+    instantiates compiled Driver objects, acting as a bridge between compiled
+    and uncompiled code.
+
+    Uses callable-based connection management to decouple from config implementation.
+    """
+
+    __slots__ = (
+        "_acquire_connection",
+        "_connection",
+        "_driver",
+        "_driver_features",
+        "_prepare_driver",
+        "_release_connection",
+        "_statement_config",
+    )
+
+    def __init__(
+        self,
+        acquire_connection: "Callable[[], Any]",
+        release_connection: "Callable[[Any], Any]",
+        statement_config: "StatementConfig",
+        driver_features: "dict[str, Any]",
+        prepare_driver: "Callable[[AiosqliteDriver], AiosqliteDriver]",
+    ) -> None:
+        self._acquire_connection = acquire_connection
+        self._release_connection = release_connection
+        self._statement_config = statement_config
+        self._driver_features = driver_features
+        self._prepare_driver = prepare_driver
+        self._connection: Any = None
+        self._driver: AiosqliteDriver | None = None
+
+    async def __aenter__(self) -> "AiosqliteDriver":
+        from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver
+
+        self._connection = await self._acquire_connection()
+        self._driver = AiosqliteDriver(
+            connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features
+        )
+        return self._prepare_driver(self._driver)
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> "bool | None":
+        if self._connection is not None:
+            await self._release_connection(self._connection)
+            self._connection = None
+        return None
+
+
+__all__ = ("AiosqliteConnection", "AiosqliteCursor", "AiosqliteRawCursor", "AiosqliteSessionContext")

sqlspec/adapters/aiosqlite/adk/__init__.py

@@ -0,0 +1,5 @@
+"""Aiosqlite ADK integration for Google Agent Development Kit."""
+
+from sqlspec.adapters.aiosqlite.adk.store import AiosqliteADKMemoryStore, AiosqliteADKStore
+
+__all__ = ("AiosqliteADKMemoryStore", "AiosqliteADKStore")