pyoq-sql 1.0.2__py3-none-any.whl

pyoq/fetching/execution.py

@@ -0,0 +1,587 @@
+"""Carrying out a resolved fetch plan against a database.
+
+A plan is decided before anything runs and says how each relation is reached.
+This runs it: one statement for the root and everything that can ride along
+with it, and one statement per batch for each relation that cannot.
+
+What comes back is uniform. However a row arrived, joined into its parent's
+row, gathered as JSON beside it, or read by a statement of its own, it is a
+row of its table in that table's own column order, carrying its own relations
+the same way. Nothing above has to know which way a relation was fetched,
+which is the reason for deciding it below them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast
+
+from pyoq.descriptors import TableDescriptor
+from pyoq.errors import RelationNotLoadedError, RowIdentityError
+from pyoq.fetching.collections import decode_collection
+from pyoq.fetching.plans import batch_statement, root_statement
+from pyoq.relations.batching import plan_key_batches
+from pyoq.relations.loading import RelationValue
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterable, Mapping, Sequence
+
+    from pyoq.descriptors import RelationshipDescriptor
+    from pyoq.fetching.plans import (
+        DeferredRelation,
+        FetchStatement,
+        JoinedRelation,
+        NestedRelation,
+        Positions,
+        RelationBlock,
+        Values,
+    )
+    from pyoq.query.execution import QueryOperations
+    from pyoq.query.execution.aio import AsyncQueryOperations
+    from pyoq.query.expressions import Condition
+    from pyoq.query.select import SelectQuery
+    from pyoq.relations.graph import RelationGraph
+    from pyoq.relations.model import TypedRelation
+    from pyoq.relations.planning import ResolvedFetch, ResolvedFetchPlan
+    from pyoq.schema.models import Identifier, ObjectReference
+
+    Loaded = RelationValue[tuple["FetchedRow", ...]]
+
+
+Row = TypeVar("Row")
+Target = TypeVar("Target")
+Held = TypeVar("Held")
+
+Work: TypeAlias = "Generator[SelectQuery[Values], Sequence[Values], Held]"
+"""Work that asks for the rows of a query and is handed them back."""
+
+
+@dataclass(frozen=True, slots=True)
+class FetchedRow:
+    """One row and everything fetched through it.
+
+    Every relation answers with rows, whatever its cardinality. A to-one holds
+    the single row it resolved to, or reads as absent. A to-many holds however
+    many there were, and none is an empty collection rather than an absence,
+    because a relation with no children was answered rather than unanswered.
+    """
+
+    table: ObjectReference
+    """The table these values came out of, which is what says how to read them.
+
+    A row carries it rather than being told it, because a caller naming a
+    table is making a claim and a row is the only thing that knows whether it
+    is true.
+    """
+
+    values: Values
+    related: Mapping[TypedRelation, RelationValue[tuple[FetchedRow, ...]]]
+
+    def read(self, table: TableDescriptor[Row, object, object], /) -> Row:
+        """This row as the class the schema generated to hold it.
+
+        The values came back in the order the table declares its columns,
+        which is the order that class names its fields, so the row it builds
+        holds each value under the name and the type the column was declared
+        with. A table written by hand names no row class and is told so.
+
+        The descriptor has to name the table this row came from. Reading a row
+        as a table it did not come from would put its values under names that
+        were never theirs, and the type checker would believe it.
+        """
+        _require_same_table(self.table, table)
+        return _built(table.row_type, table.database_name, self.values)
+
+    def read_all(
+        self,
+        relation: RelationshipDescriptor[object, Target],
+        /,
+    ) -> tuple[Target, ...]:
+        """The rows a relation resolved to, as the class that holds them.
+
+        None where a to-one resolved to nothing, and none where a to-many
+        resolved to no children, which are the two the relation itself tells
+        apart.
+        """
+        return tuple(
+            held.read(_target_of(relation))
+            for held in self.rows_of(_matching(self.related, relation))
+        )
+
+    def rows_of(self, relation: TypedRelation, /) -> tuple[FetchedRow, ...]:
+        """The rows a relation resolved to, which is none where it resolved to none.
+
+        A relation nobody fetched raises rather than reading as empty, which is
+        the difference `RelationValue` exists to keep.
+        """
+        held = self.related.get(relation)
+        if held is None:
+            return _unloaded().value or ()
+        return held.value or ()
+
+
+class FetchAssembly:
+    """Carrying out a resolved plan, apart from running anything.
+
+    It asks for the rows of a query and is handed them back, so one traversal
+    serves a database that blocks and one that does not. Nothing here awaits
+    and nothing here connects, which is why there is one of it rather than two
+    that have to be kept saying the same thing.
+
+    The parameter limit is given rather than read from a compiler, because a
+    compiler is asked to compile and nothing else. It is the same limit a bulk
+    planner is built with, and it decides how many parents one select-in
+    statement can ask about.
+    """
+
+    __slots__ = ("_graph", "_maximum_parameters")
+
+    def __init__(self, graph: RelationGraph, /, *, maximum_parameters: int) -> None:
+        self._graph = graph
+        self._maximum_parameters = maximum_parameters
+
+    def rows_for(
+        self,
+        plan: ResolvedFetchPlan,
+        where: Condition | None,
+    ) -> Work[tuple[FetchedRow, ...]]:
+        """Every row the plan's root describes, with its relations loaded."""
+        statement = root_statement(plan, self._graph, where=where)
+        rows = yield statement.query
+        return (yield from self._rows(statement.block, rows))
+
+    def _rows(
+        self,
+        block: RelationBlock,
+        rows: Sequence[Values],
+    ) -> Work[tuple[FetchedRow, ...]]:
+        """One table's rows out of a result, each with its relations loaded.
+
+        Every relation is loaded across all the rows at once rather than row by
+        row. A relation reached from a joined table is still one statement for
+        the whole result, which is the difference between this and the access
+        pattern it exists to avoid.
+        """
+        values = tuple(_picked(row, block.positions) for row in rows)
+        held: tuple[dict[TypedRelation, Loaded], ...] = tuple({} for _ in rows)
+        for joined in block.joined:
+            yield from self._load_joined(joined, rows, held)
+        for nested in block.nested:
+            self._load_nested(nested, rows, held)
+        for deferred in block.deferred:
+            yield from self._load_deferred(deferred, values, held)
+        return tuple(
+            FetchedRow(block.table, own, entry)
+            for own, entry in zip(values, held, strict=True)
+        )
+
+    def _load_joined(
+        self,
+        joined: JoinedRelation,
+        rows: Sequence[Values],
+        held: Sequence[dict[TypedRelation, Loaded]],
+    ) -> Work[None]:
+        """A to-one relation, and whatever is fetched through it in turn.
+
+        A row that matched nothing reads as absent. The rest are carried on
+        together, so a relation hanging off the joined table is one statement
+        for every row that matched rather than one for each.
+        """
+        matched = tuple(
+            index
+            for index, row in enumerate(rows)
+            if not all(row[position] is None for position in joined.key_positions)
+        )
+        for entry in held:
+            entry[joined.fetch.relation] = RelationValue.absent()
+        children = yield from self._rows(
+            joined.block, tuple(rows[index] for index in matched)
+        )
+        for index, child in zip(matched, children, strict=True):
+            held[index][joined.fetch.relation] = RelationValue.loaded((child,))
+
+    def _load_nested(
+        self,
+        nested: NestedRelation,
+        rows: Sequence[Values],
+        held: Sequence[dict[TypedRelation, Loaded]],
+    ) -> None:
+        """A collection, read back out of the JSON it arrived as.
+
+        A collection is bounded and ordered where it is built, so its rows are
+        already the ones its limit kept. It carries no deeper fetch of its own:
+        a plan that reaches through a collection is resolved to select-in,
+        because a deeper fetch is keyed on rows rather than on JSON.
+
+        What a join added is inside the same value, after the collection's own
+        columns, so it is read out of each row rather than fetched again.
+        """
+        keys = nested.keys + tuple(key for join in nested.joined for key in join.keys)
+        own = len(nested.keys)
+        for entry, row in zip(held, rows, strict=True):
+            gathered = decode_collection(row[nested.position], keys)
+            entry[nested.fetch.relation] = _resolved(
+                nested.fetch.relation,
+                tuple(
+                    FetchedRow(
+                        nested.fetch.table,
+                        values[:own],
+                        _joined_into(nested, values, own),
+                    )
+                    for values in gathered
+                ),
+            )
+
+    def _load_deferred(
+        self,
+        deferred: DeferredRelation,
+        values: Sequence[Values],
+        held: Sequence[dict[TypedRelation, Loaded]],
+    ) -> Work[None]:
+        """A relation read by statements of its own, for every parent at once."""
+        gathered = yield from self._deferred(deferred, values)
+        for entry, own in zip(held, values, strict=True):
+            entry[deferred.fetch.relation] = _resolved(
+                deferred.fetch.relation,
+                gathered.get(_picked(own, deferred.key_positions), ()),
+            )
+
+    def _deferred(
+        self,
+        deferred: DeferredRelation,
+        parents: Sequence[Values],
+    ) -> Work[dict[Values, tuple[FetchedRow, ...]]]:
+        """One relation's rows for every parent, in as few statements as allowed.
+
+        A parent whose key holds null is not asked about. Null matches no row,
+        including another null, so asking would spend a parameter to be told
+        what is already known, and the predicate refuses one rather than
+        quietly returning nothing.
+        """
+        keys = _matchable(_picked(row, deferred.key_positions) for row in parents)
+        if not keys:
+            return {}
+        gathered: dict[Values, list[FetchedRow]] = {key: [] for key in keys}
+        for batch in plan_key_batches(
+            deferred.fetch.relation,
+            keys,
+            maximum_parameters=self._maximum_parameters,
+        ):
+            yield from self._gather(deferred.fetch, batch.keys, gathered)
+        return {
+            key: _limited(rows, deferred.fetch.limit) for key, rows in gathered.items()
+        }
+
+    def _gather(
+        self,
+        fetch: ResolvedFetch,
+        keys: Sequence[Values],
+        gathered: dict[Values, list[FetchedRow]],
+    ) -> Work[None]:
+        """Run one batch, and file each row under the parent key it belongs to.
+
+        Every row a batch returns matched one of the keys it asked about, so
+        every row has somewhere to go. A row that did not would be a key that
+        read back as a different value than it was sent as, and raising says
+        so rather than dropping the row and answering with less.
+        """
+        statement = batch_statement(fetch, self._graph, keys)
+        rows = yield statement.query
+        owners = _owners(statement, fetch, rows)
+        children = yield from self._rows(statement.block, rows)
+        for owner, child in zip(owners, children, strict=True):
+            gathered[owner].append(child)
+
+
+class FetchExecutor:
+    """Runs the statements a resolved plan owes against a database.
+
+    The assembly decides what to run and what the answers mean. This hands it
+    each query and gives back the rows, which is the whole of the difference
+    between running a plan here and running one that does not block.
+    """
+
+    __slots__ = ("_assembly", "_operations")
+
+    def __init__(
+        self,
+        operations: QueryOperations,
+        graph: RelationGraph,
+        /,
+        *,
+        maximum_parameters: int,
+    ) -> None:
+        self._operations = operations
+        self._assembly = FetchAssembly(graph, maximum_parameters=maximum_parameters)
+
+    def fetch(
+        self,
+        plan: ResolvedFetchPlan,
+        /,
+        *,
+        where: Condition | None = None,
+    ) -> tuple[FetchedRow, ...]:
+        """Every row the plan's root describes, with its relations loaded."""
+        work = self._assembly.rows_for(plan, where)
+        try:
+            query = next(work)
+            while True:
+                query = work.send(self._operations.many(query))
+        except StopIteration as finished:
+            return _finished(finished)
+
+
+class AsyncFetchExecutor:
+    """The same, against a database that does not block.
+
+    It shares the assembly rather than the code, so there is one answer to
+    what a plan means and two ways of asking a database.
+    """
+
+    __slots__ = ("_assembly", "_operations")
+
+    def __init__(
+        self,
+        operations: AsyncQueryOperations,
+        graph: RelationGraph,
+        /,
+        *,
+        maximum_parameters: int,
+    ) -> None:
+        self._operations = operations
+        self._assembly = FetchAssembly(graph, maximum_parameters=maximum_parameters)
+
+    async def fetch(
+        self,
+        plan: ResolvedFetchPlan,
+        /,
+        *,
+        where: Condition | None = None,
+    ) -> tuple[FetchedRow, ...]:
+        """Every row the plan's root describes, with its relations loaded."""
+        work = self._assembly.rows_for(plan, where)
+        try:
+            query = next(work)
+            while True:
+                query = work.send(await self._operations.many(query))
+        except StopIteration as finished:
+            return _finished(finished)
+
+
+def _finished(finished: StopIteration) -> tuple[FetchedRow, ...]:
+    """What the assembly returned once it stopped asking for rows.
+
+    A generator's return value is untyped where it is caught, and this is the
+    one place that catches one.
+    """
+    return cast("tuple[FetchedRow, ...]", finished.value)
+
+
+def _owners(
+    statement: FetchStatement,
+    fetch: ResolvedFetch,
+    rows: Sequence[Values],
+) -> tuple[Values, ...]:
+    """The parent key each row of a batch belongs to."""
+    order = {column.name: index for index, column in enumerate(statement.columns)}
+    positions = tuple(
+        order[name] for name in fetch.relation.target.columns if name in order
+    )
+    return tuple(_picked(row, positions) for row in rows)
+
+
+def _unloaded() -> RelationValue[tuple[FetchedRow, ...]]:
+    return RelationValue[tuple["FetchedRow", ...]].unloaded()
+
+
+def _matching(
+    related: Mapping[TypedRelation, RelationValue[tuple[FetchedRow, ...]]],
+    descriptor: RelationshipDescriptor[object, object],
+) -> TypedRelation:
+    """The relation a generated descriptor names, among the ones fetched.
+
+    A descriptor and a relation are the same fact written for two audiences,
+    so they are matched on what identifies one: which columns of which table
+    reach which columns of which other, and in which direction.
+    """
+    for relation in related:
+        if _same(relation, descriptor):
+            return relation
+    source = _written_parts(
+        descriptor.source_catalog, descriptor.source_schema, descriptor.source_table
+    )
+    target = _written_parts(
+        descriptor.target_catalog, descriptor.target_schema, descriptor.target_table
+    )
+    message = (
+        f"the relation from {source} to {target} was not fetched; include it "
+        "in the query's fetch plan to read it"
+    )
+    raise RelationNotLoadedError(message)
+
+
+def _same(
+    relation: TypedRelation,
+    descriptor: RelationshipDescriptor[object, object],
+) -> bool:
+    """Whether a descriptor names this relation and no other.
+
+    Both tables are compared whole, schema and catalog included. Two schemas
+    may each hold a `parent` reaching a `child` over the same columns, and a
+    descriptor that named only the tables would describe both of them.
+    """
+    return (
+        _at(relation.source.table, descriptor.source_catalog, descriptor.source_schema)
+        and _at(
+            relation.target.table, descriptor.target_catalog, descriptor.target_schema
+        )
+        and relation.source.table.name.value == descriptor.source_table
+        and relation.target.table.name.value == descriptor.target_table
+        and _names(relation.source.columns) == descriptor.source_columns
+        and _names(relation.target.columns) == descriptor.target_columns
+        and relation.direction is descriptor.direction
+    )
+
+
+def _at(table: ObjectReference, catalog: str | None, schema: str | None) -> bool:
+    return _named(table.catalog) == catalog and _named(table.schema) == schema
+
+
+def _names(columns: tuple[Identifier, ...]) -> tuple[str, ...]:
+    return tuple(column.value for column in columns)
+
+
+def _require_same_table(
+    produced: ObjectReference,
+    claimed: TableDescriptor[object, object, object],
+) -> None:
+    """Refuse a descriptor that names a table other than the one this came from.
+
+    Compared whole, catalog and schema included, because two schemas may hold
+    a table of one name and reading one as the other would be exactly the
+    mistake this is here to refuse.
+    """
+    if _describes(produced, claimed):
+        return
+    named = _written_parts(
+        claimed.catalog_name, claimed.schema_name, claimed.database_name
+    )
+    message = f"this row came from {_written(produced)} and is not a row of {named}"
+    raise RowIdentityError(message)
+
+
+def _describes(
+    produced: ObjectReference,
+    claimed: TableDescriptor[object, object, object],
+) -> bool:
+    return (
+        produced.name.value == claimed.database_name
+        and _named(produced.schema) == claimed.schema_name
+        and _named(produced.catalog) == claimed.catalog_name
+    )
+
+
+def _written(reference: ObjectReference) -> str:
+    return _written_parts(
+        _named(reference.catalog), _named(reference.schema), reference.name.value
+    )
+
+
+def _written_parts(catalog: str | None, schema: str | None, table: str) -> str:
+    return ".".join(part for part in (catalog, schema, table) if part is not None)
+
+
+def _named(value: Identifier | None) -> str | None:
+    return None if value is None else value.value
+
+
+def _target_of(
+    relation: RelationshipDescriptor[object, Target],
+) -> TableDescriptor[Target, object, object]:
+    """The table a relation reaches, as the descriptor for reading its rows."""
+    return TableDescriptor(
+        relation.target_table,
+        relation.target_schema,
+        relation.target_catalog,
+        relation.target_row_type,
+    )
+
+
+def _built(row_type: type[Row] | None, table: str, values: Values) -> Row:
+    """One row as its class, where the schema generated one to hold it."""
+    if row_type is None:
+        message = (
+            f"{table!r} names no row class, so a row of it reads back as the "
+            "values it came with; generate the package to read it as a row"
+        )
+        raise RowIdentityError(message)
+    return row_type(*values)
+
+
+def _joined_into(
+    nested: NestedRelation,
+    values: Values,
+    own: int,
+) -> dict[TypedRelation, RelationValue[tuple[FetchedRow, ...]]]:
+    """What each row of a collection holds through the tables it was read with.
+
+    A join that matched nothing reads as nulls across its key, which is an
+    absence rather than a row of nulls, the same answer a join in the parent's
+    own statement gives.
+    """
+    held: dict[TypedRelation, RelationValue[tuple[FetchedRow, ...]]] = {}
+    taken = own
+    for join in nested.joined:
+        found = values[taken : taken + len(join.keys)]
+        taken += len(join.keys)
+        matched = any(found[position] is not None for position in join.key_positions)
+        held[join.relation] = _resolved(
+            join.relation,
+            (FetchedRow(join.relation.target.table, found, {}),) if matched else (),
+        )
+    return held
+
+
+def _resolved(
+    relation: TypedRelation,
+    rows: tuple[FetchedRow, ...],
+) -> RelationValue[tuple[FetchedRow, ...]]:
+    """What a relation holds, once the rows it resolved to are known.
+
+    No rows means two different things. A to-one resolved to nothing, which is
+    an absence a caller can ask about. A to-many resolved to no children,
+    which is an answer: the collection is empty. Reading either as the other
+    loses the distinction `RelationValue` exists to keep.
+    """
+    if not rows and relation.to_one:
+        return RelationValue[tuple["FetchedRow", ...]].absent()
+    return RelationValue.loaded(rows)
+
+
+def _matchable(keys: Iterable[Values]) -> tuple[Values, ...]:
+    """Every key that could match a row, once, in the order it was first seen.
+
+    A key holding null matches nothing, so it is not worth a parameter and is
+    not one the predicate accepts.
+    """
+    return _distinct(key for key in keys if all(value is not None for value in key))
+
+
+def _limited(rows: Sequence[FetchedRow], limit: int | None) -> tuple[FetchedRow, ...]:
+    """The rows a limit keeps, which are the ones the order put first."""
+    return tuple(rows) if limit is None else tuple(rows[:limit])
+
+
+def _distinct(keys: Iterable[Values]) -> tuple[Values, ...]:
+    """Every key once, in the order it was first seen."""
+    seen: dict[Values, None] = {}
+    for key in keys:
+        seen.setdefault(key, None)
+    return tuple(seen)
+
+
+def _picked(row: Values, positions: Positions) -> Values:
+    return tuple(row[index] for index in positions)
+
+
+__all__ = ("AsyncFetchExecutor", "FetchAssembly", "FetchExecutor", "FetchedRow")

