pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/postgres/__init__.py

@@ -0,0 +1,48 @@
+"""PostgreSQL query compilation and execution."""
+
+from pyoq.query.postgres.commands import PostgresCursorSession
+from pyoq.query.postgres.compiler import PostgresCompiler
+from pyoq.query.postgres.context import PostgresCompilationContext
+from pyoq.query.postgres.executor import PostgresExecutor
+from pyoq.query.postgres.expressions import PostgresExpressionCompiler
+from pyoq.query.postgres.factory import PostgresConnectionFactory
+from pyoq.query.postgres.model import PostgresCapabilities
+from pyoq.query.postgres.operations import PostgresQueryOperations
+from pyoq.query.postgres.parameters import (
+    PostgresParameter,
+    adapt_parameter,
+    adapt_parameters,
+)
+from pyoq.query.postgres.planning import PostgresBulkPlanner
+from pyoq.query.postgres.pool import (
+    PostgresConnectionFactoryProtocol,
+    PostgresConnectionLease,
+    PostgresPool,
+    PostgresPoolPolicy,
+    PostgresPoolStats,
+)
+from pyoq.query.postgres.transactions import PostgresTransaction, TransactionMode
+from pyoq.query.postgres.writes import PostgresWriteCompilation
+
+__all__ = (
+    "PostgresBulkPlanner",
+    "PostgresCapabilities",
+    "PostgresCompilationContext",
+    "PostgresCompiler",
+    "PostgresConnectionFactory",
+    "PostgresConnectionFactoryProtocol",
+    "PostgresConnectionLease",
+    "PostgresCursorSession",
+    "PostgresExecutor",
+    "PostgresExpressionCompiler",
+    "PostgresParameter",
+    "PostgresPool",
+    "PostgresPoolPolicy",
+    "PostgresPoolStats",
+    "PostgresQueryOperations",
+    "PostgresTransaction",
+    "PostgresWriteCompilation",
+    "TransactionMode",
+    "adapt_parameter",
+    "adapt_parameters",
+)

pyoq/query/postgres/aio/__init__.py

@@ -0,0 +1,25 @@
+"""Asynchronous PostgreSQL execution."""
+
+from pyoq.query.postgres.aio.commands import (
+    PostgresAsyncCursorSession,
+    PostgresAsyncExecutionMonitor,
+)
+from pyoq.query.postgres.aio.executor import PostgresAsyncExecutor
+from pyoq.query.postgres.aio.factory import PostgresAsyncConnectionFactory
+from pyoq.query.postgres.aio.operations import PostgresAsyncQueryOperations
+from pyoq.query.postgres.aio.pool import (
+    PostgresAsyncConnectionLease,
+    PostgresAsyncPool,
+)
+from pyoq.query.postgres.aio.transactions import PostgresAsyncTransaction
+
+__all__ = (
+    "PostgresAsyncConnectionFactory",
+    "PostgresAsyncConnectionLease",
+    "PostgresAsyncCursorSession",
+    "PostgresAsyncExecutionMonitor",
+    "PostgresAsyncExecutor",
+    "PostgresAsyncPool",
+    "PostgresAsyncQueryOperations",
+    "PostgresAsyncTransaction",
+)

pyoq/query/postgres/aio/bulk.py

