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

sqlspec/adapters/adbc/driver.py

@@ -7,18 +7,31 @@ 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
+from sqlspec.adapters.adbc.type_converter import ADBCTypeConverter
 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
-from sqlspec.utils.serializers import to_json
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
@@ -28,6 +41,7 @@ if TYPE_CHECKING:
     from sqlspec.adapters.adbc._types import AdbcConnection
     from sqlspec.core.result import SQLResult
     from sqlspec.driver import ExecutionResult
+    from sqlspec.driver._sync import SyncDataDictionaryBase
 
 __all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "get_adbc_statement_config")
 
@@ -293,7 +307,7 @@ def _convert_array_for_postgres_adbc(value: Any) -> Any:
 
 
 def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
-    """Get type coercion map for Arrow type handling.
+    """Get type coercion map for Arrow type handling with dialect-aware type conversion.
 
     Args:
         dialect: Database dialect name
@@ -301,7 +315,9 @@ def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
     Returns:
         Mapping of Python types to conversion functions
     """
-    type_map = {
+    tc = ADBCTypeConverter(dialect)
+
+    return {
         datetime.datetime: lambda x: x,
         datetime.date: lambda x: x,
         datetime.time: lambda x: x,
@@ -309,18 +325,13 @@ def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
         bool: lambda x: x,
         int: lambda x: x,
         float: lambda x: x,
-        str: lambda x: x,
+        str: tc.convert_if_detected,
         bytes: lambda x: x,
         tuple: _convert_array_for_postgres_adbc,
         list: _convert_array_for_postgres_adbc,
         dict: lambda x: x,
     }
 
-    if dialect in {"postgres", "postgresql"}:
-        type_map[dict] = lambda x: to_json(x) if x is not None else None
-
-    return type_map
-
 
 class AdbcCursor:
     """Context manager for cursor management."""
@@ -329,21 +340,24 @@ 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()
         return self.cursor
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)
+    def __exit__(self, *_: Any) -> None:
         if self.cursor is not None:
             with contextlib.suppress(Exception):
                 self.cursor.close()  # type: ignore[no-untyped-call]
 
 
 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__ = ()
 
@@ -351,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):
@@ -394,13 +486,13 @@ class AdbcDriver(SyncDriverAdapterBase):
     database dialects, parameter style conversion, and transaction management.
     """
 
-    __slots__ = ("_detected_dialect", "dialect")
+    __slots__ = ("_data_dictionary", "_detected_dialect", "dialect")
 
     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)
 
@@ -413,6 +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: SyncDataDictionaryBase | None = None
 
     @staticmethod
     def _ensure_pyarrow_installed() -> None:
@@ -475,6 +568,93 @@ class AdbcDriver(SyncDriverAdapterBase):
             return None
         return parameters
 
