pyoq-sql 1.0.2__py3-none-any.whl

pyoq/policies/governed.py

@@ -0,0 +1,430 @@
+"""A database that holds every statement to its policies.
+
+A compiler is the one thing every statement built from nodes passes through, so
+that is where policies are applied. Nothing reaches a driver another way except
+SQL that was already compiled, which is the one case a node cannot govern, so
+that case is refused here rather than left open.
+
+Both are one object to install, because installing half of it would be a hole
+that looked like a policy.
+"""
+
+from __future__ import annotations
+
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    asynccontextmanager,
+    contextmanager,
+)
+from typing import TYPE_CHECKING, Protocol, TypeVar, cast, runtime_checkable
+
+from pyoq.errors import OperationUnavailableError, PolicyViolationError
+from pyoq.policies.model import Policies, Policy
+from pyoq.policies.rewriting import govern
+from pyoq.query.execution import (
+    BulkResult,
+    CompiledQuery,
+    QueryOperations,
+    RowStream,
+)
+from pyoq.query.execution.aio import AsyncQueryOperations, AsyncRowStream
+from pyoq.serving import SupportsAsyncTransaction, SupportsTransaction
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        AsyncGenerator,
+        Awaitable,
+        Callable,
+        Generator,
+        Iterable,
+    )
+
+    from pyoq.query.execution import (
+        BulkPlan,
+        DatabaseCursor,
+        ExecuteResult,
+        ExecutionControl,
+        WriteProvider,
+    )
+    from pyoq.query.execution.aio import AsyncDatabaseCursor
+    from pyoq.query.select import Query
+    from pyoq.query.select_nodes import QueryNode
+    from pyoq.query.statements import (
+        StatementCompiler,
+        StatementNode,
+        StatementProvider,
+    )
+    from pyoq.query.write_nodes import WriteNode
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+
+class GovernedCompiler:
+    """A compiler that only ever sees a statement its policies allow."""
+
+    __slots__ = ("_inner", "_policies", "_raw_sql")
+
+    def __init__(
+        self,
+        inner: StatementCompiler,
+        policies: Policies,
+        *,
+        raw_sql: bool = False,
+    ) -> None:
+        self._inner = inner
+        self._policies = policies
+        self._raw_sql = raw_sql
+
+    @property
+    def capabilities(self) -> object:
+        """Whatever the dialect's own compiler says it can do.
+
+        Not every compiler declares one, so this asks rather than assumes.
+        """
+        return getattr(self._inner, "capabilities", None)
+
+    def compile(
+        self,
+        statement: StatementProvider | StatementNode,
+        /,
+    ) -> CompiledQuery:
+        node = getattr(statement, "node", statement)
+        governed = govern(node, self._policies, raw_sql=self._raw_sql)
+        return self._inner.compile(cast("StatementNode", governed))
+
+
+@runtime_checkable
+class SupportsStreaming(Protocol):
+    """A database that can read a result a piece at a time."""
+
+    def stream(self, query: Query[ResultRow], /) -> RowStream[ResultRow]: ...
+
+
+@runtime_checkable
+class SupportsAsyncStreaming(Protocol):
+    """The same, for a database that is awaited."""
+
+    def stream(self, query: Query[ResultRow], /) -> AsyncRowStream[ResultRow]: ...
+
+
+class GovernedOperations(QueryOperations):
+    """A database whose statements are held to a policy set."""
+
+    __slots__ = ("_compiler", "_inner", "_policies", "_raw_sql")
+
+    def __init__(
+        self,
+        inner: QueryOperations,
+        policies: Policies,
+        *,
+        raw_sql: bool = False,
+    ) -> None:
+        self._inner = inner
+        self._policies = policies
+        self._raw_sql = raw_sql
+        self._compiler = GovernedCompiler(inner.compiler, policies, raw_sql=raw_sql)
+
+    @property
+    def policies(self) -> Policies:
+        return self._policies
+
+    @property
+    def released(self) -> QueryOperations:
+        """The database underneath, which nothing here holds to a policy."""
+        return self._inner
+
+    @property
+    def compiler(self) -> StatementCompiler:
+        return self._compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        """A bulk write is planned from a governed statement, not the given one."""
+        return self._inner.plan_bulk(self._governed_write(statement))
+
+    def last_inserted_id(self, cursor: DatabaseCursor) -> int | None:
+        return self._inner.last_inserted_id(cursor)
+
+    def _governed_write(
+        self,
+        statement: WriteProvider | WriteNode,
+    ) -> WriteProvider | WriteNode:
+        node = getattr(statement, "node", statement)
+        governed = govern(node, self._policies, raw_sql=self._raw_sql)
+        return cast("WriteNode", governed)
+
+    def execute(
+        self,
+        statement: CompiledQuery | object,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ExecuteResult:
+        _require_governable(statement, raw_sql=self._raw_sql)
+        return super().execute(statement, control=control)  # type: ignore[arg-type]
+
+    def execute_bulk(
+        self,
+        statement: WriteProvider | WriteNode,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> BulkResult:
+        """The chunks of a plan this database made from a governed statement.
+
+        They are compiled SQL by the time they run, so they are handed to the
+        database underneath rather than back through the refusal that exists for
+        SQL a caller compiled elsewhere.
+        """
+        plan = self.plan_bulk(statement)
+        return BulkResult(self._inner.execute_chunks(plan, control), plan.row_counts)
+
+    def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        """A plan handed in from elsewhere is SQL from elsewhere."""
+        _require_governable(_first_statement(plan), raw_sql=self._raw_sql)
+        return self._inner.execute_chunks(plan, control)
+
+    def transaction(self) -> AbstractContextManager[QueryOperations]:
+        """A transaction held to the same policies the database is.
+
+        Without this a caller would reach for the database underneath to open
+        one, and everything inside it would be ungoverned.
+        """
+        return _governed_transaction(
+            _transactional(self._inner), self._policies, raw_sql=self._raw_sql
+        )
+
+    def stream(self, query: Query[ResultRow], /) -> RowStream[ResultRow]:
+        """A streamed read, governed the way any other read is.
+
+        A dialect's own streaming settings belong to the database underneath,
+        which is where it was given them.
+        """
+        return _streaming(self._inner).stream(self._governed_query(query))
+
+    def _governed_query(self, query: Query[ResultRow]) -> Query[ResultRow]:
+        governed_node = govern(query.node, self._policies, raw_sql=self._raw_sql)
+        return _rebuilt_query(query, governed_node)
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[DatabaseCursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        return self._inner._run(statement, operation, control)
+
+
+class AsyncGovernedOperations(AsyncQueryOperations):
+    """The asynchronous counterpart, governed at the same point."""
+
+    __slots__ = ("_compiler", "_inner", "_policies", "_raw_sql")
+
+    def __init__(
+        self,
+        inner: AsyncQueryOperations,
+        policies: Policies,
+        *,
+        raw_sql: bool = False,
+    ) -> None:
+        self._inner = inner
+        self._policies = policies
+        self._raw_sql = raw_sql
+        self._compiler = GovernedCompiler(inner.compiler, policies, raw_sql=raw_sql)
+
+    @property
+    def policies(self) -> Policies:
+        return self._policies
+
+    @property
+    def released(self) -> AsyncQueryOperations:
+        """The database underneath, which nothing here holds to a policy."""
+        return self._inner
+
+    @property
+    def compiler(self) -> StatementCompiler:
+        return self._compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return self._inner.plan_bulk(self._governed_write(statement))
+
+    def last_inserted_id(self, cursor: AsyncDatabaseCursor) -> int | None:
+        return self._inner.last_inserted_id(cursor)
+
+    def _governed_write(
+        self,
+        statement: WriteProvider | WriteNode,
+    ) -> WriteProvider | WriteNode:
+        node = getattr(statement, "node", statement)
+        governed = govern(node, self._policies, raw_sql=self._raw_sql)
+        return cast("WriteNode", governed)
+
+    async def execute(
+        self,
+        statement: CompiledQuery | object,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> ExecuteResult:
+        _require_governable(statement, raw_sql=self._raw_sql)
+        return await super().execute(statement, control=control)  # type: ignore[arg-type]
+
+    async def execute_bulk(
+        self,
+        statement: WriteProvider | WriteNode,
+        /,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> BulkResult:
+        """The chunks of a plan this database made from a governed statement."""
+        plan = self.plan_bulk(statement)
+        return BulkResult(
+            await self._inner.execute_chunks(plan, control), plan.row_counts
+        )
+
+    async def execute_chunks(
+        self,
+        plan: BulkPlan,
+        control: ExecutionControl | None,
+    ) -> tuple[ExecuteResult, ...]:
+        """A plan handed in from elsewhere is SQL from elsewhere."""
+        _require_governable(_first_statement(plan), raw_sql=self._raw_sql)
+        return await self._inner.execute_chunks(plan, control)
+
+    def transaction(self) -> AbstractAsyncContextManager[AsyncQueryOperations]:
+        """A transaction held to the same policies the database is."""
+        return _async_governed_transaction(
+            _async_transactional(self._inner), self._policies, raw_sql=self._raw_sql
+        )
+
+    def stream(self, query: Query[ResultRow], /) -> AsyncRowStream[ResultRow]:
+        """A streamed read, governed the way any other read is."""
+        return _async_streaming(self._inner).stream(self._governed_query(query))
+
+    def _governed_query(self, query: Query[ResultRow]) -> Query[ResultRow]:
+        governed_node = govern(query.node, self._policies, raw_sql=self._raw_sql)
+        return _rebuilt_query(query, governed_node)
+
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncDatabaseCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        return await self._inner._run(statement, operation, control)
+
+
+def governed(
+    operations: QueryOperations,
+    policies: Iterable[Policy],
+    /,
+    *,
+    raw_sql: bool = False,
+) -> GovernedOperations:
+    """A database held to these policies."""
+    return GovernedOperations(operations, _policies_of(policies), raw_sql=raw_sql)
+
+
+def async_governed(
+    operations: AsyncQueryOperations,
+    policies: Iterable[Policy],
+    /,
+    *,
+    raw_sql: bool = False,
+) -> AsyncGovernedOperations:
+    return AsyncGovernedOperations(operations, _policies_of(policies), raw_sql=raw_sql)
+
+
+def _policies_of(policies: Iterable[object]) -> Policies:
+    from typing import cast
+
+    return Policies(tuple(cast("Policy", policy) for policy in policies))
+
+
+@contextmanager
+def _governed_transaction(
+    inner: SupportsTransaction,
+    policies: Policies,
+    *,
+    raw_sql: bool,
+) -> Generator[QueryOperations]:
+    with inner.transaction() as active:
+        yield GovernedOperations(active, policies, raw_sql=raw_sql)
+
+
+@asynccontextmanager
+async def _async_governed_transaction(
+    inner: SupportsAsyncTransaction,
+    policies: Policies,
+    *,
+    raw_sql: bool,
+) -> AsyncGenerator[AsyncQueryOperations]:
+    async with inner.transaction() as active:
+        yield AsyncGovernedOperations(active, policies, raw_sql=raw_sql)
+
+
+def _rebuilt_query(query: Query[ResultRow], node: object) -> Query[ResultRow]:
+    """The same kind of query, around the statement a policy allows."""
+    return type(query)(cast("QueryNode", node))
+
+
+def _transactional(inner: QueryOperations) -> SupportsTransaction:
+    if isinstance(inner, SupportsTransaction):
+        return inner
+    raise _cannot("open a transaction")
+
+
+def _async_transactional(inner: AsyncQueryOperations) -> SupportsAsyncTransaction:
+    if isinstance(inner, SupportsAsyncTransaction):
+        return inner
+    raise _cannot("open a transaction")
+
+
+def _streaming(inner: QueryOperations) -> SupportsStreaming:
+    if isinstance(inner, SupportsStreaming):
+        return inner
+    raise _cannot("stream a result")
+
+
+def _async_streaming(inner: AsyncQueryOperations) -> SupportsAsyncStreaming:
+    if isinstance(inner, SupportsAsyncStreaming):
+        return inner
+    raise _cannot("stream a result")
+
+
+def _cannot(what: str) -> OperationUnavailableError:
+    message = f"the database underneath cannot {what}"
+    return OperationUnavailableError(message)
+
+
+def _first_statement(plan: BulkPlan) -> object:
+    """What the plan would run, judged the way one statement is.
+
+    A plan always has a chunk, because one with none is refused when it is made.
+    """
+    return plan.chunks[0].statement
+
+
+def _require_governable(statement: object, *, raw_sql: bool) -> None:
+    """SQL that was already compiled has no statement left to hold."""
+    if not isinstance(statement, CompiledQuery) or raw_sql:
+        return
+    message = (
+        "already compiled 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)
+
+
+__all__ = (
+    "AsyncGovernedOperations",
+    "GovernedCompiler",
+    "GovernedOperations",
+    "async_governed",
+    "governed",
+)

pyoq/policies/model.py

@@ -0,0 +1,242 @@
+"""What a policy says about a table a statement touches.
+
+A policy answers three questions about one table: whether a statement may touch
+it at all, what must always be true of the rows it sees, and what must always be
+set on the rows it writes. Everything else follows from those three.
+
+Policies work on the statement's own nodes, before any SQL exists. A rule
+written against SQL text can be walked around with an alias, a subquery, or a
+common table; a rule written against the node cannot, because the node is what
+the alias, the subquery, and the common table are made of.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol
+
+from pyoq.query.fields import family_for_type
+from pyoq.query.nodes import (
+    BinaryNode,
+    BinaryOperator,
+    BoundValueNode,
+    FieldNode,
+    ScalarFamily,
+    UnaryNode,
+    UnaryOperator,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from pyoq.query.nodes import ExpressionNode
+
+
+@dataclass(frozen=True, slots=True)
+class TableReference:
+    """A table a statement touches, and the name its columns answer to.
+
+    A statement that gave the table an alias refers to its columns by that
+    alias, so a condition added for it has to as well.
+    """
+
+    name: str
+    schema: str | None = None
+    catalog: str | None = None
+    alias: str | None = None
+
+    @property
+    def qualifier(self) -> str:
+        """What a column of this table is written under."""
+        return self.alias or self.name
+
+    def column(self, name: str, /) -> FieldNode:
+        """A column of this table, written the way the statement writes it."""
+        return FieldNode(name, self.qualifier)
+
+
+class Policy(Protocol):
+    """What a policy is asked about every table a statement touches.
+
+    `refuses` says why the table may not be touched at all, or nothing when it
+    may. `restricts` says what must always be true of the rows a statement may
+    see. `assigns` says what must always be set on the rows it writes.
+    """
+
+    def refuses(self, table: TableReference, /) -> str | None: ...
+
+    def restricts(self, table: TableReference, /) -> ExpressionNode | None: ...
+
+    def assigns(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TenantScope:
+    """Every row read or written belongs to one tenant.
+
+    The column is required on writes as well as reads, so a row cannot be
+    created outside the scope that would then be unable to see it.
+    """
+
+    column: str
+    value: object
+    tables: frozenset[str] | None = None
+
+    def refuses(self, table: TableReference, /) -> str | None:
+        del table
+        return None
+
+    def restricts(self, table: TableReference, /) -> ExpressionNode | None:
+        if not self._governs(table):
+            return None
+        return BinaryNode(
+            BinaryOperator.EQUAL,
+            table.column(self.column),
+            BoundValueNode(self.value, family_for_value(self.value)),
+        )
+
+    def assigns(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]:
+        if not self._governs(table):
+            return ()
+        return ((self.column, self.value),)
+
+    def _governs(self, table: TableReference) -> bool:
+        return self.tables is None or table.name in self.tables
+
+
+@dataclass(frozen=True, slots=True)
+class SoftDelete:
+    """A row marked deleted is a row nothing reads."""
+
+    column: str
+    tables: frozenset[str] | None = None
+
+    def refuses(self, table: TableReference, /) -> str | None:
+        del table
+        return None
+
+    def restricts(self, table: TableReference, /) -> ExpressionNode | None:
+        if self.tables is not None and table.name not in self.tables:
+            return None
+        return UnaryNode(UnaryOperator.IS_NULL, table.column(self.column))
+
+    def assigns(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]:
+        del table
+        return ()
+
+
+@dataclass(frozen=True, slots=True)
+class AllowedTables:
+    """Only these tables may be touched at all."""
+
+    names: frozenset[str]
+
+    def refuses(self, table: TableReference, /) -> str | None:
+        if table.name in self.names:
+            return None
+        return (
+            f"the table {table.name!r} is not one this connection may touch; "
+            f"allowed: {', '.join(sorted(self.names)) or 'none'}"
+        )
+
+    def restricts(self, table: TableReference, /) -> ExpressionNode | None:
+        del table
+        return None
+
+    def assigns(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]:
+        del table
+        return ()
+
+
+@dataclass(frozen=True, slots=True)
+class RowConstraint:
+    """One table's rows are limited to those a condition admits."""
+
+    table: str
+    build: Callable[[TableReference], ExpressionNode]
+
+    def refuses(self, table: TableReference, /) -> str | None:
+        del table
+        return None
+
+    def restricts(self, table: TableReference, /) -> ExpressionNode | None:
+        return None if table.name != self.table else self.build(table)
+
+    def assigns(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]:
+        del table
+        return ()
+
+
+@dataclass(frozen=True, slots=True)
+class Policies:
+    """Every policy a statement is held to, asked as one."""
+
+    members: tuple[Policy, ...] = ()
+
+    @classmethod
+    def of(cls, policies: Iterable[Policy], /) -> Policies:
+        return cls(tuple(policies))
+
+    def refusal(self, table: TableReference, /) -> str | None:
+        for policy in self.members:
+            reason = policy.refuses(table)
+            if reason is not None:
+                return reason
+        return None
+
+    def conditions(self, table: TableReference, /) -> tuple[ExpressionNode, ...]:
+        return tuple(
+            condition
+            for condition in (policy.restricts(table) for policy in self.members)
+            if condition is not None
+        )
+
+    def assignments(
+        self,
+        table: TableReference,
+        /,
+    ) -> tuple[tuple[str, object], ...]:
+        return tuple(
+            assignment
+            for policy in self.members
+            for assignment in policy.assigns(table)
+        )
+
+    def __bool__(self) -> bool:
+        return bool(self.members)
+
+
+__all__ = (
+    "AllowedTables",
+    "Policies",
+    "Policy",
+    "RowConstraint",
+    "SoftDelete",
+    "TableReference",
+    "TenantScope",
+)
+
+
+def family_for_value(value: object, /) -> ScalarFamily:
+    """What a policy's own value is, so it binds like any other."""
+    return ScalarFamily.NULL if value is None else family_for_type(type(value))