"""aiomysql MySQL driver implementation.
Provides MySQL/MariaDB connectivity with parameter style conversion,
type coercion, error handling, and transaction management.
aiomysql is built on top of PyMySQL, so error classes come from pymysql.err
rather than a driver-specific error module.
"""
import tempfile
from collections.abc import Sized
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, cast
from sqlspec.adapters.aiomysql._typing import (
AiomysqlCursor,
AiomysqlFieldType,
AiomysqlPymysqlError,
AiomysqlPymysqlMySQLError,
AiomysqlSessionContext,
)
from sqlspec.adapters.aiomysql.core import (
AiomysqlStreamSource,
build_insert_statement,
build_load_data_statement,
collect_rows,
create_mapped_exception,
default_statement_config,
driver_profile,
encode_records_for_local_infile,
format_identifier,
normalize_execute_many_parameters,
normalize_execute_parameters,
normalize_lastrowid,
resolve_many_rowcount,
resolve_row_plan,
resolve_rowcount,
)
from sqlspec.adapters.aiomysql.data_dictionary import AiomysqlDataDictionary
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.aiomysql._typing import AiomysqlConnection
from sqlspec.core import SQL, StatementConfig
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
__all__ = ("AiomysqlCursor", "AiomysqlDriver", "AiomysqlExceptionHandler", "AiomysqlSessionContext")
logger = get_logger(__name__)
json_type_value = (
AiomysqlFieldType.JSON if AiomysqlFieldType is not None and supports_json_type(AiomysqlFieldType) else None
)
AIOMYSQL_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 AiomysqlExceptionHandler(BaseAsyncExceptionHandler):
"""Async context manager for handling aiomysql (MySQL) database exceptions.
Maps MySQL error codes and SQLSTATE to specific SQLSpec exceptions
for better error handling in application code.
aiomysql uses pymysql.err.Error as its base exception class since
aiomysql is built on top of PyMySQL.
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, AiomysqlPymysqlError):
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 AiomysqlDriver(AsyncDriverAdapterBase):
"""MySQL/MariaDB database driver using aiomysql 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: "AiomysqlConnection",
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: AiomysqlDataDictionary | None = None