pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/execution/aio/transactions.py

@@ -0,0 +1,105 @@
+"""Asynchronous transaction scope, ownership, and savepoint state."""
+
+from __future__ import annotations
+
+from asyncio import current_task
+from typing import Generic, TypeVar, cast
+
+from pyoq.errors import TransactionStateError
+from pyoq.query.execution.aio.pooling import AsyncConnectionLease
+from pyoq.query.execution.aio.streaming import AsyncRowStream
+
+ConnectionType = TypeVar("ConnectionType")
+ScopeType = TypeVar("ScopeType")
+
+_SAVEPOINT_PREFIX = "pyoq_savepoint_"
+
+
+class AsyncTransactionState(Generic[ConnectionType, ScopeType]):
+    __slots__ = (
+        "active_stream",
+        "broken",
+        "connection",
+        "current_scope",
+        "lease",
+        "owner_task",
+        "savepoint_sequence",
+    )
+
+    def __init__(self) -> None:
+        self.lease: AsyncConnectionLease[ConnectionType] | None = None
+        self.connection: ConnectionType | None = None
+        self.owner_task: object | None = None
+        self.current_scope: ScopeType | None = None
+        self.active_stream: AsyncRowStream[object] | None = None
+        self.savepoint_sequence = 0
+        self.broken = False
+
+    def require_current(
+        self,
+        scope: ScopeType,
+        *,
+        allow_stream: bool = False,
+    ) -> ConnectionType:
+        connection = self.require_scope(scope)
+        if self.broken:
+            raise TransactionStateError("transaction connection is unusable")
+        if self.active_stream is not None and not allow_stream:
+            raise TransactionStateError("transaction has an active row stream")
+        return connection
+
+    def require_scope(self, scope: ScopeType) -> ConnectionType:
+        if self.connection is None or self.current_scope is not scope:
+            message = "only the active innermost transaction may be used"
+            raise TransactionStateError(message)
+        if self.owner_task is not current_task():
+            raise TransactionStateError("transaction used outside its owner task")
+        return self.connection
+
+    def require_owner(self) -> ConnectionType:
+        if self.owner_task is not current_task():
+            raise TransactionStateError("transaction is not active")
+        return cast("ConnectionType", self.connection)
+
+    def invalidate(self) -> None:
+        cast(
+            "AsyncConnectionLease[ConnectionType]",
+            self.lease,
+        ).invalidate()
+        self.broken = True
+
+    def next_savepoint(self) -> str:
+        self.savepoint_sequence += 1
+        return f"{_SAVEPOINT_PREFIX}{self.savepoint_sequence}"
+
+    async def close_stream(self) -> None:
+        stream = self.active_stream
+        if stream is not None:
+            await stream.close()
+
+    async def release_stream(self) -> None:
+        self.active_stream = None
+
+    async def adopt(
+        self,
+        lease: AsyncConnectionLease[ConnectionType],
+        scope: ScopeType,
+    ) -> None:
+        self.lease = lease
+        self.connection = await lease.__aenter__()
+        self.owner_task = current_task()
+        self.current_scope = scope
+
+    async def release(self) -> None:
+        lease = self.lease
+        self.current_scope = None
+        self.connection = None
+        self.lease = None
+        self.owner_task = None
+        await cast(
+            "AsyncConnectionLease[ConnectionType]",
+            lease,
+        ).__aexit__(None, None, None)
+
+
+__all__ = ("AsyncTransactionState",)

pyoq/query/execution/batch.py

