pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/mysql/model.py

@@ -0,0 +1,79 @@
+"""MySQL compilation capabilities."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pyoq.errors import QueryValidationError
+from pyoq.query.execution import ParameterStyle
+from pyoq.relations.planning import FetchCapabilities
+
+_PLACEHOLDER_LIMIT = 65535
+
+
+@dataclass(frozen=True, slots=True)
+class MySQLCapabilities:
+    maximum_parameters: int = _PLACEHOLDER_LIMIT
+    parameter_style: ParameterStyle = ParameterStyle.FORMAT
+    right_join: bool = True
+    full_join: bool = False
+    lateral_sources: bool = True
+    recursive_common_tables: bool = True
+    raw_expressions: bool = True
+    # MySQL refuses ORDER BY inside JSON_ARRAYAGG, and ordering in a
+    # derived table is discarded rather than honoured.
+    nested_collections: bool = True
+    ordered_nested_collections: bool = False
+    returning: bool = False
+    upsert: bool = True
+
+    @property
+    def fetch(self) -> FetchCapabilities:
+        """What this dialect can do when fetching a collection in one query."""
+        return FetchCapabilities(
+            self.nested_collections,
+            self.ordered_nested_collections,
+        )
+
+    def __post_init__(self) -> None:
+        if not _is_positive_integer(self.maximum_parameters):
+            message = "MySQL maximum parameters must be a positive integer"
+            raise QueryValidationError(message)
+        if self.maximum_parameters > _PLACEHOLDER_LIMIT:
+            message = (
+                "MySQL cannot bind more than "
+                f"{_PLACEHOLDER_LIMIT} placeholders in one statement"
+            )
+            raise QueryValidationError(message)
+        if not _is_parameter_style(self.parameter_style):
+            message = "MySQL parameter style must be a ParameterStyle value"
+            raise QueryValidationError(message)
+        feature_flags: tuple[object, ...] = (
+            self.right_join,
+            self.full_join,
+            self.lateral_sources,
+            self.recursive_common_tables,
+            self.raw_expressions,
+            self.nested_collections,
+            self.ordered_nested_collections,
+            self.returning,
+            self.upsert,
+        )
+        if not _all_boolean(feature_flags):
+            message = "MySQL capability flags must be boolean"
+            raise QueryValidationError(message)
+
+
+def _is_parameter_style(value: object) -> bool:
+    return isinstance(value, ParameterStyle)
+
+
+def _is_positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+def _all_boolean(values: tuple[object, ...]) -> bool:
+    return all(isinstance(value, bool) for value in values)
+
+
+__all__ = ("MySQLCapabilities",)

pyoq/query/mysql/operations.py

@@ -0,0 +1,43 @@
+"""MySQL typed result operations."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from pymysql.cursors import Cursor
+
+from pyoq.query.execution import (
+    BulkPlan,
+    DatabaseCursor,
+    QueryOperations,
+    WriteProvider,
+)
+from pyoq.query.mysql.compiler import MySQLCompiler
+from pyoq.query.mysql.planning import MySQLBulkPlanner
+from pyoq.query.write_nodes import WriteNode
+
+
+class MySQLQueryOperations(QueryOperations):
+    __slots__ = ("_compiler",)
+
+    def __init__(self, compiler: MySQLCompiler) -> None:
+        self._compiler = compiler
+
+    @property
+    def compiler(self) -> MySQLCompiler:
+        return self._compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return MySQLBulkPlanner(self.compiler).plan(statement)
+
+    def last_inserted_id(self, cursor: DatabaseCursor) -> int | None:
+        """Report the first identifier a multi-row insert generated.
+
+        MySQL reports the identifier of the first row an insert generated, not
+        the last, and reports zero when a statement generated none.
+        """
+        identifier = cast("Cursor", cursor).lastrowid
+        return identifier if identifier > 0 else None
+
+
+__all__ = ("MySQLQueryOperations",)

pyoq/query/mysql/parameters.py

