Source code for sqlspec.adapters.mysqlconnector.driver

"""MysqlConnector MySQL driver implementation.

Provides MySQL/MariaDB connectivity with parameter style conversion,
type coercion, error handling, and transaction management.
"""

import tempfile
from collections.abc import Sized
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, cast

from sqlspec.adapters.mysqlconnector._typing import (
    MysqlConnectorAsyncCursor,
    MysqlConnectorAsyncSessionContext,
    MysqlConnectorError,
    MysqlConnectorFieldType,
    MysqlConnectorSyncCursor,
    MysqlConnectorSyncSessionContext,
)
from sqlspec.adapters.mysqlconnector.core import (
    MysqlConnectorAsyncStreamSource,
    MysqlConnectorSyncStreamSource,
    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.mysqlconnector.data_dictionary import (
    MysqlConnectorAsyncDataDictionary,
    MysqlConnectorSyncDataDictionary,
)
from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile
from sqlspec.driver import (
    AsyncDriverAdapterBase,
    AsyncRowStream,
    BaseAsyncExceptionHandler,
    BaseSyncExceptionHandler,
    SyncDriverAdapterBase,
    SyncRowStream,
)
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.mysqlconnector._typing import MysqlConnectorAsyncConnection, MysqlConnectorSyncConnection
    from sqlspec.core import SQL, StatementConfig
    from sqlspec.driver import ExecutionResult
    from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry

__all__ = (
    "MysqlConnectorAsyncCursor",
    "MysqlConnectorAsyncDriver",
    "MysqlConnectorAsyncExceptionHandler",
    "MysqlConnectorAsyncSessionContext",
    "MysqlConnectorSyncCursor",
    "MysqlConnectorSyncDriver",
    "MysqlConnectorSyncExceptionHandler",
    "MysqlConnectorSyncSessionContext",
)

logger = get_logger("sqlspec.adapters.mysqlconnector")

json_type_value = MysqlConnectorFieldType.JSON if supports_json_type(MysqlConnectorFieldType) else None
MYSQLCONNECTOR_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 MysqlConnectorSyncExceptionHandler(BaseSyncExceptionHandler):
    """Context manager for handling mysql-connector sync exceptions."""

    __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, MysqlConnectorError):
            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 MysqlConnectorSyncDriver(SyncDriverAdapterBase):
    """MySQL/MariaDB database driver using mysql-connector sync library."""

    __slots__ = ("_data_dictionary",)
    dialect = "mysql"

    def __init__(
        self,
        connection: "MysqlConnectorSyncConnection",
        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: MysqlConnectorSyncDataDictionary | None = None
def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or getattr(cursor, "with_rows", False): fetched_data = cursor.fetchall() description = cursor.description or None row_plan = resolve_row_plan(description, MYSQLCONNECTOR_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) column_types = _resolve_column_types(description) return self.create_execution_result( cursor, selected_data=rows, column_names=column_names, column_types=column_types, data_row_count=len(rows), is_select_result=True, row_format=row_format, ) affected_rows = resolve_rowcount(cursor) last_id = normalize_lastrowid(cursor) return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id)