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

sqlspec/adapters/asyncmy/data_dictionary.py

@@ -1,12 +1,14 @@
 """MySQL-specific data dictionary for metadata queries via asyncmy."""
 
 import re
-from typing import TYPE_CHECKING, Callable, Optional, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.driver import AsyncDataDictionaryBase, AsyncDriverAdapterBase, VersionInfo
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from sqlspec.adapters.asyncmy.driver import AsyncmyDriver
 
 logger = get_logger("adapters.asyncmy.data_dictionary")
@@ -20,7 +22,7 @@ __all__ = ("MySQLAsyncDataDictionary",)
 class MySQLAsyncDataDictionary(AsyncDataDictionaryBase):
     """MySQL-specific async data dictionary."""
 
-    async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
         """Get MySQL database version information.
 
         Args:
@@ -102,6 +104,43 @@ class MySQLAsyncDataDictionary(AsyncDataDictionaryBase):
         }
         return type_map.get(type_category, "VARCHAR(255)")
 
+    async def get_columns(
+        self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table using information_schema.
+
+        Args:
+            driver: AsyncMy driver instance
+            table: Table name to query columns for
+            schema: Schema name (database name in MySQL)
+
+        Returns:
+            List of column metadata dictionaries with keys:
+                - column_name: Name of the column
+                - data_type: MySQL data type
+                - is_nullable: Whether column allows NULL (YES/NO)
+                - column_default: Default value if any
+        """
+        asyncmy_driver = cast("AsyncmyDriver", driver)
+
+        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 = await asyncmy_driver.execute(sql)
+        return result.data or []
+
     def list_available_features(self) -> "list[str]":
         """List available MySQL feature flags.
 

sqlspec/adapters/asyncmy/driver.py

@@ -5,20 +5,32 @@ type coercion, error handling, and transaction management.
 """
 
 import logging
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Final
 
-import asyncmy
 import asyncmy.errors  # pyright: ignore
+from asyncmy.constants import FIELD_TYPE as ASYNC_MY_FIELD_TYPE  # pyright: ignore
 from asyncmy.cursors import Cursor, DictCursor  # pyright: ignore
 
 from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import StatementConfig
 from sqlspec.driver import AsyncDriverAdapterBase
-from sqlspec.exceptions import SQLParsingError, SQLSpecError
+from sqlspec.exceptions import (
+    CheckViolationError,
+    DatabaseConnectionError,
+    DataError,
+    ForeignKeyViolationError,
+    IntegrityError,
+    NotNullViolationError,
+    SQLParsingError,
+    SQLSpecError,
+    TransactionError,
+    UniqueViolationError,
+)
 from sqlspec.utils.serializers import to_json
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from contextlib import AbstractAsyncContextManager
 
     from sqlspec.adapters.asyncmy._types import AsyncmyConnection
@@ -26,11 +38,17 @@ if TYPE_CHECKING:
     from sqlspec.core.statement import SQL
     from sqlspec.driver import ExecutionResult
     from sqlspec.driver._async import AsyncDataDictionaryBase
+__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "asyncmy_statement_config")
 
 logger = logging.getLogger(__name__)
 
