pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/execution/planning.py

@@ -0,0 +1,107 @@
+"""Bulk write planning bounded by a dialect parameter budget."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+from pyoq.errors import CompilationError, ParameterLimitError, QueryValidationError
+from pyoq.query.binding import (
+    bound_value_count,
+    projection_bound_value_count,
+    row_bound_value_count,
+)
+from pyoq.query.execution.batch import BulkChunk, BulkPlan, WriteProvider
+from pyoq.query.nodes import ExpressionNode
+from pyoq.query.statements import StatementCompiler
+from pyoq.query.write_nodes import DeleteNode, InsertNode, UpdateNode, WriteNode
+
+
+class BulkPlanner:
+    __slots__ = ("_compiler", "_dialect", "_maximum_parameters")
+
+    def __init__(
+        self,
+        compiler: StatementCompiler,
+        maximum_parameters: int,
+        dialect: str,
+    ) -> None:
+        self._compiler = compiler
+        self._maximum_parameters = maximum_parameters
+        self._dialect = dialect
+
+    def plan(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        node = _insert_node(statement)
+        if not node.rows:
+            message = "a bulk INSERT requires at least one row of values"
+            raise CompilationError(message)
+        limit = self._maximum_parameters
+        overhead = _suffix_bound_values(node)
+        if overhead > limit:
+            message = f"INSERT clauses alone exceed the {self._dialect} parameter limit"
+            raise ParameterLimitError(message)
+        return BulkPlan(tuple(self._chunks(node, limit, overhead)))
+
+    def _chunks(
+        self,
+        node: InsertNode,
+        limit: int,
+        overhead: int,
+    ) -> list[BulkChunk]:
+        chunks: list[BulkChunk] = []
+        start = 0
+        while start < len(node.rows):
+            length = _chunk_length(node.rows, start, limit, overhead, self._dialect)
+            rows = node.rows[start : start + length]
+            statement = self._compiler.compile(replace(node, rows=rows))
+            chunks.append(BulkChunk(statement, start, length))
+            start += length
+        return chunks
+
+
+def _insert_node(statement: WriteProvider | WriteNode) -> InsertNode:
+    node = (
+        statement
+        if isinstance(statement, (InsertNode, UpdateNode, DeleteNode))
+        else statement.node
+    )
+    if not isinstance(node, InsertNode):
+        message = "bulk planning requires an INSERT statement"
+        raise QueryValidationError(message)
+    return node
+
+
+def _chunk_length(
+    rows: tuple[tuple[ExpressionNode, ...], ...],
+    start: int,
+    limit: int,
+    overhead: int,
+    dialect: str,
+) -> int:
+    total = overhead
+    length = 0
+    for row in rows[start:]:
+        cost = row_bound_value_count(row)
+        if total + cost > limit:
+            if length == 0:
+                message = f"a single INSERT row exceeds the {dialect} parameter limit"
+                raise ParameterLimitError(message)
+            break
+        total += cost
+        length += 1
+    return length
+
+
+def _suffix_bound_values(node: InsertNode) -> int:
+    total = projection_bound_value_count(node.returning)
+    conflict = node.conflict
+    if conflict is None:
+        return total
+    total += sum(
+        bound_value_count(assignment.value) for assignment in conflict.assignments
+    )
+    if conflict.condition is not None:
+        total += bound_value_count(conflict.condition)
+    return total
+
+
+__all__ = ("BulkPlanner",)

pyoq/query/execution/pooling.py

@@ -0,0 +1,279 @@
+"""Exclusive synchronous connection pooling shared by every dialect."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass
+from threading import Condition, Lock
+from time import monotonic
+from typing import Generic, Self, TypeVar
+
+from pyoq.errors import (
+    ConnectionPoolClosedError,
+    ConnectionPoolError,
+    ConnectionPoolTimeoutError,
+    QueryValidationError,
+)
+
+ConnectionType = TypeVar("ConnectionType")
+ConnectionFactory = Callable[[], ConnectionType]
+
+
+@dataclass(frozen=True, slots=True)
+class PoolPolicy:
+    maximum_size: int = 5
+    minimum_size: int = 0
+    checkout_timeout: float | None = 30.0
+
+    def __post_init__(self) -> None:
+        if not _positive_integer(self.maximum_size):
+            message = "connection pool maximum size must be a positive integer"
+            raise QueryValidationError(message)
+        if not _non_negative_integer(self.minimum_size):
+            message = "connection pool minimum size must be a non-negative integer"
+            raise QueryValidationError(message)
+        if self.minimum_size > self.maximum_size:
+            message = "connection pool minimum size cannot exceed maximum size"
+            raise QueryValidationError(message)
+        if not _valid_timeout(self.checkout_timeout):
+            message = "connection pool checkout timeout must be positive or None"
+            raise QueryValidationError(message)
+
+
+@dataclass(frozen=True, slots=True)
+class PoolStats:
+    open_connections: int
+    idle_connections: int
+    checked_out_connections: int
+    closed: bool
+
+
+class ConnectionLease(Generic[ConnectionType]):
+    __slots__ = ("_connection", "_entered", "_invalid", "_pool")
+
+    def __init__(self, pool: ConnectionPool[ConnectionType]) -> None:
+        self._pool = pool
+        self._connection: ConnectionType | None = None
+        self._entered = False
+        self._invalid = False
+
+    @property
+    def connection(self) -> ConnectionType:
+        if self._connection is None or not self._entered:
+            message = "connection lease is not active"
+            raise ConnectionPoolError(message)
+        return self._connection
+
+    def invalidate(self) -> None:
+        if not self._entered:
+            message = "inactive connection lease cannot be invalidated"
+            raise ConnectionPoolError(message)
+        self._invalid = True
+
+    def __enter__(self) -> ConnectionType:
+        if self._entered:
+            message = "connection lease cannot be entered twice"
+            raise ConnectionPoolError(message)
+        connection = self._pool.acquire_for_lease()
+        self._connection = connection
+        self._entered = True
+        return connection
+
+    def __exit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        connection = self._connection
+        if connection is None or not self._entered:
+            return
+        self._entered = False
+        self._connection = None
+        self._pool.release_from_lease(connection, invalid=self._invalid)
+
+
+class ConnectionPool(ABC, Generic[ConnectionType]):
+    __slots__ = (
+        "_checked_out",
+        "_closed",
+        "_condition",
+        "_factory",
+        "_idle",
+        "_open_connections",
+        "_policy",
+    )
+
+    def __init__(
+        self,
+        factory: ConnectionFactory[ConnectionType],
+        policy: PoolPolicy | None = None,
+    ) -> None:
+        self._factory = factory
+        self._policy = policy or PoolPolicy()
+        self._condition = Condition(Lock())
+        self._idle: deque[ConnectionType] = deque()
+        self._checked_out: set[int] = set()
+        self._open_connections = 0
+        self._closed = False
+        self._initialize_minimum()
+
+    @abstractmethod
+    def reset(self, connection: ConnectionType) -> bool: ...
+
+    @abstractmethod
+    def validate(self, value: object) -> ConnectionType: ...
+
+    @abstractmethod
+    def dispose(self, connection: ConnectionType) -> None: ...
+
+    @property
+    def policy(self) -> PoolPolicy:
+        return self._policy
+
+    @property
+    def stats(self) -> PoolStats:
+        with self._condition:
+            return PoolStats(
+                self._open_connections,
+                len(self._idle),
+                len(self._checked_out),
+                self._closed,
+            )
+
+    def checkout(self) -> ConnectionLease[ConnectionType]:
+        return ConnectionLease(self)
+
+    def close(self) -> None:
+        with self._condition:
+            if self._closed:
+                return
+            self._closed = True
+            idle = tuple(self._idle)
+            self._idle.clear()
+            self._open_connections -= len(idle)
+            self._condition.notify_all()
+        for connection in idle:
+            self.dispose(connection)
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        self.close()
+
+    def acquire_for_lease(self) -> ConnectionType:
+        deadline = _deadline(self.policy.checkout_timeout)
+        with self._condition:
+            while True:
+                self._require_open()
+                if self._idle:
+                    connection = self._idle.popleft()
+                    self._checked_out.add(id(connection))
+                    return connection
+                if self._open_connections < self.policy.maximum_size:
+                    self._open_connections += 1
+                    break
+                self._wait_for_connection(deadline)
+        return self._create_reserved_connection()
+
+    def release_from_lease(
+        self,
+        connection: ConnectionType,
+        *,
+        invalid: bool,
+    ) -> None:
+        reusable = not invalid and self.reset(connection)
+        with self._condition:
+            self._checked_out.discard(id(connection))
+            if reusable and not self._closed:
+                self._idle.append(connection)
+            else:
+                self._open_connections -= 1
+            self._condition.notify()
+        if not reusable or self._closed:
+            self.dispose(connection)
+
+    def _initialize_minimum(self) -> None:
+        created: list[ConnectionType] = []
+        try:
+            for _ in range(self.policy.minimum_size):
+                created.append(self._create_connection())
+        except Exception:
+            for connection in created:
+                self.dispose(connection)
+            raise
+        self._idle.extend(created)
+        self._open_connections = len(created)
+
+    def _create_reserved_connection(self) -> ConnectionType:
+        try:
+            connection = self._create_connection()
+        except Exception:
+            with self._condition:
+                self._open_connections -= 1
+                self._condition.notify()
+            raise
+        with self._condition:
+            if self._closed:
+                self._open_connections -= 1
+                self.dispose(connection)
+                raise ConnectionPoolClosedError("connection pool is closed")
+            self._checked_out.add(id(connection))
+        return connection
+
+    def _create_connection(self) -> ConnectionType:
+        try:
+            connection_value: object = self._factory()
+        except ConnectionPoolError:
+            raise
+        except Exception as error:
+            message = f"unable to create pooled connection: {error}"
+            raise ConnectionPoolError(message) from error
+        return self.validate(connection_value)
+
+    def _wait_for_connection(self, deadline: float | None) -> None:
+        remaining = None if deadline is None else deadline - monotonic()
+        if remaining is not None and remaining <= 0:
+            message = "connection checkout exceeded its timeout"
+            raise ConnectionPoolTimeoutError(message)
+        self._condition.wait(remaining)
+
+    def _require_open(self) -> None:
+        if self._closed:
+            message = "connection pool is closed"
+            raise ConnectionPoolClosedError(message)
+
+
+def _deadline(timeout: float | None) -> float | None:
+    return None if timeout is None else monotonic() + timeout
+
+
+def _positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+def _non_negative_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value >= 0
+
+
+def _valid_timeout(value: object) -> bool:
+    return value is None or (
+        isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0
+    )
+
+
+__all__ = (
+    "ConnectionFactory",
+    "ConnectionLease",
+    "ConnectionPool",
+    "PoolPolicy",
+    "PoolStats",
+)