+    def prepare_driver_parameters(
+        self,
+        parameters: Any,
+        statement_config: "StatementConfig",
+        is_many: bool = False,
+        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
+            statement_config: Statement configuration
+            is_many: Whether this is for execute_many operation
+            prepared_statement: Prepared statement containing the original SQL statement
+
+        Returns:
+            Parameters with cast-aware type coercion applied
+        """
+        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)
+
+        return super().prepare_driver_parameters(parameters, statement_config, is_many, prepared_statement)
+
+    def _get_parameter_casts(self, statement: SQL) -> "dict[int, str]":
+        """Get parameter cast metadata from compiled statement.
+
+        Args:
+            statement: SQL statement with compiled metadata
+
+        Returns:
+            Dict mapping parameter positions to cast types
+        """
+
+        processed_state = statement.get_processed_state()
+        if processed_state is not Empty:
+            return processed_state.parameter_casts or {}
+        return {}
+
+    def _prepare_parameters_with_casts(
+        self, parameters: Any, parameter_casts: "dict[int, str]", statement_config: "StatementConfig"
+    ) -> Any:
+        """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)
+            parameter_casts: Mapping of parameter positions to cast types
+            statement_config: Statement configuration for type coercion
+
+        Returns:
+            Parameters with cast-aware type coercion applied
+        """
+        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] = []
+            for idx, param in enumerate(parameters, start=1):  # pyright: ignore
+                cast_type = parameter_casts.get(idx, "").upper()
+                if cast_type in {"JSON", "JSONB", "TYPE.JSON", "TYPE.JSONB"}:
+                    if isinstance(param, dict):
+                        result.append(json_encoder(param))
+                    else:
+                        result.append(param)
+                elif isinstance(param, dict):
+                    result.append(ADBCTypeConverter(self.dialect).convert_dict(param))  # type: ignore[arg-type]
+                else:
+                    if statement_config.parameter_config.type_coercion_map:
+                        for type_check, converter in statement_config.parameter_config.type_coercion_map.items():
+                            if type_check is not dict and isinstance(param, type_check):
+                                param = converter(param)
+                                break
+                    result.append(param)
+            return tuple(result) if isinstance(parameters, tuple) else result
+        return parameters
+
     def with_cursor(self, connection: "AdbcConnection") -> "AdbcCursor":
         """Create context manager for cursor.
 
@@ -494,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:
@@ -519,17 +699,26 @@ class AdbcDriver(SyncDriverAdapterBase):
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
+        parameter_casts = self._get_parameter_casts(statement)
+
         try:
             if not prepared_parameters:
                 cursor._rowcount = 0  # pyright: ignore[reportPrivateUsage]
                 row_count = 0
-            elif isinstance(prepared_parameters, list) and prepared_parameters:
+            elif isinstance(prepared_parameters, (list, tuple)) and prepared_parameters:
                 processed_params = []
                 for param_set in prepared_parameters:
                     postgres_compatible = self._handle_postgres_empty_parameters(param_set)
-                    formatted_params = self.prepare_driver_parameters(
-                        postgres_compatible, self.statement_config, is_many=False
-                    )
+
+                    if self.dialect in {"postgres", "postgresql"}:
+                        # For postgres, always use cast-aware parameter preparation
+                        formatted_params = self._prepare_parameters_with_casts(
+                            postgres_compatible, parameter_casts, self.statement_config
+                        )
+                    else:
+                        formatted_params = self.prepare_driver_parameters(
+                            postgres_compatible, self.statement_config, is_many=False
+                        )
                     processed_params.append(formatted_params)
 
                 cursor.executemany(sql, processed_params)
@@ -540,7 +729,6 @@ class AdbcDriver(SyncDriverAdapterBase):
 
         except Exception:
             self._handle_postgres_rollback(cursor)
-            logger.exception("Executemany failed")
             raise
 
         return self.create_execution_result(cursor, rowcount_override=row_count, is_many_result=True)
@@ -557,9 +745,18 @@ class AdbcDriver(SyncDriverAdapterBase):
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
+        parameter_casts = self._get_parameter_casts(statement)
+
         try:
             postgres_compatible_params = self._handle_postgres_empty_parameters(prepared_parameters)
-            cursor.execute(sql, parameters=postgres_compatible_params)
+
+            if self.dialect in {"postgres", "postgresql"}:
+                formatted_params = self._prepare_parameters_with_casts(
+                    postgres_compatible_params, parameter_casts, self.statement_config
+                )
+                cursor.execute(sql, parameters=formatted_params)
+            else:
+                cursor.execute(sql, parameters=postgres_compatible_params)
 
         except Exception:
             self._handle_postgres_rollback(cursor)
@@ -570,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]
 
@@ -655,3 +852,14 @@ class AdbcDriver(SyncDriverAdapterBase):
         except Exception as e:
             msg = f"Failed to commit transaction: {e}"
             raise SQLSpecError(msg) from e
+
+    @property
+    def data_dictionary(self) -> "SyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            self._data_dictionary = AdbcDataDictionary()
+        return self._data_dictionary

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",)