@@ -0,0 +1,56 @@
+"""Pipelined asynchronous PostgreSQL bulk chunk execution."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import psycopg
+
+from pyoq.query.execution import BulkChunk, ExecuteResult, ExecutionControl
+from pyoq.query.postgres.aio.commands import PostgresAsyncExecutionMonitor
+from pyoq.query.postgres.parameters import adapt_parameters
+
+AsyncConnection = psycopg.AsyncConnection[tuple[object, ...]]
+AsyncCursor = psycopg.AsyncCursor[tuple[object, ...]]
+Invalidator = Callable[[], None]
+
+
+async def run_pipeline(
+    connection: AsyncConnection,
+    chunks: tuple[BulkChunk, ...],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> tuple[ExecuteResult, ...]:
+    monitor = PostgresAsyncExecutionMonitor(control)
+    cursors: list[AsyncCursor] = []
+    try:
+        await monitor.install(connection, invalidate)
+        async with connection.pipeline():
+            for chunk in chunks:
+                monitor.checkpoint()
+                cursors.append(await _send(connection, chunk))
+        monitor.checkpoint()
+        return tuple(ExecuteResult(max(cursor.rowcount, 0)) for cursor in cursors)
+    except psycopg.Error as error:
+        monitor.translate_driver_error(error, invalidate)
+    finally:
+        for cursor in cursors:
+            await _close_quietly(cursor)
+        await monitor.close(connection, invalidate)
+
+
+async def _send(connection: AsyncConnection, chunk: BulkChunk) -> AsyncCursor:
+    cursor = connection.cursor()
+    parameters = adapt_parameters(chunk.statement.parameters)
+    await cursor.execute(chunk.statement.sql.encode(), parameters or None)
+    return cursor
+
+
+async def _close_quietly(cursor: AsyncCursor) -> None:
+    try:
+        await cursor.close()
+    except psycopg.Error:
+        return
+
+
+__all__ = ("run_pipeline",)

pyoq/query/postgres/aio/commands.py

@@ -0,0 +1,264 @@
+"""Asynchronous PostgreSQL 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 time import monotonic
+from typing import NoReturn, TypeVar
+
+import psycopg
+
+from pyoq.errors import QueryCancelledError, QueryExecutionError, QueryTimeoutError
+from pyoq.query.execution import ExecutionControl
+from pyoq.query.postgres.health import is_cancellation, is_connection_failure
+from pyoq.query.postgres.parameters import PostgresParameter
+
+OperationResult = TypeVar("OperationResult")
+Invalidator = Callable[[], None]
+AsyncConnection = psycopg.AsyncConnection[tuple[object, ...]]
+AsyncCursor = psycopg.AsyncCursor[tuple[object, ...]]
+
+_WATCHDOG_INTERVAL = 0.025
+_MILLISECONDS = 1000
+
+
+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
+
+
+class PostgresAsyncExecutionMonitor:
+    __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,
+        connection: AsyncConnection,
+        invalidate: Invalidator,
+    ) -> None:
+        self.checkpoint()
+        await self._apply_timeout(connection, invalidate)
+        if self._control.cancellation_token is None:
+            return
+        self._watchdog = create_task(self._watch(connection))
+
+    def checkpoint(self) -> None:
+        token = self._control.cancellation_token
+        if token is not None and token.cancelled:
+            raise QueryCancelledError("PostgreSQL query execution was cancelled")
+        if self._expired():
+            raise QueryTimeoutError("PostgreSQL query execution exceeded its timeout")
+
+    def translate_driver_error(
+        self,
+        error: psycopg.Error,
+        invalidate: Invalidator,
+    ) -> NoReturn:
+        if is_connection_failure(error):
+            invalidate()
+        self.checkpoint()
+        if is_cancellation(error):
+            raise QueryCancelledError(
+                "PostgreSQL query execution was cancelled"
+            ) from error
+        _raise_query_execution_error(error)
+
+    async def close(
+        self,
+        connection: AsyncConnection,
+        invalidate: Invalidator,
+    ) -> None:
+        self._stop.set()
+        watchdog = self._watchdog
+        self._watchdog = None
+        if watchdog is not None:
+            await await_watchdog(watchdog)
+        await self._clear_timeout(connection, invalidate)
+
+    async def _watch(self, connection: AsyncConnection) -> None:
+        token = self._control.cancellation_token
+        while True:
+            try:
+                await wait_for(self._stop.wait(), _WATCHDOG_INTERVAL)
+            except (TimeoutError, CancelledError):
+                if token is not None and token.cancelled:
+                    await _cancel_quietly(connection)
+                    return
+                continue
+            return
+
+    async def _apply_timeout(
+        self,
+        connection: AsyncConnection,
+        invalidate: Invalidator,
+    ) -> None:
+        timeout = self._control.timeout
+        if timeout is None:
+            return
+        milliseconds = max(int(timeout * _MILLISECONDS), 1)
+        try:
+            await connection.execute(f"SET statement_timeout = {milliseconds}".encode())
+        except psycopg.Error as error:
+            self.translate_driver_error(error, invalidate)
+        self._timeout_applied = True
+
+    async def _clear_timeout(
+        self,
+        connection: AsyncConnection,
+        invalidate: Invalidator,
+    ) -> None:
+        if not self._timeout_applied:
+            return
+        self._timeout_applied = False
+        try:
+            await connection.execute(b"SET statement_timeout = 0")
+        except psycopg.Error:
+            invalidate()
+
+    def _expired(self) -> bool:
+        timeout = self._control.timeout
+        return timeout is not None and monotonic() - self._started >= timeout
+
+
+class PostgresAsyncCursorSession:
+    __slots__ = ("_closed", "_connection", "_cursor", "_invalidate", "_monitor")
+
+    def __init__(
+        self,
+        connection: AsyncConnection,
+        cursor: AsyncCursor,
+        monitor: PostgresAsyncExecutionMonitor,
+        invalidate: Invalidator,
+    ) -> None:
+        self._connection = connection
+        self._cursor = cursor
+        self._monitor = monitor
+        self._invalidate = invalidate
+        self._closed = False
+
+    @property
+    def cursor(self) -> AsyncCursor:
+        return self._cursor
+
+    def checkpoint(self) -> None:
+        self._monitor.checkpoint()
+
+    def translate_driver_error(self, error: psycopg.Error) -> NoReturn:
+        self._monitor.translate_driver_error(error, self._invalidate)
+
+    async def fetch_batch(self, size: int) -> list[tuple[object, ...]]:
+        try:
+            self._monitor.checkpoint()
+            rows = await self._cursor.fetchmany(size)
+            self._monitor.checkpoint()
+        except psycopg.Error as error:
+            self._monitor.translate_driver_error(error, self._invalidate)
+        return rows
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await _close_cursor(self._cursor, self._invalidate)
+        await self._monitor.close(self._connection, self._invalidate)
+
+
+async def run_sql(
+    connection: AsyncConnection,
+    sql: str,
+    parameters: tuple[PostgresParameter, ...],
+    operation: Callable[[AsyncCursor], Awaitable[OperationResult]],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> OperationResult:
+    session = await open_cursor(connection, sql, parameters, control, invalidate)
+    try:
+        result = await operation(session.cursor)
+        session.checkpoint()
+        return result
+    except psycopg.Error as error:
+        session.translate_driver_error(error)
+    finally:
+        await session.close()
+
+
+async def open_cursor(
+    connection: AsyncConnection,
+    sql: str,
+    parameters: tuple[PostgresParameter, ...],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+    *,
+    server_side_name: str | None = None,
+) -> PostgresAsyncCursorSession:
+    monitor = PostgresAsyncExecutionMonitor(control)
+    cursor: AsyncCursor | None = None
+    try:
+        await monitor.install(connection, invalidate)
+        cursor = connection.cursor(name=server_side_name or "")
+        await cursor.execute(sql.encode(), parameters or None)
+        monitor.checkpoint()
+        return PostgresAsyncCursorSession(connection, cursor, monitor, invalidate)
+    except psycopg.Error as error:
+        await _release(cursor, monitor, connection, invalidate)
+        monitor.translate_driver_error(error, invalidate)
+    except BaseException:
+        await _release(cursor, monitor, connection, invalidate)
+        raise
+
+
+async def _release(
+    cursor: AsyncCursor | None,
+    monitor: PostgresAsyncExecutionMonitor,
+    connection: AsyncConnection,
+    invalidate: Invalidator,
+) -> None:
+    if cursor is not None:
+        await _close_cursor(cursor, invalidate)
+    await monitor.close(connection, invalidate)
+
+
+async def _close_cursor(cursor: AsyncCursor, invalidate: Invalidator) -> None:
+    try:
+        await cursor.close()
+    except psycopg.Error:
+        invalidate()
+
+
+async def _cancel_quietly(connection: AsyncConnection) -> None:
+    try:
+        await connection.cancel_safe()
+    except psycopg.Error:
+        return
+
+
+def _raise_query_execution_error(error: psycopg.Error) -> NoReturn:
+    message = f"PostgreSQL query execution failed: {error}"
+    raise QueryExecutionError(message) from error
+
+
+__all__ = (
+    "PostgresAsyncCursorSession",
+    "PostgresAsyncExecutionMonitor",
+    "await_watchdog",
+    "open_cursor",
+    "run_sql",
+)

pyoq/query/postgres/aio/executor.py

@@ -0,0 +1,152 @@
+"""Asynchronous typed PostgreSQL query execution."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
+import psycopg
+
+from pyoq.query.execution import (
+    BulkPlan,
+    CompiledQuery,
+    ExecuteResult,
+    ExecutionControl,
+    StreamPolicy,
+)
+from pyoq.query.execution.aio import AsyncStreamResource
+from pyoq.query.execution.aio.pooling import AsyncConnectionLease
+from pyoq.query.execution.aio.streaming import AsyncRowStream
+from pyoq.query.postgres.aio.bulk import run_pipeline
+from pyoq.query.postgres.aio.commands import open_cursor, run_sql
+from pyoq.query.postgres.aio.operations import PostgresAsyncQueryOperations
+from pyoq.query.postgres.aio.pool import PostgresAsyncPool
+from pyoq.query.postgres.aio.transactions import PostgresAsyncTransaction
+from pyoq.query.postgres.bulk import pipeline_supported
+from pyoq.query.postgres.compiler import PostgresCompiler
+from pyoq.query.postgres.parameters import adapt_parameters
+from pyoq.query.postgres.transactions import TransactionMode
+from pyoq.query.select import Query
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+AsyncConnection = psycopg.AsyncConnection[tuple[object, ...]]
+AsyncCursor = psycopg.AsyncCursor[tuple[object, ...]]
+
+_STREAM_CURSOR_PREFIX = "pyoq_stream_"
+
+
+class PostgresAsyncExecutor(PostgresAsyncQueryOperations):
+    __slots__ = ("_pool", "_stream_sequence")
+
+    def __init__(
+        self,
+        pool: PostgresAsyncPool,
+        compiler: PostgresCompiler | None = None,
+    ) -> None:
+        self._pool = pool
+        self._stream_sequence = 0
+        super().__init__(compiler or PostgresCompiler())
+
+    @property
+    def pool(self) -> PostgresAsyncPool:
+        return self._pool
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: StreamPolicy | None = None,
+    ) -> AsyncRowStream[ResultRow]:
+        selected = policy or StreamPolicy()
+        statement = self.compiler.compile(query)
+        self._stream_sequence += 1
+        name = f"{_STREAM_CURSOR_PREFIX}{self._stream_sequence}"
+        return AsyncRowStream(
+            lambda: self._open_stream(statement, selected, name),
+            selected,
+            statement.row_types,
+        )
+
+    async def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        if not pipeline_supported(self.compiler.capabilities.pipeline):
+            return await super().execute_chunks(plan, control)
+        lease = self.pool.checkout()
+        async with lease as connection:
+            return await run_pipeline(
+                connection, plan.chunks, control, lease.invalidate
+            )
+
+    def transaction(
+        self,
+        mode: TransactionMode = TransactionMode.READ_COMMITTED,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> PostgresAsyncTransaction:
+        return PostgresAsyncTransaction(self.pool, self.compiler, mode, control)
+
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        lease = self.pool.checkout()
+        async with lease as connection:
+            return await run_sql(
+                connection,
+                statement.sql,
+                parameters,
+                operation,
+                control,
+                lease.invalidate,
+            )
+
+    async def _open_stream(
+        self,
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+        name: str,
+    ) -> AsyncStreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        lease = self.pool.checkout()
+        connection = await lease.__aenter__()
+        try:
+            await connection.execute(b"BEGIN")
+            session = await open_cursor(
+                connection,
+                statement.sql,
+                parameters,
+                policy.control,
+                lease.invalidate,
+                server_side_name=name,
+            )
+        except BaseException:
+            await _end_stream_transaction(connection, lease)
+            raise
+        return AsyncStreamResource(
+            session,
+            lambda: _end_stream_transaction(connection, lease),
+        )
+
+
+async def _end_stream_transaction(
+    connection: AsyncConnection,
+    lease: AsyncConnectionLease[AsyncConnection],
+) -> None:
+    try:
+        await connection.execute(b"ROLLBACK")
+    except psycopg.Error:
+        lease.invalidate()
+    finally:
+        await lease.__aexit__(None, None, None)
+
+
+__all__ = ("PostgresAsyncExecutor",)

pyoq/query/postgres/aio/factory.py

@@ -0,0 +1,26 @@
+"""Asynchronous PostgreSQL connection creation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import psycopg
+
+from pyoq.errors import ConnectionPoolError
+
+AsyncConnection = psycopg.AsyncConnection[tuple[object, ...]]
+
+
+@dataclass(frozen=True, slots=True)
+class PostgresAsyncConnectionFactory:
+    dsn: str
+
+    async def __call__(self) -> AsyncConnection:
+        try:
+            return await psycopg.AsyncConnection.connect(self.dsn, autocommit=True)
+        except psycopg.Error as error:
+            message = f"unable to open PostgreSQL connection: {error}"
+            raise ConnectionPoolError(message) from error
+
+
+__all__ = ("PostgresAsyncConnectionFactory",)

pyoq/query/postgres/aio/operations.py

@@ -0,0 +1,26 @@
+"""Asynchronous PostgreSQL typed result operations."""
+
+from __future__ import annotations
+
+from pyoq.query.execution import BulkPlan, WriteProvider
+from pyoq.query.execution.aio import AsyncQueryOperations
+from pyoq.query.postgres.compiler import PostgresCompiler
+from pyoq.query.postgres.planning import PostgresBulkPlanner
+from pyoq.query.write_nodes import WriteNode
+
+
+class PostgresAsyncQueryOperations(AsyncQueryOperations):
+    __slots__ = ("_compiler",)
+
+    def __init__(self, compiler: PostgresCompiler) -> None:
+        self._compiler = compiler
+
+    @property
+    def compiler(self) -> PostgresCompiler:
+        return self._compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return PostgresBulkPlanner(self.compiler).plan(statement)
+
+
+__all__ = ("PostgresAsyncQueryOperations",)

pyoq/query/postgres/aio/pool.py

@@ -0,0 +1,40 @@
+"""Exclusive asynchronous PostgreSQL connection pool."""
+
+from __future__ import annotations
+
+from typing import TypeAlias, cast
+
+import psycopg
+
+from pyoq.errors import ConnectionPoolError
+from pyoq.query.execution.aio.pooling import (
+    AsyncConnectionLease,
+    AsyncConnectionPool,
+)
+
+AsyncConnection = psycopg.AsyncConnection[tuple[object, ...]]
+PostgresAsyncConnectionLease: TypeAlias = AsyncConnectionLease[AsyncConnection]
+
+
+class PostgresAsyncPool(AsyncConnectionPool[AsyncConnection]):
+    __slots__ = ()
+
+    async def reset(self, connection: AsyncConnection) -> bool:
+        try:
+            await connection.rollback()
+            await connection.execute(b"SET statement_timeout = 0")
+        except psycopg.Error:
+            return False
+        return not connection.closed
+
+    def validate(self, value: object) -> AsyncConnection:
+        if not isinstance(value, psycopg.AsyncConnection):
+            message = "PostgreSQL connection factory returned an invalid connection"
+            raise ConnectionPoolError(message)
+        return cast("AsyncConnection", value)
+
+    async def dispose(self, connection: AsyncConnection) -> None:
+        await connection.close()
+
+
+__all__ = ("PostgresAsyncConnectionLease", "PostgresAsyncPool")