sqlspec 0.16.2__py3-none-any.whl → 0.17.0__py3-none-any.whl

sqlspec/adapters/aiosqlite/config.py

@@ -1,116 +1,33 @@
 """Aiosqlite database configuration with optimized connection management."""
 
-import asyncio
 import logging
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, TypedDict
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union
 
-import aiosqlite
 from typing_extensions import NotRequired
 
+from sqlspec.adapters.aiosqlite._types import AiosqliteConnection
 from sqlspec.adapters.aiosqlite.driver import AiosqliteCursor, AiosqliteDriver, aiosqlite_statement_config
+from sqlspec.adapters.aiosqlite.pool import (
+    AiosqliteConnectionPool,
+    AiosqliteConnectTimeoutError,
+    AiosqlitePoolClosedError,
+    AiosqlitePoolConnection,
+)
 from sqlspec.config import AsyncDatabaseConfig
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
 
-    from sqlspec.adapters.aiosqlite._types import AiosqliteConnection
     from sqlspec.core.statement import StatementConfig
 
-__all__ = ("AiosqliteConfig", "AiosqliteConnectionParams", "AiosqliteConnectionPool")
+__all__ = ("AiosqliteConfig", "AiosqliteConnectionParams", "AiosqlitePoolParams")
 
 logger = logging.getLogger(__name__)
 
-# Core PRAGMAs for SQLite performance optimization
-WAL_PRAGMA_SQL: Final[str] = "PRAGMA journal_mode = WAL"
-FOREIGN_KEYS_SQL: Final[str] = "PRAGMA foreign_keys = ON"
-SYNC_NORMAL_SQL: Final[str] = "PRAGMA synchronous = NORMAL"
-BUSY_TIMEOUT_SQL: Final[str] = "PRAGMA busy_timeout = 5000"  # 5 seconds
-
-
-class AiosqliteConnectionPool:
-    """Connection pool for Aiosqlite using a single shared connection approach.
-
-    Uses a single shared connection per database file since aiosqlite internally
-    handles queuing and serialization of operations.
-    """
-
-    __slots__ = ("_closed", "_connection", "_connection_parameters", "_lock")
-
-    def __init__(self, connection_parameters: "dict[str, Any]") -> None:
-        """Initialize connection manager.
-
-        Args:
-            connection_parameters: SQLite connection parameters
-        """
-        self._connection: Optional[AiosqliteConnection] = None
-        self._connection_parameters = connection_parameters
-        self._lock = asyncio.Lock()
-        self._closed = False
-
-    async def _ensure_connection(self) -> "AiosqliteConnection":
-        """Ensure we have a valid connection, creating one if needed."""
-        async with self._lock:
-            if self._connection is None or self._closed:
-                self._connection = await aiosqlite.connect(**self._connection_parameters)
-
-                await self._connection.execute(WAL_PRAGMA_SQL)
-                await self._connection.execute(FOREIGN_KEYS_SQL)
-                await self._connection.execute(SYNC_NORMAL_SQL)
-                await self._connection.execute(BUSY_TIMEOUT_SQL)
-                await self._connection.commit()
-
-                self._closed = False
-                logger.debug("Created new aiosqlite connection")
-
-            return self._connection
-
-    @asynccontextmanager
-    async def get_connection(self) -> "AsyncGenerator[AiosqliteConnection, None]":
-        """Get the shared connection.
-
-        Yields:
-            The shared Aiosqlite connection instance.
-        """
-        connection = await self._ensure_connection()
-        yield connection
-
-    async def close(self) -> None:
-        """Close the shared connection."""
-        async with self._lock:
-            if self._connection is not None and not self._closed:
-                await self._connection.close()
-                self._connection = None
-                self._closed = True
-                logger.debug("Closed aiosqlite connection")
-
-    def size(self) -> int:
-        """Get connection count."""
-        return 0 if self._closed or self._connection is None else 1
-
-    def checked_out(self) -> int:
-        """Get number of checked out connections."""
-        return 0
-
-    async def acquire(self) -> "AiosqliteConnection":
-        """Get the shared connection directly.
-
-        Returns:
-            The shared connection instance.
-        """
-        return await self._ensure_connection()
-
-    async def release(self, connection: "AiosqliteConnection") -> None:
-        """No-op release for compatibility.
-
-        Args:
-            connection: Connection to release (ignored)
-        """
-        _ = connection
-
 
 class AiosqliteConnectionParams(TypedDict, total=False):
