"""AsyncMy MySQL driver implementation.
Provides MySQL/MariaDB connectivity with parameter style conversion,
type coercion, error handling, and transaction management.
"""
from collections.abc import Sized
from typing import TYPE_CHECKING, Any, Final, cast
from sqlspec.adapters.asyncmy._typing import (
AsyncmyCursor,
AsyncmyError,
AsyncmyFieldType,
AsyncmyMySQLError,
AsyncmySessionContext,
)
from sqlspec.adapters.asyncmy.core import (
AsyncmyStreamSource,
build_insert_statement,
collect_rows,
create_mapped_exception,
default_statement_config,
driver_profile,
format_identifier,
normalize_execute_many_parameters,
normalize_execute_parameters,
normalize_lastrowid,
resolve_many_rowcount,
resolve_row_plan,
resolve_rowcount,
)
from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary
from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile
from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler
from sqlspec.exceptions import SQLSpecError
from sqlspec.utils.logging import get_logger
from sqlspec.utils.serializers import from_json
from sqlspec.utils.type_guards import supports_json_type
if TYPE_CHECKING:
from collections.abc import Callable
from sqlspec.adapters.asyncmy._typing import AsyncmyConnection
from sqlspec.core import SQL, StatementConfig
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "AsyncmySessionContext")
logger = get_logger(__name__)
json_type_value = (
AsyncmyFieldType.JSON if AsyncmyFieldType is not None and supports_json_type(AsyncmyFieldType) else None
)
ASYNCMY_JSON_TYPE_CODES: Final[set[int]] = {json_type_value} if json_type_value is not None else set()
_MYSQL_TYPE_CODE_TOKENS: Final[dict[int, str]] = {
0: "decimal",
1: "int32",
2: "int32",
3: "int64",
4: "float32",
5: "float64",
7: "timestamp",
8: "int64",
10: "date",
11: "time",
12: "timestamp",
246: "decimal",
252: "binary",
253: "string",
254: "string",
}
def _resolve_column_types(description: Any) -> "dict[str, str] | None":
"""Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens.
Returns ``None`` when the cursor has no description or reports no
recognizable type codes.
"""
if not description:
return None
column_types: dict[str, str] = {}
for col in description:
token = _MYSQL_TYPE_CODE_TOKENS.get(col[1])
if token is not None:
column_types[col[0]] = token
return column_types or None
class AsyncmyExceptionHandler(BaseAsyncExceptionHandler):
"""Async context manager for handling asyncmy (MySQL) database exceptions.
Maps MySQL error codes and SQLSTATE 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:
if exc_type is None:
return False
if issubclass(exc_type, cast("type[BaseException]", AsyncmyError)):
result = create_mapped_exception(exc_val, logger=logger)
if result is True:
return True
self.pending_exception = cast("Exception", result)
return True
return False
class AsyncmyDriver(AsyncDriverAdapterBase):
"""MySQL/MariaDB database driver using AsyncMy client library.
Implements asynchronous database operations for MySQL and MariaDB servers
with support for parameter style conversion, type coercion, error handling,
and transaction management.
"""
__slots__ = ("_data_dictionary",)
dialect = "mysql"
def __init__(
self,
connection: "AsyncmyConnection",
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: AsyncmyDataDictionary | None = None