pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/postgres/context.py

@@ -0,0 +1,20 @@
+"""PostgreSQL compilation context."""
+
+from __future__ import annotations
+
+from pyoq.query.execution import CompilationContext
+from pyoq.query.postgres.model import PostgresCapabilities
+
+
+class PostgresCompilationContext(CompilationContext):
+    __slots__ = ("capabilities",)
+
+    def __init__(self, capabilities: PostgresCapabilities) -> None:
+        super().__init__(
+            capabilities.maximum_parameters,
+            capabilities.parameter_style,
+        )
+        self.capabilities = capabilities
+
+
+__all__ = ("PostgresCompilationContext",)

pyoq/query/postgres/executor.py

@@ -0,0 +1,147 @@
+"""Synchronous typed PostgreSQL query execution."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TypeVar
+
+import psycopg
+
+from pyoq.query.execution import (
+    BulkPlan,
+    CompiledQuery,
+    ExecuteResult,
+    ExecutionControl,
+    RowStream,
+    StreamPolicy,
+    StreamResource,
+)
+from pyoq.query.execution.pooling import ConnectionLease
+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.postgres.transactions import PostgresTransaction, TransactionMode
+from pyoq.query.select import Query
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+Connection = psycopg.Connection[tuple[object, ...]]
+
+_STREAM_CURSOR_PREFIX = "pyoq_stream_"
+
+
+class PostgresExecutor(PostgresQueryOperations):
+    __slots__ = ("_pool", "_stream_sequence")
+
+    def __init__(
+        self,
+        pool: PostgresPool,
+        compiler: PostgresCompiler | None = None,
+    ) -> None:
+        self._pool = pool
+        self._stream_sequence = 0
+        super().__init__(compiler or PostgresCompiler())
+
+    @property
+    def pool(self) -> PostgresPool:
+        return self._pool
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: StreamPolicy | None = None,
+    ) -> RowStream[ResultRow]:
+        selected = policy or StreamPolicy()
+        statement = self.compiler.compile(query)
+        self._stream_sequence += 1
+        name = f"{_STREAM_CURSOR_PREFIX}{self._stream_sequence}"
+        return RowStream(
+            lambda: self._open_stream(statement, selected, name),
+            selected,
+            statement.row_types,
+        )
+
+    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)
+        lease = self.pool.checkout()
+        with lease as connection:
+            return run_pipeline(connection, plan.chunks, control, lease.invalidate)
+
+    def transaction(
+        self,
+        mode: TransactionMode = TransactionMode.READ_COMMITTED,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> PostgresTransaction:
+        return PostgresTransaction(self.pool, self.compiler, mode, control)
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[psycopg.Cursor[tuple[object, ...]]], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        lease = self.pool.checkout()
+        with lease as connection:
+            return run_sql(
+                connection,
+                statement.sql,
+                parameters,
+                operation,
+                control,
+                lease.invalidate,
+            )
+
+    def _open_stream(
+        self,
+        statement: CompiledQuery,
+        policy: StreamPolicy,
+        name: str,
+    ) -> StreamResource:
+        parameters = adapt_parameters(statement.parameters)
+        lease = self.pool.checkout()
+        connection = lease.__enter__()
+        try:
+            connection.execute(b"BEGIN")
+            session = open_cursor(
+                connection,
+                statement.sql,
+                parameters,
+                policy.control,
+                lease.invalidate,
+                server_side_name=name,
+            )
+        except BaseException:
+            _end_stream_transaction(connection, lease)
+            raise
+        return StreamResource(
+            session,
+            lambda: _end_stream_transaction(connection, lease),
+        )
+
+
+def _end_stream_transaction(
+    connection: Connection,
+    lease: ConnectionLease[Connection],
+) -> None:
+    try:
+        connection.execute(b"ROLLBACK")
+    except psycopg.Error:
+        lease.invalidate()
+    finally:
+        lease.__exit__(None, None, None)
+
+
+__all__ = ("PostgresExecutor",)

pyoq/query/postgres/expressions.py

@@ -0,0 +1,311 @@
+"""PostgreSQL scalar expression compilation."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+
+from pyoq.config import DatabaseDialect
+from pyoq.errors import UnsupportedQueryError
+from pyoq.query.nodes import (
+    ArrayConstructNode,
+    ArrayDimensionsNode,
+    ArrayElementNode,
+    ArrayLengthNode,
+    ArrayNode,
+    ArrayOperation,
+    BinaryNode,
+    BinaryOperator,
+    BoundValueNode,
+    DatePart,
+    ExpressionNode,
+    ExtractNode,
+    FieldNode,
+    FunctionName,
+    FunctionNode,
+    JsonBuild,
+    JsonBuildNode,
+    JsonContainsNode,
+    JsonKeysNode,
+    JsonNode,
+    JsonOperation,
+    JsonWrite,
+    JsonWriteNode,
+    ScalarFamily,
+)
+from pyoq.query.postgres.context import PostgresCompilationContext
+from pyoq.query.postgres.identifiers import QUOTING
+from pyoq.query.rendering.expressions import (
+    CollectionFunctions,
+    ExpressionRenderer,
+)
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.results import expression_family
+from pyoq.query.select_nodes import ELEMENT_COLUMN, JsonElementsSourceNode
+
+_POSTGRES_BINARY_SQL = {
+    BinaryOperator.IS_DISTINCT_FROM: "IS DISTINCT FROM",
+    BinaryOperator.IS_NOT_DISTINCT_FROM: "IS NOT DISTINCT FROM",
+}
+
+_DATE_FIELD = {
+    DatePart.YEAR: "YEAR",
+    DatePart.MONTH: "MONTH",
+    DatePart.DAY: "DAY",
+    DatePart.HOUR: "HOUR",
+    DatePart.MINUTE: "MINUTE",
+    DatePart.SECOND: "SECOND",
+}
+
+
+_ARRAY_FUNCTIONS: Mapping[ArrayOperation, str] = {
+    ArrayOperation.APPEND: "ARRAY_APPEND",
+    ArrayOperation.PREPEND: "ARRAY_PREPEND",
+    ArrayOperation.CONCAT: "ARRAY_CAT",
+    ArrayOperation.REMOVE: "ARRAY_REMOVE",
+    ArrayOperation.REPLACE: "ARRAY_REPLACE",
+}
+"""What PostgreSQL calls each way of making one array out of another."""
+
+_ARRAY_OPERATORS: Mapping[ArrayOperation, str] = {
+    ArrayOperation.CONTAINS: "@>",
+    ArrayOperation.CONTAINED_BY: "<@",
+    ArrayOperation.OVERLAPS: "&&",
+}
+"""Verified against a running server, which is where these came from."""
+
+
+_NO_MERGE_PATCH = (
+    "PostgreSQL has no merge-patch; its concatenation replaces a member whole "
+    "and keeps a null rather than removing one, which json_concat says plainly"
+)
+"""Refused rather than emulated, because one level is not all the way down."""
+
+
+class PostgresExpressionCompiler(ExpressionRenderer[PostgresCompilationContext]):
+    __slots__ = ()
+
+    @property
+    def collection_functions(self) -> CollectionFunctions:
+        return CollectionFunctions(
+            "jsonb_agg",
+            "jsonb_build_object",
+            "'[]'::jsonb",
+            "jsonb_build_array",
+            "jsonb_object_agg",
+            ordered=True,
+        )
+
+    @property
+    def dialect(self) -> str:
+        return "PostgreSQL"
+
+    def qualified_field(self, node: FieldNode) -> str:
+        return QUOTING.qualified(
+            node.name,
+            table=node.table,
+            schema=node.schema,
+            catalog=node.catalog,
+        )
+
+    def binary_operator_sql(self, operator: BinaryOperator) -> str | None:
+        override = _POSTGRES_BINARY_SQL.get(operator)
+        if override is not None:
+            return override
+        return super().binary_operator_sql(operator)
+
+    def validate_binary(self, node: BinaryNode) -> None:
+        if node.operator is not BinaryOperator.ADD:
+            return
+        families = (
+            expression_family(node.left),
+            expression_family(node.right),
+        )
+        if all(family is ScalarFamily.TEMPORAL for family in families):
+            message = "PostgreSQL cannot add two temporal values"
+            raise UnsupportedQueryError(message)
+
+    def render_array(self, node: ArrayNode) -> str:
+        subject = self.compile(node.expression)
+        written_operands = [self.compile(item) for item in node.operands]
+        operand = written_operands[0]
+        if node.operation is ArrayOperation.HAS:
+            return f"({operand} = ANY({subject}))"
+        if node.operation is ArrayOperation.LACKS:
+            return f"({operand} <> ALL({subject}))"
+        if node.operation in _ARRAY_FUNCTIONS:
+            written = _ARRAY_FUNCTIONS[node.operation]
+            if node.operation is ArrayOperation.PREPEND:
+                return f"{written}({operand}, {subject})"
+            values = ", ".join(written_operands)
+            return f"{written}({subject}, {values})"
+        return f"({subject} {_ARRAY_OPERATORS[node.operation]} {operand})"
+
+    def render_array_element(self, node: ArrayElementNode) -> str:
+        """PostgreSQL counts an array from one, and a tuple counts from zero.
+
+        The array is bracketed because PostgreSQL will not subscript the
+        answer of a function without it, and the array here is often one.
+        Every dimension is subscripted at once, because PostgreSQL reads one
+        subscript of an array of arrays as nothing at all.
+        """
+        subject = self.compile(node.expression)
+        positions = "".join(f"[{self.compile(index)} + 1]" for index in node.indexes)
+        return f"({subject}){positions}"
+
+    def render_array_length(self, node: ArrayLengthNode) -> str:
+        subject = self.compile(node.expression)
+        if node.dimension is None:
+            return f"CARDINALITY({subject})"
+        return f"ARRAY_LENGTH({subject}, {node.dimension})"
+
+    def render_array_construct(self, node: ArrayConstructNode) -> str:
+        values = ", ".join(self.compile(element) for element in node.elements)
+        return f"ARRAY[{values}]"
+
+    def render_array_dimensions(self, node: ArrayDimensionsNode) -> str:
+        return f"ARRAY_NDIMS({self.compile(node.expression)})"
+
+    def render_json_elements(self, node: JsonElementsSourceNode) -> str:
+        """PostgreSQL expands an array into rows and lets the column be named."""
+        subject = self.compile(node.expression)
+        alias = QUOTING.quote(node.alias)
+        return f"JSONB_ARRAY_ELEMENTS({subject}) AS {alias}({ELEMENT_COLUMN})"
+
+    def render_json_keys(self, node: JsonKeysNode) -> str:
+        """PostgreSQL gives back one row per key, and nothing for an empty object."""
+        subject = self.compile(node.expression)
+        return f"JSONB_PATH_QUERY_ARRAY({subject}, '$.keyvalue().key')"
+
+    def render_json(self, node: JsonNode) -> str:
+        """PostgreSQL takes a path as steps, and reads a written one as a key.
+
+        `#>` and `#>>` walk a list of step names, which is why the steps are
+        bound as an array rather than joined into the `$.a.b` text the other
+        two dialects want. Given that text, PostgreSQL would look for a member
+        actually called `$.a.b` and answer null without complaining.
+        """
+        subject = self.compile(node.expression)
+        if node.operation is JsonOperation.REMOVE:
+            return f"({subject} #- {self._json_steps(node.path)})"
+        if node.operation is JsonOperation.EXISTS:
+            return f"({subject} #> {self._json_steps(node.path)} IS NOT NULL)"
+        if node.operation is JsonOperation.LENGTH:
+            return f"JSONB_ARRAY_LENGTH({subject} #> {self._json_steps(node.path)})"
+        operator = "#>>" if node.operation is JsonOperation.TEXT else "#>"
+        return f"({subject} {operator} {self._json_steps(node.path)})"
+
+    def _json_steps(self, path: tuple[str | int, ...]) -> str:
+        steps = tuple(str(step) for step in path)
+        return self._context.bind(BoundValueNode(steps, ScalarFamily.OTHER))
+
+    def render_json_write(self, node: JsonWriteNode) -> str:
+        """Compiled in the order it is written, because binding follows that."""
+        subject = self.compile(node.expression)
+        if node.operation is JsonWrite.MERGE:
+            raise UnsupportedQueryError(_NO_MERGE_PATCH)
+        if node.operation is JsonWrite.CONCAT:
+            return f"({subject} || {self._as_json(node.value)})"
+        if node.operation is JsonWrite.INSERT:
+            return self._json_inserted(node, subject)
+        steps = self._json_steps(node.path)
+        value = self._as_json(node.value)
+        created = "false" if node.operation is JsonWrite.REPLACE else "true"
+        return f"JSONB_SET({subject}, {steps}, {value}, {created})"
+
+    def _json_inserted(self, node: JsonWriteNode, subject: str) -> str:
+        """Written only where nothing is, the way the other dialects insert.
+
+        PostgreSQL's own insert raises where a member already exists, while
+        SQLite and MySQL leave it. Asking first is what makes one word mean
+        one thing on all three.
+        """
+        present = self._json_steps(node.path)
+        steps = self._json_steps(node.path)
+        value = self._as_json(node.value)
+        return (
+            f"CASE WHEN ({subject} #> {present}) IS NULL "
+            f"THEN JSONB_SET({subject}, {steps}, {value}, true) "
+            f"ELSE {subject} END"
+        )
+
+    def _as_json(self, node: ExpressionNode) -> str:
+        """A value as JSON, whatever it arrived as.
+
+        A cast will not turn a bound integer into JSON, because the driver
+        sends it as a smallint and PostgreSQL has no cast from one to the
+        other, and a bare parameter has no type to infer from. A value is
+        therefore written as the JSON text of itself and cast from there,
+        which every JSON value has.
+        """
+        if not isinstance(node, BoundValueNode):
+            return f"({self.compile(node)})::jsonb"
+        written = json.dumps(node.value, ensure_ascii=False, separators=(",", ":"))
+        placeholder = self._context.bind(
+            BoundValueNode(written, ScalarFamily.STRING, node.sensitive)
+        )
+        return f"{placeholder}::jsonb"
+
+    def render_json_build(self, node: JsonBuildNode) -> str:
+        """PostgreSQL infers nothing from a bare parameter, so each one says.
+
+        A builder takes any type at all, which leaves a bound value with
+        nothing to be inferred from. Members are written as text and values
+        as the JSON of themselves, which is what they are being built into.
+        """
+        functions = self.collection_functions
+        if node.kind is JsonBuild.ARRAY:
+            values = ", ".join(self._as_json(item) for item in node.arguments)
+            return f"{functions.array_constructor}({values})"
+        members = ", ".join(
+            self._as_text(item) if position % 2 == 0 else self._as_json(item)
+            for position, item in enumerate(node.arguments)
+        )
+        return f"{functions.object_function}({members})"
+
+    def _as_text(self, node: ExpressionNode) -> str:
+        """A member name, which PostgreSQL takes as text and nothing else."""
+        return f"({self.compile(node)})::text"
+
+    def render_json_contains(self, node: JsonContainsNode) -> str:
+        subject = self.compile(node.expression)
+        return f"({subject} @> {self.compile(node.value)}::jsonb)"
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    @property
+    def database_dialect(self) -> DatabaseDialect:
+        return DatabaseDialect.POSTGRES
+
+    def render_extract(self, node: ExtractNode) -> str:
+        field = _DATE_FIELD.get(node.part)
+        if field is None:
+            message = f"unsupported PostgreSQL date part: {node.part}"
+            raise UnsupportedQueryError(message)
+        expression = self.compile(node.expression)
+        return f"CAST(EXTRACT({field} FROM {expression}) AS INTEGER)"
+
+    def render_string_predicate(self, node: FunctionNode) -> str:
+        left_node, right_node = node.arguments
+        left = self.compile(left_node)
+        right = self.compile(right_node)
+        if node.name is FunctionName.CONTAINS:
+            return f"(POSITION({right} IN {left}) > 0)"
+        if node.name is FunctionName.STARTS_WITH:
+            comparison = self.compile(right_node)
+            return f"(LEFT({left}, LENGTH({right})) = {comparison})"
+        if node.name is FunctionName.ENDS_WITH:
+            comparison = self.compile(right_node)
+            return f"(RIGHT({left}, LENGTH({right})) = {comparison})"
+        message = f"unsupported PostgreSQL function: {node.name}"
+        raise UnsupportedQueryError(message)
+
+    def require_raw_expressions(self) -> None:
+        if not self._context.capabilities.raw_expressions:
+            message = "PostgreSQL raw expressions are disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+
+__all__ = ("PostgresExpressionCompiler",)

