sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/adapters/pymysql/driver.py

@@ -0,0 +1,259 @@
+"""PyMySQL MySQL driver implementation."""
+
+from typing import TYPE_CHECKING, Any, Final, cast
+
+import pymysql
+from pymysql.constants import FIELD_TYPE
+
+from sqlspec.adapters.pymysql.core import (
+    build_insert_statement,
+    collect_rows,
+    create_mapped_exception,
+    default_statement_config,
+    detect_json_columns,
+    driver_profile,
+    format_identifier,
+    normalize_execute_many_parameters,
+    normalize_execute_parameters,
+    normalize_lastrowid,
+    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 SyncDriverAdapterBase
+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
+
+from sqlspec.adapters.pymysql._typing import PyMysqlSessionContext
+
+__all__ = ("PyMysqlCursor", "PyMysqlDriver", "PyMysqlExceptionHandler", "PyMysqlSessionContext")
+
+logger = get_logger("sqlspec.adapters.pymysql")
+
+json_type_value = FIELD_TYPE.JSON if supports_json_type(FIELD_TYPE) else None
+PYMYSQL_JSON_TYPE_CODES: Final[set[int]] = {json_type_value} if json_type_value is not None else set()
+
+
+class PyMysqlCursor:
+    """Context manager for PyMySQL cursor operations."""
+
+    __slots__ = ("connection", "cursor")
+
+    def __init__(self, connection: "PyMysqlConnection") -> None:
+        self.connection = connection
+        self.cursor: Any | None = None
+
+    def __enter__(self) -> Any:
+        self.cursor = self.connection.cursor(pymysql.cursors.DictCursor)
+        return self.cursor
+
+    def __exit__(self, *_: Any) -> None:
+        if self.cursor is not None:
+            self.cursor.close()
+
+
+class PyMysqlExceptionHandler:
+    """Context manager for handling PyMySQL exceptions."""
+
+    __slots__ = ("pending_exception",)
+
+    def __init__(self) -> None:
+        self.pending_exception: Exception | None = None
+
+    def __enter__(self) -> "PyMysqlExceptionHandler":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        if exc_type is None:
+            return False
+        if issubclass(exc_type, pymysql.MySQLError):
+            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
+
+    def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        cursor.execute(sql, normalize_execute_parameters(prepared_parameters))
+
+        if statement.returns_rows():
+            fetched_data = cursor.fetchall()
+            fetched_rows = list(fetched_data) if fetched_data else None
+            description = list(cursor.description) if cursor.description else None
+            json_indexes = detect_json_columns(cursor, PYMYSQL_JSON_TYPE_CODES)
+            deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json))
+            rows, column_names = collect_rows(fetched_rows, description, json_indexes, deserializer, logger=logger)
+
+            return self.create_execution_result(
+                cursor, selected_data=rows, column_names=column_names, data_row_count=len(rows), is_select_result=True
+            )
+
+        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)
+
+    def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+
+        prepared_parameters = normalize_execute_many_parameters(prepared_parameters)
+        cursor.executemany(sql, prepared_parameters)
+
+        affected_rows = len(prepared_parameters)
+        return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True)
+
+    def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
+
+        successful_count = 0
+        last_cursor = cursor
+
+        for stmt in statements:
+            cursor.execute(stmt, normalize_execute_parameters(prepared_parameters))
+            successful_count += 1
+
+        return self.create_execution_result(
+            last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True
+        )
+
+    def begin(self) -> None:
+        try:
+            with PyMysqlCursor(self.connection) as cursor:
+                cursor.execute("BEGIN")
+        except pymysql.MySQLError as exc:
+            msg = f"Failed to begin MySQL transaction: {exc}"
+            raise SQLSpecError(msg) from exc
+
+    def commit(self) -> None:
+        try:
+            self.connection.commit()
+        except pymysql.MySQLError as exc:
+            msg = f"Failed to commit MySQL transaction: {exc}"
+            raise SQLSpecError(msg) from exc
+
+    def rollback(self) -> None:
+        try:
+            self.connection.rollback()
+        except pymysql.MySQLError as exc:
+            msg = f"Failed to rollback MySQL transaction: {exc}"
+            raise SQLSpecError(msg) from exc
+
+    def with_cursor(self, connection: "PyMysqlConnection") -> "PyMysqlCursor":
+        return PyMysqlCursor(connection)
+
+    def handle_database_exceptions(self) -> "PyMysqlExceptionHandler":
+        return PyMysqlExceptionHandler()
+
+    def select_to_storage(
+        self,
+        statement: "SQL | str",
+        destination: "StorageDestination",
+        /,
+        *parameters: Any,
+        statement_config: "StatementConfig | None" = None,
+        partitioner: "dict[str, object] | None" = None,
+        format_hint: "StorageFormat | None" = None,
+        telemetry: "StorageTelemetry | None" = None,
+        **kwargs: Any,
+    ) -> "StorageBridgeJob":
+        self._require_capability("arrow_export_enabled")
+        arrow_result = self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs)
+        pipeline = self._storage_pipeline()
+        telemetry_payload = self._write_result_to_storage_sync(
+            arrow_result, destination, format_hint=format_hint, pipeline=pipeline
+        )
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    def load_from_arrow(
+        self,
+        table: str,
+        source: "ArrowResult | Any",
+        *,
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+        telemetry: "StorageTelemetry | None" = None,
+    ) -> "StorageBridgeJob":
+        self._require_capability("arrow_import_enabled")
+        arrow_table = self._coerce_arrow_table(source)
+        if overwrite:
+            statement = f"TRUNCATE TABLE {format_identifier(table)}"
+            with self.handle_database_exceptions(), self.with_cursor(self.connection) as cursor:
+                cursor.execute(statement)
+
+        columns, records = self._arrow_table_to_rows(arrow_table)
+        if records:
+            insert_sql = build_insert_statement(table, columns)
+            with self.handle_database_exceptions(), self.with_cursor(self.connection) as cursor:
+                cursor.executemany(insert_sql, records)
+
+        telemetry_payload = self._build_ingest_telemetry(arrow_table)
+        telemetry_payload["destination"] = table
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    def load_from_storage(
+        self,
+        table: str,
+        source: "StorageDestination",
+        *,
+        file_format: "StorageFormat",
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+    ) -> "StorageBridgeJob":
+        arrow_table, inbound = self._read_arrow_from_storage_sync(source, file_format=file_format)
+        return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound)
+
+    @property
+    def data_dictionary(self) -> "PyMysqlDataDictionary":
+        if self._data_dictionary is None:
+            self._data_dictionary = PyMysqlDataDictionary()
+        return self._data_dictionary
+
+    def _connection_in_transaction(self) -> bool:
+        get_autocommit = getattr(self.connection, "get_autocommit", None)
+        if callable(get_autocommit):
+            return not bool(get_autocommit())
+        autocommit = getattr(self.connection, "autocommit", None)
+        if autocommit is not None:
+            try:
+                return not bool(autocommit)
+            except Exception:
+                return False
+        return False
+
+
+register_driver_profile("pymysql", driver_profile)