-__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "asyncmy_statement_config")
-
+json_type_value = (
+    ASYNC_MY_FIELD_TYPE.JSON if ASYNC_MY_FIELD_TYPE is not None and hasattr(ASYNC_MY_FIELD_TYPE, "JSON") else None
+)
+ASYNCMY_JSON_TYPE_CODES: Final[set[int]] = {json_type_value} if json_type_value is not None else set()
+MYSQL_ER_DUP_ENTRY = 1062
+MYSQL_ER_NO_DEFAULT_FOR_FIELD = 1364
+MYSQL_ER_CHECK_CONSTRAINT_VIOLATED = 3819
 
 asyncmy_statement_config = StatementConfig(
     dialect="mysql",
@@ -61,9 +79,9 @@ class AsyncmyCursor:
 
     def __init__(self, connection: "AsyncmyConnection") -> None:
         self.connection = connection
-        self.cursor: Optional[Union[Cursor, DictCursor]] = None
+        self.cursor: Cursor | DictCursor | None = None
 
-    async def __aenter__(self) -> Union[Cursor, DictCursor]:
+    async def __aenter__(self) -> Cursor | DictCursor:
         self.cursor = self.connection.cursor()
         return self.cursor
 
@@ -73,10 +91,10 @@ class AsyncmyCursor:
 
 
 class AsyncmyExceptionHandler:
-    """Context manager for AsyncMy database exception handling.
+    """Async context manager for handling asyncmy (MySQL) database exceptions.
 
-    Converts AsyncMy-specific exceptions to SQLSpec exceptions with appropriate
-    error categorization and context preservation.
+    Maps MySQL error codes and SQLSTATE to specific SQLSpec exceptions
+    for better error handling in application code.
     """
 
     __slots__ = ()
@@ -84,53 +102,119 @@ class AsyncmyExceptionHandler:
     async def __aenter__(self) -> None:
         return None
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> "Optional[bool]":
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> "bool | None":
         if exc_type is None:
             return None
-
-        if issubclass(exc_type, asyncmy.errors.IntegrityError):
-            e = exc_val
-            msg = f"AsyncMy MySQL integrity constraint violation: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, asyncmy.errors.ProgrammingError):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "syntax" in error_msg or "parse" in error_msg:
-                msg = f"AsyncMy MySQL SQL syntax error: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"AsyncMy MySQL programming error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, asyncmy.errors.OperationalError):
-            e = exc_val
-            # Handle specific MySQL errors that are expected in migrations
-            if hasattr(e, "args") and len(e.args) >= 1 and isinstance(e.args[0], int):
-                error_code = e.args[0]
-                # Error 1061: Duplicate key name (index already exists)
-                # Error 1091: Can't DROP index that doesn't exist
-                if error_code in {1061, 1091}:
-                    # These are acceptable during migrations - log and continue
-                    logger.warning("AsyncMy MySQL expected migration error (ignoring): %s", e)
-                    return True  # Suppress the exception by returning True
-            msg = f"AsyncMy MySQL operational error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, asyncmy.errors.DatabaseError):
-            e = exc_val
-            msg = f"AsyncMy MySQL database error: {e}"
-            raise SQLSpecError(msg) from e
         if issubclass(exc_type, asyncmy.errors.Error):
-            e = exc_val
-            msg = f"AsyncMy MySQL error: {e}"
-            raise SQLSpecError(msg) from e
-        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 async database operation error: {e}"
-            raise SQLSpecError(msg) from e
+            return self._map_mysql_exception(exc_val)
         return None
 
+    def _map_mysql_exception(self, e: Any) -> "bool | None":
+        """Map MySQL exception to SQLSpec exception.
+
+        Args:
+            e: MySQL error instance
+
+        Returns:
+            True to suppress migration-related errors, None otherwise
+
+        Raises:
+            Specific SQLSpec exception based on error code
+        """
+        error_code = None
+        sqlstate = None
+
+        if hasattr(e, "args") and len(e.args) >= 1 and isinstance(e.args[0], int):
+            error_code = e.args[0]
+
+        sqlstate = getattr(e, "sqlstate", None)
+
+        if error_code in {1061, 1091}:
+            logger.warning("AsyncMy MySQL expected migration error (ignoring): %s", e)
+            return True
+
+        if sqlstate == "23505" or error_code == MYSQL_ER_DUP_ENTRY:
+            self._raise_unique_violation(e, sqlstate, error_code)
+        elif sqlstate == "23503" or error_code in (1216, 1217, 1451, 1452):
+            self._raise_foreign_key_violation(e, sqlstate, error_code)
+        elif sqlstate == "23502" or error_code in (1048, MYSQL_ER_NO_DEFAULT_FOR_FIELD):
+            self._raise_not_null_violation(e, sqlstate, error_code)
+        elif sqlstate == "23514" or error_code == MYSQL_ER_CHECK_CONSTRAINT_VIOLATED:
+            self._raise_check_violation(e, sqlstate, error_code)
+        elif sqlstate and sqlstate.startswith("23"):
+            self._raise_integrity_error(e, sqlstate, error_code)
+        elif sqlstate and sqlstate.startswith("42"):
+            self._raise_parsing_error(e, sqlstate, error_code)
+        elif sqlstate and sqlstate.startswith("08"):
+            self._raise_connection_error(e, sqlstate, error_code)
+        elif sqlstate and sqlstate.startswith("40"):
+            self._raise_transaction_error(e, sqlstate, error_code)
+        elif sqlstate and sqlstate.startswith("22"):
+            self._raise_data_error(e, sqlstate, error_code)
+        elif error_code in {2002, 2003, 2005, 2006, 2013}:
+            self._raise_connection_error(e, sqlstate, error_code)
+        elif error_code in {1205, 1213}:
+            self._raise_transaction_error(e, sqlstate, error_code)
+        elif error_code in range(1064, 1100):
+            self._raise_parsing_error(e, sqlstate, error_code)
+        else:
+            self._raise_generic_error(e, sqlstate, error_code)
+        return None
+
+    def _raise_unique_violation(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL unique constraint violation {code_str}: {e}"
+        raise UniqueViolationError(msg) from e
+
+    def _raise_foreign_key_violation(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL foreign key constraint violation {code_str}: {e}"
+        raise ForeignKeyViolationError(msg) from e
+
+    def _raise_not_null_violation(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL not-null constraint violation {code_str}: {e}"
+        raise NotNullViolationError(msg) from e
+
+    def _raise_check_violation(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL check constraint violation {code_str}: {e}"
+        raise CheckViolationError(msg) from e
+
+    def _raise_integrity_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL integrity constraint violation {code_str}: {e}"
+        raise IntegrityError(msg) from e
+
+    def _raise_parsing_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL SQL syntax error {code_str}: {e}"
+        raise SQLParsingError(msg) from e
+
+    def _raise_connection_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL connection error {code_str}: {e}"
+        raise DatabaseConnectionError(msg) from e
+
+    def _raise_transaction_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL transaction error {code_str}: {e}"
+        raise TransactionError(msg) from e
+
+    def _raise_data_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        code_str = f"[{sqlstate or code}]"
+        msg = f"MySQL data error {code_str}: {e}"
+        raise DataError(msg) from e
+
+    def _raise_generic_error(self, e: Any, sqlstate: "str | None", code: "int | None") -> None:
+        if sqlstate and code:
+            msg = f"MySQL database error [{sqlstate}:{code}]: {e}"
+        elif sqlstate or code:
+            msg = f"MySQL database error [{sqlstate or code}]: {e}"
+        else:
+            msg = f"MySQL database error: {e}"
+        raise SQLSpecError(msg) from e
+
 
 class AsyncmyDriver(AsyncDriverAdapterBase):
     """MySQL/MariaDB database driver using AsyncMy client library.
@@ -146,20 +230,94 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
     def __init__(
         self,
         connection: "AsyncmyConnection",
-        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:
-        if statement_config is None:
+        final_statement_config = statement_config
+        if final_statement_config is None:
             cache_config = get_cache_config()
-            statement_config = asyncmy_statement_config.replace(
+            final_statement_config = asyncmy_statement_config.replace(
                 enable_caching=cache_config.compiled_cache_enabled,
                 enable_parsing=True,
                 enable_validation=True,
                 dialect="mysql",
             )
 
-        super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
-        self._data_dictionary: Optional[AsyncDataDictionaryBase] = None
+        final_statement_config = self._apply_json_serializer_feature(final_statement_config, driver_features)
+
+        super().__init__(
+            connection=connection, statement_config=final_statement_config, driver_features=driver_features
+        )
+        self._data_dictionary: AsyncDataDictionaryBase | None = None
+
+    @staticmethod
+    def _clone_parameter_config(
+        parameter_config: ParameterStyleConfig, type_coercion_map: "dict[type[Any], Callable[[Any], Any]]"
+    ) -> ParameterStyleConfig:
+        """Create a copy of the parameter configuration with updated coercion map.
+
+        Args:
+            parameter_config: Existing parameter configuration to copy.
+            type_coercion_map: Updated coercion mapping for parameter serialization.
+
+        Returns:
+            ParameterStyleConfig with the updated type coercion map applied.
+        """
+
+        supported_execution_styles = (
+            set(parameter_config.supported_execution_parameter_styles)
+            if parameter_config.supported_execution_parameter_styles is not None
+            else None
+        )
+
+        return ParameterStyleConfig(
+            default_parameter_style=parameter_config.default_parameter_style,
+            supported_parameter_styles=set(parameter_config.supported_parameter_styles),
+            supported_execution_parameter_styles=supported_execution_styles,
+            default_execution_parameter_style=parameter_config.default_execution_parameter_style,
+            type_coercion_map=type_coercion_map,
+            has_native_list_expansion=parameter_config.has_native_list_expansion,
+            needs_static_script_compilation=parameter_config.needs_static_script_compilation,
+            allow_mixed_parameter_styles=parameter_config.allow_mixed_parameter_styles,
+            preserve_parameter_format=parameter_config.preserve_parameter_format,
+            preserve_original_params_for_many=parameter_config.preserve_original_params_for_many,
+            output_transformer=parameter_config.output_transformer,
+            ast_transformer=parameter_config.ast_transformer,
+        )
+
+    @staticmethod
+    def _apply_json_serializer_feature(
+        statement_config: "StatementConfig", driver_features: "dict[str, Any] | None"
+    ) -> "StatementConfig":
+        """Apply driver-level JSON serializer customization to the statement config.
+
+        Args:
+            statement_config: Base statement configuration for the driver.
+            driver_features: Driver feature mapping provided via configuration.
+
+        Returns:
+            StatementConfig with serializer adjustments applied when configured.
+        """
+
+        if not driver_features:
+            return statement_config
+
+        serializer = driver_features.get("json_serializer")
+        if serializer is None:
+            return statement_config
+
+        parameter_config = statement_config.parameter_config
+        type_coercion_map = dict(parameter_config.type_coercion_map)
+
+        def serialize_tuple(value: Any) -> Any:
+            return serializer(list(value))
+
+        type_coercion_map[dict] = serializer
+        type_coercion_map[list] = serializer
+        type_coercion_map[tuple] = serialize_tuple
+
+        updated_parameter_config = AsyncmyDriver._clone_parameter_config(parameter_config, type_coercion_map)
+        return statement_config.replace(parameter_config=updated_parameter_config)
 
     def with_cursor(self, connection: "AsyncmyConnection") -> "AsyncmyCursor":
         """Create cursor context manager for the connection.
@@ -180,7 +338,7 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         """
         return AsyncmyExceptionHandler()
 
-    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
+    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None":
         """Handle AsyncMy-specific operations before standard execution.
 
         Args:
@@ -193,6 +351,75 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         _ = (cursor, statement)
         return None
 
+    def _detect_json_columns(self, cursor: Any) -> "list[int]":
+        """Identify JSON column indexes from cursor metadata.
+
+        Args:
+            cursor: Database cursor with description metadata available.
+
+        Returns:
+            List of index positions where JSON values are present.
+        """
+
+        description = getattr(cursor, "description", None)
+        if not description or not ASYNCMY_JSON_TYPE_CODES:
+            return []
+
+        json_indexes: list[int] = []
+        for index, column in enumerate(description):
+            type_code = getattr(column, "type_code", None)
+            if type_code is None and isinstance(column, (tuple, list)) and len(column) > 1:
+                type_code = column[1]
+            if type_code in ASYNCMY_JSON_TYPE_CODES:
+                json_indexes.append(index)
+        return json_indexes
+
+    def _deserialize_json_columns(
+        self, cursor: Any, column_names: "list[str]", rows: "list[dict[str, Any]]"
+    ) -> "list[dict[str, Any]]":
+        """Apply configured JSON deserializer to result rows.
+
+        Args:
+            cursor: Database cursor used for the current result set.
+            column_names: Ordered column names from the cursor description.
+            rows: Result rows represented as dictionaries.
+
+        Returns:
+            Rows with JSON columns decoded when a deserializer is configured.
+        """
+
+        if not rows or not column_names:
+            return rows
+
+        deserializer = self.driver_features.get("json_deserializer")
+        if deserializer is None:
+            return rows
+
+        json_indexes = self._detect_json_columns(cursor)
+        if not json_indexes:
+            return rows
+
+        target_columns = [column_names[index] for index in json_indexes if index < len(column_names)]
+        if not target_columns:
+            return rows
+
+        for row in rows:
+            for column in target_columns:
+                if column not in row:
+                    continue
+                raw_value = row[column]
+                if raw_value is None:
+                    continue
+                if isinstance(raw_value, bytearray):
+                    raw_value = bytes(raw_value)
+                if not isinstance(raw_value, (str, bytes)):
+                    continue
+                try:
+                    row[column] = deserializer(raw_value)
+                except Exception:
+                    logger.debug("Failed to deserialize JSON column %s", column, exc_info=True)
+        return rows
+
     async def _execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
         """Execute SQL script with statement splitting and parameter handling.
 
@@ -269,12 +496,16 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
             column_names = [desc[0] for desc in cursor.description or []]
 
             if fetched_data and not isinstance(fetched_data[0], dict):
-                data = [dict(zip(column_names, row)) for row in fetched_data]
+                rows = [dict(zip(column_names, row, strict=False)) for row in fetched_data]
+            elif fetched_data:
+                rows = [dict(row) for row in fetched_data]
             else:
-                data = fetched_data
+                rows = []
+
+            rows = self._deserialize_json_columns(cursor, column_names, rows)
 
             return self.create_execution_result(
-                cursor, selected_data=data, column_names=column_names, data_row_count=len(data), is_select_result=True
+                cursor, selected_data=rows, column_names=column_names, data_row_count=len(rows), is_select_result=True
             )
 
         affected_rows = cursor.rowcount if cursor.rowcount is not None else -1

sqlspec/adapters/asyncmy/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for AsyncMy adapter."""
+
+from sqlspec.adapters.asyncmy.litestar.store import AsyncmyStore
+
+__all__ = ("AsyncmyStore",)