pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/execution/aio/operations.py

@@ -0,0 +1,228 @@
+"""Dialect-neutral asynchronous typed result operations."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, 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
+
+
+class AsyncDatabaseCursor(Protocol):
+    @property
+    def rowcount(self) -> int: ...
+
+    async def fetchall(self) -> Sequence[object]: ...
+
+    async def fetchmany(self, size: int, /) -> Sequence[object]: ...
+
+
+ResultRow = TypeVar("ResultRow")
+ScalarValue = TypeVar("ScalarValue")
+OperationResult = TypeVar("OperationResult")
+
+
+class AsyncQueryOperations(ABC):
+    __slots__ = ()
+
+    @property
+    @abstractmethod
+    def compiler(self) -> StatementCompiler: ...
+
+    @abstractmethod
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan: ...
+
+    @abstractmethod
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncDatabaseCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult: ...
+
+    def last_inserted_id(self, cursor: AsyncDatabaseCursor) -> int | None:
+        del cursor
+        return None
+
+    async def one(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ResultRow:
+        compiled = self.compiler.compile(query)
+        rows = await 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))
+
+    async def one_or_none(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ResultRow | None:
+        compiled = self.compiler.compile(query)
+        rows = await 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))
+
+    async def many(
+        self,
+        query: Query[ResultRow] | ReturningStatement[ResultRow] | Call[ResultRow],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> list[ResultRow]:
+        compiled = self.compiler.compile(query)
+        rows = await self._run(compiled, fetch_all, control)
+        return cast("list[ResultRow]", decode_rows(rows, compiled.row_types))
+
+    async def scalar(
+        self,
+        query: Query[tuple[ScalarValue]] | ReturningStatement[tuple[ScalarValue]],
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ScalarValue:
+        row = cast("tuple[object, ...]", await self.one(query, control=control))
+        return cast("ScalarValue", _scalar_value(row))
+
+    async 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)
+        )
+
+        async def operation(cursor: AsyncDatabaseCursor) -> ExecuteResult:
+            return self._execution_result(cursor, compiled.statement_kind)
+
+        return await self._run(compiled, operation, control)
+
+    async def execute_bulk(
+        self,
+        statement: WriteProvider | WriteNode,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> BulkResult:
+        plan = self.plan_bulk(statement)
+        return BulkResult(await self.execute_chunks(plan, control), plan.row_counts)
+
+    async def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        results = [
+            await self.execute(chunk.statement, control=control)
+            for chunk in plan.chunks
+        ]
+        return tuple(results)
+
+    async 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 = await self._run(chunk.statement, fetch_all, control)
+            rows.extend(decode_rows(answered, chunk.statement.row_types))
+        return cast("list[ResultRow]", rows)
+
+    async 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)
+        results = [
+            await self.execute(operation, control=control) for operation in operations
+        ]
+        return tuple(results)
+
+    def _execution_result(
+        self,
+        cursor: AsyncDatabaseCursor,
+        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)
+
+    async def _limited_rows(
+        self,
+        statement: CompiledQuery,
+        control: ExecutionControl | None,
+    ) -> list[object]:
+        return await self._run(statement, _fetch_up_to_two, control)
+
+
+async def fetch_all(cursor: AsyncDatabaseCursor) -> list[object]:
+    return list(await cursor.fetchall())
+
+
+async def _fetch_up_to_two(cursor: AsyncDatabaseCursor) -> list[object]:
+    return list(await 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__ = ("AsyncDatabaseCursor", "AsyncQueryOperations", "fetch_all")

pyoq/query/execution/aio/pooling.py

@@ -0,0 +1,233 @@
+"""Exclusive asynchronous connection pooling shared by every dialect."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from asyncio import Condition, wait_for
+from collections import deque
+from collections.abc import Awaitable, Callable
+from time import monotonic
+from typing import Generic, Self, TypeVar
+
+from pyoq.errors import (
+    ConnectionPoolClosedError,
+    ConnectionPoolError,
+    ConnectionPoolTimeoutError,
+)
+from pyoq.query.execution.pooling import PoolPolicy, PoolStats
+
+ConnectionType = TypeVar("ConnectionType")
+AsyncConnectionFactory = Callable[[], Awaitable[ConnectionType]]
+
+
+class AsyncConnectionLease(Generic[ConnectionType]):
+    __slots__ = ("_connection", "_entered", "_invalid", "_pool")
+
+    def __init__(self, pool: AsyncConnectionPool[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
+
+    async def __aenter__(self) -> ConnectionType:
+        if self._entered:
+            message = "connection lease cannot be entered twice"
+            raise ConnectionPoolError(message)
+        connection = await self._pool.acquire_for_lease()
+        self._connection = connection
+        self._entered = True
+        return connection
+
+    async def __aexit__(
+        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
+        await self._pool.release_from_lease(connection, invalid=self._invalid)
+
+
+class AsyncConnectionPool(ABC, Generic[ConnectionType]):
+    __slots__ = (
+        "_checked_out",
+        "_closed",
+        "_condition",
+        "_factory",
+        "_idle",
+        "_open_connections",
+        "_policy",
+    )
+
+    def __init__(
+        self,
+        factory: AsyncConnectionFactory[ConnectionType],
+        policy: PoolPolicy | None = None,
+    ) -> None:
+        self._factory = factory
+        self._policy = policy or PoolPolicy()
+        self._condition = Condition()
+        self._idle: deque[ConnectionType] = deque()
+        self._checked_out: set[int] = set()
+        self._open_connections = 0
+        self._closed = False
+
+    @abstractmethod
+    async def reset(self, connection: ConnectionType) -> bool: ...
+
+    @abstractmethod
+    def validate(self, value: object) -> ConnectionType: ...
+
+    @abstractmethod
+    async def dispose(self, connection: ConnectionType) -> None: ...
+
+    @property
+    def policy(self) -> PoolPolicy:
+        return self._policy
+
+    @property
+    def stats(self) -> PoolStats:
+        return PoolStats(
+            self._open_connections,
+            len(self._idle),
+            len(self._checked_out),
+            self._closed,
+        )
+
+    def checkout(self) -> AsyncConnectionLease[ConnectionType]:
+        return AsyncConnectionLease(self)
+
+    async def open(self) -> Self:
+        created: list[ConnectionType] = []
+        try:
+            for _ in range(self.policy.minimum_size):
+                created.append(await self._create_connection())
+        except Exception:
+            for connection in created:
+                await self.dispose(connection)
+            raise
+        self._idle.extend(created)
+        self._open_connections += len(created)
+        return self
+
+    async def close(self) -> None:
+        async 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:
+            await self.dispose(connection)
+
+    async def __aenter__(self) -> Self:
+        return await self.open()
+
+    async def __aexit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        await self.close()
+
+    async def acquire_for_lease(self) -> ConnectionType:
+        deadline = _deadline(self.policy.checkout_timeout)
+        async 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
+                await self._wait_for_connection(deadline)
+        return await self._create_reserved_connection()
+
+    async def release_from_lease(
+        self,
+        connection: ConnectionType,
+        *,
+        invalid: bool,
+    ) -> None:
+        reusable = not invalid and await self.reset(connection)
+        async 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:
+            await self.dispose(connection)
+
+    async def _create_reserved_connection(self) -> ConnectionType:
+        try:
+            connection = await self._create_connection()
+        except Exception:
+            async with self._condition:
+                self._open_connections -= 1
+                self._condition.notify()
+            raise
+        async with self._condition:
+            if self._closed:
+                self._open_connections -= 1
+                await self.dispose(connection)
+                raise ConnectionPoolClosedError("connection pool is closed")
+            self._checked_out.add(id(connection))
+        return connection
+
+    async def _create_connection(self) -> ConnectionType:
+        try:
+            connection_value: object = await 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)
+
+    async def _wait_for_connection(self, deadline: float | None) -> None:
+        remaining = None if deadline is None else deadline - monotonic()
+        try:
+            await wait_for(self._condition.wait(), remaining)
+        except TimeoutError as error:
+            message = "connection checkout exceeded its timeout"
+            raise ConnectionPoolTimeoutError(message) from error
+
+    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
+
+
+__all__ = (
+    "AsyncConnectionFactory",
+    "AsyncConnectionLease",
+    "AsyncConnectionPool",
+)

pyoq/query/execution/aio/streaming.py

@@ -0,0 +1,161 @@
+"""Bounded asynchronous row streaming shared by every dialect."""
+
+from __future__ import annotations
+
+from asyncio import current_task
+from collections import deque
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Generic, Protocol, Self, TypeVar, cast
+
+from pyoq.errors import StreamStateError
+from pyoq.query.decoding import decode_rows
+from pyoq.query.execution.streaming import StreamPolicy
+
+ResultRow = TypeVar("ResultRow")
+
+
+class AsyncCursorSession(Protocol):
+    def checkpoint(self) -> None: ...
+
+    async def fetch_batch(self, size: int) -> list[tuple[object, ...]]: ...
+
+    async def close(self) -> None: ...
+
+
+class AsyncStreamResource:
+    __slots__ = ("_closed", "_release", "session")
+
+    def __init__(
+        self,
+        session: AsyncCursorSession,
+        release: Callable[[], Awaitable[None]],
+    ) -> None:
+        self.session = session
+        self._release = release
+        self._closed = False
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            await self.session.close()
+        finally:
+            await self._release()
+
+
+class AsyncRowStream(AsyncIterator[ResultRow], Generic[ResultRow]):
+    __slots__ = (
+        "_buffer",
+        "_closed",
+        "_opened",
+        "_opener",
+        "_owner_task",
+        "_policy",
+        "_resource",
+        "_row_types",
+    )
+
+    def __init__(
+        self,
+        opener: Callable[[], Awaitable[AsyncStreamResource]],
+        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: AsyncStreamResource | None = None
+        self._owner_task: object | 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)
+
+    async def __aenter__(self) -> Self:
+        if self._closed:
+            raise StreamStateError("closed row stream cannot be entered")
+        await self._ensure_open()
+        return self
+
+    async def __aexit__(
+        self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: object,
+    ) -> None:
+        await self.close()
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> ResultRow:
+        if self._closed:
+            raise StopAsyncIteration
+        await self._ensure_open()
+        self._require_owner()
+        try:
+            resource = self._require_resource()
+            resource.session.checkpoint()
+            if not self._buffer:
+                await self._load_batch()
+            return cast("ResultRow", self._buffer.popleft())
+        except BaseException:
+            await self.close()
+            raise
+
+    async 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:
+            await resource.close()
+
+    async def _ensure_open(self) -> None:
+        if self._opened:
+            self._require_owner()
+            return
+        self._owner_task = current_task()
+        try:
+            self._resource = await self._opener()
+        except BaseException:
+            self._closed = True
+            raise
+        self._opened = True
+
+    async def _load_batch(self) -> None:
+        resource = self._require_resource()
+        rows = await resource.session.fetch_batch(self._policy.batch_size)
+        if not rows:
+            raise StopAsyncIteration
+        self._buffer.extend(
+            cast("list[tuple[object, ...]]", decode_rows(rows, self._row_types))
+        )
+
+    def _require_resource(self) -> AsyncStreamResource:
+        if self._resource is None:
+            raise StreamStateError("row stream has no active cursor")
+        return self._resource
+
+    def _require_owner(self) -> None:
+        if self._owner_task is not current_task():
+            raise StreamStateError("row stream used outside its owner task")
+
+
+__all__ = (
+    "AsyncCursorSession",
+    "AsyncRowStream",
+    "AsyncStreamResource",
+)