pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/postgres/parameters.py

@@ -0,0 +1,71 @@
+"""PostgreSQL driver parameter adaptation."""
+
+from __future__ import annotations
+
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from enum import Enum
+from typing import cast
+from uuid import UUID
+
+from psycopg.types.json import Jsonb
+
+from pyoq.errors import ParameterBindingError
+
+PostgresParameter = (
+    str
+    | int
+    | float
+    | bool
+    | bytes
+    | Decimal
+    | date
+    | datetime
+    | time
+    | timedelta
+    | UUID
+    | list[object]
+    | Jsonb
+    | None
+)
+
+_SCALAR_TYPES = (
+    str,
+    int,
+    float,
+    bool,
+    bytes,
+    Decimal,
+    datetime,
+    date,
+    time,
+    timedelta,
+    UUID,
+)
+
+
+def adapt_parameters(values: tuple[object, ...]) -> tuple[PostgresParameter, ...]:
+    return tuple(adapt_parameter(value) for value in values)
+
+
+def adapt_parameter(value: object) -> PostgresParameter:
+    """One value as its driver carries it.
+
+    An array is a tuple and JSON is a dict or a list, which is the mapping
+    generation writes and the one the expression families agree with. The two
+    are told apart here because the driver sends a sequence as an array and
+    has no way to guess that a list was meant as JSON.
+    """
+    if isinstance(value, Enum):
+        return adapt_parameter(value.value)
+    if value is None or isinstance(value, _SCALAR_TYPES):
+        return value
+    if isinstance(value, tuple):
+        return [adapt_parameter(item) for item in cast("tuple[object, ...]", value)]
+    if isinstance(value, (dict, list)):
+        return Jsonb(value)
+    message = f"unsupported PostgreSQL parameter type: {type(value).__name__}"
+    raise ParameterBindingError(message)
+
+
+__all__ = ("PostgresParameter", "adapt_parameter", "adapt_parameters")

pyoq/query/postgres/planning.py

@@ -0,0 +1,20 @@
+"""PostgreSQL bulk write planning."""
+
+from __future__ import annotations
+
+from pyoq.query.execution.planning import BulkPlanner
+from pyoq.query.postgres.compiler import PostgresCompiler
+
+
+class PostgresBulkPlanner(BulkPlanner):
+    __slots__ = ()
+
+    def __init__(self, compiler: PostgresCompiler) -> None:
+        super().__init__(
+            compiler,
+            compiler.capabilities.maximum_parameters,
+            "PostgreSQL",
+        )
+
+
+__all__ = ("PostgresBulkPlanner",)

pyoq/query/postgres/pool.py

