pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/fields.py

@@ -0,0 +1,60 @@
+"""Typed structural field expressions."""
+
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+from pyoq.query.expressions import Expression
+from pyoq.query.nodes import FieldNode
+from pyoq.query.results import family_for_type
+
+Value = TypeVar("Value")
+FieldValue = TypeVar("FieldValue")
+
+
+class Field(Expression[Value], Generic[Value]):
+    __slots__ = ()
+
+    def __init__(
+        self,
+        database_name: str,
+        /,
+        *,
+        value_type: type[object],
+        table_name: str | None = None,
+        schema_name: str | None = None,
+        catalog_name: str | None = None,
+    ) -> None:
+        family = family_for_type(value_type)
+        super().__init__(
+            FieldNode(
+                database_name,
+                table_name,
+                schema_name,
+                catalog_name,
+                family,
+                value_type,
+            ),
+            family,
+        )
+
+
+def field(
+    value_type: type[FieldValue],
+    database_name: str,
+    /,
+    *,
+    table_name: str | None = None,
+    schema_name: str | None = None,
+    catalog_name: str | None = None,
+) -> Field[FieldValue]:
+    return Field(
+        database_name,
+        value_type=value_type,
+        table_name=table_name,
+        schema_name=schema_name,
+        catalog_name=catalog_name,
+    )
+
+
+__all__ = ("Field", "family_for_type", "field")

pyoq/query/mysql/__init__.py

@@ -0,0 +1,59 @@
+"""MySQL query compilation and execution."""
+
+from pyoq.query.mysql.commands import (
+    ExecutionEnvironment,
+    FailureHandler,
+    MySQLCursorSession,
+)
+from pyoq.query.mysql.compiler import MySQLCompiler
+from pyoq.query.mysql.context import MySQLCompilationContext
+from pyoq.query.mysql.executor import MySQLExecutor
+from pyoq.query.mysql.expressions import MySQLExpressionCompiler
+from pyoq.query.mysql.factory import MySQLConnectionFactory
+from pyoq.query.mysql.model import MySQLCapabilities
+from pyoq.query.mysql.operations import MySQLQueryOperations
+from pyoq.query.mysql.parameters import (
+    MySQLParameter,
+    adapt_parameter,
+    adapt_parameters,
+)
+from pyoq.query.mysql.planning import MySQLBulkPlanner
+from pyoq.query.mysql.pool import (
+    MySQLConnectionFactoryProtocol,
+    MySQLConnectionLease,
+    MySQLPool,
+    MySQLPoolPolicy,
+    MySQLPoolStats,
+)
+from pyoq.query.mysql.transactions import (
+    MySQLTransaction,
+    MySQLTransactionState,
+    TransactionMode,
+)
+from pyoq.query.mysql.writes import MySQLWriteCompilation
+
+__all__ = (
+    "ExecutionEnvironment",
+    "FailureHandler",
+    "MySQLBulkPlanner",
+    "MySQLCapabilities",
+    "MySQLCompilationContext",
+    "MySQLCompiler",
+    "MySQLConnectionFactory",
+    "MySQLConnectionFactoryProtocol",
+    "MySQLConnectionLease",
+    "MySQLCursorSession",
+    "MySQLExecutor",
+    "MySQLExpressionCompiler",
+    "MySQLParameter",
+    "MySQLPool",
+    "MySQLPoolPolicy",
+    "MySQLPoolStats",
+    "MySQLQueryOperations",
+    "MySQLTransaction",
+    "MySQLTransactionState",
+    "MySQLWriteCompilation",
+    "TransactionMode",
+    "adapt_parameter",
+    "adapt_parameters",
+)

pyoq/query/mysql/aio/__init__.py