sqlspec/adapters/pymysql/events/__init__.py

@@ -0,0 +1,5 @@
+"""Events helpers for the PyMySQL adapter."""
+
+from sqlspec.adapters.pymysql.events.store import PyMysqlEventQueueStore
+
+__all__ = ("PyMysqlEventQueueStore",)

sqlspec/adapters/pymysql/events/store.py

@@ -0,0 +1,50 @@
+"""PyMySQL event queue store with MySQL-specific DDL."""
+
+from typing import Final
+
+from sqlspec.adapters.pymysql.config import PyMysqlConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("PyMysqlEventQueueStore",)
+
+SCHEMA_QUALIFIED_SEGMENTS: Final[int] = 2
+
+
+class PyMysqlEventQueueStore(BaseEventQueueStore[PyMysqlConfig]):
+    """Queue DDL for PyMySQL configs."""
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        return "JSON", "JSON", "DATETIME(6)"
+
+    def _timestamp_default(self) -> str:
+        return "CURRENT_TIMESTAMP(6)"
+
+    def _build_index_sql(self) -> str | None:
+        table_name = self.table_name
+        segments = table_name.split(".", 1)
+
+        if len(segments) == SCHEMA_QUALIFIED_SEGMENTS:
+            schema = segments[0]
+            table = segments[1]
+            schema_selector = f"'{schema}'"
+        else:
+            table = segments[0]
+            schema_selector = "DATABASE()"
+
+        index_name = self._index_name()
+
+        return (
+            "SET @sqlspec_events_idx_exists := ("
+            "SELECT COUNT(1) FROM information_schema.statistics "
+            f"WHERE table_schema = {schema_selector} "
+            f"AND table_name = '{table}' "
+            f"AND index_name = '{index_name}');"
+            "SET @sqlspec_events_idx_stmt := IF(@sqlspec_events_idx_exists = 0, "
+            f"'ALTER TABLE {table_name} ADD INDEX {index_name} (channel, status, available_at)', "
+            "'SELECT 1');"
+            "PREPARE sqlspec_events_stmt FROM @sqlspec_events_idx_stmt;"
+            "EXECUTE sqlspec_events_stmt;"
+            "DEALLOCATE PREPARE sqlspec_events_stmt;"
+        )

