sqlspec 0.26.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/adapters/adbc/config.py

@@ -1,8 +1,9 @@
 """ADBC database configuration."""
 
 import logging
+from collections.abc import Callable
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, TypedDict, Union
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
 
 from typing_extensions import NotRequired
 
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger("sqlspec.adapters.adbc")
 
 
-class AdbcConnectionParams(TypedDict, total=False):
+class AdbcConnectionParams(TypedDict):
     """ADBC connection parameters."""
 
     uri: NotRequired[str]
@@ -54,7 +55,37 @@ class AdbcConnectionParams(TypedDict, total=False):
     extra: NotRequired[dict[str, Any]]
 
 
-__all__ = ("AdbcConfig", "AdbcConnectionParams")
+class AdbcDriverFeatures(TypedDict):
+    """ADBC driver feature configuration.
+
+    Controls optional type handling and serialization behavior for the ADBC adapter.
+    These features configure how data is converted between Python and Arrow types.
+
+    Attributes:
+        json_serializer: JSON serialization function to use.
+            Callable that takes Any and returns str (JSON string).
+            Default: sqlspec.utils.serializers.to_json
+        enable_cast_detection: Enable cast-aware parameter processing.
+            When True, detects SQL casts (e.g., ::JSONB) and applies appropriate
+            serialization. Currently used for PostgreSQL JSONB handling.
+            Default: True
+        strict_type_coercion: Enforce strict type coercion rules.
+            When True, raises errors for unsupported type conversions.
+            When False, attempts best-effort conversion.
+            Default: False
+        arrow_extension_types: Enable PyArrow extension type support.
+            When True, preserves Arrow extension type metadata when reading data.
+            When False, falls back to storage types.
+            Default: True
+    """
+
+    json_serializer: "NotRequired[Callable[[Any], str]]"
+    enable_cast_detection: NotRequired[bool]
+    strict_type_coercion: NotRequired[bool]
+    arrow_extension_types: NotRequired[bool]
+
+
+__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures")
 
 
 class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
@@ -69,15 +100,17 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
 
     driver_type: ClassVar[type[AdbcDriver]] = AdbcDriver
     connection_type: "ClassVar[type[AdbcConnection]]" = AdbcConnection
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
         *,
-        connection_config: Optional[Union[AdbcConnectionParams, dict[str, Any]]] = None,
-        migration_config: Optional[dict[str, Any]] = None,
-        statement_config: Optional[StatementConfig] = None,
-        driver_features: Optional[dict[str, Any]] = None,
-        bind_key: Optional[str] = None,
+        connection_config: AdbcConnectionParams | dict[str, Any] | None = None,
+        migration_config: dict[str, Any] | None = None,
+        statement_config: StatementConfig | None = None,
+        driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: str | None = None,
+        extension_config: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize configuration.
 
@@ -85,8 +118,9 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
             connection_config: Connection configuration parameters
             migration_config: Migration configuration
             statement_config: Default SQL statement configuration
-            driver_features: Driver feature configuration
+            driver_features: Driver feature configuration (AdbcDriverFeatures)
             bind_key: Optional unique identifier for this configuration
+            extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
         """
         if connection_config is None:
             connection_config = {}
@@ -101,12 +135,26 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
             detected_dialect = str(self._get_dialect() or "sqlite")
             statement_config = get_adbc_statement_config(detected_dialect)
 
+        from sqlspec.utils.serializers import to_json
+
+        if driver_features is None:
+            driver_features = {}
+        if "json_serializer" not in driver_features:
+            driver_features["json_serializer"] = to_json
+        if "enable_cast_detection" not in driver_features:
+            driver_features["enable_cast_detection"] = True
+        if "strict_type_coercion" not in driver_features:
+            driver_features["strict_type_coercion"] = False
+        if "arrow_extension_types" not in driver_features:
+            driver_features["arrow_extension_types"] = True
+
         super().__init__(
             connection_config=self.connection_config,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features or {},
+            driver_features=dict(driver_features),
             bind_key=bind_key,
+            extension_config=extension_config,
         )
 
     def _resolve_driver_name(self) -> str:
@@ -284,7 +332,7 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
             connection.close()
 
     def provide_session(
-        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
     ) -> "AbstractContextManager[AdbcDriver]":
         """Provide a driver session context manager.
 
@@ -305,7 +353,9 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
                     or self.statement_config
                     or get_adbc_statement_config(str(self._get_dialect() or "sqlite"))
                 )
-                yield self.driver_type(connection=connection, statement_config=final_statement_config)
+                yield self.driver_type(
+                    connection=connection, statement_config=final_statement_config, driver_features=self.driver_features
+                )
 
         return session_manager()
 

