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

sqlspec/adapters/aiosqlite/data_dictionary.py

@@ -0,0 +1,208 @@
+"""SQLite-specific data dictionary for metadata queries via aiosqlite."""
+
+from typing import TYPE_CHECKING, ClassVar
+
+from mypy_extensions import mypyc_attr
+
+from sqlspec.adapters.aiosqlite.core import format_identifier
+from sqlspec.data_dictionary import get_dialect_config
+from sqlspec.driver import AsyncDataDictionaryBase
+from sqlspec.typing import ColumnMetadata, ForeignKeyMetadata, IndexMetadata, TableMetadata, VersionInfo
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver
+    from sqlspec.data_dictionary._types import DialectConfig
+
+__all__ = ("AiosqliteDataDictionary",)
+
+
+@mypyc_attr(allow_interpreted_subclasses=True, native_class=False)
+class AiosqliteDataDictionary(AsyncDataDictionaryBase):
+    """SQLite-specific async data dictionary."""
+
+    dialect: ClassVar[str] = "sqlite"
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get_dialect_config(self) -> "DialectConfig":
+        """Return the dialect configuration for this data dictionary."""
+        return get_dialect_config(type(self).dialect)
+
+    def resolve_schema(self, schema: "str | None") -> "str | None":
+        """Return a schema name using dialect defaults when missing."""
+        if schema is not None:
+            return schema
+        return self.get_dialect_config().default_schema
+
+    def list_available_features(self) -> "list[str]":
+        """List available feature flags for this dialect."""
+        config = self.get_dialect_config()
+        features = {"supports_transactions", "supports_prepared_statements"}
+        features.update(config.feature_flags.keys())
+        features.update(config.feature_versions.keys())
+        return sorted(features)
+
+    async def get_version(self, driver: "AiosqliteDriver") -> "VersionInfo | None":
+        """Get SQLite database version information.
+
+        Args:
+            driver: Async database driver instance.
+
+        Returns:
+            SQLite version information or None if detection fails.
+
+        """
+        driver_id = id(driver)
+        # Inline cache check to avoid cross-module method call that causes mypyc segfault
+        if driver_id in self._version_fetch_attempted:
+            return self._version_cache.get(driver_id)
+        # Not cached, fetch from database
+
+        version_value = await driver.select_value_or_none(self.get_query("version"))
+        if not version_value:
+            self._log_version_unavailable(type(self).dialect, "missing")
+            self.cache_version(driver_id, None)
+            return None
+
+        version_info = self.parse_version_with_pattern(self.get_dialect_config().version_pattern, str(version_value))
+        if version_info is None:
+            self._log_version_unavailable(type(self).dialect, "parse_failed")
+            self.cache_version(driver_id, None)
+            return None
+
+        self._log_version_detected(type(self).dialect, version_info)
+        self.cache_version(driver_id, version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: "AiosqliteDriver", feature: str) -> bool:
+        """Check if SQLite database supports a specific feature.
+
+        Args:
+            driver: Async database driver instance.
+            feature: Feature name to check.
+
+        Returns:
+            True if feature is supported, False otherwise.
+
+        """
+        version_info = await self.get_version(driver)
+        return self.resolve_feature_flag(feature, version_info)
+
+    async def get_optimal_type(self, driver: "AiosqliteDriver", type_category: str) -> str:
+        """Get optimal SQLite type for a category.
+
+        Args:
+            driver: Async database driver instance.
+            type_category: Type category.
+
+        Returns:
+            SQLite-specific type name.
+
+        """
+        config = self.get_dialect_config()
+        version_info = await self.get_version(driver)
+
+        if type_category == "json":
+            json_version = config.get_feature_version("supports_json")
+            if version_info and json_version and version_info >= json_version:
+                return "JSON"
+            return "TEXT"
+
+        return config.get_optimal_type(type_category)
+
+    async def get_tables(self, driver: "AiosqliteDriver", schema: "str | None" = None) -> "list[TableMetadata]":
+        """Get tables sorted by topological dependency order using SQLite catalog."""
+        schema_name = self.resolve_schema(schema)
+        self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables")
+        schema_prefix = f"{format_identifier(schema_name)}." if schema_name else ""
+        query_text = self.get_query_text("tables_by_schema").format(schema_prefix=schema_prefix)
+        return await driver.select(query_text, schema_type=TableMetadata)
+
+    async def get_columns(
+        self, driver: "AiosqliteDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ColumnMetadata]":
+        """Get column information for a table or schema."""
+        schema_name = self.resolve_schema(schema)
+        schema_prefix = f"{format_identifier(schema_name)}." if schema_name else ""
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns")
+            query_text = self.get_query_text("columns_by_schema").format(schema_prefix=schema_prefix)
+            return await driver.select(query_text, schema_type=ColumnMetadata)
+
+        assert table is not None
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns")
+        table_name = table
+        table_identifier = f"{schema_name}.{table_name}" if schema_name else table_name
+        query_text = self.get_query_text("columns_by_table").format(table_name=format_identifier(table_identifier))
+        return await driver.select(query_text, schema_type=ColumnMetadata)
+
+    async def get_indexes(
+        self, driver: "AiosqliteDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[IndexMetadata]":
+        """Get index metadata for a table or schema."""
+        schema_name = self.resolve_schema(schema)
+        indexes: list[IndexMetadata] = []
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes")
+            for table_info in await self.get_tables(driver, schema=schema_name):
+                table_name = table_info.get("table_name")
+                if not table_name:
+                    continue
+                indexes.extend(await self.get_indexes(driver, table=table_name, schema=schema_name))
+            return indexes
+
+        assert table is not None
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes")
+        table_name = table
+        table_identifier = f"{schema_name}.{table_name}" if schema_name else table_name
+        index_list_sql = self.get_query_text("indexes_by_table").format(table_name=format_identifier(table_identifier))
+        index_rows = await driver.select(index_list_sql)
+        for row in index_rows:
+            index_name = row.get("name")
+            if not index_name:
+                continue
+            index_identifier = f"{schema_name}.{index_name}" if schema_name else index_name
+            columns_sql = self.get_query_text("index_columns_by_index").format(
+                index_name=format_identifier(index_identifier)
+            )
+            columns_rows = await driver.select(columns_sql)
+            columns: list[str] = []
+            for col in columns_rows:
+                column_name = col.get("name")
+                if column_name is None:
+                    continue
+                columns.append(str(column_name))
+            is_primary = row.get("origin") == "pk"
+            index_metadata: IndexMetadata = {
+                "index_name": index_name,
+                "table_name": table_name,
+                "columns": columns,
+                "is_primary": is_primary,
+            }
+            if schema_name is not None:
+                index_metadata["schema_name"] = schema_name
+            unique_value = row.get("unique")
+            if unique_value is not None:
+                index_metadata["is_unique"] = unique_value
+            indexes.append(index_metadata)
+        return indexes
+
+    async def get_foreign_keys(
+        self, driver: "AiosqliteDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ForeignKeyMetadata]":
+        """Get foreign key metadata."""
+        schema_name = self.resolve_schema(schema)
+        schema_prefix = f"{format_identifier(schema_name)}." if schema_name else ""
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys")
+            query_text = self.get_query_text("foreign_keys_by_schema").format(schema_prefix=schema_prefix)
+            return await driver.select(query_text, schema_type=ForeignKeyMetadata)
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys")
+        table_label = table.replace("'", "''")
+        table_identifier = f"{schema_name}.{table}" if schema_name else table
+        query_text = self.get_query_text("foreign_keys_by_table").format(
+            table_name=format_identifier(table_identifier), table_label=table_label
+        )
+        return await driver.select(query_text, schema_type=ForeignKeyMetadata)

sqlspec/adapters/aiosqlite/driver.py

@@ -0,0 +1,313 @@
+"""AIOSQLite driver implementation for async SQLite operations."""
+
+import asyncio
+import contextlib
+import random
+import sqlite3
+from typing import TYPE_CHECKING, Any, cast
+
+import aiosqlite
+
+from sqlspec.adapters.aiosqlite.core import (
+    build_insert_statement,
+    collect_rows,
+    create_mapped_exception,
+    default_statement_config,
+    driver_profile,
+    format_identifier,
+    normalize_execute_many_parameters,
+    normalize_execute_parameters,
+    resolve_rowcount,
+)
+from sqlspec.adapters.aiosqlite.data_dictionary import AiosqliteDataDictionary
+from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile
+from sqlspec.driver import AsyncDriverAdapterBase
+from sqlspec.exceptions import SQLSpecError
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection
+    from sqlspec.core import SQL, StatementConfig
+    from sqlspec.driver import ExecutionResult
+    from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
+
+from sqlspec.adapters.aiosqlite._typing import AiosqliteSessionContext
+
+__all__ = ("AiosqliteCursor", "AiosqliteDriver", "AiosqliteExceptionHandler", "AiosqliteSessionContext")
+
+SQLITE_CONSTRAINT_UNIQUE_CODE = 2067
+SQLITE_CONSTRAINT_FOREIGNKEY_CODE = 787
+SQLITE_CONSTRAINT_NOTNULL_CODE = 1811
+SQLITE_CONSTRAINT_CHECK_CODE = 531
+SQLITE_CONSTRAINT_CODE = 19
+SQLITE_CANTOPEN_CODE = 14
+SQLITE_IOERR_CODE = 10
+SQLITE_MISMATCH_CODE = 20
+
+
+class AiosqliteCursor:
+    """Async context manager for AIOSQLite cursors."""
+
+    __slots__ = ("connection", "cursor")
+
+    def __init__(self, connection: "AiosqliteConnection") -> None:
+        self.connection = connection
+        self.cursor: aiosqlite.Cursor | None = None
+
+    async def __aenter__(self) -> "aiosqlite.Cursor":
+        self.cursor = await self.connection.cursor()
+        return self.cursor
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
+        if exc_type is not None:
+            return
+        if self.cursor is not None:
+            with contextlib.suppress(Exception):
+                await self.cursor.close()
+
+
+class AiosqliteExceptionHandler:
+    """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__ = ("pending_exception",)
+
+    def __init__(self) -> None:
+        self.pending_exception: Exception | None = None
+
+    async def __aenter__(self) -> "AiosqliteExceptionHandler":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        if exc_val is None:
+            return False
+        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",)
+    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
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # CORE DISPATCH METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def dispatch_execute(self, cursor: "aiosqlite.Cursor", statement: "SQL") -> "ExecutionResult":
+        """Execute single SQL statement."""
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        await cursor.execute(sql, normalize_execute_parameters(prepared_parameters))
+
+        if statement.returns_rows():
+            fetched_data = await cursor.fetchall()
+
+            # aiosqlite returns Iterable[Row], core helper expects Iterable[Any]
+            # Use cast to satisfy mypy and pyright
+            data, column_names, row_count = collect_rows(cast("list[Any]", fetched_data), cursor.description)
+
+            return self.create_execution_result(
+                cursor, selected_data=data, column_names=column_names, data_row_count=row_count, is_select_result=True
+            )
+
+        affected_rows = resolve_rowcount(cursor)
+        return self.create_execution_result(cursor, rowcount_override=affected_rows)
+
+    async def dispatch_execute_many(self, cursor: "aiosqlite.Cursor", statement: "SQL") -> "ExecutionResult":
+        """Execute SQL with multiple parameter sets."""
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+
+        await cursor.executemany(sql, normalize_execute_many_parameters(prepared_parameters))
+
+        affected_rows = resolve_rowcount(cursor)
+
+        return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True)
+
+    async def dispatch_execute_script(self, cursor: "aiosqlite.Cursor", statement: "SQL") -> "ExecutionResult":
+        """Execute SQL script."""
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
+
+        successful_count = 0
+        last_cursor = cursor
+
+        for stmt in statements:
+            await cursor.execute(stmt, normalize_execute_parameters(prepared_parameters))
+            successful_count += 1
+
+        return self.create_execution_result(
+            last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True
+        )
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # TRANSACTION MANAGEMENT
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def begin(self) -> None:
+        """Begin a database transaction."""
+        try:
+            if not self.connection.in_transaction:
+                await self.connection.execute("BEGIN IMMEDIATE")
+        except aiosqlite.Error as e:
+            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:
+                        break
+                else:
+                    return
+            msg = f"Failed to begin transaction after retries: {e}"
+            raise SQLSpecError(msg) from e
+
+    async def commit(self) -> None:
+        """Commit the current transaction."""
+        try:
+            await self.connection.commit()
+        except aiosqlite.Error as e:
+            msg = f"Failed to commit transaction: {e}"
+            raise SQLSpecError(msg) from e
+
+    async def rollback(self) -> None:
+        """Rollback the current transaction."""
+        try:
+            await self.connection.rollback()
+        except aiosqlite.Error as e:
+            msg = f"Failed to rollback transaction: {e}"
+            raise SQLSpecError(msg) from e
+
+    def with_cursor(self, connection: "AiosqliteConnection") -> "AiosqliteCursor":
+        """Create async context manager for AIOSQLite cursor."""
+        return AiosqliteCursor(connection)
+
+    def handle_database_exceptions(self) -> "AiosqliteExceptionHandler":
+        """Handle AIOSQLite-specific exceptions."""
+        return AiosqliteExceptionHandler()
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # STORAGE API METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def select_to_storage(
+        self,
+        statement: "SQL | str",
+        destination: "StorageDestination",
+        /,
+        *parameters: Any,
+        statement_config: "StatementConfig | None" = None,
+        partitioner: "dict[str, object] | None" = None,
+        format_hint: "StorageFormat | None" = None,
+        telemetry: "StorageTelemetry | None" = None,
+        **kwargs: Any,
+    ) -> "StorageBridgeJob":
+        """Execute a query and stream Arrow results into storage."""
+
+        self._require_capability("arrow_export_enabled")
+        arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs)
+        async_pipeline = self._storage_pipeline()
+        telemetry_payload = await self._write_result_to_storage_async(
+            arrow_result, destination, format_hint=format_hint, pipeline=async_pipeline
+        )
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    async def load_from_arrow(
+        self,
+        table: str,
+        source: "ArrowResult | Any",
+        *,
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+        telemetry: "StorageTelemetry | None" = None,
+    ) -> "StorageBridgeJob":
+        """Load Arrow data into SQLite using batched inserts."""
+
+        self._require_capability("arrow_import_enabled")
+        arrow_table = self._coerce_arrow_table(source)
+        if overwrite:
+            statement = f"DELETE FROM {format_identifier(table)}"
+            async with self.handle_database_exceptions(), self.with_cursor(self.connection) as cursor:
+                await cursor.execute(statement)
+
+        columns, records = self._arrow_table_to_rows(arrow_table)
+        if records:
+            insert_sql = build_insert_statement(table, columns)
+            async with self.handle_database_exceptions(), self.with_cursor(self.connection) as cursor:
+                await cursor.executemany(insert_sql, records)
+
+        telemetry_payload = self._build_ingest_telemetry(arrow_table)
+        telemetry_payload["destination"] = table
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    async def load_from_storage(
+        self,
+        table: str,
+        source: "StorageDestination",
+        *,
+        file_format: "StorageFormat",
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+    ) -> "StorageBridgeJob":
+        """Load staged artifacts from storage into SQLite."""
+
+        arrow_table, inbound = await self._read_arrow_from_storage_async(source, file_format=file_format)
+        return await self.load_from_arrow(
+            table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound
+        )
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # UTILITY METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    @property
+    def data_dictionary(self) -> "AiosqliteDataDictionary":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            self._data_dictionary = AiosqliteDataDictionary()
+        return self._data_dictionary
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # PRIVATE/INTERNAL METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    def _connection_in_transaction(self) -> bool:
+        """Check if connection is in transaction.
+
+        Returns:
+            True if connection is in an active transaction.
+        """
+        return bool(self.connection.in_transaction)
+
+
+register_driver_profile("aiosqlite", driver_profile)

sqlspec/adapters/aiosqlite/events/__init__.py

@@ -0,0 +1,5 @@
+"""Events helpers for the aiosqlite adapter."""
+
+from sqlspec.adapters.aiosqlite.events.store import AiosqliteEventQueueStore
+
+__all__ = ("AiosqliteEventQueueStore",)

sqlspec/adapters/aiosqlite/events/store.py

@@ -0,0 +1,20 @@
+"""AioSQLite event queue store."""
+
+from sqlspec.adapters.aiosqlite.config import AiosqliteConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("AiosqliteEventQueueStore",)
+
+
+class AiosqliteEventQueueStore(BaseEventQueueStore[AiosqliteConfig]):
+    """Provide column definitions for the async SQLite adapter.
+
+    SQLite stores JSON as TEXT since it lacks a native JSON column type.
+    JSON functions can still operate on TEXT columns containing valid JSON.
+    """
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        """Return SQLite-compatible column types for the event queue."""
+        return "TEXT", "TEXT", "TIMESTAMP"

sqlspec/adapters/aiosqlite/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for AioSQLite adapter."""
+
+from sqlspec.adapters.aiosqlite.litestar.store import AiosqliteStore
+
+__all__ = ("AiosqliteStore",)