-    """aiosqlite connection parameters."""
+    """TypedDict for aiosqlite connection parameters."""
 
     database: NotRequired[str]
     timeout: NotRequired[float]
@@ -121,17 +38,27 @@ class AiosqliteConnectionParams(TypedDict, total=False):
     uri: NotRequired[bool]
 
 
-class AiosqliteConfig(AsyncDatabaseConfig):
+class AiosqlitePoolParams(AiosqliteConnectionParams, total=False):
+    """TypedDict for aiosqlite pool parameters, inheriting connection parameters."""
+
+    pool_size: NotRequired[int]
+    connect_timeout: NotRequired[float]
+    idle_timeout: NotRequired[float]
+    operation_timeout: NotRequired[float]
+    extra: NotRequired[dict[str, Any]]
+
+
+class AiosqliteConfig(AsyncDatabaseConfig["AiosqliteConnection", AiosqliteConnectionPool, AiosqliteDriver]):
     """Database configuration for AioSQLite engine."""
 
-    driver_type: ClassVar[type[AiosqliteDriver]] = AiosqliteDriver
-    cursor_type: ClassVar[type[AiosqliteCursor]] = AiosqliteCursor
+    driver_type: "ClassVar[type[AiosqliteDriver]]" = AiosqliteDriver
+    connection_type: "ClassVar[type[AiosqliteConnection]]" = AiosqliteConnection
 
     def __init__(
         self,
         *,
+        pool_config: "Optional[Union[AiosqlitePoolParams, dict[str, Any]]]" = None,
         pool_instance: "Optional[AiosqliteConnectionPool]" = None,
-        pool_config: "Optional[dict[str, Any]]" = None,
         migration_config: "Optional[dict[str, Any]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
         **kwargs: Any,
@@ -139,115 +66,160 @@ class AiosqliteConfig(AsyncDatabaseConfig):
         """Initialize AioSQLite configuration.
 
         Args:
+            pool_config: Pool configuration parameters (TypedDict or dict)
             pool_instance: Optional pre-configured connection pool instance.
-            pool_config: Optional pool configuration dict (AiosqliteConnectionParams).
             migration_config: Optional migration configuration.
             statement_config: Optional statement configuration.
-            **kwargs: Additional connection parameters.
+            **kwargs: Additional connection parameters that override pool_config.
         """
-        connection_params = {}
-        if pool_config:
-            connection_params.update(pool_config)
-        connection_params.update(kwargs)
+        config_dict = dict(pool_config) if pool_config else {}
+        config_dict.update(kwargs)  # Allow kwargs to override pool_config values
+
+        # Handle memory database URI conversion - test expectation is different than sqlite pattern
+        if "database" not in config_dict or config_dict["database"] == ":memory:":
+            config_dict["database"] = "file::memory:?cache=shared"
+            config_dict["uri"] = True
 
         super().__init__(
-            pool_config=connection_params,
+            pool_config=config_dict,
             pool_instance=pool_instance,
-            migration_config=migration_config or {},
+            migration_config=migration_config,
             statement_config=statement_config or aiosqlite_statement_config,
+            driver_features={},
         )
 
-        self._connection_parameters = self._parse_connection_parameters(connection_params)
-
-        if pool_instance is None:
-            self.pool_instance: AiosqliteConnectionPool = AiosqliteConnectionPool(self._connection_parameters)
-
-    def _parse_connection_parameters(self, params: "dict[str, Any]") -> "dict[str, Any]":
-        """Parse connection parameters for AioSQLite.
-
-        Args:
-            params: Connection parameters dict.
+    def _get_pool_config_dict(self) -> "dict[str, Any]":
+        """Get pool configuration as plain dict for external library.
 
         Returns:
-            Processed connection parameters dict.
+            Dictionary with pool parameters, filtering out None values.
         """
-        result = params.copy()
-
-        if "database" not in result:
-            # Default to in-memory database
-            result["database"] = ":memory:"
-
-        # Convert regular :memory: to shared memory for multi-connection access
-        if result.get("database") == ":memory:":
-            result["database"] = "file::memory:?cache=shared"
-            result["uri"] = True
+        config: dict[str, Any] = dict(self.pool_config)
+        extras = config.pop("extra", {})
+        config.update(extras)
+        return {k: v for k, v in config.items() if v is not None}
 
-        for pool_param in ["pool_min_size", "pool_max_size", "pool_timeout", "pool_recycle_seconds"]:
-            result.pop(pool_param, None)
+    def _get_connection_config_dict(self) -> "dict[str, Any]":
+        """Get connection configuration as plain dict for pool creation.
 
-        return result
+        Returns:
+            Dictionary with connection parameters for creating connections.
+        """
+        # Filter out all pool-specific parameters that aiosqlite.connect() doesn't accept
+        excluded_keys = {
+            "pool_size",
+            "connect_timeout",
+            "idle_timeout",
+            "operation_timeout",
+            "extra",
+            "pool_min_size",
+            "pool_max_size",
+            "pool_timeout",
+            "pool_recycle_seconds",
+        }
+        return {k: v for k, v in self.pool_config.items() if k not in excluded_keys}
 
     @asynccontextmanager
-    async def provide_connection(self) -> "AsyncGenerator[AiosqliteConnection, None]":
-        """Provide a database connection.
+    async def provide_connection(self, *args: Any, **kwargs: Any) -> "AsyncGenerator[AiosqliteConnection, None]":
+        """Provide an async connection context manager.
+
+        Args:
+            *args: Additional arguments.
+            **kwargs: Additional keyword arguments.
 
         Yields:
-            AiosqliteConnection: Database connection instance.
+            An aiosqlite connection instance.
         """
+        if self.pool_instance is None:
+            self.pool_instance = await self._create_pool()
         async with self.pool_instance.get_connection() as connection:
             yield connection
 
     @asynccontextmanager
     async def provide_session(
-        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+        self, *_args: Any, statement_config: "Optional[StatementConfig]" = None, **_kwargs: Any
     ) -> "AsyncGenerator[AiosqliteDriver, None]":
-        """Provide an async database session.
+        """Provide an async driver session context manager.
 
         Args:
-            *args: Additional positional arguments.
+            *_args: Additional arguments.
             statement_config: Optional statement configuration override.
-            **kwargs: Additional keyword arguments.
+            **_kwargs: Additional keyword arguments.
 
         Yields:
-            AiosqliteDriver: Database session instance.
+            An AiosqliteDriver instance.
         """
-        _ = args, kwargs
-        effective_statement_config = statement_config or self.statement_config
-        async with self.pool_instance.get_connection() as connection:
-            session = self.driver_type(connection, statement_config=effective_statement_config)
-            try:
-                yield session
-            finally:
-                pass
-
-    async def close(self) -> None:
-        """Close the connection manager."""
-        if self.pool_instance:
-            await self.pool_instance.close()
+        async with self.provide_connection(*_args, **_kwargs) as connection:
+            yield self.driver_type(connection=connection, statement_config=statement_config or self.statement_config)
 
-    def _get_connection_config_dict(self) -> "dict[str, Any]":
-        """Get connection configuration dictionary.
+    async def _create_pool(self) -> AiosqliteConnectionPool:
+        """Create the connection pool instance.
 
         Returns:
-            Connection parameters for creating connections.
+            AiosqliteConnectionPool: The connection pool instance.
         """
-        return self._connection_parameters.copy()
+        config = self._get_pool_config_dict()
+        pool_size = config.pop("pool_size", 5)
+        connect_timeout = config.pop("connect_timeout", 30.0)
+        idle_timeout = config.pop("idle_timeout", 24 * 60 * 60)
+        operation_timeout = config.pop("operation_timeout", 10.0)
+
+        return AiosqliteConnectionPool(
+            connection_parameters=self._get_connection_config_dict(),
+            pool_size=pool_size,
+            connect_timeout=connect_timeout,
+            idle_timeout=idle_timeout,
+            operation_timeout=operation_timeout,
+        )
 
-    async def _create_pool(self) -> "AiosqliteConnectionPool":
-        """Create the connection manager instance.
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        if self.pool_instance and not self.pool_instance.is_closed:
+            await self.pool_instance.close()
+
+    async def create_connection(self) -> "AiosqliteConnection":
+        """Create a single async connection from the pool.
 
         Returns:
-            AiosqliteConnectionPool: The connection manager instance.
+            An aiosqlite connection instance.
         """
         if self.pool_instance is None:
-            self.pool_instance = AiosqliteConnectionPool(self._connection_parameters)
+            self.pool_instance = await self._create_pool()
+        pool_connection = await self.pool_instance.acquire()
+        return pool_connection.connection
+
+    async def provide_pool(self) -> AiosqliteConnectionPool:
+        """Provide async pool instance.
+
+        Returns:
+            The async connection pool.
+        """
+        if not self.pool_instance:
+            self.pool_instance = await self.create_pool()
         return self.pool_instance
 
-    async def _close_pool(self) -> None:
-        """Close the connection manager."""
-        if self.pool_instance:
-            await self.pool_instance.close()
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for aiosqlite types.
 
-    async def close_pool(self) -> None:
-        """Close the connection pool (delegates to _close_pool)."""
-        await self._close_pool()
+        This provides all aiosqlite-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update(
+            {
+                "AiosqliteConnection": AiosqliteConnection,
+                "AiosqliteConnectionPool": AiosqliteConnectionPool,
+                "AiosqliteConnectTimeoutError": AiosqliteConnectTimeoutError,
+                "AiosqliteCursor": AiosqliteCursor,
+                "AiosqlitePoolClosedError": AiosqlitePoolClosedError,
+                "AiosqlitePoolConnection": AiosqlitePoolConnection,
+            }
+        )
+        return namespace
+
+    async def _close_pool(self) -> None:
+        """Close the connection pool."""
+        await self.close_pool()

sqlspec/adapters/aiosqlite/driver.py

@@ -7,6 +7,7 @@ Provides async SQLite database connectivity with:
 - SQLite-specific optimizations
 """
 
+import asyncio
 import contextlib
 import datetime
 from decimal import Decimal
@@ -143,13 +144,7 @@ class AiosqliteDriver(AsyncDriverAdapterBase):
     ) -> None:
         if statement_config is None:
             cache_config = get_cache_config()