@@ -0,0 +1,96 @@
+"""Immutable bulk and multi-operation execution contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from pyoq.errors import QueryValidationError
+from pyoq.query.execution.compilation import CompiledQuery
+from pyoq.query.execution.results import ExecuteResult
+from pyoq.query.write_nodes import WriteNode
+
+
+@dataclass(frozen=True, slots=True)
+class OperationBudget:
+    maximum_operations: int = 256
+
+    def __post_init__(self) -> None:
+        if not _is_positive_integer(self.maximum_operations):
+            message = "maximum operations must be a positive integer"
+            raise QueryValidationError(message)
+
+
+@dataclass(frozen=True, slots=True)
+class BulkResult:
+    results: tuple[ExecuteResult, ...]
+    row_counts: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        if len(self.results) != len(self.row_counts):
+            message = "bulk results must align with their planned row counts"
+            raise QueryValidationError(message)
+
+    @property
+    def rows_affected(self) -> int:
+        return sum(result.rows_affected for result in self.results)
+
+    @property
+    def statement_count(self) -> int:
+        return len(self.results)
+
+    @property
+    def input_rows(self) -> int:
+        return sum(self.row_counts)
+
+
+def _is_positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+class WriteProvider(Protocol):
+    @property
+    def node(self) -> WriteNode: ...
+
+
+@dataclass(frozen=True, slots=True)
+class BulkChunk:
+    statement: CompiledQuery
+    start: int
+    length: int
+
+    def __post_init__(self) -> None:
+        if self.start < 0 or self.length <= 0:
+            message = "a bulk chunk must cover at least one input row"
+            raise QueryValidationError(message)
+
+
+@dataclass(frozen=True, slots=True)
+class BulkPlan:
+    chunks: tuple[BulkChunk, ...]
+
+    def __post_init__(self) -> None:
+        if not self.chunks:
+            message = "a bulk plan requires at least one chunk"
+            raise QueryValidationError(message)
+
+    @property
+    def statement_count(self) -> int:
+        return len(self.chunks)
+
+    @property
+    def input_rows(self) -> int:
+        return sum(chunk.length for chunk in self.chunks)
+
+    @property
+    def row_counts(self) -> tuple[int, ...]:
+        return tuple(chunk.length for chunk in self.chunks)
+
+
+__all__ = (
+    "BulkChunk",
+    "BulkPlan",
+    "BulkResult",
+    "OperationBudget",
+    "WriteProvider",
+)

pyoq/query/execution/binding_style.py

@@ -0,0 +1,30 @@
+"""Driver parameter placeholder styles."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class ParameterStyle(StrEnum):
+    QMARK = "qmark"
+    NUMERIC = "numeric"
+    FORMAT = "format"
+
+    def placeholder(self, index: int, /) -> str:
+        if self is ParameterStyle.QMARK:
+            return "?"
+        if self is ParameterStyle.NUMERIC:
+            return f"${index + 1}"
+        return "%s"
+
+    @property
+    def escapes_percent(self) -> bool:
+        return self is ParameterStyle.FORMAT
+
+    def structural_text(self, text: str, /) -> str:
+        if not self.escapes_percent:
+            return text
+        return text.replace("%", "%%")
+
+
+__all__ = ("ParameterStyle",)

pyoq/query/execution/compilation.py