sqlspec/adapters/adbc/data_dictionary.py

@@ -1,7 +1,7 @@
 """ADBC multi-dialect data dictionary for metadata queries."""
 
 import re
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, VersionInfo
 from sqlspec.utils.logging import get_logger
@@ -38,7 +38,7 @@ class AdbcDataDictionary(SyncDataDictionaryBase):
         """
         return str(cast("AdbcDriver", driver).dialect)
 
-    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
         """Get database version information based on detected dialect.
 
         Args:
@@ -268,6 +268,56 @@ class AdbcDataDictionary(SyncDataDictionaryBase):
 
         return type_map.get(type_category, "TEXT")
 
+    def get_columns(
+        self, driver: SyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table based on detected dialect.
+
+        Args:
+            driver: ADBC driver instance
+            table: Table name to query columns for
+            schema: Schema name (None for default)
+
+        Returns:
+            List of column metadata dictionaries with keys:
+                - column_name: Name of the column
+                - data_type: Database data type
+                - is_nullable or nullable: Whether column allows NULL
+                - column_default or default_value: Default value if any
+        """
+        dialect = self._get_dialect(driver)
+        adbc_driver = cast("AdbcDriver", driver)
+
+        if dialect == "sqlite":
+            result = adbc_driver.execute(f"PRAGMA table_info({table})")
+            return [
+                {
+                    "column_name": row["name"] if isinstance(row, dict) else row[1],
+                    "data_type": row["type"] if isinstance(row, dict) else row[2],
+                    "nullable": not (row["notnull"] if isinstance(row, dict) else row[3]),
+                    "default_value": row["dflt_value"] if isinstance(row, dict) else row[4],
+                }
+                for row in result.data or []
+            ]
+
+        if schema:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                ORDER BY ordinal_position
+            """
+        else:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}'
+                ORDER BY ordinal_position
+            """
+
+        result = adbc_driver.execute(sql)
+        return result.data or []
+
     def list_available_features(self) -> "list[str]":
         """List available feature flags across all supported dialects.
 

sqlspec/adapters/adbc/driver.py

@@ -7,9 +7,8 @@ database dialects, parameter style conversion, and transaction management.
 import contextlib
 import datetime
 import decimal
-from typing import TYPE_CHECKING, Any, Optional, cast
+from typing import TYPE_CHECKING, Any, cast
 
-from adbc_driver_manager.dbapi import DatabaseError, IntegrityError, OperationalError, ProgrammingError
 from sqlglot import exp
 
 from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary
@@ -18,7 +17,19 @@ from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import SyncDriverAdapterBase
-from sqlspec.exceptions import MissingDependencyError, SQLParsingError, SQLSpecError
+from sqlspec.exceptions import (
+    CheckViolationError,
+    DatabaseConnectionError,
+    DataError,
+    ForeignKeyViolationError,
+    IntegrityError,
+    MissingDependencyError,
+    NotNullViolationError,
+    SQLParsingError,
+    SQLSpecError,
+    TransactionError,
+    UniqueViolationError,
+)
 from sqlspec.typing import Empty
 from sqlspec.utils.logging import get_logger
 
@@ -329,7 +340,7 @@ class AdbcCursor:
 
     def __init__(self, connection: "AdbcConnection") -> None:
         self.connection = connection
-        self.cursor: Optional[Cursor] = None
+        self.cursor: Cursor | None = None
 
     def __enter__(self) -> "Cursor":
         self.cursor = self.connection.cursor()
@@ -342,7 +353,11 @@ class AdbcCursor:
 
 
 class AdbcExceptionHandler:
-    """Context manager for handling database exceptions."""
+    """Context manager for handling ADBC database exceptions.
+
+    ADBC propagates underlying database errors. Exception mapping
+    depends on the specific ADBC driver being used.
+    """
 
     __slots__ = ()
 
@@ -350,40 +365,118 @@ class AdbcExceptionHandler:
         return None
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        _ = exc_tb
         if exc_type is None:
             return
+        self._map_adbc_exception(exc_val)
 