-            enhanced_config = aiosqlite_statement_config.replace(
-                enable_caching=cache_config.compiled_cache_enabled,
-                enable_parsing=True,
-                enable_validation=True,
-                dialect="sqlite",
-            )
-            statement_config = enhanced_config
+            statement_config = aiosqlite_statement_config.replace(enable_caching=cache_config.compiled_cache_enabled)
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
 
@@ -223,12 +218,28 @@ class AiosqliteDriver(AsyncDriverAdapterBase):
         return self.create_execution_result(cursor, rowcount_override=affected_rows)
 
     async def begin(self) -> None:
-        """Begin a database transaction."""
+        """Begin a database transaction with appropriate locking strategy."""
         try:
             if not self.connection.in_transaction:
-                await self.connection.execute("BEGIN")
+                # For shared cache databases, use IMMEDIATE to reduce lock contention
+                # For other databases, BEGIN IMMEDIATE is also safer for concurrent access
+                await self.connection.execute("BEGIN IMMEDIATE")
         except aiosqlite.Error as e:
-            msg = f"Failed to begin transaction: {e}"
+            # If IMMEDIATE fails due to lock, try with exponential backoff
+            import random
+
+            max_retries = 3
+            for attempt in range(max_retries):
+                delay = 0.01 * (2**attempt) + random.uniform(0, 0.01)  # noqa: S311
+                await asyncio.sleep(delay)
+                try:
+                    await self.connection.execute("BEGIN IMMEDIATE")
+                except aiosqlite.Error:
+                    if attempt == max_retries - 1:  # Last attempt
+                        break
+                else:
+                    return
+            msg = f"Failed to begin transaction after retries: {e}"
             raise SQLSpecError(msg) from e
 
     async def rollback(self) -> None: