pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/mysql/aio/transactions.py

@@ -0,0 +1,313 @@
+"""Owned asynchronous MySQL transactions and nested savepoints."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TypeVar, cast
+
+from pyoq.errors import ExecutionError, QueryValidationError, TransactionStateError
+from pyoq.query.execution import (
+    CompiledQuery,
+    ExecutionControl,
+    StatementKind,
+    StreamPolicy,
+)
+from pyoq.query.execution.aio import (
+    AsyncRowStream,
+    AsyncStreamResource,
+    AsyncTransactionState,
+)
+from pyoq.query.mysql.aio.commands import (
+    AsyncExecutionEnvironment,
+    AsyncFailureHandler,
+    open_cursor,
+    run_sql,
+)
+from pyoq.query.mysql.aio.driver import AsyncMySQLConnection, AsyncMySQLCursor
+from pyoq.query.mysql.aio.operations import MySQLAsyncQueryOperations
+from pyoq.query.mysql.aio.pool import MySQLAsyncPool
+from pyoq.query.mysql.compiler import MySQLCompiler
+from pyoq.query.mysql.parameters import adapt_parameters
+from pyoq.query.mysql.transactions import TransactionMode
+from pyoq.query.select import Query
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+
+class MySQLAsyncTransactionState(
+    AsyncTransactionState[AsyncMySQLConnection, "MySQLAsyncTransaction"]
+):
+    """Adds the one transaction fact MySQL reports that other dialects do not.
+
+    A deadlock rolls the whole transaction back on the server while leaving the
+    connection healthy, and a later ``COMMIT`` then succeeds having committed
+    nothing.
+    """
+
+    __slots__ = ("server_rolled_back",)
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.server_rolled_back = False
+
+    def mark_server_rollback(self) -> None:
+        self.server_rolled_back = True
+
+
+class MySQLAsyncTransaction(MySQLAsyncQueryOperations):
+    __slots__ = (
+        "_control",
+        "_entered",
+        "_finished",
+        "_mode",
+        "_parent",
+        "_pool",
+        "_savepoint_name",
+        "_state",
+    )
+
+    def __init__(
+        self,
+        pool: MySQLAsyncPool,
+        compiler: MySQLCompiler,
+        mode: TransactionMode = TransactionMode.REPEATABLE_READ,
+        control: ExecutionControl | None = None,
+        *,
+        parent: MySQLAsyncTransaction | None = None,
+        state: MySQLAsyncTransactionState | None = None,
+    ) -> None:
+        super().__init__(compiler)
+        if not _is_transaction_mode(mode):
+            raise QueryValidationError("MySQL transaction mode is invalid")
+        self._pool = pool
+        self._mode = mode
+        self._control = control
+        self._parent = parent
+        self._state = state or MySQLAsyncTransactionState()
+        self._savepoint_name: str | None = None
+        self._entered = False
+        self._finished = False
+
+    async def __aenter__(self) -> MySQLAsyncTransaction:
+        if self._entered or self._finished:
+            raise TransactionStateError("transaction cannot be entered twice")
+        if self._parent is None:
+            await self._enter_outer()
+        else:
+            await self._enter_savepoint()
+        self._entered = True
+        return self
+
+    async def __aexit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        if not self._entered or self._finished:
+            return
+        if self._parent is None:
+            await self._exit_outer(exception_type)
+        else:
+            await self._exit_savepoint(exception_type)
+
+    def savepoint(
+        self,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> MySQLAsyncTransaction:
+        self._state.require_current(self)
+        return MySQLAsyncTransaction(
+            self._pool,
+            self.compiler,
+            self._mode,
+            control or self._control,
+            parent=self,
+            state=self._state,
+        )
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: StreamPolicy | None = None,
+    ) -> AsyncRowStream[ResultRow]:
+        self._state.require_current(self)
+        selected = policy or StreamPolicy(control=self._control)
+        if selected.control is None and self._control is not None:
+            selected = StreamPolicy(selected.batch_size, self._control)
+        statement = self.compiler.compile(query)
+        stream: AsyncRowStream[ResultRow]
+        stream = AsyncRowStream(
+            lambda: self._open_stream(stream, statement, selected),
+            selected,
+        )
+        return stream
+
+    def _environment(
+        self,
+        connection: AsyncMySQLConnection,
+    ) -> AsyncExecutionEnvironment:
+        return AsyncExecutionEnvironment(
+            connection,
+            self._pool.connection_factory,
+            AsyncFailureHandler(
+                self._state.invalidate,
+                self._state.mark_server_rollback,
+            ),
+        )
+
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncMySQLCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        return await run_sql(
+            self._environment(connection),
+            statement.sql,
+            parameters,
+            operation,
+            control or self._control,
+            one_result_set=statement.statement_kind is StatementKind.CALL,
+        )
+
+    async def _enter_outer(self) -> None:
+        await self._state.adopt(self._pool.checkout(), self)
+        try:
+            await self._run_transaction_sql(self._mode.isolation_sql)
+            await self._run_transaction_sql(self._mode.start_sql)
+        except BaseException:
+            await self._release_outer()
+            raise
+
+    async def _enter_savepoint(self) -> None:
+        self._state.require_current(cast("MySQLAsyncTransaction", self._parent))
+        name = self._state.next_savepoint()
+        await self._run_transaction_sql(f"SAVEPOINT `{name}`", scope=self._parent)
+        self._savepoint_name = name
+        self._state.current_scope = self
+
+    async def _exit_outer(self, exception_type: type[BaseException] | None) -> None:
+        try:
+            self._state.require_scope(self)
+        except TransactionStateError:
+            await self._abort_misnested_outer(exception_type)
+            return
+        await self._state.close_stream()
+        try:
+            await self._finish_outer(exception_type)
+        except ExecutionError:
+            self._state.invalidate()
+            if exception_type is None:
+                raise
+        finally:
+            await self._release_outer()
+
+    async def _finish_outer(self, exception_type: type[BaseException] | None) -> None:
+        if self._state.server_rolled_back:
+            if exception_type is None:
+                message = (
+                    "MySQL rolled this transaction back on the server; its work "
+                    "was discarded and it cannot be committed"
+                )
+                raise TransactionStateError(message)
+            return
+        if self._state.broken:
+            return
+        await self._run_transaction_sql("ROLLBACK" if exception_type else "COMMIT")
+
+    async def _exit_savepoint(self, exception_type: type[BaseException] | None) -> None:
+        self._state.require_scope(self)
+        await self._state.close_stream()
+        name = cast("str", self._savepoint_name)
+        try:
+            await self._finish_savepoint(name, exception_type)
+        except ExecutionError:
+            self._state.invalidate()
+            if exception_type is None:
+                raise
+        finally:
+            self._state.current_scope = self._parent
+            self._finished = True
+
+    async def _finish_savepoint(
+        self,
+        name: str,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        if self._state.broken or self._state.server_rolled_back:
+            return
+        if exception_type is not None:
+            await self._run_transaction_sql(f"ROLLBACK TO SAVEPOINT `{name}`")
+        await self._run_transaction_sql(f"RELEASE SAVEPOINT `{name}`")
+
+    async def _run_transaction_sql(
+        self,
+        sql: str,
+        *,
+        scope: MySQLAsyncTransaction | None = None,
+    ) -> None:
+        connection = self._state.require_current(scope or self, allow_stream=True)
+        await run_sql(
+            self._environment(connection),
+            sql,
+            (),
+            _discard_cursor,
+            self._control,
+        )
+
+    async def _open_stream(
+        self,
+        stream: AsyncRowStream[ResultRow],
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+    ) -> AsyncStreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        session = await open_cursor(
+            self._environment(connection),
+            statement.sql,
+            parameters,
+            policy.control,
+            unbuffered=True,
+        )
+        self._state.active_stream = cast("AsyncRowStream[object]", stream)
+        return AsyncStreamResource(session, self._state.release_stream)
+
+    async def _release_outer(self) -> None:
+        self._finished = True
+        await self._state.release()
+
+    async def _abort_misnested_outer(
+        self,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        connection = self._state.require_owner()
+        await self._state.close_stream()
+        try:
+            cursor = await connection.cursor()
+            await cursor.execute("ROLLBACK", ())
+            await cursor.close()
+        except Exception:
+            self._state.invalidate()
+        finally:
+            await self._release_outer()
+        if exception_type is None:
+            message = "transaction exited with an active nested savepoint"
+            raise TransactionStateError(message)
+
+
+def _is_transaction_mode(value: object) -> bool:
+    return isinstance(value, TransactionMode)
+
+
+async def _discard_cursor(cursor: AsyncMySQLCursor) -> None:
+    del cursor
+
+
+__all__ = ("MySQLAsyncTransaction", "MySQLAsyncTransactionState", "TransactionMode")

pyoq/query/mysql/commands.py

@@ -0,0 +1,354 @@
+"""MySQL cursor execution, timeouts, and cancellation."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from threading import Event, Thread
+from time import monotonic
+from typing import TYPE_CHECKING, NoReturn, TypeAlias, TypeVar, cast
+
+import pymysql
+from pymysql.cursors import Cursor, SSCursor
+
+from pyoq.errors import QueryCancelledError, QueryExecutionError, QueryTimeoutError
+from pyoq.query.execution import ExecutionControl
+from pyoq.query.mysql.health import (
+    is_cancellation,
+    is_connection_failure,
+    is_timeout,
+    is_transaction_abort,
+)
+from pyoq.query.mysql.parameters import MySQLParameter
+
+if TYPE_CHECKING:
+    from pymysql.connections import Connection as DriverConnection
+
+OperationResult = TypeVar("OperationResult")
+
+Connection: TypeAlias = "DriverConnection[Cursor]"
+ConnectionOpener: TypeAlias = "Callable[[], Connection]"
+
+_WATCHDOG_INTERVAL = 0.025
+_MILLISECONDS = 1000
+_CONNECTION_ID = "SELECT CONNECTION_ID()"
+_CLEAR_TIMEOUT = "SET @@max_execution_time = 0"
+
+
+def _is_connection_failure(error: pymysql.Error) -> bool:
+    """PyMySQL reports a locally closed connection as an interface error."""
+    return isinstance(error, pymysql.InterfaceError) or is_connection_failure(error)
+
+
+def _ignore_transaction_abort() -> None:
+    return
+
+
+@dataclass(frozen=True, slots=True)
+class FailureHandler:
+    """Where a driver failure is reported once it has been classified."""
+
+    invalidate: Callable[[], None]
+    abort_transaction: Callable[[], None] = field(default=_ignore_transaction_abort)
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionEnvironment:
+    connection: Connection
+    opener: ConnectionOpener
+    handler: FailureHandler
+
+
+class MySQLExecutionMonitor:
+    """Server-side SELECT timeout plus a watchdog that kills the running query.
+
+    MySQL applies ``max_execution_time`` to read-only SELECT statements only, so
+    the watchdog enforces both cancellation and the deadline for every other
+    statement by issuing ``KILL QUERY`` from a separate connection.
+    """
+
+    __slots__ = ("_control", "_started", "_stop", "_timeout_applied", "_watchdog")
+
+    def __init__(self, control: ExecutionControl | None) -> None:
+        self._control = control or ExecutionControl()
+        self._started = monotonic() if self._control.timeout is not None else 0.0
+        self._stop = Event()
+        self._watchdog: Thread | None = None
+        self._timeout_applied = False
+
+    def install(self, environment: ExecutionEnvironment) -> None:
+        self.checkpoint()
+        self._apply_timeout(environment)
+        if not self._watchdog_required():
+            return
+        watchdog = Thread(
+            target=self._watch,
+            args=(environment, _read_connection_id(environment.connection)),
+            daemon=True,
+            name="pyoq-mysql-cancel",
+        )
+        self._watchdog = watchdog
+        watchdog.start()
+
+    def checkpoint(self) -> None:
+        if self._cancel_requested():
+            raise QueryCancelledError("MySQL query execution was cancelled")
+        if self._expired():
+            raise QueryTimeoutError("MySQL query execution exceeded its timeout")
+
+    def translate_driver_error(
+        self,
+        error: pymysql.Error,
+        environment: ExecutionEnvironment,
+    ) -> NoReturn:
+        if _is_connection_failure(error):
+            environment.handler.invalidate()
+        elif is_transaction_abort(error):
+            environment.handler.abort_transaction()
+        self.checkpoint()
+        if is_timeout(error):
+            raise QueryTimeoutError(
+                "MySQL query execution exceeded its timeout"
+            ) from error
+        if is_cancellation(error):
+            raise QueryCancelledError("MySQL query execution was cancelled") from error
+        message = f"MySQL query execution failed: {error}"
+        raise QueryExecutionError(message) from error
+
+    def close(self, environment: ExecutionEnvironment) -> None:
+        self._stop.set()
+        watchdog = self._watchdog
+        self._watchdog = None
+        if watchdog is not None:
+            watchdog.join()
+        self._clear_timeout(environment)
+
+    def _watch(self, environment: ExecutionEnvironment, victim: int) -> None:
+        while not self._stop.wait(_WATCHDOG_INTERVAL):
+            if self._cancel_requested() or self._expired():
+                kill_query(environment.opener, victim)
+                return
+
+    def _watchdog_required(self) -> bool:
+        return (
+            self._control.cancellation_token is not None
+            or self._control.timeout is not None
+        )
+
+    def _apply_timeout(self, environment: ExecutionEnvironment) -> None:
+        timeout = self._control.timeout
+        if timeout is None:
+            return
+        milliseconds = max(int(timeout * _MILLISECONDS), 1)
+        with environment.connection.cursor() as cursor:
+            cursor.execute(f"SET @@max_execution_time = {milliseconds:d}", ())
+        self._timeout_applied = True
+
+    def _clear_timeout(self, environment: ExecutionEnvironment) -> None:
+        if not self._timeout_applied:
+            return
+        self._timeout_applied = False
+        try:
+            with environment.connection.cursor() as cursor:
+                cursor.execute(_CLEAR_TIMEOUT, ())
+        except pymysql.Error:
+            environment.handler.invalidate()
+
+    def _cancel_requested(self) -> bool:
+        token = self._control.cancellation_token
+        return token is not None and token.cancelled
+
+    def _expired(self) -> bool:
+        timeout = self._control.timeout
+        return timeout is not None and monotonic() - self._started >= timeout
+
+
+class MySQLCursorSession:
+    __slots__ = ("_closed", "_cursor", "_environment", "_monitor")
+
+    def __init__(
+        self,
+        environment: ExecutionEnvironment,
+        cursor: Cursor,
+        monitor: MySQLExecutionMonitor,
+    ) -> None:
+        self._environment = environment
+        self._cursor = cursor
+        self._monitor = monitor
+        self._closed = False
+
+    @property
+    def cursor(self) -> Cursor:
+        return self._cursor
+
+    def checkpoint(self) -> None:
+        self._monitor.checkpoint()
+
+    def translate_driver_error(self, error: pymysql.Error) -> NoReturn:
+        self._monitor.translate_driver_error(error, self._environment)
+
+    def fetch_batch(self, size: int) -> list[tuple[object, ...]]:
+        _require_live_connection(self._environment)
+        try:
+            self._monitor.checkpoint()
+            rows = cast("tuple[tuple[object, ...], ...]", self._cursor.fetchmany(size))
+            self._monitor.checkpoint()
+        except pymysql.Error as error:
+            self._monitor.translate_driver_error(error, self._environment)
+        return list(rows)
+
+    def answered_with_more_than_one_result_set(self) -> bool:
+        """Whether the procedure left result sets behind that nothing would read.
+
+        MySQL follows the sets of a procedure with a packet that carries no
+        columns, so a set that has columns is one that holds rows.
+        """
+        try:
+            while self._cursor.nextset():
+                if self._cursor.description:
+                    return True
+        except pymysql.Error as error:
+            self._monitor.translate_driver_error(error, self._environment)
+        return False
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        _close_cursor(self._cursor, self._environment)
+        self._monitor.close(self._environment)
+
+
+def run_sql(
+    environment: ExecutionEnvironment,
+    sql: str,
+    parameters: tuple[MySQLParameter, ...],
+    operation: Callable[[Cursor], OperationResult],
+    control: ExecutionControl | None,
+    *,
+    one_result_set: bool = False,
+) -> OperationResult:
+    session = open_cursor(environment, sql, parameters, control)
+    try:
+        result = operation(session.cursor)
+        if one_result_set and session.answered_with_more_than_one_result_set():
+            raise QueryExecutionError(MANY_RESULT_SETS)
+        session.checkpoint()
+        return result
+    except pymysql.Error as error:
+        session.translate_driver_error(error)
+    finally:
+        session.close()
+
+
+MANY_RESULT_SETS = (
+    "the procedure answered with more than one result set, and only the first "
+    "would be read; have it answer with one, or read the rest in statements of "
+    "their own"
+)
+"""Refused rather than dropped, because a dropped result set is silent."""
+
+
+def open_cursor(
+    environment: ExecutionEnvironment,
+    sql: str,
+    parameters: tuple[MySQLParameter, ...],
+    control: ExecutionControl | None,
+    *,
+    unbuffered: bool = False,
+) -> MySQLCursorSession:
+    monitor = MySQLExecutionMonitor(control)
+    cursor: Cursor | None = None
+    try:
+        monitor.install(environment)
+        cursor = _open_driver_cursor(environment.connection, unbuffered=unbuffered)
+        cursor.execute(sql, parameters)
+        monitor.checkpoint()
+        return MySQLCursorSession(environment, cursor, monitor)
+    except pymysql.Error as error:
+        _release(cursor, monitor, environment)
+        monitor.translate_driver_error(error, environment)
+    except BaseException:
+        _release(cursor, monitor, environment)
+        raise
+
+
+def kill_query(opener: ConnectionOpener, victim: int) -> None:
+    """Interrupt a running statement from a second connection.
+
+    MySQL has no client-side cancel, so the only way to stop a statement is to
+    ask the server to kill it from somewhere else. A kill that cannot be
+    delivered leaves the statement running, which the caller detects at its own
+    checkpoint.
+    """
+    try:
+        killer = opener()
+    except Exception:
+        return
+    try:
+        with killer.cursor() as cursor:
+            cursor.execute(f"KILL QUERY {victim:d}", ())
+    except pymysql.Error:
+        return
+    finally:
+        killer.close()
+
+
+def _open_driver_cursor(connection: Connection, *, unbuffered: bool) -> Cursor:
+    return connection.cursor(SSCursor) if unbuffered else connection.cursor()
+
+
+def _read_connection_id(connection: Connection) -> int:
+    """Read the identifier a kill has to name.
+
+    ``SELECT CONNECTION_ID()`` always answers with one integer row, so a reply
+    of any other shape means the peer is not a MySQL server and the resulting
+    error belongs to the caller.
+    """
+    with connection.cursor() as cursor:
+        cursor.execute(_CONNECTION_ID, ())
+        row = cast("tuple[int, ...]", cursor.fetchone())
+    return row[0]
+
+
+def _release(
+    cursor: Cursor | None,
+    monitor: MySQLExecutionMonitor,
+    environment: ExecutionEnvironment,
+) -> None:
+    if cursor is not None:
+        _close_cursor(cursor, environment)
+    monitor.close(environment)
+
+
+def _close_cursor(cursor: Cursor, environment: ExecutionEnvironment) -> None:
+    if not environment.connection.open:
+        environment.handler.invalidate()
+        return
+    cursor.close()
+
+
+def _require_live_connection(environment: ExecutionEnvironment) -> None:
+    """Refuse to read through a connection whose socket is already gone.
+
+    PyMySQL reads an unbuffered result straight from the socket, so a closed
+    connection fails inside the driver with an attribute error rather than a
+    driver error that could be classified.
+    """
+    if environment.connection.open:
+        return
+    environment.handler.invalidate()
+    message = "MySQL connection was closed while a result was open"
+    raise QueryExecutionError(message)
+
+
+__all__ = (
+    "ConnectionOpener",
+    "ExecutionEnvironment",
+    "FailureHandler",
+    "MySQLCursorSession",
+    "MySQLExecutionMonitor",
+    "kill_query",
+    "open_cursor",
+    "run_sql",
+)