"""PyMySQL MySQL driver implementation."""
import tempfile
from collections.abc import Sized
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, cast
from sqlspec.adapters.pymysql._typing import (
PyMysqlCursor,
PyMysqlFieldType,
PyMysqlMySQLError,
PyMysqlServerStatus,
PyMysqlSessionContext,
)
from sqlspec.adapters.pymysql.core import (
PymysqlStreamSource,
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.pymysql.data_dictionary import PyMysqlDataDictionary
from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile
from sqlspec.driver import 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.pymysql._typing import PyMysqlConnection
from sqlspec.core import SQL, StatementConfig
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
__all__ = ("PyMysqlCursor", "PyMysqlDriver", "PyMysqlExceptionHandler", "PyMysqlSessionContext")
logger = get_logger("sqlspec.adapters.pymysql")
json_type_value = PyMysqlFieldType.JSON if supports_json_type(PyMysqlFieldType) else None
PYMYSQL_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 PyMysqlExceptionHandler(BaseSyncExceptionHandler):
"""Context manager for handling PyMySQL 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, PyMysqlMySQLError):
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 PyMysqlDriver(SyncDriverAdapterBase):
"""MySQL/MariaDB database driver using PyMySQL."""
__slots__ = ("_data_dictionary",)
dialect = "mysql"
def __init__(
self,
connection: "PyMysqlConnection",
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: PyMysqlDataDictionary | None = None