pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/writes.py

@@ -0,0 +1,628 @@
+"""Typed immutable write statement construction."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, replace
+from typing import (
+    Generic,
+    Protocol,
+    Self,
+    TypeVar,
+    TypeVarTuple,
+    overload,
+    runtime_checkable,
+)
+
+from pyoq.errors import QueryStateError, QueryValidationError
+from pyoq.query.expressions import Expression, require_boolean
+from pyoq.query.nodes import ExpressionNode, FieldNode, ScalarFamily
+from pyoq.query.projections import Projection, projection_node
+from pyoq.query.select_nodes import ProjectionNode
+from pyoq.query.sources import TableLike, write_target
+from pyoq.query.values import bind
+from pyoq.query.write_nodes import (
+    AssignmentNode,
+    ConflictAction,
+    ConflictNode,
+    DeleteNode,
+    InsertNode,
+    UpdateNode,
+    WriteNode,
+)
+from pyoq.unset import UnsetType
+
+Values = TypeVarTuple("Values")
+NodeType = TypeVar("NodeType", bound=WriteNode, covariant=True)
+ScopedNode = TypeVar("ScopedNode", bound="UpdateNode | DeleteNode")
+AssignedValue = TypeVar("AssignedValue")
+_EXCLUDED_ROW = "excluded"
+ResultRow = TypeVar("ResultRow")
+Value1 = TypeVar("Value1")
+Value2 = TypeVar("Value2")
+Value3 = TypeVar("Value3")
+Value4 = TypeVar("Value4")
+Value5 = TypeVar("Value5")
+Value6 = TypeVar("Value6")
+Value7 = TypeVar("Value7")
+Value8 = TypeVar("Value8")
+
+
+@runtime_checkable
+class ColumnReference(Protocol):
+    @property
+    def node(self) -> ExpressionNode: ...
+
+    @property
+    def family(self) -> ScalarFamily: ...
+
+
+@runtime_checkable
+class WriteColumn(Protocol):
+    @property
+    def writable(self) -> bool: ...
+
+    @property
+    def generated(self) -> bool: ...
+
+    @property
+    def field_name(self) -> str | None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class ColumnBinding:
+    database_name: str
+    field_name: str | None = None
+
+
+class WriteStatement(Generic[NodeType]):
+    __slots__ = ("_node",)
+    _node: NodeType
+
+    def __init__(self, node: NodeType) -> None:
+        object.__setattr__(self, "_node", node)
+
+    def __setattr__(self, name: str, value: object) -> None:
+        message = "write statements are immutable"
+        raise AttributeError(message)
+
+    @property
+    def node(self) -> NodeType:
+        return self._node
+
+
+class ReturningStatement(WriteStatement[WriteNode], Generic[ResultRow]):
+    __slots__ = ()
+
+
+class MutationStatement(WriteStatement[NodeType], Generic[NodeType]):
+    __slots__ = ()
+
+    def returning_all(self) -> ReturningStatement[tuple[object, ...]]:
+        """Every column the table declares, without listing them.
+
+        Generation writes the column list onto the table, so the statement
+        already knows what they are. A table written by hand declares none,
+        and is asked to name them.
+        """
+        target = self.node.target
+        if not target.columns:
+            message = (
+                "table names no columns to return; generated tables declare "
+                "them, and any other table has to name them"
+            )
+            raise QueryValidationError(message)
+        projections = tuple(
+            ProjectionNode(FieldNode(name, target.name, target.schema, target.catalog))
+            for name in target.columns
+        )
+        return ReturningStatement(replace(self.node, returning=projections))
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        /,
+    ) -> ReturningStatement[tuple[Value1]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        /,
+    ) -> ReturningStatement[tuple[Value1, Value2]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        /,
+    ) -> ReturningStatement[tuple[Value1, Value2, Value3]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        value4: Expression[Value4] | Projection[Value4],
+        /,
+    ) -> ReturningStatement[tuple[Value1, Value2, Value3, Value4]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        value4: Expression[Value4] | Projection[Value4],
+        value5: Expression[Value5] | Projection[Value5],
+        /,
+    ) -> ReturningStatement[tuple[Value1, Value2, Value3, Value4, Value5]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        value4: Expression[Value4] | Projection[Value4],
+        value5: Expression[Value5] | Projection[Value5],
+        value6: Expression[Value6] | Projection[Value6],
+        /,
+    ) -> ReturningStatement[tuple[Value1, Value2, Value3, Value4, Value5, Value6]]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        value4: Expression[Value4] | Projection[Value4],
+        value5: Expression[Value5] | Projection[Value5],
+        value6: Expression[Value6] | Projection[Value6],
+        value7: Expression[Value7] | Projection[Value7],
+        /,
+    ) -> ReturningStatement[
+        tuple[Value1, Value2, Value3, Value4, Value5, Value6, Value7]
+    ]: ...
+
+    @overload
+    def returning(
+        self,
+        value1: Expression[Value1] | Projection[Value1],
+        value2: Expression[Value2] | Projection[Value2],
+        value3: Expression[Value3] | Projection[Value3],
+        value4: Expression[Value4] | Projection[Value4],
+        value5: Expression[Value5] | Projection[Value5],
+        value6: Expression[Value6] | Projection[Value6],
+        value7: Expression[Value7] | Projection[Value7],
+        value8: Expression[Value8] | Projection[Value8],
+        /,
+    ) -> ReturningStatement[
+        tuple[Value1, Value2, Value3, Value4, Value5, Value6, Value7, Value8]
+    ]: ...
+
+    def returning(  # type: ignore[misc]
+        self,
+        *values: object,
+    ) -> ReturningStatement[object]:
+        if not values:
+            message = "RETURNING requires at least one projection"
+            raise QueryValidationError(message)
+        projections = tuple(projection_node(value) for value in values)
+        return ReturningStatement(replace(self.node, returning=projections))
+
+
+class ScopedStatement(MutationStatement[ScopedNode], Generic[ScopedNode]):
+    __slots__ = ()
+
+    def where(self, condition: Expression[bool], /) -> Self:
+        require_boolean(condition, "WHERE")
+        node = self.node
+        if node.condition is not None:
+            message = "WHERE condition is already defined"
+            raise QueryStateError(message)
+        if node.unconditional:
+            message = "a full-table statement cannot also define a WHERE condition"
+            raise QueryStateError(message)
+        return type(self)(replace(node, condition=condition.node))
+
+    def all_rows(self) -> Self:
+        node = self.node
+        if node.condition is not None:
+            message = "a conditional statement cannot also target every row"
+            raise QueryStateError(message)
+        return type(self)(replace(node, unconditional=True))
+
+
+class InsertStatement(MutationStatement[InsertNode], Generic[*Values]):
+    __slots__ = ("_columns",)
+    _columns: tuple[ColumnBinding, ...]
+
+    def __init__(
+        self,
+        node: InsertNode,
+        columns: tuple[ColumnBinding, ...],
+    ) -> None:
+        super().__init__(node)
+        object.__setattr__(self, "_columns", columns)
+
+    @property
+    def columns(self) -> tuple[ColumnBinding, ...]:
+        return self._columns
+
+    def values(self, *values: *Values) -> InsertStatement[*Values]:
+        row = _value_row(values, len(self._columns))
+        return self._with_rows((row,))
+
+    def values_many(self, rows: Iterable[object], /) -> InsertStatement[*Values]:
+        additions = tuple(_generated_row(row, self._columns) for row in rows)
+        if not additions:
+            message = "values_many requires at least one insert value"
+            raise QueryValidationError(message)
+        return self._with_rows(additions)
+
+    def on_conflict_do_nothing(
+        self,
+        *targets: ColumnReference,
+    ) -> InsertStatement[*Values]:
+        conflict = ConflictNode(
+            ConflictAction.NOTHING,
+            self._conflict_targets(targets),
+        )
+        return InsertStatement[*Values](self._with_conflict(conflict), self._columns)
+
+    def on_conflict_do_update(
+        self,
+        *targets: ColumnReference,
+    ) -> ConflictUpdate[*Values]:
+        conflict = ConflictNode(
+            ConflictAction.UPDATE,
+            self._conflict_targets(targets),
+        )
+        return ConflictUpdate[*Values](self._with_conflict(conflict), self._columns)
+
+    def _with_conflict(self, conflict: ConflictNode) -> InsertNode:
+        node = self.node
+        if node.conflict is not None:
+            message = "conflict resolution is already defined"
+            raise QueryStateError(message)
+        return replace(node, conflict=conflict)
+
+    def _conflict_targets(
+        self,
+        targets: tuple[ColumnReference, ...],
+    ) -> tuple[str, ...]:
+        selected = frozenset(binding.database_name for binding in self._columns)
+        names: list[str] = []
+        for target in targets:
+            name = _write_column(target, "ON CONFLICT").database_name
+            if name not in selected:
+                message = f"conflict target is not an inserted column: {name}"
+                raise QueryValidationError(message)
+            names.append(name)
+        return tuple(names)
+
+    def _with_rows(
+        self,
+        rows: tuple[tuple[ExpressionNode, ...], ...],
+    ) -> InsertStatement[*Values]:
+        node = self.node
+        return InsertStatement[*Values](
+            replace(node, rows=(*node.rows, *rows)),
+            self._columns,
+        )
+
+
+class ConflictUpdate(InsertStatement[*Values], Generic[*Values]):
+    __slots__ = ()
+
+    def set(
+        self,
+        column: Expression[AssignedValue],
+        value: AssignedValue | Expression[AssignedValue],
+        /,
+    ) -> ConflictUpdate[*Values]:
+        conflict = self._conflict
+        return self._with_resolution(
+            replace(
+                conflict,
+                assignments=(*conflict.assignments, _assignment(column, value)),
+            )
+        )
+
+    def where(self, condition: Expression[bool], /) -> ConflictUpdate[*Values]:
+        require_boolean(condition, "ON CONFLICT")
+        conflict = self._conflict
+        if conflict.condition is not None:
+            message = "conflict update condition is already defined"
+            raise QueryStateError(message)
+        return self._with_resolution(replace(conflict, condition=condition.node))
+
+    @property
+    def _conflict(self) -> ConflictNode:
+        conflict = self.node.conflict
+        if conflict is None:
+            message = "conflict update statement has no conflict resolution"
+            raise QueryStateError(message)
+        return conflict
+
+    def _with_resolution(self, conflict: ConflictNode) -> ConflictUpdate[*Values]:
+        return ConflictUpdate[*Values](
+            replace(self.node, conflict=conflict),
+            self._columns,
+        )
+
+
+class UpdateStatement(ScopedStatement[UpdateNode]):
+    __slots__ = ()
+
+    def set(
+        self,
+        column: Expression[AssignedValue],
+        value: AssignedValue | Expression[AssignedValue],
+        /,
+    ) -> UpdateStatement:
+        return self._with_assignments((_assignment(column, value),))
+
+    def set_values(
+        self,
+        values: object,
+        /,
+        *columns: ColumnReference,
+    ) -> UpdateStatement:
+        if not columns:
+            message = "set_values requires at least one column"
+            raise QueryValidationError(message)
+        assignments = _generated_assignments(values, columns)
+        if not assignments:
+            message = "update value leaves every selected column unset"
+            raise QueryValidationError(message)
+        return self._with_assignments(assignments)
+
+    def _with_assignments(
+        self,
+        assignments: tuple[AssignmentNode, ...],
+    ) -> UpdateStatement:
+        node = self.node
+        return UpdateStatement(
+            replace(node, assignments=(*node.assignments, *assignments))
+        )
+
+
+class DeleteStatement(ScopedStatement[DeleteNode]):
+    __slots__ = ()
+
+
+def update(table: TableLike, /) -> UpdateStatement:
+    return UpdateStatement(UpdateNode(write_target(table)))
+
+
+def delete_from(table: TableLike, /) -> DeleteStatement:
+    return DeleteStatement(DeleteNode(write_target(table)))
+
+
+def _assignment(
+    column: object,
+    value: object,
+) -> AssignmentNode:
+    binding = _write_column(column, "UPDATE")
+    return AssignmentNode(binding.database_name, _value_node(value))
+
+
+def _generated_assignments(
+    values: object,
+    columns: tuple[ColumnReference, ...],
+) -> tuple[AssignmentNode, ...]:
+    assignments: list[AssignmentNode] = []
+    for column in columns:
+        binding = _write_column(column, "UPDATE")
+        attribute = _generated_attribute(values, binding)
+        if isinstance(attribute, UnsetType):
+            continue
+        assignments.append(
+            AssignmentNode(binding.database_name, _value_node(attribute))
+        )
+    return tuple(assignments)
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    /,
+) -> InsertStatement[Value1]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    /,
+) -> InsertStatement[Value1, Value2]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    /,
+) -> InsertStatement[Value1, Value2, Value3]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    column4: Expression[Value4],
+    /,
+) -> InsertStatement[Value1, Value2, Value3, Value4]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    column4: Expression[Value4],
+    column5: Expression[Value5],
+    /,
+) -> InsertStatement[Value1, Value2, Value3, Value4, Value5]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    column4: Expression[Value4],
+    column5: Expression[Value5],
+    column6: Expression[Value6],
+    /,
+) -> InsertStatement[Value1, Value2, Value3, Value4, Value5, Value6]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    column4: Expression[Value4],
+    column5: Expression[Value5],
+    column6: Expression[Value6],
+    column7: Expression[Value7],
+    /,
+) -> InsertStatement[Value1, Value2, Value3, Value4, Value5, Value6, Value7]: ...
+
+
+@overload
+def insert_into(
+    table: TableLike,
+    column1: Expression[Value1],
+    column2: Expression[Value2],
+    column3: Expression[Value3],
+    column4: Expression[Value4],
+    column5: Expression[Value5],
+    column6: Expression[Value6],
+    column7: Expression[Value7],
+    column8: Expression[Value8],
+    /,
+) -> InsertStatement[
+    Value1, Value2, Value3, Value4, Value5, Value6, Value7, Value8
+]: ...
+
+
+def insert_into(  # type: ignore[misc]
+    table: TableLike,
+    *columns: object,
+) -> InsertStatement[object]:
+    target = write_target(table)
+    if not columns:
+        message = "INSERT requires at least one column"
+        raise QueryValidationError(message)
+    bindings = tuple(_write_column(column, "INSERT") for column in columns)
+    names = tuple(binding.database_name for binding in bindings)
+    return InsertStatement(InsertNode(target, names), bindings)
+
+
+def excluded(column: Expression[AssignedValue], /) -> Expression[AssignedValue]:
+    binding = _write_column(column, "ON CONFLICT")
+    return Expression(
+        FieldNode(binding.database_name, _EXCLUDED_ROW, family=column.family),
+        column.family,
+    )
+
+
+def _write_column(value: object, clause: str) -> ColumnBinding:
+    if not isinstance(value, Expression):
+        message = f"{clause} columns must be typed column expressions"
+        raise QueryValidationError(message)
+    node = value.node
+    if not isinstance(node, FieldNode):
+        message = f"{clause} columns must reference table columns"
+        raise QueryValidationError(message)
+    if not isinstance(value, WriteColumn):
+        return ColumnBinding(node.name)
+    _require_writable_column(value, node.name)
+    return ColumnBinding(node.name, value.field_name)
+
+
+def _require_writable_column(column: WriteColumn, name: str) -> None:
+    if column.generated:
+        message = f"generated column cannot be written: {name}"
+        raise QueryValidationError(message)
+    if not column.writable:
+        message = f"column is not writable: {name}"
+        raise QueryValidationError(message)
+
+
+def _value_row(
+    values: tuple[object, ...],
+    expected: int,
+) -> tuple[ExpressionNode, ...]:
+    if len(values) != expected:
+        message = "INSERT values must match the selected column count"
+        raise QueryValidationError(message)
+    return tuple(_value_node(value) for value in values)
+
+
+def _generated_row(
+    value: object,
+    columns: tuple[ColumnBinding, ...],
+) -> tuple[ExpressionNode, ...]:
+    return tuple(_generated_value(value, column) for column in columns)
+
+
+def _generated_attribute(value: object, column: ColumnBinding) -> object:
+    if column.field_name is None:
+        message = "generated write values require generated column descriptors"
+        raise QueryValidationError(message)
+    if not hasattr(value, column.field_name):
+        message = f"write value does not define column: {column.database_name}"
+        raise QueryValidationError(message)
+    return getattr(value, column.field_name)
+
+
+def _generated_value(value: object, column: ColumnBinding) -> ExpressionNode:
+    attribute = _generated_attribute(value, column)
+    if isinstance(attribute, UnsetType):
+        message = f"insert value leaves column unset: {column.database_name}"
+        raise QueryValidationError(message)
+    return _value_node(attribute)
+
+
+def _value_node(value: object) -> ExpressionNode:
+    if isinstance(value, Expression):
+        return value.node
+    return bind(value).node
+
+
+__all__ = (
+    "ColumnBinding",
+    "ColumnReference",
+    "ConflictUpdate",
+    "DeleteStatement",
+    "InsertStatement",
+    "MutationStatement",
+    "ReturningStatement",
+    "ScopedStatement",
+    "UpdateStatement",
+    "WriteColumn",
+    "WriteStatement",
+    "delete_from",
+    "excluded",
+    "insert_into",
+    "update",
+)

pyoq/relations/__init__.py

@@ -0,0 +1,66 @@
+"""Typed relation facts, navigation, fetch plans, and loading state."""
+
+from pyoq.relations.batching import (
+    DEFAULT_KEY_LIMIT,
+    KeyBatch,
+    RelationBatch,
+    RelationBatchLoader,
+    plan_key_batches,
+)
+from pyoq.relations.derivation import derive_relations, table_relations
+from pyoq.relations.fetching import (
+    MAXIMUM_FETCH_DEPTH,
+    FetchJoin,
+    FetchOrder,
+    FetchPlan,
+    FetchRequest,
+    FetchStrategy,
+    by_column,
+    by_expression,
+    validate_fetch_plan,
+)
+from pyoq.relations.graph import RelationGraph, build_relation_graph
+from pyoq.relations.loading import RelationState, RelationValue
+from pyoq.relations.model import (
+    RelationCardinality,
+    RelationDirection,
+    RelationEndpoint,
+    TypedRelation,
+)
+from pyoq.relations.planning import (
+    FetchCapabilities,
+    ResolvedFetch,
+    ResolvedFetchPlan,
+    resolve_fetch_plan,
+)
+
+__all__ = (
+    "DEFAULT_KEY_LIMIT",
+    "MAXIMUM_FETCH_DEPTH",
+    "FetchCapabilities",
+    "FetchJoin",
+    "FetchOrder",
+    "FetchPlan",
+    "FetchRequest",
+    "FetchStrategy",
+    "KeyBatch",
+    "RelationBatch",
+    "RelationBatchLoader",
+    "RelationCardinality",
+    "RelationDirection",
+    "RelationEndpoint",
+    "RelationGraph",
+    "RelationState",
+    "RelationValue",
+    "ResolvedFetch",
+    "ResolvedFetchPlan",
+    "TypedRelation",
+    "build_relation_graph",
+    "by_column",
+    "by_expression",
+    "derive_relations",
+    "plan_key_batches",
+    "resolve_fetch_plan",
+    "table_relations",
+    "validate_fetch_plan",
+)