pyoq/query/execution/results.py

@@ -0,0 +1,36 @@
+"""Immutable execution result contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from pyoq.errors import QueryValidationError
+
+
+class StatementKind(Enum):
+    UNKNOWN = "unknown"
+    SELECT = "select"
+    INSERT = "insert"
+    UPDATE = "update"
+    DELETE = "delete"
+    DEFINITION = "definition"
+    CALL = "call"
+
+
+@dataclass(frozen=True, slots=True)
+class ExecuteResult:
+    rows_affected: int
+    last_inserted_id: int | None = None
+
+    def __post_init__(self) -> None:
+        if not _is_non_negative_integer(self.rows_affected):
+            message = "affected row count must be a non-negative integer"
+            raise QueryValidationError(message)
+
+
+def _is_non_negative_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value >= 0
+
+
+__all__ = ("ExecuteResult", "StatementKind")

pyoq/query/execution/streaming.py

@@ -0,0 +1,178 @@
+"""Bounded synchronous row streaming shared by every dialect."""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from threading import get_ident
+from typing import Generic, Protocol, Self, TypeVar, cast
+
+from pyoq.errors import QueryValidationError, StreamStateError
+from pyoq.query.decoding import decode_rows
+from pyoq.query.execution.control import ExecutionControl
+
+ResultRow = TypeVar("ResultRow")
+
+
+class CursorSession(Protocol):
+    def checkpoint(self) -> None: ...
+
+    def fetch_batch(self, size: int) -> list[tuple[object, ...]]: ...
+
+    def close(self) -> None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class StreamPolicy:
+    batch_size: int = 256
+    control: ExecutionControl | None = None
+
+    def __post_init__(self) -> None:
+        if not _positive_integer(self.batch_size):
+            message = "stream batch size must be a positive integer"
+            raise QueryValidationError(message)
+
+
+class StreamResource:
+    __slots__ = ("_closed", "_release", "session")
+
+    def __init__(
+        self,
+        session: CursorSession,
+        release: Callable[[], None],
+    ) -> None:
+        self.session = session
+        self._release = release
+        self._closed = False
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            self.session.close()
+        finally:
+            self._release()
+
+
+class RowStream(Iterator[ResultRow], Generic[ResultRow]):
+    __slots__ = (
+        "_buffer",
+        "_closed",
+        "_opened",
+        "_opener",
+        "_owner_thread",
+        "_policy",
+        "_resource",
+        "_row_types",
+    )
+
+    def __init__(
+        self,
+        opener: Callable[[], StreamResource],
+        policy: StreamPolicy,
+        row_types: tuple[type[object] | None, ...] = (),
+    ) -> None:
+        self._opener = opener
+        self._policy = policy
+        self._row_types = row_types
+        self._buffer: deque[tuple[object, ...]] = deque()
+        self._resource: StreamResource | None = None
+        self._owner_thread: int | None = None
+        self._opened = False
+        self._closed = False
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    @property
+    def buffered_rows(self) -> int:
+        return len(self._buffer)
+
+    def __enter__(self) -> Self:
+        if self._closed:
+            raise StreamStateError("closed row stream cannot be entered")
+        self._ensure_open()
+        return self
+
+    def __exit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        self.close()
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> ResultRow:
+        if self._closed:
+            raise StopIteration
+        self._ensure_open()
+        self._require_owner()
+        try:
+            resource = self._require_resource()
+            resource.session.checkpoint()
+            if not self._buffer:
+                self._load_batch()
+            return cast(ResultRow, self._buffer.popleft())
+        except BaseException:
+            self.close()
+            raise
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        if self._opened:
+            self._require_owner()
+        self._closed = True
+        self._buffer.clear()
+        resource = self._resource
+        self._resource = None
+        if resource is not None:
+            resource.close()
+
+    def _ensure_open(self) -> None:
+        if self._opened:
+            self._require_owner()
+            return
+        self._owner_thread = get_ident()
+        try:
+            self._resource = self._opener()
+        except BaseException:
+            self._closed = True
+            raise
+        self._opened = True
+
+    def _load_batch(self) -> None:
+        resource = self._require_resource()
+        rows = resource.session.fetch_batch(self._policy.batch_size)
+        if not rows:
+            raise StopIteration
+        self._buffer.extend(
+            cast("list[tuple[object, ...]]", decode_rows(rows, self._row_types))
+        )
+
+    def _require_resource(self) -> StreamResource:
+        if self._resource is None:
+            raise StreamStateError("row stream has no active cursor")
+        return self._resource
+
+    def _require_owner(self) -> None:
+        if self._owner_thread != get_ident():
+            raise StreamStateError("row stream used outside its owner thread")
+
+
+def _positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+__all__ = (
+    "CursorSession",
+    "RowStream",
+    "StreamPolicy",
+    "StreamResource",
+)

pyoq/query/execution/transactions.py

@@ -0,0 +1,95 @@
+"""Transaction scope, ownership, and savepoint state shared by every dialect."""
+
+from __future__ import annotations
+
+from threading import get_ident
+from typing import Generic, TypeVar, cast
+
+from pyoq.errors import TransactionStateError
+from pyoq.query.execution.pooling import ConnectionLease
+from pyoq.query.execution.streaming import RowStream
+
+ConnectionType = TypeVar("ConnectionType")
+ScopeType = TypeVar("ScopeType")
+
+_SAVEPOINT_PREFIX = "pyoq_savepoint_"
+
+
+class TransactionState(Generic[ConnectionType, ScopeType]):
+    __slots__ = (
+        "active_stream",
+        "broken",
+        "connection",
+        "current_scope",
+        "lease",
+        "owner_thread",
+        "savepoint_sequence",
+    )
+
+    def __init__(self) -> None:
+        self.lease: ConnectionLease[ConnectionType] | None = None
+        self.connection: ConnectionType | None = None
+        self.owner_thread: int | None = None
+        self.current_scope: ScopeType | None = None
+        self.active_stream: RowStream[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_thread != get_ident():
+            raise TransactionStateError("transaction used outside its owner thread")
+        return self.connection
+
+    def require_owner(self) -> ConnectionType:
+        if self.owner_thread != get_ident():
+            raise TransactionStateError("transaction is not active")
+        return cast(ConnectionType, self.connection)
+
+    def invalidate(self) -> None:
+        cast(ConnectionLease[ConnectionType], self.lease).invalidate()
+        self.broken = True
+
+    def next_savepoint(self) -> str:
+        self.savepoint_sequence += 1
+        return f"{_SAVEPOINT_PREFIX}{self.savepoint_sequence}"
+
+    def close_stream(self) -> None:
+        stream = self.active_stream
+        if stream is not None:
+            stream.close()
+
+    def release_stream(self) -> None:
+        self.active_stream = None
+
+    def adopt(self, lease: ConnectionLease[ConnectionType], scope: ScopeType) -> None:
+        self.lease = lease
+        self.connection = lease.__enter__()
+        self.owner_thread = get_ident()
+        self.current_scope = scope
+
+    def release(self) -> None:
+        lease = self.lease
+        self.current_scope = None
+        self.connection = None
+        self.lease = None
+        self.owner_thread = None
+        cast(ConnectionLease[ConnectionType], lease).__exit__(None, None, None)
+
+
+__all__ = ("TransactionState",)