@@ -0,0 +1,38 @@
+"""Asynchronous MySQL query execution."""
+
+from pyoq.query.mysql.aio.commands import (
+    AsyncExecutionEnvironment,
+    AsyncFailureHandler,
+    MySQLAsyncCursorSession,
+)
+from pyoq.query.mysql.aio.driver import (
+    AsyncMySQLConnection,
+    AsyncMySQLCursor,
+    preferred_tls_context,
+)
+from pyoq.query.mysql.aio.executor import MySQLAsyncExecutor
+from pyoq.query.mysql.aio.factory import MySQLAsyncConnectionFactory
+from pyoq.query.mysql.aio.operations import MySQLAsyncQueryOperations
+from pyoq.query.mysql.aio.pool import MySQLAsyncConnectionLease, MySQLAsyncPool
+from pyoq.query.mysql.aio.transactions import (
+    MySQLAsyncTransaction,
+    MySQLAsyncTransactionState,
+    TransactionMode,
+)
+
+__all__ = (
+    "AsyncExecutionEnvironment",
+    "AsyncFailureHandler",
+    "AsyncMySQLConnection",
+    "AsyncMySQLCursor",
+    "MySQLAsyncConnectionFactory",
+    "MySQLAsyncConnectionLease",
+    "MySQLAsyncCursorSession",
+    "MySQLAsyncExecutor",
+    "MySQLAsyncPool",
+    "MySQLAsyncQueryOperations",
+    "MySQLAsyncTransaction",
+    "MySQLAsyncTransactionState",
+    "TransactionMode",
+    "preferred_tls_context",
+)

pyoq/query/mysql/aio/commands.py