@@ -0,0 +1,52 @@
+"""Exclusive synchronous PostgreSQL connection pool."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TypeAlias, cast
+
+import psycopg
+
+from pyoq.errors import ConnectionPoolError
+from pyoq.query.execution.pooling import (
+    ConnectionLease,
+    ConnectionPool,
+    PoolPolicy,
+    PoolStats,
+)
+
+Connection = psycopg.Connection[tuple[object, ...]]
+PostgresConnectionFactoryProtocol = Callable[[], Connection]
+PostgresConnectionLease: TypeAlias = ConnectionLease[Connection]
+PostgresPoolPolicy: TypeAlias = PoolPolicy
+PostgresPoolStats: TypeAlias = PoolStats
+
+
+class PostgresPool(ConnectionPool[Connection]):
+    __slots__ = ()
+
+    def reset(self, connection: Connection) -> bool:
+        try:
+            connection.rollback()
+            connection.execute(b"SET statement_timeout = 0")
+        except psycopg.Error:
+            return False
+        return not connection.closed
+
+    def validate(self, value: object) -> Connection:
+        if not isinstance(value, psycopg.Connection):
+            message = "PostgreSQL connection factory returned an invalid connection"
+            raise ConnectionPoolError(message)
+        return cast("Connection", value)
+
+    def dispose(self, connection: Connection) -> None:
+        connection.close()
+
+
+__all__ = (
+    "PostgresConnectionFactoryProtocol",
+    "PostgresConnectionLease",
+    "PostgresPool",
+    "PostgresPoolPolicy",
+    "PostgresPoolStats",
+)

pyoq/query/postgres/transactions.py

@@ -0,0 +1,295 @@
+"""Owned PostgreSQL transactions and nested savepoints."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+from typing import TypeVar, cast
+
+import psycopg
+
+from pyoq.errors import ExecutionError, QueryValidationError, TransactionStateError
+from pyoq.query.execution import (
+    BulkPlan,
+    CompiledQuery,
+    ExecuteResult,
+    ExecutionControl,
+    RowStream,
+    StreamPolicy,
+    StreamResource,
+    TransactionState,
+)
+from pyoq.query.postgres.bulk import pipeline_supported, run_pipeline
+from pyoq.query.postgres.commands import open_cursor, run_sql
+from pyoq.query.postgres.compiler import PostgresCompiler
+from pyoq.query.postgres.operations import PostgresQueryOperations
+from pyoq.query.postgres.parameters import adapt_parameters
+from pyoq.query.postgres.pool import PostgresPool
+from pyoq.query.select import Query
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+Connection = psycopg.Connection[tuple[object, ...]]
+_PostgresTransactionState = TransactionState[Connection, "PostgresTransaction"]
+
+_STREAM_CURSOR_PREFIX = "pyoq_stream_"
+
+
+class TransactionMode(Enum):
+    READ_COMMITTED = "ISOLATION LEVEL READ COMMITTED"
+    REPEATABLE_READ = "ISOLATION LEVEL REPEATABLE READ"
+    SERIALIZABLE = "ISOLATION LEVEL SERIALIZABLE"
+    READ_ONLY = "ISOLATION LEVEL READ COMMITTED READ ONLY"
+
+
+class PostgresTransaction(PostgresQueryOperations):
+    __slots__ = (
+        "_control",
+        "_entered",
+        "_finished",
+        "_mode",
+        "_parent",
+        "_pool",
+        "_savepoint_name",
+        "_state",
+        "_stream_sequence",
+    )
+
+    def __init__(
+        self,
+        pool: PostgresPool,
+        compiler: PostgresCompiler,
+        mode: TransactionMode = TransactionMode.READ_COMMITTED,
+        control: ExecutionControl | None = None,
+        *,
+        parent: PostgresTransaction | None = None,
+        state: _PostgresTransactionState | 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 TransactionState()
+        self._savepoint_name: str | None = None
+        self._stream_sequence = 0
+        self._entered = False
+        self._finished = False
+
+    def __enter__(self) -> PostgresTransaction:
+        if self._entered or self._finished:
+            raise TransactionStateError("transaction cannot be entered twice")
+        if self._parent is None:
+            self._enter_outer()
+        else:
+            self._enter_savepoint()
+        self._entered = True
+        return self
+
+    def __exit__(
+        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:
+            self._exit_outer(exception_type)
+        else:
+            self._exit_savepoint(exception_type)
+
+    def savepoint(
+        self,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> PostgresTransaction:
+        self._state.require_current(self)
+        return PostgresTransaction(
+            self._pool,
+            self.compiler,
+            self._mode,
+            control or self._control,
+            parent=self,
+            state=self._state,
+        )
+
+    def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        if not pipeline_supported(self.compiler.capabilities.pipeline):
+            return super().execute_chunks(plan, control)
+        connection = self._state.require_current(self)
+        return run_pipeline(
+            connection,
+            plan.chunks,
+            control or self._control,
+            self._state.invalidate,
+        )
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: StreamPolicy | None = None,
+    ) -> RowStream[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: RowStream[ResultRow]
+        stream = RowStream(
+            lambda: self._open_stream(stream, statement, selected),
+            selected,
+            statement.row_types,
+        )
+        return stream
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[psycopg.Cursor[tuple[object, ...]]], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        return run_sql(
+            connection,
+            statement.sql,
+            parameters,
+            operation,
+            control or self._control,
+            self._state.invalidate,
+        )
+
+    def _enter_outer(self) -> None:
+        self._state.adopt(self._pool.checkout(), self)
+        try:
+            self._run_transaction_sql(f"BEGIN {self._mode.value}")
+        except BaseException:
+            self._release_outer()
+            raise
+
+    def _enter_savepoint(self) -> None:
+        self._state.require_current(cast("PostgresTransaction", self._parent))
+        name = self._state.next_savepoint()
+        self._run_transaction_sql(f'SAVEPOINT "{name}"', scope=self._parent)
+        self._savepoint_name = name
+        self._state.current_scope = self
+
+    def _exit_outer(self, exception_type: type[BaseException] | None) -> None:
+        try:
+            self._state.require_scope(self)
+        except TransactionStateError:
+            self._abort_misnested_outer(exception_type)
+            return
+        self._state.close_stream()
+        command = "ROLLBACK" if exception_type is not None else "COMMIT"
+        try:
+            if not self._state.broken:
+                self._run_transaction_sql(command)
+        except ExecutionError:
+            self._state.invalidate()
+            if exception_type is None:
+                raise
+        finally:
+            self._release_outer()
+
+    def _exit_savepoint(self, exception_type: type[BaseException] | None) -> None:
+        self._state.require_scope(self)
+        self._state.close_stream()
+        name = cast("str", self._savepoint_name)
+        try:
+            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
+
+    def _finish_savepoint(
+        self,
+        name: str,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        if self._state.broken:
+            return
+        if exception_type is not None:
+            self._run_transaction_sql(f'ROLLBACK TO SAVEPOINT "{name}"')
+        self._run_transaction_sql(f'RELEASE SAVEPOINT "{name}"')
+
+    def _run_transaction_sql(
+        self,
+        sql: str,
+        *,
+        scope: PostgresTransaction | None = None,
+    ) -> None:
+        connection = self._state.require_current(scope or self, allow_stream=True)
+        run_sql(
+            connection,
+            sql,
+            (),
+            _discard_cursor,
+            self._control,
+            self._state.invalidate,
+        )
+
+    def _open_stream(
+        self,
+        stream: RowStream[ResultRow],
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+    ) -> StreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        connection = self._state.require_current(self)
+        self._stream_sequence += 1
+        session = 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("RowStream[object]", stream)
+        return StreamResource(session, self._state.release_stream)
+
+    def _release_outer(self) -> None:
+        self._finished = True
+        self._state.release()
+
+    def _abort_misnested_outer(
+        self,
+        exception_type: type[BaseException] | None,
+    ) -> None:
+        connection = self._state.require_owner()
+        self._state.close_stream()
+        try:
+            connection.rollback()
+        except psycopg.Error:
+            self._state.invalidate()
+        finally:
+            self._release_outer()
+        if exception_type is None:
+            message = "transaction exited with an active nested savepoint"
+            raise TransactionStateError(message)
+
+
+def _is_transaction_mode(value: object) -> bool:
+    return isinstance(value, TransactionMode)
+
+
+def _discard_cursor(cursor: psycopg.Cursor[tuple[object, ...]]) -> None:
+    del cursor
+
+
+__all__ = ("PostgresTransaction", "TransactionMode")

pyoq/query/postgres/writes.py

@@ -0,0 +1,37 @@
+"""PostgreSQL write statement compilation."""
+
+from __future__ import annotations
+
+from pyoq.errors import UnsupportedQueryError
+from pyoq.query.postgres.context import PostgresCompilationContext
+from pyoq.query.postgres.expressions import PostgresExpressionCompiler
+from pyoq.query.postgres.identifiers import QUOTING
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.rendering.writes import WriteRenderer
+
+
+class PostgresWriteCompilation(WriteRenderer):
+    __slots__ = ("_context",)
+
+    def __init__(self, context: PostgresCompilationContext) -> None:
+        self._context = context
+        super().__init__(PostgresExpressionCompiler(context))
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    def require_returning(self) -> None:
+        if not self._context.capabilities.returning:
+            message = "PostgreSQL RETURNING is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+    def require_conflict_resolution(self) -> None:
+        if not self._context.capabilities.upsert:
+            message = (
+                "PostgreSQL conflict resolution is disabled by compiler capabilities"
+            )
+            raise UnsupportedQueryError(message)
+
+
+__all__ = ("PostgresWriteCompilation",)

pyoq/query/projections.py

@@ -0,0 +1,105 @@
+"""Typed SELECT projections and ordering terms."""
+
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+from pyoq.query.expressions import Expression
+from pyoq.query.select_nodes import (
+    NullOrder,
+    OrderDirection,
+    OrderNode,
+    ProjectionNode,
+)
+
+Value = TypeVar("Value")
+
+
+class Projection(Generic[Value]):
+    __slots__ = ("_node",)
+    _node: ProjectionNode
+
+    def __init__(self, node: ProjectionNode) -> None:
+        object.__setattr__(self, "_node", node)
+
+    def __setattr__(self, name: str, value: object) -> None:
+        message = "projections are immutable"
+        raise AttributeError(message)
+
+    @property
+    def node(self) -> ProjectionNode:
+        return self._node
+
+
+class OrderTerm:
+    __slots__ = ("_node",)
+    _node: OrderNode
+
+    def __init__(self, node: OrderNode) -> None:
+        object.__setattr__(self, "_node", node)
+
+    def __setattr__(self, name: str, value: object) -> None:
+        message = "order terms are immutable"
+        raise AttributeError(message)
+
+    @property
+    def node(self) -> OrderNode:
+        return self._node
+
+    def nulls_first(self) -> OrderTerm:
+        return OrderTerm(
+            OrderNode(self.node.expression, self.node.direction, NullOrder.FIRST)
+        )
+
+    def nulls_last(self) -> OrderTerm:
+        return OrderTerm(
+            OrderNode(self.node.expression, self.node.direction, NullOrder.LAST)
+        )
+
+
+def alias(expression: Expression[Value], name: str, /) -> Projection[Value]:
+    _require_identifier(name, "projection alias")
+    return Projection(ProjectionNode(expression.node, name))
+
+
+def ascending(expression: Expression[Value], /) -> OrderTerm:
+    return OrderTerm(OrderNode(expression.node, OrderDirection.ASCENDING))
+
+
+def descending(expression: Expression[Value], /) -> OrderTerm:
+    return OrderTerm(OrderNode(expression.node, OrderDirection.DESCENDING))
+
+
+def projection_node(value: object) -> ProjectionNode:
+    if isinstance(value, Projection):
+        return value.node
+    if isinstance(value, Expression):
+        return ProjectionNode(value.node)
+    message = "SELECT projections must be typed expressions"
+    raise TypeError(message)
+
+
+def order_node(value: object) -> OrderNode:
+    if isinstance(value, OrderTerm):
+        return value.node
+    if isinstance(value, Expression):
+        return OrderNode(value.node, OrderDirection.ASCENDING)
+    message = "ORDER BY terms must be typed expressions or order terms"
+    raise TypeError(message)
+
+
+def _require_identifier(value: str, label: str) -> None:
+    if not value or "\x00" in value:
+        message = f"{label} cannot be empty or contain a null character"
+        raise ValueError(message)
+
+
+__all__ = (
+    "OrderTerm",
+    "Projection",
+    "alias",
+    "ascending",
+    "descending",
+    "order_node",
+    "projection_node",
+)

pyoq/query/raw.py

@@ -0,0 +1,90 @@
+"""Explicit typed raw expression templates."""
+
+from __future__ import annotations
+
+from string import Formatter
+from typing import TypeVar
+
+from pyoq.errors import RawExpressionError
+from pyoq.query.expressions import ComputedExpression, Condition
+from pyoq.query.fields import family_for_type
+from pyoq.query.nodes import (
+    NodeProvider,
+    RawNode,
+    RawPart,
+    ScalarFamily,
+)
+
+Value = TypeVar("Value")
+
+
+def raw(
+    value_type: type[Value],
+    template: str,
+    /,
+    **expressions: NodeProvider,
+) -> ComputedExpression[Value]:
+    parts = _template_parts(template, expressions)
+    family = family_for_type(value_type)
+    node = RawNode(parts, family)
+    return ComputedExpression(node, family)
+
+
+def raw_condition(template: str, /, **expressions: NodeProvider) -> Condition:
+    return Condition(
+        RawNode(_template_parts(template, expressions), ScalarFamily.BOOLEAN)
+    )
+
+
+def _template_parts(
+    template: str,
+    expressions: dict[str, NodeProvider],
+) -> tuple[RawPart, ...]:
+    if not template or "\x00" in template:
+        message = "raw expression template cannot be empty or contain a null character"
+        raise RawExpressionError(message)
+    try:
+        parsed = tuple(Formatter().parse(template))
+    except ValueError as error:
+        message = "raw expression template has invalid braces"
+        raise RawExpressionError(message) from error
+    parts: list[RawPart] = []
+    referenced: set[str] = set()
+    for literal, name, format_spec, conversion in parsed:
+        _append_literal(parts, literal)
+        if name is None:
+            continue
+        _require_plain_name(name, format_spec or "", conversion)
+        expression: object = expressions.get(name)
+        if not isinstance(expression, NodeProvider):
+            message = f"raw expression is missing expression argument: {name}"
+            raise RawExpressionError(message)
+        parts.append(expression.node)
+        referenced.add(name)
+    unused = sorted(expressions.keys() - referenced)
+    if unused:
+        message = f"raw expression has unused arguments: {', '.join(unused)}"
+        raise RawExpressionError(message)
+    return tuple(parts)
+
+
+def _append_literal(parts: list[RawPart], literal: str) -> None:
+    if not literal:
+        return
+    if parts and isinstance(parts[-1], str):
+        parts[-1] += literal
+        return
+    parts.append(literal)
+
+
+def _require_plain_name(
+    name: str,
+    format_spec: str,
+    conversion: str | None,
+) -> None:
+    if not name.isidentifier() or format_spec or conversion is not None:
+        message = "raw placeholders must be plain Python identifiers"
+        raise RawExpressionError(message)
+
+
+__all__ = ("NodeProvider", "raw", "raw_condition")