-        try:
-            if issubclass(exc_type, IntegrityError):
-                e = exc_val
-                msg = f"Integrity constraint violation: {e}"
-                raise SQLSpecError(msg) from e
-            if issubclass(exc_type, ProgrammingError):
-                e = exc_val
-                error_msg = str(e).lower()
-                if "syntax" in error_msg or "parse" in error_msg:
-                    msg = f"SQL syntax error: {e}"
-                    raise SQLParsingError(msg) from e
-                msg = f"Programming error: {e}"
-                raise SQLSpecError(msg) from e
-            if issubclass(exc_type, OperationalError):
-                e = exc_val
-                msg = f"Operational error: {e}"
-                raise SQLSpecError(msg) from e
-            if issubclass(exc_type, DatabaseError):
-                e = exc_val
-                msg = f"Database error: {e}"
-                raise SQLSpecError(msg) from e
-        except ImportError:
-            pass
-        if issubclass(exc_type, Exception):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "parse" in error_msg or "syntax" in error_msg:
-                msg = f"SQL parsing failed: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"Unexpected database operation error: {e}"
-            raise SQLSpecError(msg) from e
+    def _map_adbc_exception(self, e: Any) -> None:
+        """Map ADBC exception to SQLSpec exception.
+
+        ADBC drivers may expose SQLSTATE codes or driver-specific codes.
+
+        Args:
+            e: ADBC exception instance
+        """
+        sqlstate = getattr(e, "sqlstate", None)
+
+        if sqlstate:
+            self._map_sqlstate_exception(e, sqlstate)
+        else:
+            self._map_message_based_exception(e)
+
+    def _map_sqlstate_exception(self, e: Any, sqlstate: str) -> None:
+        """Map SQLSTATE code to exception.
+
+        Args:
+            e: Exception instance
+            sqlstate: SQLSTATE error code
+        """
+        if sqlstate == "23505":
+            self._raise_unique_violation(e)
+        elif sqlstate == "23503":
+            self._raise_foreign_key_violation(e)
+        elif sqlstate == "23502":
+            self._raise_not_null_violation(e)
+        elif sqlstate == "23514":
+            self._raise_check_violation(e)
+        elif sqlstate.startswith("23"):
+            self._raise_integrity_error(e)
+        elif sqlstate.startswith("42"):
+            self._raise_parsing_error(e)
+        elif sqlstate.startswith("08"):
+            self._raise_connection_error(e)
+        elif sqlstate.startswith("40"):
+            self._raise_transaction_error(e)
+        elif sqlstate.startswith("22"):
+            self._raise_data_error(e)
+        else:
+            self._raise_generic_error(e)
+
+    def _map_message_based_exception(self, e: Any) -> None:
+        """Map exception using message-based detection.
+
+        Args:
+            e: Exception instance
+        """
+        error_msg = str(e).lower()
+
+        if "unique" in error_msg or "duplicate" in error_msg:
+            self._raise_unique_violation(e)
+        elif "foreign key" in error_msg:
+            self._raise_foreign_key_violation(e)
+        elif "not null" in error_msg or "null value" in error_msg:
+            self._raise_not_null_violation(e)
+        elif "check constraint" in error_msg:
+            self._raise_check_violation(e)
+        elif "constraint" in error_msg:
+            self._raise_integrity_error(e)
+        elif "syntax" in error_msg:
+            self._raise_parsing_error(e)
+        elif "connection" in error_msg or "connect" in error_msg:
+            self._raise_connection_error(e)
+        else:
+            self._raise_generic_error(e)
+
+    def _raise_unique_violation(self, e: Any) -> None:
+        msg = f"ADBC unique constraint violation: {e}"
+        raise UniqueViolationError(msg) from e
+
+    def _raise_foreign_key_violation(self, e: Any) -> None:
+        msg = f"ADBC foreign key constraint violation: {e}"
+        raise ForeignKeyViolationError(msg) from e
+
+    def _raise_not_null_violation(self, e: Any) -> None:
+        msg = f"ADBC not-null constraint violation: {e}"
+        raise NotNullViolationError(msg) from e
+
+    def _raise_check_violation(self, e: Any) -> None:
+        msg = f"ADBC check constraint violation: {e}"
+        raise CheckViolationError(msg) from e
+
+    def _raise_integrity_error(self, e: Any) -> None:
+        msg = f"ADBC integrity constraint violation: {e}"
+        raise IntegrityError(msg) from e
+
+    def _raise_parsing_error(self, e: Any) -> None:
+        msg = f"ADBC SQL parsing error: {e}"
+        raise SQLParsingError(msg) from e
+
+    def _raise_connection_error(self, e: Any) -> None:
+        msg = f"ADBC connection error: {e}"
+        raise DatabaseConnectionError(msg) from e
+
+    def _raise_transaction_error(self, e: Any) -> None:
+        msg = f"ADBC transaction error: {e}"
+        raise TransactionError(msg) from e
+
+    def _raise_data_error(self, e: Any) -> None:
+        msg = f"ADBC data error: {e}"
+        raise DataError(msg) from e
+
+    def _raise_generic_error(self, e: Any) -> None:
+        msg = f"ADBC database error: {e}"
+        raise SQLSpecError(msg) from e
 
 
 class AdbcDriver(SyncDriverAdapterBase):