@@ -0,0 +1,48 @@
+"""Dialect-neutral compiled statement output."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pyoq.errors import QueryValidationError
+from pyoq.query.execution.results import StatementKind
+
+
+@dataclass(frozen=True, slots=True)
+class CompiledQuery:
+    sql: str
+    parameters: tuple[object, ...] = ()
+    sensitive_parameter_indexes: frozenset[int] = frozenset()
+    statement_kind: StatementKind = StatementKind.UNKNOWN
+    # What each projected column was declared to hold, so a row can be given
+    # back as it was described rather than as a driver happened to answer.
+    row_types: tuple[type[object] | None, ...] = ()
+
+    def __post_init__(self) -> None:
+        if not self.sql or "\x00" in self.sql:
+            message = "compiled SQL cannot be empty or contain a null character"
+            raise QueryValidationError(message)
+        if any(
+            not _is_valid_parameter_index(index, len(self.parameters))
+            for index in self.sensitive_parameter_indexes
+        ):
+            message = "sensitive parameter index is outside the parameter set"
+            raise QueryValidationError(message)
+        if not _is_statement_kind(self.statement_kind):
+            message = "compiled statement kind is invalid"
+            raise QueryValidationError(message)
+
+
+def _is_valid_parameter_index(value: object, parameter_count: int) -> bool:
+    return (
+        isinstance(value, int)
+        and not isinstance(value, bool)
+        and 0 <= value < parameter_count
+    )
+
+
+def _is_statement_kind(value: object) -> bool:
+    return isinstance(value, StatementKind)
+
+
+__all__ = ("CompiledQuery",)

pyoq/query/execution/context.py

@@ -0,0 +1,61 @@
+"""Per-compilation parameter state shared by every dialect."""
+
+from __future__ import annotations
+
+from pyoq.errors import ParameterLimitError
+from pyoq.query.decoding import projected_types
+from pyoq.query.execution.binding_style import ParameterStyle
+from pyoq.query.execution.compilation import CompiledQuery
+from pyoq.query.execution.results import StatementKind
+from pyoq.query.nodes import BoundValueNode, ScalarFamily
+
+
+class CompilationContext:
+    __slots__ = ("_maximum_parameters", "_style", "parameters", "sensitive_indexes")
+
+    def __init__(
+        self,
+        maximum_parameters: int,
+        parameter_style: ParameterStyle,
+    ) -> None:
+        self._maximum_parameters = maximum_parameters
+        self._style = parameter_style
+        self.parameters: list[object] = []
+        self.sensitive_indexes: set[int] = set()
+
+    @property
+    def parameter_style(self) -> ParameterStyle:
+        return self._style
+
+    def bind(self, node: BoundValueNode) -> str:
+        index = len(self.parameters)
+        if index >= self._maximum_parameters:
+            message = "statement exceeds the configured parameter limit"
+            raise ParameterLimitError(message)
+        self.parameters.append(node.value)
+        if node.sensitive:
+            self.sensitive_indexes.add(index)
+        return self._style.placeholder(index)
+
+    def bind_page_value(self, value: int) -> str:
+        return self.bind(BoundValueNode(value, family=ScalarFamily.NUMERIC))
+
+    def structural_text(self, text: str, /) -> str:
+        return self._style.structural_text(text)
+
+    def result(
+        self,
+        sql: str,
+        statement_kind: StatementKind,
+        node: object = None,
+    ) -> CompiledQuery:
+        return CompiledQuery(
+            sql,
+            tuple(self.parameters),
+            frozenset(self.sensitive_indexes),
+            statement_kind,
+            projected_types(node) if node is not None else (),
+        )
+
+
+__all__ = ("CompilationContext",)

pyoq/query/execution/control.py

@@ -0,0 +1,50 @@
+"""Cooperative execution timeout and cancellation contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from threading import Event
+
+from pyoq.errors import QueryValidationError
+
+
+class CancellationToken:
+    __slots__ = ("_event",)
+
+    def __init__(self) -> None:
+        self._event = Event()
+
+    @property
+    def cancelled(self) -> bool:
+        return self._event.is_set()
+
+    def cancel(self) -> None:
+        self._event.set()
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionControl:
+    timeout: float | None = None
+    cancellation_token: CancellationToken | None = None
+    progress_steps: int = 1_000
+
+    def __post_init__(self) -> None:
+        if not _valid_timeout(self.timeout):
+            message = "execution timeout must be positive or None"
+            raise QueryValidationError(message)
+        if not _positive_integer(self.progress_steps):
+            message = "execution progress steps must be a positive integer"
+            raise QueryValidationError(message)
+
+
+def _valid_timeout(value: object) -> bool:
+    return value is None or (
+        isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0
+    )
+
+
+def _positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+__all__ = ("CancellationToken", "ExecutionControl")

pyoq/query/execution/operations.py

