"""AIOSQLite driver implementation for async SQLite operations."""
import asyncio
import random
import sqlite3
from typing import TYPE_CHECKING, Any, cast
import aiosqlite
from sqlspec.adapters.aiosqlite._typing import AiosqliteCursor, AiosqliteRawCursor, AiosqliteSessionContext
from sqlspec.adapters.aiosqlite.core import (
AiosqliteStreamSource,
_execute_and_resolve_metadata,
_execute_fetchall_with_metadata,
build_insert_statement,
collect_rows,
create_mapped_exception,
default_statement_config,
driver_profile,
format_identifier,
normalize_execute_many_parameters,
normalize_execute_parameters,
resolve_rowcount,
run_on_worker_thread,
)
from sqlspec.adapters.aiosqlite.data_dictionary import AiosqliteDataDictionary
from sqlspec.core import ArrowResult, ParameterStyle, TypedParameter, get_cache_config, register_driver_profile
from sqlspec.core.result import DMLResult
from sqlspec.driver import (
AsyncDriverAdapterBase,
AsyncRowStream,
BaseAsyncExceptionHandler,
parameter_value_needs_processing,
type_coercion_fallbacks,
)
from sqlspec.exceptions import SQLSpecError
from sqlspec.utils.type_guards import resolve_row_format
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection
from sqlspec.builder import QueryBuilder
from sqlspec.core import SQL, SQLResult, Statement, StatementConfig, StatementFilter
from sqlspec.core.compiler import OperationType
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
from sqlspec.typing import StatementParameters
__all__ = (
"AiosqliteCursor",
"AiosqliteDriver",
"AiosqliteExceptionHandler",
"AiosqliteRawCursor",
"AiosqliteSessionContext",
)
class AiosqliteExceptionHandler(BaseAsyncExceptionHandler):
"""Async context manager for handling aiosqlite database exceptions.
Maps SQLite extended result codes to specific SQLSpec exceptions
for better error handling in application code.
Uses deferred exception pattern for mypyc compatibility: exceptions
are stored in pending_exception rather than raised from __aexit__
to avoid ABI boundary violations with compiled code.
"""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
_ = exc_type
if isinstance(exc_val, (aiosqlite.Error, sqlite3.Error)):
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
class AiosqliteDriver(AsyncDriverAdapterBase):
"""AIOSQLite driver for async SQLite database operations."""
__slots__ = ("_data_dictionary", "_rowid_target_cache")
dialect = "sqlite"
def __init__(
self,
connection: "AiosqliteConnection",
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = default_statement_config.replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self._data_dictionary: AiosqliteDataDictionary | None = None
self._rowid_target_cache: dict[tuple[str | None, str], bool] = {}