@@ -0,0 +1,69 @@
+"""MySQL driver parameter adaptation."""
+
+from __future__ import annotations
+
+import json
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from enum import Enum
+from typing import cast
+from uuid import UUID
+
+from pyoq.errors import ParameterBindingError
+
+MySQLParameter = (
+    str
+    | int
+    | float
+    | bool
+    | bytes
+    | Decimal
+    | date
+    | datetime
+    | time
+    | timedelta
+    | None
+)
+
+_NATIVE_TYPES = (
+    str,
+    int,
+    float,
+    bool,
+    bytes,
+    Decimal,
+    datetime,
+    date,
+    time,
+    timedelta,
+)
+
+
+def adapt_parameters(values: tuple[object, ...]) -> tuple[MySQLParameter, ...]:
+    return tuple(adapt_parameter(value) for value in values)
+
+
+def adapt_parameter(value: object) -> MySQLParameter:
+    if isinstance(value, Enum):
+        return adapt_parameter(value.value)
+    if value is None or isinstance(value, _NATIVE_TYPES):
+        return value
+    if isinstance(value, UUID):
+        return str(value)
+    if isinstance(value, memoryview):
+        return value.tobytes()
+    if isinstance(value, (list, tuple, dict)):
+        return _json_parameter(cast("object", value))
+    message = f"unsupported MySQL parameter type: {type(value).__name__}"
+    raise ParameterBindingError(message)
+
+
+def _json_parameter(value: object) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+    except (TypeError, ValueError) as error:
+        message = "MySQL JSON parameter is not serializable"
+        raise ParameterBindingError(message) from error
+
+
+__all__ = ("MySQLParameter", "adapt_parameter", "adapt_parameters")

pyoq/query/mysql/planning.py

@@ -0,0 +1,20 @@
+"""MySQL bulk write planning."""
+
+from __future__ import annotations
+
+from pyoq.query.execution.planning import BulkPlanner
+from pyoq.query.mysql.compiler import MySQLCompiler
+
+
+class MySQLBulkPlanner(BulkPlanner):
+    __slots__ = ()
+
+    def __init__(self, compiler: MySQLCompiler) -> None:
+        super().__init__(
+            compiler,
+            compiler.capabilities.maximum_parameters,
+            "MySQL",
+        )
+
+
+__all__ = ("MySQLBulkPlanner",)

pyoq/query/mysql/pool.py