sqlspec/adapters/pymysql/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar helpers for the PyMySQL adapter."""
+
+from sqlspec.adapters.pymysql.litestar.store import PyMysqlStore
+
+__all__ = ("PyMysqlStore",)

sqlspec/adapters/pymysql/litestar/store.py

@@ -0,0 +1,232 @@
+"""PyMySQL session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Final
+
+import pymysql
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.pymysql.config import PyMysqlConfig
+
+logger = get_logger("sqlspec.adapters.pymysql.litestar.store")
+
+__all__ = ("PyMysqlStore",)
+
+MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146
+
+
+class PyMysqlStore(BaseSQLSpecStore["PyMysqlConfig"]):
+    """MySQL/MariaDB session store using PyMySQL sync driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "PyMysqlConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id VARCHAR(255) PRIMARY KEY,
+            data LONGBLOB NOT NULL,
+            expires_at DATETIME(6),
+            created_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6),
+            updated_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+            INDEX idx_{self._table_name}_expires_at (expires_at)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        return [
+            f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}",
+            f"DROP TABLE IF EXISTS {self._table_name}",
+        ]
+
+    def _create_table(self) -> None:
+        sql = self._get_create_table_sql()
+        with self._config.provide_session() as driver:
+            driver.execute_script(sql)
+            driver.commit()
+        self._log_table_created()
+
+    async def create_table(self) -> None:
+        await async_(self._create_table)()
+
+    def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor(pymysql.cursors.DictCursor)
+                try:
+                    cursor.execute(sql, (key,))
+                    row = cursor.fetchone()
+                finally:
+                    cursor.close()
+
+                if row is None:
+                    return None
+
+                if renew_for is not None and row["expires_at"] is not None:
+                    new_expires_at = self._calculate_expires_at(renew_for)
+                    if new_expires_at is not None:
+                        naive_expires_at = new_expires_at.replace(tzinfo=None)
+                        update_sql = f"""
+                        UPDATE {self._table_name}
+                        SET expires_at = %s, updated_at = UTC_TIMESTAMP(6)
+                        WHERE session_id = %s
+                        """
+                        update_cursor = conn.cursor()
+                        try:
+                            update_cursor.execute(update_sql, (naive_expires_at, key))
+                        finally:
+                            update_cursor.close()
+                        conn.commit()
+
+                return bytes(row["data"])
+        except pymysql.MySQLError as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "args", [None])[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return None
+            raise
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        return await async_(self._get)(key, renew_for)
+
+    def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+        naive_expires_at = expires_at.replace(tzinfo=None) if expires_at else None
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES (%s, %s, %s) AS new
+        ON DUPLICATE KEY UPDATE
+            data = new.data,
+            expires_at = new.expires_at,
+            updated_at = UTC_TIMESTAMP(6)
+        """
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key, data, naive_expires_at))
+            finally:
+                cursor.close()
+            conn.commit()
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        await async_(self._set)(key, value, expires_in)
+
+    def _delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key,))
+            finally:
+                cursor.close()
+            conn.commit()
+
+    async def delete(self, key: str) -> None:
+        await async_(self._delete)(key)
+
+    def _delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor()
+                try:
+                    cursor.execute(sql)
+                finally:
+                    cursor.close()
+                conn.commit()
+            self._log_delete_all()
+        except pymysql.MySQLError as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "args", [None])[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                logger.debug("Table %s does not exist, skipping delete_all", self._table_name)
+                return
+            raise
+
+    async def delete_all(self) -> None:
+        await async_(self._delete_all)()
+
+    def _exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor()
+                try:
+                    cursor.execute(sql, (key,))
+                    result = cursor.fetchone()
+                finally:
+                    cursor.close()
+                return result is not None
+        except pymysql.MySQLError as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "args", [None])[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return False
+            raise
+
+    async def exists(self, key: str) -> bool:
+        return await async_(self._exists)(key)
+
+    def _expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key,))
+                row = cursor.fetchone()
+            finally:
+                cursor.close()
+
+            if row is None or row[0] is None:
+                return None
+
+            expires_at_naive = row[0]
+            expires_at_utc = expires_at_naive.replace(tzinfo=timezone.utc)
+            now = datetime.now(timezone.utc)
+
+            if expires_at_utc <= now:
+                return 0
+
+            delta = expires_at_utc - now
+            return int(delta.total_seconds())
+
+    async def expires_in(self, key: str) -> "int | None":
+        return await async_(self._expires_in)(key)
+
+    def _delete_expired(self) -> int:
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= UTC_TIMESTAMP(6)"
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql)
+                conn.commit()
+                count: int = cursor.rowcount
+            finally:
+                cursor.close()
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+    async def delete_expired(self) -> int:
+        return await async_(self._delete_expired)()

