"""Multi-connection pool for aiosqlite."""
import asyncio
import logging
import sqlite3
import time
from contextlib import suppress
from inspect import isawaitable
from threading import Thread
from typing import TYPE_CHECKING, Any, Final
import aiosqlite
from sqlspec.adapters.aiosqlite.core import run_on_worker_thread
from sqlspec.exceptions import SQLSpecError
from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context
from sqlspec.utils.uuids import uuid4
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from types import TracebackType
from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection
__all__ = (
"AiosqliteConnectTimeoutError",
"AiosqliteConnectionPool",
"AiosqlitePoolClosedError",
"AiosqlitePoolConnection",
"AiosqlitePoolConnectionContext",
)
logger = get_logger(POOL_LOGGER_NAME)
_ADAPTER_NAME = "aiosqlite"
SQLITE_BUSY_TIMEOUT: Final = 5000
SQLITE_DEFAULT_ENABLE_FOREIGN_KEYS: Final = False
SQLITE_DEFAULT_ENABLE_OPTIMIZATIONS: Final = True
SQLITE_MEMORY_CACHE_SIZE: Final = -16000
def _dict_row_factory(cursor: Any, row: "tuple[Any, ...]") -> "dict[str, Any]":
return {description[0]: row[index] for index, description in enumerate(cursor.description)}
def _resolve_row_factory(row_factory: Any) -> Any:
if row_factory == "row":
return sqlite3.Row
if row_factory == "dict":
return _dict_row_factory
if row_factory == "tuple":
return None
return row_factory
def _has_active_transaction(connection: "AiosqliteConnection") -> bool:
return bool(getattr(connection, "in_transaction", False))
async def _apply_runtime_setup(connection: "AiosqliteConnection", runtime_setup: "dict[str, Any]") -> None:
pragmas = runtime_setup.get("pragmas", ())
if pragmas:
pragma_script = "\n".join(f"PRAGMA {pragma_name} = {pragma_value};" for pragma_name, pragma_value in pragmas)
await connection.executescript(pragma_script)
extensions = runtime_setup.get("extensions")
if extensions:
await connection.enable_load_extension(True)
try:
for extension_path in extensions:
await connection.load_extension(extension_path)
finally:
await connection.enable_load_extension(False)
for function_config in runtime_setup.get("custom_functions", ()):
await connection.create_function(
function_config["name"],
function_config["narg"],
function_config["func"],
deterministic=function_config.get("deterministic", False),
)
raw_connection = connection._conn # pyright: ignore[reportPrivateUsage]
for aggregate_config in runtime_setup.get("custom_aggregates", ()):
await run_on_worker_thread(
connection,
raw_connection.create_aggregate,
aggregate_config["name"],
aggregate_config["narg"],
aggregate_config["aggregate_class"],
)
for collation_config in runtime_setup.get("custom_collations", ()):
await run_on_worker_thread(
connection, raw_connection.create_collation, collation_config["name"], collation_config["func"]
)
authorizer_callback = runtime_setup.get("authorizer_callback")
if authorizer_callback is not None:
await connection.set_authorizer(authorizer_callback)
trace_callback = runtime_setup.get("trace_callback")
if trace_callback is not None:
await connection.set_trace_callback(trace_callback)
progress_handler = runtime_setup.get("progress_handler")
if progress_handler is not None:
await connection.set_progress_handler(progress_handler, runtime_setup.get("progress_handler_interval", 1000))
if "row_factory" in runtime_setup:
connection.row_factory = _resolve_row_factory(runtime_setup["row_factory"])
if "text_factory" in runtime_setup:
connection.text_factory = runtime_setup["text_factory"]
class AiosqlitePoolClosedError(SQLSpecError):
"""Pool has been closed and cannot accept new operations."""
class AiosqliteConnectTimeoutError(SQLSpecError):
"""Connection could not be established within the specified timeout period."""
class AiosqlitePoolConnection:
"""Wrapper for database connections in the pool."""
__slots__ = ("_closed", "_healthy", "connection", "id", "idle_since")
def __init__(self, connection: "AiosqliteConnection") -> None:
"""Initialize pool connection wrapper.
Args:
connection: The raw aiosqlite connection
"""
self.id = uuid4().hex
self.connection = connection
self.idle_since: float | None = None
self._closed = False
self._healthy = True
@property
def idle_time(self) -> float:
"""Get idle time in seconds.
Returns:
Idle time in seconds, 0.0 if connection is in use
"""
if self.idle_since is None:
return 0.0
return time.time() - self.idle_since
@property
def is_closed(self) -> bool:
"""Check if connection is closed.
Returns:
True if connection is closed
"""
return self._closed
@property
def is_healthy(self) -> bool:
"""Check if connection was healthy on last check.
Returns:
True if connection is presumed healthy
"""
return self._healthy and not self._closed
def mark_as_in_use(self) -> None:
"""Mark connection as in use."""
self.idle_since = None
def mark_as_idle(self) -> None:
"""Mark connection as idle."""
self.idle_since = time.time()
def mark_unhealthy(self) -> None:
"""Mark connection as unhealthy."""
self._healthy = False
async def is_alive(self) -> bool:
"""Check if connection is alive and functional.
Returns:
True if connection is healthy
"""
if self._closed:
self._healthy = False
return False
try:
await self.connection.execute("SELECT 1")
except Exception:
self._healthy = False
return False
else:
self._healthy = True
return True
async def reset(self) -> None:
"""Reset connection to clean state."""
if self._closed:
return
if not _has_active_transaction(self.connection):
return
with suppress(Exception):
await self.connection.rollback()
async def close(self) -> None:
"""Close the connection."""
if self._closed:
return
try:
if _has_active_transaction(self.connection):
with suppress(Exception):
await self.connection.rollback()
await self.connection.close()
except Exception:
# Note: No pool context available at connection level
log_with_context(
logger, logging.DEBUG, "pool.connection.close.error", adapter=_ADAPTER_NAME, connection_id=self.id
)
finally:
self._closed = True
class AiosqlitePoolConnectionContext:
"""Async context manager for pooled aiosqlite connections."""
__slots__ = ("_connection", "_pool")
def __init__(self, pool: "AiosqliteConnectionPool") -> None:
"""Initialize the context manager.
Args:
pool: Connection pool instance.
"""
self._pool = pool
self._connection: AiosqlitePoolConnection | None = None
async def __aenter__(self) -> "AiosqliteConnection":
self._connection = await self._pool.acquire()
return self._connection.connection
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> "bool | None":
if self._connection is None:
return False
await self._pool.release(self._connection)
self._connection = None
return False
class AiosqliteConnectionPool:
"""Multi-connection pool for aiosqlite."""
__slots__ = (
"_closed_event_instance",
"_connect_timeout",
"_connection_parameters",
"_connection_registry",
"_enable_foreign_keys",
"_enable_optimizations",
"_health_check_interval",
"_idle_timeout",
"_lock_instance",
"_min_size",
"_on_connection_create",
"_operation_timeout",
"_pool_id",
"_pool_size",
"_queue_instance",
"_runtime_setup",
"_warmed",
)
def __init__(
self,
connection_parameters: "dict[str, Any]",
pool_size: int = 5,
min_size: int = 0,
connect_timeout: float = 30.0,
idle_timeout: float = 24 * 60 * 60,
operation_timeout: float = 10.0,
health_check_interval: float = 30.0,
enable_optimizations: bool = SQLITE_DEFAULT_ENABLE_OPTIMIZATIONS,
enable_foreign_keys: bool = SQLITE_DEFAULT_ENABLE_FOREIGN_KEYS,
on_connection_create: "Callable[[AiosqliteConnection], Awaitable[None]] | None" = None,
runtime_setup: "dict[str, Any] | None" = None,
) -> None:
"""Initialize connection pool.
Args:
connection_parameters: SQLite connection parameters
pool_size: Maximum number of connections in the pool
min_size: Minimum connections to pre-create (pool warming)
connect_timeout: Maximum time to wait for connection acquisition
idle_timeout: Maximum time a connection can remain idle
operation_timeout: Maximum time for connection operations
health_check_interval: Seconds of idle time before running health check
enable_optimizations: Whether to apply performance PRAGMAs
enable_foreign_keys: Whether to enable foreign-key enforcement
on_connection_create: Async callback executed when connection is created
runtime_setup: Runtime feature setup to apply to new connections
"""
self._connection_parameters = connection_parameters
self._pool_size = pool_size
self._min_size = min(min_size, pool_size)
self._connect_timeout = connect_timeout
self._idle_timeout = idle_timeout
self._operation_timeout = operation_timeout
self._health_check_interval = health_check_interval
self._enable_optimizations = enable_optimizations
self._enable_foreign_keys = enable_foreign_keys
self._on_connection_create = on_connection_create
self._runtime_setup = runtime_setup
self._connection_registry: dict[str, AiosqlitePoolConnection] = {}
self._warmed = False
self._pool_id = uuid4().hex[:8] # Short ID for logging
self._queue_instance: asyncio.Queue[AiosqlitePoolConnection] | None = None
self._lock_instance: asyncio.Lock | None = None
self._closed_event_instance: asyncio.Event | None = None