pyoq/query/postgres/factory.py

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

pyoq/query/postgres/health.py

@@ -0,0 +1,24 @@
+"""PostgreSQL connection failure classification."""
+
+from __future__ import annotations
+
+import psycopg
+
+_FATAL_STATE_CLASSES = frozenset({"08", "53", "57P", "58", "XX"})
+_CANCELLED_STATE = "57014"
+
+
+def is_connection_failure(error: psycopg.Error) -> bool:
+    if isinstance(error, (psycopg.OperationalError, psycopg.InterfaceError)):
+        return not is_cancellation(error)
+    state = error.sqlstate
+    if state is None:
+        return False
+    return any(state.startswith(prefix) for prefix in _FATAL_STATE_CLASSES)
+
+
+def is_cancellation(error: psycopg.Error) -> bool:
+    return error.sqlstate == _CANCELLED_STATE
+
+
+__all__ = ("is_cancellation", "is_connection_failure")

pyoq/query/postgres/identifiers.py

@@ -0,0 +1,9 @@
+"""PostgreSQL identifier rendering."""
+
+from __future__ import annotations
+
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+
+QUOTING = IdentifierQuoting("PostgreSQL")
+
+__all__ = ("QUOTING",)

pyoq/query/postgres/model.py

@@ -0,0 +1,81 @@
+"""PostgreSQL compilation capabilities."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pyoq.errors import QueryValidationError
+from pyoq.query.execution import ParameterStyle
+from pyoq.relations.planning import FetchCapabilities
+
+_EXTENDED_PROTOCOL_PARAMETERS = 65535
+
+
+@dataclass(frozen=True, slots=True)
+class PostgresCapabilities:
+    maximum_parameters: int = _EXTENDED_PROTOCOL_PARAMETERS
+    parameter_style: ParameterStyle = ParameterStyle.FORMAT
+    right_join: bool = True
+    full_join: bool = True
+    null_ordering: bool = True
+    lateral_sources: bool = True
+    recursive_common_tables: bool = True
+    raw_expressions: bool = True
+    nested_collections: bool = True
+    ordered_nested_collections: bool = True
+    returning: bool = True
+    upsert: bool = True
+    pipeline: bool = True
+
+    @property
+    def fetch(self) -> FetchCapabilities:
+        """What this dialect can do when fetching a collection in one query."""
+        return FetchCapabilities(
+            self.nested_collections,
+            self.ordered_nested_collections,
+        )
+
+    def __post_init__(self) -> None:
+        if not _is_positive_integer(self.maximum_parameters):
+            message = "PostgreSQL maximum parameters must be a positive integer"
+            raise QueryValidationError(message)
+        if self.maximum_parameters > _EXTENDED_PROTOCOL_PARAMETERS:
+            message = (
+                "PostgreSQL cannot bind more than "
+                f"{_EXTENDED_PROTOCOL_PARAMETERS} parameters in one statement"
+            )
+            raise QueryValidationError(message)
+        if not _is_parameter_style(self.parameter_style):
+            message = "PostgreSQL parameter style must be a ParameterStyle value"
+            raise QueryValidationError(message)
+        feature_flags: tuple[object, ...] = (
+            self.right_join,
+            self.full_join,
+            self.null_ordering,
+            self.lateral_sources,
+            self.recursive_common_tables,
+            self.raw_expressions,
+            self.nested_collections,
+            self.ordered_nested_collections,
+            self.returning,
+            self.upsert,
+            self.pipeline,
+        )
+        if not _all_boolean(feature_flags):
+            message = "PostgreSQL capability flags must be boolean"
+            raise QueryValidationError(message)
+
+
+def _is_parameter_style(value: object) -> bool:
+    return isinstance(value, ParameterStyle)
+
+
+def _is_positive_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+def _all_boolean(values: tuple[object, ...]) -> bool:
+    return all(isinstance(value, bool) for value in values)
+
+
+__all__ = ("PostgresCapabilities",)

pyoq/query/postgres/operations.py

@@ -0,0 +1,25 @@
+"""PostgreSQL typed result operations."""
+
+from __future__ import annotations
+
+from pyoq.query.execution import BulkPlan, QueryOperations, WriteProvider
+from pyoq.query.postgres.compiler import PostgresCompiler
+from pyoq.query.postgres.planning import PostgresBulkPlanner
+from pyoq.query.write_nodes import WriteNode
+
+
+class PostgresQueryOperations(QueryOperations):
+    __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__ = ("PostgresQueryOperations",)