sqlspec/adapters/pymysql/pool.py

@@ -0,0 +1,137 @@
+"""PyMySQL database configuration with thread-local connections."""
+
+import contextlib
+import threading
+import time
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, TypedDict, cast
+
+import pymysql
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.pymysql._typing import PyMysqlConnection
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+class PyMysqlConnectionParams(TypedDict):
+    """PyMySQL connection parameters."""
+
+    host: NotRequired[str]
+    user: NotRequired[str]
+    password: NotRequired[str]
+    database: NotRequired[str]
+    port: NotRequired[int]
+    unix_socket: NotRequired[str]
+    charset: NotRequired[str]
+    connect_timeout: NotRequired[int]
+    read_timeout: NotRequired[int]
+    write_timeout: NotRequired[int]
+    autocommit: NotRequired[bool]
+    ssl: NotRequired["dict[str, Any]"]
+    client_flag: NotRequired[int]
+    cursorclass: NotRequired[type]
+    init_command: NotRequired[str]
+    sql_mode: NotRequired[str]
+    extra: NotRequired["dict[str, Any]"]
+
+
+__all__ = ("PyMysqlConnectionPool",)
+
+logger = get_logger(__name__)
+
+
+class PyMysqlConnectionPool:
+    """Thread-local connection manager for PyMySQL."""
+
+    __slots__ = ("_connection_parameters", "_health_check_interval", "_recycle_seconds", "_thread_local")
+
+    def __init__(
+        self, connection_parameters: "dict[str, Any]", recycle_seconds: int = 86400, health_check_interval: float = 30.0
+    ) -> None:
+        self._connection_parameters = connection_parameters
+        self._thread_local = threading.local()
+        self._recycle_seconds = recycle_seconds
+        self._health_check_interval = health_check_interval
+
+    def _create_connection(self) -> PyMysqlConnection:
+        connection = pymysql.connect(**self._connection_parameters)
+        return cast("PyMysqlConnection", connection)
+
+    def _is_connection_alive(self, connection: PyMysqlConnection) -> bool:
+        try:
+            connection.ping(reconnect=False)
+        except Exception:
+            return False
+        return True
+
+    def _get_thread_connection(self) -> PyMysqlConnection:
+        thread_state = self._thread_local.__dict__
+        if "connection" not in thread_state:
+            self._thread_local.connection = self._create_connection()
+            self._thread_local.created_at = time.time()
+            self._thread_local.last_used = time.time()
+            return cast("PyMysqlConnection", self._thread_local.connection)
+
+        if self._recycle_seconds > 0 and time.time() - self._thread_local.created_at > self._recycle_seconds:
+            logger.debug("PyMySQL connection exceeded recycle time, recreating")
+            with contextlib.suppress(Exception):
+                self._thread_local.connection.close()
+            self._thread_local.connection = self._create_connection()
+            self._thread_local.created_at = time.time()
+            self._thread_local.last_used = time.time()
+            return cast("PyMysqlConnection", self._thread_local.connection)
+
+        idle_time = time.time() - thread_state.get("last_used", 0)
+        if idle_time > self._health_check_interval and not self._is_connection_alive(self._thread_local.connection):
+            logger.debug("PyMySQL connection failed health check after %.1fs idle, recreating", idle_time)
+            with contextlib.suppress(Exception):
+                self._thread_local.connection.close()
+            self._thread_local.connection = self._create_connection()
+            self._thread_local.created_at = time.time()
+
+        self._thread_local.last_used = time.time()
+        return cast("PyMysqlConnection", self._thread_local.connection)
+
+    def _close_thread_connection(self) -> None:
+        thread_state = self._thread_local.__dict__
+        if "connection" in thread_state:
+            with contextlib.suppress(Exception):
+                self._thread_local.connection.close()
+            del self._thread_local.connection
+            if "created_at" in thread_state:
+                del self._thread_local.created_at
+            if "last_used" in thread_state:
+                del self._thread_local.last_used
+
+    @contextmanager
+    def get_connection(self) -> "Generator[PyMysqlConnection, None, None]":
+        connection = self._get_thread_connection()
+        try:
+            yield connection
+        finally:
+            with contextlib.suppress(Exception):
+                if connection.open and connection.get_autocommit() is False:
+                    connection.commit()
+
+    def close(self) -> None:
+        self._close_thread_connection()
+
+    def acquire(self) -> PyMysqlConnection:
+        return self._get_thread_connection()
+
+    def release(self, connection: PyMysqlConnection) -> None:
+        _ = connection
+
+    def size(self) -> int:
+        try:
+            _ = self._thread_local.connection
+        except AttributeError:
+            return 0
+        else:
+            return 1
+
+    def checked_out(self) -> int:
+        return 0