pyoq-sql 1.0.2__py3-none-any.whl

pyoq/policies/rewriting.py

@@ -0,0 +1,263 @@
+"""Holding a statement to its policies before any SQL exists.
+
+Every way a statement can reach a table is walked: the table it selects from,
+the tables it joins, the queries inside its subqueries and common tables, the
+halves of a set operation, the table a collection gathers, and the target of a
+write. A rule that only looked at the first of those would be walked around by
+any of the rest.
+
+Raw SQL is the one thing that cannot be walked, because there is no node to
+read. It is refused unless a policy set says otherwise, which is a decision a
+project makes once rather than a hole that is always open.
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+from pyoq.errors import PolicyViolationError
+from pyoq.policies.model import Policies, TableReference
+from pyoq.query.nodes import (
+    AggregateNode,
+    BinaryNode,
+    BoundValueNode,
+    CollectionNode,
+    ExtractNode,
+    FunctionNode,
+    RawNode,
+    UnaryNode,
+    VariadicNode,
+    VariadicOperator,
+)
+from pyoq.query.select_nodes import (
+    JoinNode,
+    SelectNode,
+    SetNode,
+    SubquerySourceNode,
+    TableSourceNode,
+)
+from pyoq.query.write_nodes import DeleteNode, InsertNode, UpdateNode
+
+if TYPE_CHECKING:
+    from pyoq.query.nodes import ExpressionNode
+    from pyoq.query.select_nodes import SourceNode
+    from pyoq.query.write_nodes import WriteTargetNode
+
+
+def govern(node: object, policies: Policies, /, *, raw_sql: bool = False) -> object:
+    """The statement a policy set allows, built from the one it was given."""
+    if not policies:
+        return node
+    return _Governing(policies, raw_sql=raw_sql).statement(node)
+
+
+class _Governing:
+    """One pass over one statement, carrying the policies it is held to."""
+
+    __slots__ = ("_policies", "_raw_sql")
+
+    def __init__(self, policies: Policies, *, raw_sql: bool) -> None:
+        self._policies = policies
+        self._raw_sql = raw_sql
+
+    def statement(self, node: object) -> object:
+        if isinstance(node, (SelectNode, SetNode)):
+            return self.query(node)
+        if isinstance(node, InsertNode):
+            return self._insert(node)
+        if isinstance(node, UpdateNode):
+            return self._scoped_write(node)
+        if isinstance(node, DeleteNode):
+            return self._scoped_write(node)
+        return node
+
+    def query(self, node: SelectNode | SetNode) -> SelectNode | SetNode:
+        if isinstance(node, SetNode):
+            return replace(
+                node,
+                left=self.query(node.left),
+                right=self.query(node.right),
+            )
+        return self._select(node)
+
+    def _select(self, node: SelectNode) -> SelectNode:
+        common = tuple(
+            replace(table, query=self.query(table.query))
+            for table in node.common_tables
+        )
+        source = self._source(node.source)
+        joins = tuple(
+            replace(join, source=self._joined(join.source)) for join in node.joins
+        )
+        tables = self._tables_of(node.source, node.joins)
+        condition = _all_of(
+            (
+                *(() if node.condition is None else (self.expression(node.condition),)),
+                *self._conditions(tables),
+            )
+        )
+        return replace(
+            node,
+            common_tables=common,
+            source=source,
+            joins=joins,
+            condition=condition,
+            projections=tuple(
+                replace(projection, expression=self.expression(projection.expression))
+                for projection in node.projections
+            ),
+        )
+
+    def _source(self, source: SourceNode | None) -> SourceNode | None:
+        if isinstance(source, TableSourceNode):
+            self._require_allowed(_reference_of(source))
+            return source
+        if isinstance(source, SubquerySourceNode):
+            return replace(source, query=self.query(source.query))
+        return source
+
+    def _joined(self, source: SourceNode) -> SourceNode:
+        """A join always has a source, so it always gets one back."""
+        governed = self._source(source)
+        return source if governed is None else governed
+
+    def _tables_of(
+        self,
+        source: SourceNode | None,
+        joins: tuple[JoinNode, ...],
+    ) -> tuple[TableReference, ...]:
+        """Only the tables named here; a subquery answers for its own."""
+        sources = (source, *(join.source for join in joins))
+        return tuple(
+            _reference_of(one) for one in sources if isinstance(one, TableSourceNode)
+        )
+
+    def _conditions(
+        self,
+        tables: tuple[TableReference, ...],
+    ) -> tuple[ExpressionNode, ...]:
+        return tuple(
+            condition
+            for table in tables
+            for condition in self._policies.conditions(table)
+        )
+
+    def _insert(self, node: InsertNode) -> InsertNode:
+        table = _target_reference(node.target)
+        self._require_allowed(table)
+        assignments = self._policies.assignments(table)
+        if not assignments:
+            return node
+        return _with_assignments(node, assignments)
+
+    def _scoped_write(self, node: UpdateNode | DeleteNode) -> UpdateNode | DeleteNode:
+        table = _target_reference(node.target)
+        self._require_allowed(table)
+        conditions = self._policies.conditions(table)
+        if not conditions:
+            return node
+        existing = () if node.condition is None else (self.expression(node.condition),)
+        governed = _all_of((*existing, *conditions))
+        # A statement that said it meant every row now means every row a policy
+        # admits, which is what a policy is for.
+        return replace(node, condition=governed, unconditional=False)
+
+    def expression(self, node: ExpressionNode) -> ExpressionNode:
+        """A collection reads a table of its own, and raw SQL reads anything."""
+        if isinstance(node, RawNode):
+            self._require_raw_sql()
+            return node
+        if isinstance(node, CollectionNode):
+            return self._collection(node)
+        return _rebuilt(node, self.expression)
+
+    def _collection(self, node: CollectionNode) -> CollectionNode:
+        table = TableReference(node.table, node.schema, node.catalog, node.alias)
+        self._require_allowed(table)
+        conditions = self._policies.conditions(table)
+        governed = _all_of((self.expression(node.condition), *conditions))
+        if governed is None:  # pragma: no cover - a collection always correlates
+            return node
+        return replace(node, condition=governed)
+
+    def _require_allowed(self, table: TableReference) -> None:
+        reason = self._policies.refusal(table)
+        if reason is not None:
+            raise PolicyViolationError(reason)
+
+    def _require_raw_sql(self) -> None:
+        if self._raw_sql:
+            return
+        message = (
+            "raw SQL cannot be held to a policy, because there is no statement "
+            "to read; allow it deliberately if this connection may run it"
+        )
+        raise PolicyViolationError(message)
+
+
+def _reference_of(source: TableSourceNode) -> TableReference:
+    return TableReference(source.name, source.schema, source.catalog, source.alias)
+
+
+def _target_reference(target: WriteTargetNode) -> TableReference:
+    return TableReference(target.name, target.schema, target.catalog)
+
+
+def _all_of(conditions: tuple[ExpressionNode, ...]) -> ExpressionNode | None:
+    if not conditions:
+        return None
+    if len(conditions) == 1:
+        return conditions[0]
+    return VariadicNode(VariadicOperator.AND, conditions)
+
+
+def _with_assignments(
+    node: InsertNode,
+    assignments: tuple[tuple[str, object], ...],
+) -> InsertNode:
+    """Set what a policy requires, replacing whatever a caller gave for it."""
+    from pyoq.policies.model import family_for_value
+
+    columns = list(node.columns)
+    rows = [list(row) for row in node.rows]
+    for name, value in assignments:
+        bound = BoundValueNode(value, family_for_value(value))
+        if name in columns:
+            index = columns.index(name)
+            for row in rows:
+                row[index] = bound
+            continue
+        columns.append(name)
+        for row in rows:
+            row.append(bound)
+    return replace(
+        node,
+        columns=tuple(columns),
+        rows=tuple(tuple(row) for row in rows),
+    )
+
+
+def _rebuilt(
+    node: ExpressionNode,
+    visit: object,
+) -> ExpressionNode:
+    """Every expression that holds another one, rebuilt around the visited."""
+    if not callable(visit):  # pragma: no cover - called with a method
+        return node
+    if isinstance(node, UnaryNode):
+        return replace(node, operand=visit(node.operand))
+    if isinstance(node, BinaryNode):
+        return replace(node, left=visit(node.left), right=visit(node.right))
+    if isinstance(node, VariadicNode):
+        return replace(node, operands=tuple(visit(one) for one in node.operands))
+    if isinstance(node, FunctionNode):
+        return replace(node, arguments=tuple(visit(one) for one in node.arguments))
+    if isinstance(node, AggregateNode):
+        return replace(node, arguments=tuple(visit(one) for one in node.arguments))
+    if isinstance(node, ExtractNode):
+        return replace(node, expression=visit(node.expression))
+    return node
+
+
+__all__ = ("govern",)

pyoq/py.typed

@@ -0,0 +1 @@
+

pyoq/query/__init__.py

@@ -0,0 +1,312 @@
+"""Public typed query construction interface."""
+
+from pyoq.query.aggregates import (
+    AggregateExpression,
+    average,
+    count,
+    json_agg,
+    json_object_agg,
+    maximum,
+    minimum,
+    sum_,
+)
+from pyoq.query.arrays import array_of
+from pyoq.query.casts import cast
+from pyoq.query.choices import (
+    CaseBranches,
+    CaseBuilder,
+    case,
+    coalesce,
+    greatest,
+    least,
+    nullif,
+)
+from pyoq.query.documents import json_array, json_object
+from pyoq.query.expressions import ComputedExpression, Condition, Expression
+from pyoq.query.fields import Field, field
+from pyoq.query.nodes import (
+    AggregateName,
+    AggregateNode,
+    BinaryNode,
+    BinaryOperator,
+    BoundValueNode,
+    CaseNode,
+    CastNode,
+    DatePart,
+    ExpressionNode,
+    ExtractNode,
+    FieldNode,
+    FunctionName,
+    FunctionNode,
+    RawNode,
+    RowValueNode,
+    ScalarFamily,
+    UnaryNode,
+    UnaryOperator,
+    VariadicNode,
+    VariadicOperator,
+    VendorFunctionNode,
+    WindowFunctionNode,
+    WindowName,
+    WindowNode,
+)
+from pyoq.query.projections import (
+    OrderTerm,
+    Projection,
+    alias,
+    ascending,
+    descending,
+)
+from pyoq.query.raw import NodeProvider, raw, raw_condition
+from pyoq.query.recursion import (
+    RecursiveColumn,
+    RecursiveTable,
+    column,
+    recursive_table,
+)
+from pyoq.query.routines import Call, call, routine_output
+from pyoq.query.rows import RowValue, row
+from pyoq.query.select import (
+    JoinStep,
+    Query,
+    SelectQuery,
+    SemiJoinStep,
+    SetQuery,
+    exists,
+    not_exists,
+    select,
+    select_from,
+)
+from pyoq.query.select_nodes import (
+    CommonTableNode,
+    CommonTableSourceNode,
+    JoinKind,
+    JoinNode,
+    NullOrder,
+    OrderDirection,
+    OrderNode,
+    ProjectionNode,
+    SelectNode,
+    SetNode,
+    SetOperator,
+    SubquerySourceNode,
+    TableSourceNode,
+)
+from pyoq.query.sources import (
+    ELEMENT_COLUMN,
+    CommonTable,
+    Source,
+    Subquery,
+    common_table,
+    json_elements,
+    table_source,
+)
+from pyoq.query.values import BoundValue, bind
+from pyoq.query.vendor import (
+    VendorFunction,
+    VendorPredicate,
+    vendor_function,
+    vendor_predicate,
+)
+from pyoq.query.windows import (
+    FrameEdge,
+    FrameExclusion,
+    Window,
+    WindowFunction,
+    WindowSpecification,
+    cume_dist,
+    current_row,
+    dense_rank,
+    first_value,
+    following,
+    lag,
+    last_value,
+    lead,
+    named_window,
+    nth_value,
+    ntile,
+    percent_rank,
+    preceding,
+    rank,
+    row_number,
+    unbounded_following,
+    unbounded_preceding,
+)
+from pyoq.query.write_nodes import (
+    AssignmentNode,
+    ConflictAction,
+    ConflictNode,
+    DeleteNode,
+    InsertNode,
+    UpdateNode,
+    WriteNode,
+    WriteTargetNode,
+)
+from pyoq.query.writes import (
+    ColumnBinding,
+    ColumnReference,
+    ConflictUpdate,
+    DeleteStatement,
+    InsertStatement,
+    MutationStatement,
+    ReturningStatement,
+    ScopedStatement,
+    UpdateStatement,
+    WriteColumn,
+    WriteStatement,
+    delete_from,
+    excluded,
+    insert_into,
+    update,
+)
+from pyoq.schema.models import JsonScalar, JsonValue
+
+__all__ = (
+    "ELEMENT_COLUMN",
+    "AggregateExpression",
+    "AggregateName",
+    "AggregateNode",
+    "AssignmentNode",
+    "BinaryNode",
+    "BinaryOperator",
+    "BoundValue",
+    "BoundValueNode",
+    "Call",
+    "CaseBranches",
+    "CaseBuilder",
+    "CaseNode",
+    "CastNode",
+    "ColumnBinding",
+    "ColumnReference",
+    "CommonTable",
+    "CommonTableNode",
+    "CommonTableSourceNode",
+    "ComputedExpression",
+    "Condition",
+    "ConflictAction",
+    "ConflictNode",
+    "ConflictUpdate",
+    "DatePart",
+    "DeleteNode",
+    "DeleteStatement",
+    "Expression",
+    "ExpressionNode",
+    "ExtractNode",
+    "Field",
+    "FieldNode",
+    "FrameEdge",
+    "FrameExclusion",
+    "FunctionName",
+    "FunctionNode",
+    "InsertNode",
+    "InsertStatement",
+    "JoinKind",
+    "JoinNode",
+    "JoinStep",
+    "JsonScalar",
+    "JsonValue",
+    "MutationStatement",
+    "NodeProvider",
+    "NullOrder",
+    "OrderDirection",
+    "OrderNode",
+    "OrderTerm",
+    "Projection",
+    "ProjectionNode",
+    "Query",
+    "RawNode",
+    "RecursiveColumn",
+    "RecursiveTable",
+    "ReturningStatement",
+    "RowValue",
+    "RowValueNode",
+    "ScalarFamily",
+    "ScopedStatement",
+    "SelectNode",
+    "SelectQuery",
+    "SemiJoinStep",
+    "SetNode",
+    "SetOperator",
+    "SetQuery",
+    "Source",
+    "Subquery",
+    "SubquerySourceNode",
+    "TableSourceNode",
+    "UnaryNode",
+    "UnaryOperator",
+    "UpdateNode",
+    "UpdateStatement",
+    "VariadicNode",
+    "VariadicOperator",
+    "VendorFunction",
+    "VendorFunctionNode",
+    "VendorPredicate",
+    "Window",
+    "WindowFunction",
+    "WindowFunctionNode",
+    "WindowName",
+    "WindowNode",
+    "WindowSpecification",
+    "WriteColumn",
+    "WriteNode",
+    "WriteStatement",
+    "WriteTargetNode",
+    "alias",
+    "array_of",
+    "ascending",
+    "average",
+    "bind",
+    "call",
+    "case",
+    "cast",
+    "coalesce",
+    "column",
+    "common_table",
+    "count",
+    "cume_dist",
+    "current_row",
+    "delete_from",
+    "dense_rank",
+    "descending",
+    "excluded",
+    "exists",
+    "field",
+    "first_value",
+    "following",
+    "greatest",
+    "insert_into",
+    "json_agg",
+    "json_array",
+    "json_elements",
+    "json_object",
+    "json_object_agg",
+    "lag",
+    "last_value",
+    "lead",
+    "least",
+    "maximum",
+    "minimum",
+    "named_window",
+    "not_exists",
+    "nth_value",
+    "ntile",
+    "nullif",
+    "percent_rank",
+    "preceding",
+    "rank",
+    "raw",
+    "raw_condition",
+    "recursive_table",
+    "routine_output",
+    "row",
+    "row_number",
+    "select",
+    "select_from",
+    "sum_",
+    "table_source",
+    "unbounded_following",
+    "unbounded_preceding",
+    "update",
+    "vendor_function",
+    "vendor_predicate",
+)