@@ -0,0 +1,389 @@
+"""Asynchronous MySQL cursor execution, timeouts, and cancellation."""
+
+from __future__ import annotations
+
+from asyncio import CancelledError, Event, Task, create_task, shield, wait_for
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from time import monotonic
+from typing import NoReturn, TypeAlias, TypeVar, cast
+
+from pyoq.errors import QueryCancelledError, QueryExecutionError, QueryTimeoutError
+from pyoq.query.execution import ExecutionControl
+from pyoq.query.mysql.aio.driver import (
+    AsyncMySQLConnection,
+    AsyncMySQLCursor,
+    driver_errors,
+    interface_errors,
+    unbuffered_cursor_class,
+)
+from pyoq.query.mysql.commands import MANY_RESULT_SETS
+from pyoq.query.mysql.health import (
+    is_cancellation,
+    is_connection_failure,
+    is_timeout,
+    is_transaction_abort,
+)
+from pyoq.query.mysql.parameters import MySQLParameter
+
+OperationResult = TypeVar("OperationResult")
+
+AsyncConnectionOpener: TypeAlias = "Callable[[], Awaitable[AsyncMySQLConnection]]"
+
+_WATCHDOG_INTERVAL = 0.025
+_MILLISECONDS = 1000
+_CONNECTION_ID = "SELECT CONNECTION_ID()"
+_CLEAR_TIMEOUT = "SET @@max_execution_time = 0"
+
+
+def _is_connection_failure(error: BaseException) -> bool:
+    """The driver reports a locally closed connection as an interface error."""
+    return isinstance(error, interface_errors()) or is_connection_failure(error)
+
+
+async def await_watchdog(watchdog: Task[None]) -> None:
+    """Wait for a watchdog this monitor owns, without inheriting its fate.
+
+    Cleanup runs while an error is already in flight, so a watchdog that some
+    outer shutdown cancelled must not turn a release into a cancellation. Being
+    cancelled ourselves is a different matter and still propagates, which is
+    what the shield distinguishes: awaiting a task directly would cancel it
+    along with us and leave the two cases indistinguishable.
+    """
+    try:
+        await shield(watchdog)
+    except CancelledError:
+        if not watchdog.cancelled():
+            raise
+
+
+def _ignore_transaction_abort() -> None:
+    return
+
+
+@dataclass(frozen=True, slots=True)
+class AsyncFailureHandler:
+    """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 AsyncExecutionEnvironment:
+    connection: AsyncMySQLConnection
+    opener: AsyncConnectionOpener
+    handler: AsyncFailureHandler
+
+
+async def kill_query(opener: AsyncConnectionOpener, 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 = await opener()
+    except Exception:
+        return
+    try:
+        cursor = await killer.cursor()
+        await cursor.execute(f"KILL QUERY {victim:d}", ())
+    except driver_errors():
+        return
+    finally:
+        killer.close()
+
+
+class MySQLAsyncExecutionMonitor:
+    """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.
+    """
+
+    __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: Task[None] | None = None
+        self._timeout_applied = False
+
+    async def install(self, environment: AsyncExecutionEnvironment) -> None:
+        self.checkpoint()
+        await self._apply_timeout(environment)
+        if not self._watchdog_required():
+            return
+        victim = await read_connection_id(environment.connection)
+        self._watchdog = create_task(self._watch(environment, victim))
+
+    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: BaseException,
+        environment: AsyncExecutionEnvironment,
+    ) -> 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
+
+    async def close(self, environment: AsyncExecutionEnvironment) -> None:
+        self._stop.set()
+        watchdog = self._watchdog
+        self._watchdog = None
+        if watchdog is not None:
+            await await_watchdog(watchdog)
+        await self._clear_timeout(environment)
+
+    async def _watch(
+        self,
+        environment: AsyncExecutionEnvironment,
+        victim: int,
+    ) -> None:
+        while True:
+            try:
+                await wait_for(self._stop.wait(), _WATCHDOG_INTERVAL)
+            except (TimeoutError, CancelledError):
+                if self._cancel_requested() or self._expired():
+                    await kill_query(environment.opener, victim)
+                    return
+                continue
+            return
+
+    def _watchdog_required(self) -> bool:
+        return (
+            self._control.cancellation_token is not None
+            or self._control.timeout is not None
+        )
+
+    async def _apply_timeout(self, environment: AsyncExecutionEnvironment) -> None:
+        timeout = self._control.timeout
+        if timeout is None:
+            return
+        milliseconds = max(int(timeout * _MILLISECONDS), 1)
+        cursor = await environment.connection.cursor()
+        await cursor.execute(f"SET @@max_execution_time = {milliseconds:d}", ())
+        await cursor.close()
+        self._timeout_applied = True
+
+    async def _clear_timeout(self, environment: AsyncExecutionEnvironment) -> None:
+        if not self._timeout_applied:
+            return
+        self._timeout_applied = False
+        if environment.connection.closed:
+            environment.handler.invalidate()
+            return
+        try:
+            cursor = await environment.connection.cursor()
+            await cursor.execute(_CLEAR_TIMEOUT, ())
+            await cursor.close()
+        except driver_errors():
+            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 MySQLAsyncCursorSession:
+    __slots__ = ("_closed", "_cursor", "_environment", "_monitor")
+
+    def __init__(
+        self,
+        environment: AsyncExecutionEnvironment,
+        cursor: AsyncMySQLCursor,
+        monitor: MySQLAsyncExecutionMonitor,
+    ) -> None:
+        self._environment = environment
+        self._cursor = cursor
+        self._monitor = monitor
+        self._closed = False
+
+    @property
+    def cursor(self) -> AsyncMySQLCursor:
+        return self._cursor
+
+    def checkpoint(self) -> None:
+        self._monitor.checkpoint()
+
+    def translate_driver_error(self, error: BaseException) -> NoReturn:
+        self._monitor.translate_driver_error(error, self._environment)
+
+    async def fetch_batch(self, size: int) -> list[tuple[object, ...]]:
+        require_live_connection(self._environment)
+        try:
+            self._monitor.checkpoint()
+            rows = await self._cursor.fetchmany(size)
+            self._monitor.checkpoint()
+        except driver_errors() as error:
+            self._monitor.translate_driver_error(error, self._environment)
+        return cast("list[tuple[object, ...]]", list(rows))
+
+    async 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 await self._cursor.nextset():
+                if self._cursor.description:
+                    return True
+        except driver_errors() as error:
+            self._monitor.translate_driver_error(error, self._environment)
+        return False
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await _close_cursor(self._cursor, self._environment)
+        await self._monitor.close(self._environment)
+
+
+async def run_sql(
+    environment: AsyncExecutionEnvironment,
+    sql: str,
+    parameters: tuple[MySQLParameter, ...],
+    operation: Callable[[AsyncMySQLCursor], Awaitable[OperationResult]],
+    control: ExecutionControl | None,
+    *,
+    one_result_set: bool = False,
+) -> OperationResult:
+    session = await open_cursor(environment, sql, parameters, control)
+    try:
+        result = await operation(session.cursor)
+        if one_result_set and await session.answered_with_more_than_one_result_set():
+            raise QueryExecutionError(MANY_RESULT_SETS)
+        session.checkpoint()
+        return result
+    except driver_errors() as error:
+        session.translate_driver_error(error)
+    finally:
+        await session.close()
+
+
+async def open_cursor(
+    environment: AsyncExecutionEnvironment,
+    sql: str,
+    parameters: tuple[MySQLParameter, ...],
+    control: ExecutionControl | None,
+    *,
+    unbuffered: bool = False,
+) -> MySQLAsyncCursorSession:
+    monitor = MySQLAsyncExecutionMonitor(control)
+    cursor: AsyncMySQLCursor | None = None
+    try:
+        await monitor.install(environment)
+        cursor = await _open_driver_cursor(
+            environment.connection,
+            unbuffered=unbuffered,
+        )
+        await cursor.execute(sql, parameters)
+        monitor.checkpoint()
+        return MySQLAsyncCursorSession(environment, cursor, monitor)
+    except driver_errors() as error:
+        await _release(cursor, monitor, environment)
+        monitor.translate_driver_error(error, environment)
+    except BaseException:
+        await _release(cursor, monitor, environment)
+        raise
+
+
+async def read_connection_id(connection: AsyncMySQLConnection) -> 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.
+    """
+    cursor = await connection.cursor()
+    try:
+        await cursor.execute(_CONNECTION_ID, ())
+        row = cast("tuple[int, ...]", await cursor.fetchone())
+    finally:
+        await cursor.close()
+    return row[0]
+
+
+def require_live_connection(environment: AsyncExecutionEnvironment) -> None:
+    """Refuse to read through a connection whose socket is already gone.
+
+    The driver reads an unbuffered result straight from the socket, so a closed
+    connection fails inside it rather than raising a driver error that could be
+    classified.
+    """
+    if not environment.connection.closed:
+        return
+    environment.handler.invalidate()
+    message = "MySQL connection was closed while a result was open"
+    raise QueryExecutionError(message)
+
+
+async def _open_driver_cursor(
+    connection: AsyncMySQLConnection,
+    *,
+    unbuffered: bool,
+) -> AsyncMySQLCursor:
+    if unbuffered:
+        return await connection.cursor(unbuffered_cursor_class())
+    return await connection.cursor()
+
+
+async def _release(
+    cursor: AsyncMySQLCursor | None,
+    monitor: MySQLAsyncExecutionMonitor,
+    environment: AsyncExecutionEnvironment,
+) -> None:
+    if cursor is not None:
+        await _close_cursor(cursor, environment)
+    await monitor.close(environment)
+
+
+async def _close_cursor(
+    cursor: AsyncMySQLCursor,
+    environment: AsyncExecutionEnvironment,
+) -> None:
+    if environment.connection.closed:
+        environment.handler.invalidate()
+        return
+    await cursor.close()
+
+
+__all__ = (
+    "AsyncConnectionOpener",
+    "AsyncExecutionEnvironment",
+    "AsyncFailureHandler",
+    "MySQLAsyncCursorSession",
+    "MySQLAsyncExecutionMonitor",
+    "await_watchdog",
+    "kill_query",
+    "open_cursor",
+    "read_connection_id",
+    "require_live_connection",
+    "run_sql",
+)