@@ -398,8 +491,8 @@ class AdbcDriver(SyncDriverAdapterBase):
     def __init__(
         self,
         connection: "AdbcConnection",
-        statement_config: "Optional[StatementConfig]" = None,
-        driver_features: "Optional[dict[str, Any]]" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "dict[str, Any] | None" = None,
     ) -> None:
         self._detected_dialect = self._get_dialect(connection)
 
@@ -412,7 +505,7 @@ class AdbcDriver(SyncDriverAdapterBase):
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
         self.dialect = statement_config.dialect
-        self._data_dictionary: Optional[SyncDataDictionaryBase] = None
+        self._data_dictionary: SyncDataDictionaryBase | None = None
 
     @staticmethod
     def _ensure_pyarrow_installed() -> None:
@@ -480,12 +573,13 @@ class AdbcDriver(SyncDriverAdapterBase):
         parameters: Any,
         statement_config: "StatementConfig",
         is_many: bool = False,
-        prepared_statement: Optional[Any] = None,
+        prepared_statement: Any | None = None,
     ) -> Any:
         """Prepare parameters with cast-aware type coercion for ADBC.
 
         For PostgreSQL, applies cast-aware parameter processing using metadata from the compiled statement.
         This allows proper handling of JSONB casts and other type conversions.
+        Respects driver_features['enable_cast_detection'] configuration.
 
         Args:
             parameters: Parameters in any format
@@ -496,7 +590,9 @@ class AdbcDriver(SyncDriverAdapterBase):
         Returns:
             Parameters with cast-aware type coercion applied
         """
-        if prepared_statement and self.dialect in {"postgres", "postgresql"} and not is_many:
+        enable_cast_detection = self.driver_features.get("enable_cast_detection", True)
+
+        if enable_cast_detection and prepared_statement and self.dialect in {"postgres", "postgresql"} and not is_many:
             parameter_casts = self._get_parameter_casts(prepared_statement)
             postgres_compatible = self._handle_postgres_empty_parameters(parameters)
             return self._prepare_parameters_with_casts(postgres_compatible, parameter_casts, statement_config)
@@ -524,6 +620,7 @@ class AdbcDriver(SyncDriverAdapterBase):
         """Prepare parameters with cast-aware type coercion.
 
         Uses type coercion map for non-dict types and dialect-aware dict handling.
+        Respects driver_features configuration for JSON serialization backend.
 
         Args:
             parameters: Parameter values (list, tuple, or scalar)
@@ -533,7 +630,9 @@ class AdbcDriver(SyncDriverAdapterBase):
         Returns:
             Parameters with cast-aware type coercion applied
         """
-        from sqlspec._serialization import encode_json
+        from sqlspec.utils.serializers import to_json
+
+        json_encoder = self.driver_features.get("json_serializer", to_json)
 
         if isinstance(parameters, (list, tuple)):
             result: list[Any] = []
@@ -541,7 +640,7 @@ class AdbcDriver(SyncDriverAdapterBase):
                 cast_type = parameter_casts.get(idx, "").upper()
                 if cast_type in {"JSON", "JSONB", "TYPE.JSON", "TYPE.JSONB"}:
                     if isinstance(param, dict):
-                        result.append(encode_json(param))
+                        result.append(json_encoder(param))
                     else:
                         result.append(param)
                 elif isinstance(param, dict):
@@ -575,7 +674,7 @@ class AdbcDriver(SyncDriverAdapterBase):
         """
         return AdbcExceptionHandler()
 
-    def _try_special_handling(self, cursor: "Cursor", statement: SQL) -> "Optional[SQLResult]":
+    def _try_special_handling(self, cursor: "Cursor", statement: SQL) -> "SQLResult | None":
         """Handle special operations.
 
         Args:
@@ -668,7 +767,7 @@ class AdbcDriver(SyncDriverAdapterBase):
             column_names = [col[0] for col in cursor.description or []]
 
             if fetched_data and isinstance(fetched_data[0], tuple):
-                dict_data: list[dict[Any, Any]] = [dict(zip(column_names, row)) for row in fetched_data]
+                dict_data: list[dict[Any, Any]] = [dict(zip(column_names, row, strict=False)) for row in fetched_data]
             else:
                 dict_data = fetched_data  # type: ignore[assignment]
 

sqlspec/adapters/adbc/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for ADBC adapter."""
+
+from sqlspec.adapters.adbc.litestar.store import ADBCStore
+
+__all__ = ("ADBCStore",)