pyoq-sql 1.0.0__py3-none-any.whl

pyoq/query/postgres/aio/transactions.py

@@ -0,0 +1,295 @@
+"""Owned asynchronous PostgreSQL transactions and nested savepoints."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TypeVar, cast
+
+import psycopg
+
+from pyoq.errors import ExecutionError, QueryValidationError, TransactionStateError
+from pyoq.query.execution import (
+    BulkPlan,
+    CompiledQuery,
+    ExecuteResult,
+    ExecutionControl,
+    StreamPolicy,
+)
+from pyoq.query.execution.aio import (
+    AsyncRowStream,
+    AsyncStreamResource,
+    AsyncTransactionState,
+)
+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.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, ...]]
+_AsyncState = AsyncTransactionState[AsyncConnection, "PostgresAsyncTransaction"]
+
+_STREAM_CURSOR_PREFIX = "pyoq_stream_"
+
+
+class PostgresAsyncTransaction(PostgresAsyncQueryOperations):
+    __slots__ = (
+        "_control",
+        "_entered",
+        "_finished",
+        "_mode",
+        "_parent",
+        "_pool",
+        "_savepoint_name",
+        "_state",
+        "_stream_sequence",
+    )
+
+    def __init__(
+        self,
+        pool: PostgresAsyncPool,
+        compiler: PostgresCompiler,
+        mode: TransactionMode = TransactionMode.READ_COMMITTED,
+        control: ExecutionControl | None = None,
+        *,
+        parent: PostgresAsyncTransaction | None = None,
+        state: _AsyncState | None = None,
+    ) -> None:
+        super().__init__(compiler)
+        if not _is_transaction_mode(mode):
+            raise QueryValidationError("PostgreSQL transaction mode is invalid")
+        self._pool = pool
+        self._mode = mode
+        self._control = control
+        self._parent = parent
+        self._state = state or AsyncTransactionState()
+        self._savepoint_name: str | None = None
+        self._stream_sequence = 0
+        self._entered = False
+        self._finished = False
+
+    async def __aenter__(self) -> PostgresAsyncTransaction:
+        if self._entered or self._finished:
+            raise TransactionStateError("transaction cannot be entered twice")
+        if self._parent is None:
+            await self._enter_outer()
+        else:
+            await self._enter_savepoint()
+        self._entered = True
+        return self
+
+    async def __aexit__(
+        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:
+            await self._exit_outer(exception_type)
+        else:
+            await self._exit_savepoint(exception_type)
+
+    def savepoint(
+        self,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> PostgresAsyncTransaction:
+        self._state.require_current(self)
+        return PostgresAsyncTransaction(
+            self._pool,
+            self.compiler,
+            self._mode,
+            control or self._control,
+            parent=self,
+            state=self._state,
+        )
+
+    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)
+        connection = self._state.require_current(self)
+        return await run_pipeline(
+            connection,
+            plan.chunks,
+            control or self._control,
+            self._state.invalidate,
+        )
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: StreamPolicy | None = None,
+    ) -> AsyncRowStream[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: AsyncRowStream[ResultRow]
+        stream = AsyncRowStream(
+            lambda: self._open_stream(stream, statement, selected),
+            selected,
+            statement.row_types,
+        )
+        return stream
+
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        return await run_sql(
+            connection,
+            statement.sql,
+            parameters,
+            operation,
+            control or self._control,
+            self._state.invalidate,
+        )
+
+    async def _enter_outer(self) -> None:
+        await self._state.adopt(self._pool.checkout(), self)
+        try:
+            await self._run_transaction_sql(f"BEGIN {self._mode.value}")
+        except BaseException:
+            await self._release_outer()
+            raise
+
+    async def _enter_savepoint(self) -> None:
+        self._state.require_current(cast("PostgresAsyncTransaction", self._parent))
+        name = self._state.next_savepoint()
+        await self._run_transaction_sql(f'SAVEPOINT "{name}"', scope=self._parent)
+        self._savepoint_name = name
+        self._state.current_scope = self
+
+    async def _exit_outer(self, exception_type: type[BaseException] | None) -> None:
+        try:
+            self._state.require_scope(self)
+        except TransactionStateError:
+            await self._abort_misnested_outer(exception_type)
+            return
+        await self._state.close_stream()
+        command = "ROLLBACK" if exception_type is not None else "COMMIT"
+        try:
+            if not self._state.broken:
+                await self._run_transaction_sql(command)
+        except ExecutionError:
+            self._state.invalidate()
+            if exception_type is None:
+                raise
+        finally:
+            await self._release_outer()
+
+    async def _exit_savepoint(
+        self,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        self._state.require_scope(self)
+        await self._state.close_stream()
+        name = cast("str", self._savepoint_name)
+        try:
+            await 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
+
+    async def _finish_savepoint(
+        self,
+        name: str,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        if self._state.broken:
+            return
+        if exception_type is not None:
+            await self._run_transaction_sql(f'ROLLBACK TO SAVEPOINT "{name}"')
+        await self._run_transaction_sql(f'RELEASE SAVEPOINT "{name}"')
+
+    async def _run_transaction_sql(
+        self,
+        sql: str,
+        *,
+        scope: PostgresAsyncTransaction | None = None,
+    ) -> None:
+        connection = self._state.require_current(scope or self, allow_stream=True)
+        await run_sql(
+            connection,
+            sql,
+            (),
+            _discard_cursor,
+            self._control,
+            self._state.invalidate,
+        )
+
+    async def _open_stream(
+        self,
+        stream: AsyncRowStream[ResultRow],
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+    ) -> AsyncStreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        self._stream_sequence += 1
+        session = await open_cursor(
+            connection,
+            statement.sql,
+            parameters,
+            policy.control,
+            self._state.invalidate,
+            server_side_name=f"{_STREAM_CURSOR_PREFIX}{self._stream_sequence}",
+        )
+        self._state.active_stream = cast("AsyncRowStream[object]", stream)
+        return AsyncStreamResource(session, self._state.release_stream)
+
+    async def _release_outer(self) -> None:
+        self._finished = True
+        await self._state.release()
+
+    async def _abort_misnested_outer(
+        self,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        connection = self._state.require_owner()
+        await self._state.close_stream()
+        try:
+            await connection.rollback()
+        except psycopg.Error:
+            self._state.invalidate()
+        finally:
+            await self._release_outer()
+        if exception_type is None:
+            message = "transaction exited with an active nested savepoint"
+            raise TransactionStateError(message)
+
+
+async def _discard_cursor(cursor: AsyncCursor) -> None:
+    del cursor
+
+
+def _is_transaction_mode(value: object) -> bool:
+    return isinstance(value, TransactionMode)
+
+
+__all__ = ("PostgresAsyncTransaction",)

pyoq/query/postgres/bulk.py

@@ -0,0 +1,62 @@
+"""Pipelined 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.commands import PostgresExecutionMonitor
+from pyoq.query.postgres.parameters import adapt_parameters
+
+Connection = psycopg.Connection[tuple[object, ...]]
+Invalidator = Callable[[], None]
+
+
+def pipeline_supported(enabled: bool) -> bool:
+    return enabled and psycopg.capabilities.has_pipeline()
+
+
+def run_pipeline(
+    connection: Connection,
+    chunks: tuple[BulkChunk, ...],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> tuple[ExecuteResult, ...]:
+    monitor = PostgresExecutionMonitor(control)
+    cursors: list[psycopg.Cursor[tuple[object, ...]]] = []
+    try:
+        monitor.install(connection, invalidate)
+        with connection.pipeline():
+            for chunk in chunks:
+                monitor.checkpoint()
+                cursors.append(_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:
+            _close_quietly(cursor)
+        monitor.close(connection, invalidate)
+
+
+def _send(
+    connection: Connection,
+    chunk: BulkChunk,
+) -> psycopg.Cursor[tuple[object, ...]]:
+    cursor = connection.cursor()
+    parameters = adapt_parameters(chunk.statement.parameters)
+    cursor.execute(chunk.statement.sql.encode(), parameters or None)
+    return cursor
+
+
+def _close_quietly(cursor: psycopg.Cursor[tuple[object, ...]]) -> None:
+    try:
+        cursor.close()
+    except psycopg.Error:
+        return
+
+
+__all__ = ("pipeline_supported", "run_pipeline")

pyoq/query/postgres/commands.py

@@ -0,0 +1,238 @@
+"""PostgreSQL cursor execution, timeouts, and cancellation."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from threading import Event, Thread
+from time import monotonic
+from typing import NoReturn, TypeVar, cast
+
+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]
+Connection = psycopg.Connection[tuple[object, ...]]
+
+_WATCHDOG_INTERVAL = 0.025
+_MILLISECONDS = 1000
+
+
+class PostgresExecutionMonitor:
+    """Server-side statement timeout plus a cancellation watchdog."""
+
+    __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: Thread | None = None
+        self._timeout_applied = False
+
+    def install(self, connection: Connection, invalidate: Invalidator) -> None:
+        self.checkpoint()
+        self._apply_timeout(connection, invalidate)
+        token = self._control.cancellation_token
+        if token is None:
+            return
+        watchdog = Thread(
+            target=self._watch,
+            args=(connection,),
+            daemon=True,
+            name="pyoq-postgres-cancel",
+        )
+        self._watchdog = watchdog
+        watchdog.start()
+
+    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)
+
+    def close(self, connection: Connection, invalidate: Invalidator) -> None:
+        self._stop.set()
+        watchdog = self._watchdog
+        self._watchdog = None
+        if watchdog is not None:
+            watchdog.join()
+        self._clear_timeout(connection, invalidate)
+
+    def _watch(self, connection: Connection) -> None:
+        token = cast("object", self._control.cancellation_token)
+        while not self._stop.wait(_WATCHDOG_INTERVAL):
+            if getattr(token, "cancelled", False):
+                _cancel_quietly(connection)
+                return
+
+    def _apply_timeout(self, connection: Connection, invalidate: Invalidator) -> None:
+        timeout = self._control.timeout
+        if timeout is None:
+            return
+        milliseconds = max(int(timeout * _MILLISECONDS), 1)
+        try:
+            connection.execute(f"SET statement_timeout = {milliseconds}".encode())
+        except psycopg.Error as error:
+            self.translate_driver_error(error, invalidate)
+        self._timeout_applied = True
+
+    def _clear_timeout(self, connection: Connection, invalidate: Invalidator) -> None:
+        if not self._timeout_applied:
+            return
+        self._timeout_applied = False
+        try:
+            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 PostgresCursorSession:
+    __slots__ = ("_closed", "_connection", "_cursor", "_invalidate", "_monitor")
+
+    def __init__(
+        self,
+        connection: Connection,
+        cursor: psycopg.Cursor[tuple[object, ...]],
+        monitor: PostgresExecutionMonitor,
+        invalidate: Invalidator,
+    ) -> None:
+        self._connection = connection
+        self._cursor = cursor
+        self._monitor = monitor
+        self._invalidate = invalidate
+        self._closed = False
+
+    @property
+    def cursor(self) -> psycopg.Cursor[tuple[object, ...]]:
+        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)
+
+    def fetch_batch(self, size: int) -> list[tuple[object, ...]]:
+        try:
+            self._monitor.checkpoint()
+            rows = self._cursor.fetchmany(size)
+            self._monitor.checkpoint()
+        except psycopg.Error as error:
+            self._monitor.translate_driver_error(error, self._invalidate)
+        return rows
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        _close_cursor(self._cursor, self._invalidate)
+        self._monitor.close(self._connection, self._invalidate)
+
+
+def run_sql(
+    connection: Connection,
+    sql: str,
+    parameters: tuple[PostgresParameter, ...],
+    operation: Callable[[psycopg.Cursor[tuple[object, ...]]], OperationResult],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> OperationResult:
+    session = open_cursor(connection, sql, parameters, control, invalidate)
+    try:
+        result = operation(session.cursor)
+        session.checkpoint()
+        return result
+    except psycopg.Error as error:
+        session.translate_driver_error(error)
+    finally:
+        session.close()
+
+
+def open_cursor(
+    connection: Connection,
+    sql: str,
+    parameters: tuple[PostgresParameter, ...],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+    *,
+    server_side_name: str | None = None,
+) -> PostgresCursorSession:
+    monitor = PostgresExecutionMonitor(control)
+    cursor: psycopg.Cursor[tuple[object, ...]] | None = None
+    try:
+        monitor.install(connection, invalidate)
+        cursor = connection.cursor(name=server_side_name or "")
+        cursor.execute(sql.encode(), parameters or None)
+        monitor.checkpoint()
+        return PostgresCursorSession(connection, cursor, monitor, invalidate)
+    except psycopg.Error as error:
+        _release(cursor, monitor, connection, invalidate)
+        monitor.translate_driver_error(error, invalidate)
+    except BaseException:
+        _release(cursor, monitor, connection, invalidate)
+        raise
+
+
+def _release(
+    cursor: psycopg.Cursor[tuple[object, ...]] | None,
+    monitor: PostgresExecutionMonitor,
+    connection: Connection,
+    invalidate: Invalidator,
+) -> None:
+    if cursor is not None:
+        _close_cursor(cursor, invalidate)
+    monitor.close(connection, invalidate)
+
+
+def _close_cursor(
+    cursor: psycopg.Cursor[tuple[object, ...]],
+    invalidate: Invalidator,
+) -> None:
+    try:
+        cursor.close()
+    except psycopg.Error:
+        invalidate()
+
+
+def _cancel_quietly(connection: Connection) -> None:
+    try:
+        connection.cancel()
+    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__ = (
+    "PostgresCursorSession",
+    "PostgresExecutionMonitor",
+    "open_cursor",
+    "run_sql",
+)

pyoq/query/postgres/compiler.py

@@ -0,0 +1,114 @@
+"""PostgreSQL query and statement compiler."""
+
+from __future__ import annotations
+
+from pyoq.errors import UnsupportedQueryError
+from pyoq.query.execution import CompiledQuery, StatementKind
+from pyoq.query.postgres.context import PostgresCompilationContext
+from pyoq.query.postgres.expressions import PostgresExpressionCompiler
+from pyoq.query.postgres.identifiers import QUOTING
+from pyoq.query.postgres.model import PostgresCapabilities
+from pyoq.query.postgres.writes import PostgresWriteCompilation
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.rendering.queries import QueryRenderer
+from pyoq.query.select_nodes import JoinKind, NullOrder, SelectNode, SetNode
+from pyoq.query.statements import StatementNode, StatementProvider
+from pyoq.query.write_nodes import CallNode, DeleteNode, InsertNode, UpdateNode
+
+_STATEMENT_NODES = (
+    SelectNode,
+    SetNode,
+    InsertNode,
+    UpdateNode,
+    DeleteNode,
+    CallNode,
+)
+
+
+class PostgresCompiler:
+    __slots__ = ("_capabilities",)
+
+    def __init__(
+        self,
+        capabilities: PostgresCapabilities | None = None,
+    ) -> None:
+        self._capabilities = capabilities or PostgresCapabilities()
+
+    @property
+    def capabilities(self) -> PostgresCapabilities:
+        return self._capabilities
+
+    def compile(
+        self,
+        statement: StatementProvider | StatementNode,
+        /,
+    ) -> CompiledQuery:
+        node = statement if isinstance(statement, _STATEMENT_NODES) else statement.node
+        context = PostgresCompilationContext(self.capabilities)
+        if isinstance(node, CallNode):
+            sql = PostgresWriteCompilation(context).compile_call(node)
+            return context.result(sql, StatementKind.CALL, node)
+        if isinstance(node, (InsertNode, UpdateNode, DeleteNode)):
+            sql, statement_kind = PostgresWriteCompilation(context).compile(node)
+            return context.result(sql, statement_kind, node)
+        sql = _PostgresQueryCompilation(context).compile(node)
+        return context.result(sql, StatementKind.SELECT, node)
+
+
+class _PostgresQueryCompilation(QueryRenderer):
+    __slots__ = ("_context",)
+
+    def __init__(self, context: PostgresCompilationContext) -> None:
+        self._context = context
+        super().__init__(PostgresExpressionCompiler(context))
+
+    @property
+    def dialect(self) -> str:
+        return "PostgreSQL"
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    def require_join_capability(self, kind: JoinKind) -> None:
+        if kind is JoinKind.RIGHT and not self._context.capabilities.right_join:
+            message = "PostgreSQL RIGHT JOIN is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+        if kind is JoinKind.FULL and not self._context.capabilities.full_join:
+            message = "PostgreSQL FULL JOIN is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+    def null_ordered(
+        self,
+        expression: str,
+        direction: str,
+        nulls: NullOrder,
+    ) -> str:
+        if not self._context.capabilities.null_ordering:
+            message = "PostgreSQL null ordering is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+        return super().null_ordered(expression, direction, nulls)
+
+    def require_recursive_common_tables(self) -> None:
+        if not self._context.capabilities.recursive_common_tables:
+            message = "PostgreSQL recursive common tables are disabled"
+            raise UnsupportedQueryError(message)
+
+    def require_lateral_sources(self) -> None:
+        if not self._context.capabilities.lateral_sources:
+            message = "PostgreSQL has no lateral sources"
+            raise UnsupportedQueryError(message)
+
+    def append_pagination(
+        self,
+        fragments: list[str],
+        limit: int | None,
+        offset: int | None,
+    ) -> None:
+        if limit is not None:
+            fragments.append(f"LIMIT {self._context.bind_page_value(limit)}")
+        if offset is not None:
+            fragments.append(f"OFFSET {self._context.bind_page_value(offset)}")
+
+
+__all__ = ("PostgresCompiler",)