pyoq/fetching/joined.py

@@ -0,0 +1,79 @@
+"""Fetching a to-one relation in the query that already reads its parent.
+
+A to-one relation is one row the join carries anyway, so it costs no extra
+statement. What it does cost is knowing where each table's values sit in the
+result, which is what a projection layout records.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from pyoq.errors import QueryValidationError
+from pyoq.query.expressions import Condition
+from pyoq.query.nodes import (
+    BinaryNode,
+    BinaryOperator,
+    ExpressionNode,
+    VariadicNode,
+    VariadicOperator,
+)
+
+
+def join_condition(
+    source: Sequence[ExpressionNode],
+    target: Sequence[ExpressionNode],
+    /,
+) -> Condition:
+    """Match a row to the one it points at, across every column of the key."""
+    _require_aligned(source, target)
+    equalities = tuple(
+        BinaryNode(BinaryOperator.EQUAL, left, right)
+        for left, right in zip(source, target, strict=True)
+    )
+    if len(equalities) == 1:
+        return Condition(equalities[0])
+    return Condition(VariadicNode(VariadicOperator.AND, equalities))
+
+
+def projection_layout(widths: Sequence[int], /) -> tuple[tuple[int, ...], ...]:
+    """Where each table's values sit once its columns are laid end to end.
+
+    Hydration reads a result by position, so laying the projection out and
+    reading it back have to agree. Deriving one from the other keeps them from
+    drifting apart as a plan grows.
+    """
+    _require_widths(widths)
+    positions: list[tuple[int, ...]] = []
+    start = 0
+    for width in widths:
+        positions.append(tuple(range(start, start + width)))
+        start += width
+    return tuple(positions)
+
+
+def _require_aligned(
+    source: Sequence[ExpressionNode],
+    target: Sequence[ExpressionNode],
+) -> None:
+    if not source:
+        message = "a join needs at least one column on each side"
+        raise QueryValidationError(message)
+    if len(source) != len(target):
+        message = (
+            f"a join matches {len(source)} columns against {len(target)}, "
+            f"which cannot line up"
+        )
+        raise QueryValidationError(message)
+
+
+def _require_widths(widths: Sequence[int]) -> None:
+    if not widths:
+        message = "a projection layout needs at least one table"
+        raise QueryValidationError(message)
+    if any(width < 1 for width in widths):
+        message = "every table in a projection contributes at least one column"
+        raise QueryValidationError(message)
+
+
+__all__ = ("join_condition", "projection_layout")