@@ -0,0 +1,224 @@
+"""Dialect-neutral typed result operations."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable, Sequence
+from typing import Protocol, TypeVar, cast
+
+from pyoq.errors import (
+    MultipleResultsError,
+    NoResultError,
+    QueryValidationError,
+    ResultCardinalityError,
+)
+from pyoq.query.decoding import decode_row, decode_rows
+from pyoq.query.execution.batch import (
+    BulkPlan,
+    BulkResult,
+    OperationBudget,
+    WriteProvider,
+)
+from pyoq.query.execution.compilation import CompiledQuery
+from pyoq.query.execution.control import ExecutionControl
+from pyoq.query.execution.results import ExecuteResult, StatementKind
+from pyoq.query.routines import Call
+from pyoq.query.select import Query
+from pyoq.query.statements import StatementCompiler
+from pyoq.query.write_nodes import WriteNode
+from pyoq.query.writes import ReturningStatement, WriteStatement
+
+ResultRow = TypeVar("ResultRow")
+ScalarValue = TypeVar("ScalarValue")
+OperationResult = TypeVar("OperationResult")
+
+
+class DatabaseCursor(Protocol):
+    @property
+    def rowcount(self) -> int: ...
+
+    def fetchall(self) -> Sequence[object]: ...
+
+    def fetchmany(self, size: int, /) -> Sequence[object]: ...
+
+
+class QueryOperations(ABC):
+    __slots__ = ()
+
+    @property
+    @abstractmethod
+    def compiler(self) -> StatementCompiler: ...
+
+    def one(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ResultRow:
+        compiled = self.compiler.compile(query)
+        rows = self._limited_rows(compiled, control)
+        if not rows:
+            raise NoResultError("query returned no rows")
+        if len(rows) > 1:
+            raise MultipleResultsError("query returned more than one row")
+        return cast(ResultRow, decode_row(rows[0], compiled.row_types))
+
+    def one_or_none(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ResultRow | None:
+        compiled = self.compiler.compile(query)
+        rows = self._limited_rows(compiled, control)
+        if len(rows) > 1:
+            raise MultipleResultsError("query returned more than one row")
+        if not rows:
+            return None
+        return cast(ResultRow, decode_row(rows[0], compiled.row_types))
+
+    def many(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> list[ResultRow]:
+        compiled = self.compiler.compile(query)
+        rows = self._run(compiled, fetch_all, control)
+        return cast(list[ResultRow], decode_rows(rows, compiled.row_types))
+
+    def scalar(
+        self,
+        query: Query[tuple[ScalarValue]] | ReturningStatement[tuple[ScalarValue]],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ScalarValue:
+        row = cast(tuple[object, ...], self.one(query, control=control))
+        return cast(ScalarValue, _scalar_value(row))
+
+    def execute(
+        self,
+        statement: CompiledQuery | WriteStatement[WriteNode] | Call[tuple[object, ...]],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ExecuteResult:
+        compiled = (
+            statement
+            if isinstance(statement, CompiledQuery)
+            else self.compiler.compile(statement)
+        )
+        return self._run(
+            compiled,
+            lambda cursor: self._execution_result(cursor, compiled.statement_kind),
+            control,
+        )
+
+    def last_inserted_id(self, cursor: DatabaseCursor) -> int | None:
+        del cursor
+        return None
+
+    @abstractmethod
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan: ...
+
+    def execute_bulk(
+        self,
+        statement: WriteProvider | WriteNode,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> BulkResult:
+        plan = self.plan_bulk(statement)
+        return BulkResult(self.execute_chunks(plan, control), plan.row_counts)
+
+    def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        return tuple(
+            self.execute(chunk.statement, control=control) for chunk in plan.chunks
+        )
+
+    def many_bulk(
+        self,
+        statement: ReturningStatement[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> list[ResultRow]:
+        plan = self.plan_bulk(statement)
+        rows: list[object] = []
+        for chunk in plan.chunks:
+            answered = self._run(chunk.statement, fetch_all, control)
+            rows.extend(decode_rows(answered, chunk.statement.row_types))
+        return cast(list[ResultRow], rows)
+
+    def execute_all(
+        self,
+        statements: Iterable[CompiledQuery | WriteStatement[WriteNode]],
+        /,
+        *,
+        budget: OperationBudget | None = None,
+        control: ExecutionControl | None = None,
+    ) -> tuple[ExecuteResult, ...]:
+        operations: Sequence[CompiledQuery | WriteStatement[WriteNode]] = tuple(
+            statements
+        )
+        limit = (budget or OperationBudget()).maximum_operations
+        if len(operations) > limit:
+            message = f"operation batch exceeds its budget of {limit} operations"
+            raise QueryValidationError(message)
+        return tuple(
+            self.execute(operation, control=control) for operation in operations
+        )
+
+    def _execution_result(
+        self,
+        cursor: DatabaseCursor,
+        statement_kind: StatementKind,
+    ) -> ExecuteResult:
+        row_count = max(cursor.rowcount, 0)
+        inserted = (
+            self.last_inserted_id(cursor)
+            if statement_kind is StatementKind.INSERT
+            else None
+        )
+        return ExecuteResult(row_count, inserted)
+
+    def _limited_rows(
+        self,
+        statement: CompiledQuery,
+        control: ExecutionControl | None,
+    ) -> list[object]:
+        return self._run(statement, _fetch_up_to_two, control)
+
+    @abstractmethod
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[DatabaseCursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult: ...
+
+
+def fetch_all(cursor: DatabaseCursor) -> list[object]:
+    return list(cursor.fetchall())
+
+
+def _fetch_up_to_two(cursor: DatabaseCursor) -> list[object]:
+    return list(cursor.fetchmany(2))
+
+
+def _scalar_value(row: tuple[object, ...]) -> object:
+    if len(row) != 1:
+        message = "scalar query projection must contain exactly one value"
+        raise ResultCardinalityError(message)
+    return row[0]
+
+
+__all__ = ("DatabaseCursor", "QueryOperations", "fetch_all")