@@ -0,0 +1,67 @@
+"""Exclusive synchronous MySQL connection pool."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TypeAlias, cast
+
+import pymysql
+import pymysql.connections
+
+from pyoq.errors import ConnectionPoolError
+from pyoq.query.execution.pooling import (
+    ConnectionLease,
+    ConnectionPool,
+    PoolPolicy,
+    PoolStats,
+)
+from pyoq.query.mysql.commands import Connection
+
+MySQLConnectionFactoryProtocol: TypeAlias = "Callable[[], Connection]"
+MySQLConnectionLease: TypeAlias = ConnectionLease[Connection]
+MySQLPoolPolicy: TypeAlias = PoolPolicy
+MySQLPoolStats: TypeAlias = PoolStats
+
+_RESET_TIMEOUT = "SET @@max_execution_time = 0"
+
+
+class MySQLPool(ConnectionPool[Connection]):
+    __slots__ = ()
+
+    @property
+    def connection_factory(self) -> MySQLConnectionFactoryProtocol:
+        return self._factory
+
+    def reset(self, connection: Connection) -> bool:
+        try:
+            connection.rollback()
+            with connection.cursor() as cursor:
+                cursor.execute(_RESET_TIMEOUT, ())
+        except pymysql.Error:
+            return False
+        return connection.open
+
+    def validate(self, value: object) -> Connection:
+        if not isinstance(value, pymysql.connections.Connection):
+            message = "MySQL connection factory returned an invalid connection"
+            raise ConnectionPoolError(message)
+        return cast("Connection", value)
+
+    def dispose(self, connection: Connection) -> None:
+        """Close a connection the driver has not already closed itself.
+
+        PyMySQL refuses a second close, and it force-closes the socket itself
+        after a fatal error, so an open socket is the only thing left to send a
+        quit message on.
+        """
+        if connection.open:
+            connection.close()
+
+
+__all__ = (
+    "MySQLConnectionFactoryProtocol",
+    "MySQLConnectionLease",
+    "MySQLPool",
+    "MySQLPoolPolicy",
+    "MySQLPoolStats",
+)

pyoq/query/mysql/transactions.py

@@ -0,0 +1,331 @@
+"""Owned MySQL transactions and nested savepoints."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+from typing import TypeVar, cast
+
+import pymysql
+from pymysql.cursors import Cursor
+
+from pyoq.errors import ExecutionError, QueryValidationError, TransactionStateError
+from pyoq.query.execution import (
+    CompiledQuery,
+    ExecutionControl,
+    RowStream,
+    StatementKind,
+    StreamPolicy,
+    StreamResource,
+    TransactionState,
+)
+from pyoq.query.mysql.commands import (
+    Connection,
+    ExecutionEnvironment,
+    FailureHandler,
+    open_cursor,
+    run_sql,
+)
+from pyoq.query.mysql.compiler import MySQLCompiler
+from pyoq.query.mysql.operations import MySQLQueryOperations
+from pyoq.query.mysql.parameters import adapt_parameters
+from pyoq.query.mysql.pool import MySQLPool
+from pyoq.query.select import Query
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+_ISOLATION_STATEMENTS = {
+    "READ_COMMITTED": "SET TRANSACTION ISOLATION LEVEL READ COMMITTED",
+    "REPEATABLE_READ": "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ",
+    "SERIALIZABLE": "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE",
+    "READ_ONLY": "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ",
+}
+_READ_ONLY_START = "START TRANSACTION READ ONLY"
+_START = "START TRANSACTION"
+
+
+class TransactionMode(Enum):
+    """MySQL sets isolation before starting, not as part of starting."""
+
+    READ_COMMITTED = "READ_COMMITTED"
+    REPEATABLE_READ = "REPEATABLE_READ"
+    SERIALIZABLE = "SERIALIZABLE"
+    READ_ONLY = "READ_ONLY"
+
+    @property
+    def isolation_sql(self) -> str:
+        return _ISOLATION_STATEMENTS[self.value]
+
+    @property
+    def start_sql(self) -> str:
+        return _READ_ONLY_START if self is TransactionMode.READ_ONLY else _START
+
+
+class MySQLTransactionState(TransactionState[Connection, "MySQLTransaction"]):
+    """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 MySQLTransaction(MySQLQueryOperations):
+    __slots__ = (
+        "_control",
+        "_entered",
+        "_finished",
+        "_mode",
+        "_parent",
+        "_pool",
+        "_savepoint_name",
+        "_state",
+    )
+
+    def __init__(
+        self,
+        pool: MySQLPool,
+        compiler: MySQLCompiler,
+        mode: TransactionMode = TransactionMode.REPEATABLE_READ,
+        control: ExecutionControl | None = None,
+        *,
+        parent: MySQLTransaction | None = None,
+        state: MySQLTransactionState | 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 MySQLTransactionState()
+        self._savepoint_name: str | None = None
+        self._entered = False
+        self._finished = False
+
+    def __enter__(self) -> MySQLTransaction:
+        if self._entered or self._finished:
+            raise TransactionStateError("transaction cannot be entered twice")
+        if self._parent is None:
+            self._enter_outer()
+        else:
+            self._enter_savepoint()
+        self._entered = True
+        return self
+
+    def __exit__(
+        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:
+            self._exit_outer(exception_type)
+        else:
+            self._exit_savepoint(exception_type)
+
+    def savepoint(
+        self,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> MySQLTransaction:
+        self._state.require_current(self)
+        return MySQLTransaction(
+            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,
+    ) -> RowStream[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: RowStream[ResultRow]
+        stream = RowStream(
+            lambda: self._open_stream(stream, statement, selected),
+            selected,
+            statement.row_types,
+        )
+        return stream
+
+    def _environment(self, connection: Connection) -> ExecutionEnvironment:
+        return ExecutionEnvironment(
+            connection,
+            self._pool.connection_factory,
+            FailureHandler(self._state.invalidate, self._state.mark_server_rollback),
+        )
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[Cursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        return run_sql(
+            self._environment(connection),
+            statement.sql,
+            parameters,
+            operation,
+            control or self._control,
+            one_result_set=statement.statement_kind is StatementKind.CALL,
+        )
+
+    def _enter_outer(self) -> None:
+        self._state.adopt(self._pool.checkout(), self)
+        try:
+            self._run_transaction_sql(self._mode.isolation_sql)
+            self._run_transaction_sql(self._mode.start_sql)
+        except BaseException:
+            self._release_outer()
+            raise
+
+    def _enter_savepoint(self) -> None:
+        self._state.require_current(cast("MySQLTransaction", self._parent))
+        name = self._state.next_savepoint()
+        self._run_transaction_sql(f"SAVEPOINT `{name}`", scope=self._parent)
+        self._savepoint_name = name
+        self._state.current_scope = self
+
+    def _exit_outer(self, exception_type: type[BaseException] | None) -> None:
+        try:
+            self._state.require_scope(self)
+        except TransactionStateError:
+            self._abort_misnested_outer(exception_type)
+            return
+        self._state.close_stream()
+        try:
+            self._finish_outer(exception_type)
+        except ExecutionError:
+            self._state.invalidate()
+            if exception_type is None:
+                raise
+        finally:
+            self._release_outer()
+
+    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
+        self._run_transaction_sql("ROLLBACK" if exception_type else "COMMIT")
+
+    def _exit_savepoint(self, exception_type: type[BaseException] | None) -> None:
+        self._state.require_scope(self)
+        self._state.close_stream()
+        name = cast("str", self._savepoint_name)
+        try:
+            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
+
+    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:
+            self._run_transaction_sql(f"ROLLBACK TO SAVEPOINT `{name}`")
+        self._run_transaction_sql(f"RELEASE SAVEPOINT `{name}`")
+
+    def _run_transaction_sql(
+        self,
+        sql: str,
+        *,
+        scope: MySQLTransaction | None = None,
+    ) -> None:
+        connection = self._state.require_current(scope or self, allow_stream=True)
+        run_sql(
+            self._environment(connection),
+            sql,
+            (),
+            _discard_cursor,
+            self._control,
+        )
+
+    def _open_stream(
+        self,
+        stream: RowStream[ResultRow],
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+    ) -> StreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        session = open_cursor(
+            self._environment(connection),
+            statement.sql,
+            parameters,
+            policy.control,
+            unbuffered=True,
+        )
+        self._state.active_stream = cast("RowStream[object]", stream)
+        return StreamResource(session, self._state.release_stream)
+
+    def _release_outer(self) -> None:
+        self._finished = True
+        self._state.release()
+
+    def _abort_misnested_outer(
+        self,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        connection = self._state.require_owner()
+        self._state.close_stream()
+        try:
+            connection.rollback()
+        except pymysql.Error:
+            self._state.invalidate()
+        finally:
+            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)
+
+
+def _discard_cursor(cursor: Cursor) -> None:
+    del cursor
+
+
+__all__ = ("MySQLTransaction", "MySQLTransactionState", "TransactionMode")

pyoq/query/mysql/writes.py

@@ -0,0 +1,73 @@
+"""MySQL write statement compilation."""
+
+from __future__ import annotations
+
+from pyoq.errors import CompilationError, UnsupportedQueryError
+from pyoq.query.mysql.context import MySQLCompilationContext
+from pyoq.query.mysql.expressions import MySQLExpressionCompiler
+from pyoq.query.mysql.identifiers import QUOTING
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.rendering.writes import WriteRenderer
+from pyoq.query.write_nodes import ConflictAction, ConflictNode, InsertNode
+
+ROW_ALIAS = "new"
+
+
+class MySQLWriteCompilation(WriteRenderer):
+    __slots__ = ("_context",)
+
+    def __init__(self, context: MySQLCompilationContext) -> None:
+        self._context = context
+        super().__init__(MySQLExpressionCompiler(context))
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    def insert_row_alias(self, node: InsertNode) -> str | None:
+        if node.conflict is None:
+            return None
+        return ROW_ALIAS
+
+    def render_conflict(self, node: ConflictNode) -> str:
+        self._require_conflict_support(node)
+        assignments = node.assignments if node.action is ConflictAction.UPDATE else ()
+        if not assignments:
+            message = "MySQL conflict resolution requires an assignment"
+            raise CompilationError(message)
+        rendered = ", ".join(
+            self.compile_assignment(assignment) for assignment in assignments
+        )
+        return f"ON DUPLICATE KEY UPDATE {rendered}"
+
+    def require_returning(self) -> None:
+        message = (
+            "MySQL does not support RETURNING; read the written rows with a "
+            "separate query"
+        )
+        raise UnsupportedQueryError(message)
+
+    def require_conflict_resolution(self) -> None:
+        if not self._context.capabilities.upsert:
+            message = "MySQL conflict resolution is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+    def _require_conflict_support(self, node: ConflictNode) -> None:
+        self.require_conflict_resolution()
+        if node.targets:
+            message = (
+                "MySQL infers the conflicting key and cannot accept a conflict target"
+            )
+            raise UnsupportedQueryError(message)
+        if node.condition is not None:
+            message = "MySQL conflict resolution cannot filter with a condition"
+            raise UnsupportedQueryError(message)
+        if node.action is ConflictAction.NOTHING:
+            message = (
+                "MySQL has no safe do-nothing conflict resolution; assign the "
+                "conflicting columns explicitly"
+            )
+            raise UnsupportedQueryError(message)
+
+
+__all__ = ("ROW_ALIAS", "MySQLWriteCompilation")