dagrun 0.1.0__py3-none-any.whl

dagrun/__init__.py

@@ -0,0 +1,6 @@
+from dagrun import model
+from dagrun.control import Defer
+from dagrun.dag import Dag, DagValidationError
+from dagrun.dag_runner import DagRunner
+
+__all__ = ["Dag", "DagRunner", "DagValidationError", "Defer", "model"]

dagrun/chain.py

@@ -0,0 +1,71 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+def member(fn: Callable[..., Any]) -> Callable[..., Any]:
+    fn.__dagrun_member__ = True  # type: ignore[attr-defined]
+    return fn
+
+
+class Combine(Protocol):
+    def reduce(self, column: str, values: list[Any]) -> Any: ...
+
+
+@dataclass(frozen=True)
+class LastWins:
+    def reduce(self, column: str, values: list[Any]) -> Any:
+        return values[-1]
+
+
+@dataclass(frozen=True)
+class FirstWins:
+    def reduce(self, column: str, values: list[Any]) -> Any:
+        return values[0]
+
+
+@dataclass(frozen=True)
+class Error:
+    def reduce(self, column: str, values: list[Any]) -> Any:
+        distinct = {_hash(v) for v in values}
+        if len(distinct) > 1:
+            raise CombineConflict(column, values)
+        return values[0]
+
+
+class CombineConflict(RuntimeError):
+    def __init__(self, column: str, values: list[Any]) -> None:
+        super().__init__(f"Conflicting values for column {column!r}: {values!r}")
+        self.column = column
+        self.values = values
+
+
+def _hash(value: Any) -> Any:
+    try:
+        hash(value)
+    except TypeError:
+        return repr(value)
+    return value
+
+
+@dataclass(frozen=True)
+class PerColumn:
+    """Per-column combine override map. '*' means default for unlisted columns."""
+
+    rules: tuple[tuple[str, Combine], ...]
+
+    def __init__(self, rules: Mapping[str, Combine] | "tuple[tuple[str, Combine], ...]") -> None:
+        if isinstance(rules, Mapping):
+            items = tuple(rules.items())
+        else:
+            items = tuple(rules)
+        object.__setattr__(self, "rules", items)
+
+    def reduce(self, column: str, values: list[Any]) -> Any:
+        lookup = dict(self.rules)
+        rule = lookup.get(column, lookup.get("*"))
+        if rule is None:
+            return values[-1]
+        return rule.reduce(column, values)
+
+

dagrun/control.py

@@ -0,0 +1,9 @@
+class Defer(Exception):
+    """Signal from a node that the current input should not be processed now.
+
+    When a node raises Defer the runner discards any output produced for that
+    input and records no invocation, so the input is evaluated again from
+    scratch on the next run rather than being treated as done. Use it for
+    decisions that depend on runtime configuration (e.g. a token budget)
+    rather than on an intrinsic, persistable property of the input.
+    """

dagrun/dag.py

@@ -0,0 +1,396 @@
+from __future__ import annotations
+
+import inspect
+import types
+import typing
+from collections.abc import (
+    AsyncIterable,
+    AsyncIterator,
+    Callable,
+    Generator,
+    Iterable,
+    Iterator,
+    Mapping,
+)
+from dataclasses import dataclass, field
+from typing import Any, Union, get_args, get_origin
+
+from dagrun.chain import Combine, LastWins, PerColumn
+from dagrun.model import (
+    Entity,
+    SYSTEM_COLUMNS,
+    _entity_specs,
+    _EntityView,
+    _find_pk,
+    _find_unique_fk_to,
+    _required_columns,
+)
+from dagrun.plan import CompiledChain, CompiledInput, CompiledMember, Plan
+
+
+class DagValidationError(Exception):
+    pass
+
+
+@dataclass
+class _Chain:
+    members: list[Callable[..., Any]]
+    combine: Combine
+    labels: set[str] = field(default_factory=set)
+    cost: Mapping[str, int] = field(default_factory=dict)
+
+
+_ITERABLE_ORIGINS: frozenset[Any] = frozenset(
+    {Iterable, Iterator, Generator, AsyncIterable, AsyncIterator, list, tuple, set, frozenset}
+)
+
+
+def _unwrap_input_param(
+    ann: Any,
+) -> tuple[type | None, tuple[str, ...], bool]:
+    """Like _unwrap_entity, but additionally reports whether the outermost
+    wrapper was a collection (list/Iterable/etc.). A collection annotation on
+    an input parameter signals an aggregate chain: dagrun groups child rows
+    by their FK to the produced parent entity and calls the fn once per
+    parent with all matching children."""
+    if ann is inspect.Parameter.empty or ann is inspect.Signature.empty:
+        return None, (), False
+    origin = get_origin(ann)
+    if origin in _ITERABLE_ORIGINS:
+        args = get_args(ann)
+        if not args:
+            return None, (), False
+        entity, cols = _unwrap_entity(args[0])
+        if entity is None:
+            return None, (), False
+        return entity, cols, True
+    entity, cols = _unwrap_entity(ann)
+    return entity, cols, False
+
+
+def _unwrap_entity(ann: Any) -> tuple[type | None, tuple[str, ...]]:
+    """Peel Iterable/Union/Annotated wrappers and return (entity, columns).
+
+    Returns (None, ()) for annotations that don't resolve to an Entity subclass.
+    For a bare Entity, returns all required column names.
+    """
+    if ann is inspect.Parameter.empty or ann is inspect.Signature.empty:
+        return None, ()
+
+    metadata = getattr(ann, "__metadata__", None)
+    if metadata is not None:
+        for m in metadata:
+            if isinstance(m, _EntityView):
+                return m.entity, m.columns
+        return _unwrap_entity(ann.__origin__)
+
+    origin = get_origin(ann)
+    if origin is Union or origin is types.UnionType:
+        non_none = [a for a in get_args(ann) if a is not type(None)]
+        if len(non_none) == 1:
+            return _unwrap_entity(non_none[0])
+        return None, ()
+
+    if origin in _ITERABLE_ORIGINS:
+        args = get_args(ann)
+        if args:
+            return _unwrap_entity(args[0])
+        return None, ()
+
+    if isinstance(ann, type) and issubclass(ann, Entity):
+        return ann, _required_columns(ann)
+
+    return None, ()
+
+
+def _detect_cycle(
+    edges: dict[CompiledChain, list[CompiledChain]],
+) -> list[CompiledChain] | None:
+    color: dict[CompiledChain, int] = {}
+    stack: list[CompiledChain] = []
+
+    def dfs(node: CompiledChain) -> list[CompiledChain] | None:
+        color[node] = 1
+        stack.append(node)
+        for nxt in edges.get(node, ()):
+            c = color.get(nxt, 0)
+            if c == 1:
+                i = stack.index(nxt)
+                return stack[i:] + [nxt]
+            if c == 0:
+                cyc = dfs(nxt)
+                if cyc is not None:
+                    return cyc
+        stack.pop()
+        color[node] = 2
+        return None
+
+    for n in edges:
+        if color.get(n, 0) == 0:
+            cyc = dfs(n)
+            if cyc is not None:
+                return cyc
+    return None
+
+
+class Dag:
+    def __init__(self) -> None:
+        self._chains: list[_Chain] = []
+
+    def fn(
+        self,
+        target: Any = None,
+        *,
+        labels: set[str] | None = None,
+        cost: Mapping[str, int] | None = None,
+    ) -> Any:
+        return self._register(
+            target, labels, combine=LastWins(), is_group=False, cost=cost
+        )
+
+    def group(
+        self,
+        target: Any = None,
+        *,
+        labels: set[str] | None = None,
+        combine: Combine | Mapping[str, Combine] | None = None,
+        cost: Mapping[str, int] | None = None,
+    ) -> Any:
+        resolved = combine
+        if isinstance(resolved, Mapping):
+            resolved = PerColumn(rules=dict(resolved))
+        return self._register(
+            target, labels, combine=resolved or LastWins(), is_group=True, cost=cost
+        )
+
+    def compile(self) -> Plan:
+        parsed: list[_ParsedChain] = [_parse_chain(raw) for raw in self._chains]
+
+        compiled_chains: list[CompiledChain] = []
+        parsed_to_compiled: dict[int, CompiledChain] = {}
+        for p in parsed:
+            labels = set(p.raw.labels)
+            if p.produces is not None:
+                labels.add(p.produces.__name__.lower())
+            for e, _cols, is_coll in p.inputs:
+                if not is_coll:
+                    continue
+                if p.produces is None:
+                    raise DagValidationError(
+                        f"Aggregate chain reading {e.__name__} must produce "
+                        f"an entity (none declared)."
+                    )
+                try:
+                    _find_unique_fk_to(e, p.produces)
+                except LookupError as exc:
+                    raise DagValidationError(str(exc)) from exc
+
+            compiled = CompiledChain(
+                produces=p.produces if p.produces is not None else object,
+                inputs=tuple(
+                    CompiledInput(entity=e, columns=cols, is_collection=is_coll)
+                    for e, cols, is_coll in p.inputs
+                ),
+                members=tuple(p.members),
+                combine=p.raw.combine,
+                labels=frozenset(labels),
+                cost=dict(p.raw.cost),
+            )
+            compiled_chains.append(compiled)
+            parsed_to_compiled[id(p)] = compiled
+
+        column_owners: dict[tuple[type, str], CompiledChain] = {}
+        column_kinds: dict[tuple[type, str], type] = {}
+
+        for p in parsed:
+            if p.produces is None:
+                continue
+            specs = _entity_specs(p.produces)
+            compiled = parsed_to_compiled[id(p)]
+            for col in p.produces_cols:
+                if col in SYSTEM_COLUMNS:
+                    continue
+                key = (p.produces, col)
+                if key in column_owners:
+                    raise DagValidationError(
+                        f"Column {p.produces.__name__}.{col} produced by multiple chains"
+                    )
+                column_owners[key] = compiled
+                if col in specs:
+                    column_kinds[key] = specs[col].kind
+
+        edges: dict[CompiledChain, list[CompiledChain]] = {c: [] for c in compiled_chains}
+        for p in parsed:
+            compiled = parsed_to_compiled[id(p)]
+            for m in p.members:
+                for entity, cols in m.needs_entities:
+                    needed_raw = list(cols) if cols else [_find_pk(entity)]
+                    needed = [c for c in needed_raw if c not in SYSTEM_COLUMNS] or [
+                        _find_pk(entity)
+                    ]
+                    for col in needed:
+                        owner = column_owners.get((entity, col))
+                        if owner is None:
+                            raise DagValidationError(
+                                f"No chain produces {entity.__name__}.{col} "
+                                f"(required by {_chain_label(p)})"
+                            )
+                        if owner is compiled:
+                            continue
+                        if compiled not in edges[owner]:
+                            edges[owner].append(compiled)
+
+        cycle = _detect_cycle(edges)
+        if cycle is not None:
+            names = " -> ".join(c.produces.__name__ for c in cycle)
+            raise DagValidationError(f"Cycle detected: {names}")
+
+        return Plan(
+            chains=tuple(compiled_chains),
+            column_owners=column_owners,
+            column_kinds=column_kinds,
+            edges={k: tuple(v) for k, v in edges.items()},
+        )
+
+    def _register(
+        self,
+        target: Any,
+        labels: set[str] | None,
+        combine: Combine,
+        is_group: bool,
+        cost: Mapping[str, int] | None = None,
+    ) -> Any:
+        def register(t: Any) -> Any:
+            if is_group:
+                if not isinstance(t, type):
+                    raise DagValidationError(
+                        "@dag.group expects a class with @member functions; "
+                        "use @dag.fn for a solo function"
+                    )
+                members = [
+                    v
+                    for v in vars(t).values()
+                    if callable(v) and getattr(v, "__dagrun_member__", False)
+                ]
+                if not members:
+                    raise DagValidationError(
+                        f"@dag.group {t.__name__} has no @member functions"
+                    )
+            else:
+                if isinstance(t, type):
+                    raise DagValidationError(
+                        "@dag.fn expects a function; use @dag.group for a class with members"
+                    )
+                members = [t]
+            auto_labels = {t.__name__}
+            merged_cost: dict[str, int] = {"concurrency": 1}
+            if cost is not None:
+                merged_cost.update(cost)
+            self._chains.append(
+                _Chain(
+                    members=members,
+                    combine=combine,
+                    labels=auto_labels | (labels or set()),
+                    cost=merged_cost,
+                )
+            )
+            return t
+
+        if target is None:
+            return register
+        return register(target)
+
+
+@dataclass
+class _ParsedChain:
+    raw: _Chain
+    produces: type | None
+    produces_cols: frozenset[str]
+    inputs: tuple[tuple[type, frozenset[str], bool], ...]
+    members: list[CompiledMember]
+
+
+def _chain_label(p: _ParsedChain) -> str:
+    fn = p.raw.members[0] if p.raw.members else None
+    return getattr(fn, "__qualname__", repr(p.raw))
+
+
+def _parse_chain(raw: _Chain) -> _ParsedChain:
+    members: list[CompiledMember] = []
+    produces: type | None = None
+    produces_cols_union: set[str] = set()
+    chain_inputs: tuple[tuple[type, frozenset[str], bool], ...] | None = None
+    any_member_has_same_entity_input = False
+
+    for fn in raw.members:
+        try:
+            hints = typing.get_type_hints(fn, include_extras=True)
+        except Exception:
+            hints = {}
+        sig = inspect.signature(fn)
+
+        needs_entities: list[tuple[type, frozenset[str], bool]] = []
+        needs_providers: list[type] = []
+        for pname, param in sig.parameters.items():
+            ann = hints.get(pname, param.annotation)
+            entity, cols, is_collection = _unwrap_input_param(ann)
+            if entity is not None:
+                needs_entities.append((entity, frozenset(cols), is_collection))
+            elif ann is not inspect.Parameter.empty and isinstance(ann, type):
+                needs_providers.append(ann)
+
+        ret_ann = hints.get("return", sig.return_annotation)
+        ret_entity, ret_cols = _unwrap_entity(ret_ann)
+
+        if ret_entity is not None:
+            if produces is None:
+                produces = ret_entity
+            elif produces is not ret_entity:
+                raise DagValidationError(
+                    f"Members of chain disagree on produced entity: "
+                    f"{produces.__name__} vs {ret_entity.__name__}"
+                )
+
+        if ret_entity is not None and any(
+            e is ret_entity for e, _, _ in needs_entities
+        ):
+            any_member_has_same_entity_input = True
+
+        if ret_entity is not None and any(
+            is_coll for _, _, is_coll in needs_entities
+        ):
+            any_member_has_same_entity_input = True
+
+        produced_cols = frozenset(ret_cols)
+        produces_cols_union |= set(ret_cols)
+
+        member_inputs = tuple(needs_entities)
+        if chain_inputs is None:
+            chain_inputs = member_inputs
+        elif {e for e, _, _ in chain_inputs} != {e for e, _, _ in member_inputs}:
+            raise DagValidationError(
+                "Members of group disagree on input entities"
+            )
+
+        members.append(
+            CompiledMember(
+                fn=fn,
+                needs_entities=tuple(
+                    (e, c) for e, c, _ in needs_entities
+                ),
+                needs_providers=tuple(needs_providers),
+                produces_columns=produced_cols,
+            )
+        )
+
+    if produces is not None and not any_member_has_same_entity_input:
+        pk_col = _find_pk(produces)
+        produces_cols_union.add(pk_col)
+
+    return _ParsedChain(
+        raw=raw,
+        produces=produces,
+        produces_cols=frozenset(produces_cols_union),
+        inputs=chain_inputs or (),
+        members=members,
+    )

dagrun/dag_runner.py

@@ -0,0 +1,59 @@
+from collections.abc import Callable, Iterable, Mapping
+from datetime import timedelta
+from typing import Any
+
+from dagrun.dag import Dag
+from dagrun.events import Observer
+from dagrun.model import _internalize_columns
+from dagrun.store import MemoryStore, Store
+from dagrun.strategy import SingleThread, Strategy
+
+
+class DagRunner:
+    def __init__(
+        self,
+        store: Store | None = None,
+        *,
+        strategy: Strategy | None = None,
+        observers: list[Observer] | None = None,
+    ) -> None:
+        self._store: Store = store if store is not None else MemoryStore()
+        self._strategy: Strategy = strategy if strategy is not None else SingleThread()
+        self._observers: list[Observer] = list(observers) if observers is not None else []
+        self._providers: dict[type, Callable[[], Any]] = {}
+        self._external_backends: dict[type, Any] = {}
+
+    def provide(self, kind: type, factory: Callable[[], Any]) -> None:
+        self._providers[kind] = factory
+
+    def register(self, column_kind: type, backend: Any) -> None:
+        self._external_backends[column_kind] = backend
+
+    def add_observer(self, observer: Observer) -> None:
+        self._observers.append(observer)
+
+    def execute(
+        self,
+        dag: Dag,
+        *,
+        labels: set[str] | None = None,
+        max_age: timedelta | None = None,
+        pools: Mapping[str, int] | None = None,
+    ) -> None:
+        plan = dag.compile()
+        if labels is not None:
+            plan = plan.filter(labels=labels)
+        self._strategy.run(
+            plan,
+            store=self._store,
+            external_backends=self._external_backends,
+            providers=self._providers,
+            observers=self._observers,
+            max_age=max_age,
+            pools=pools,
+        )
+
+    def get_entities(self, entity: type) -> Iterable[Any]:
+        for row in self._store.get_entities(entity):
+            hydrated = _internalize_columns(entity, row, self._external_backends)
+            yield entity(**hydrated)

dagrun/events.py

@@ -0,0 +1,96 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+@dataclass(frozen=True, kw_only=True)
+class Event:
+    timestamp: float
+
+
+@dataclass(frozen=True, kw_only=True)
+class ChainStarted(Event):
+    chain: Any
+
+
+@dataclass(frozen=True, kw_only=True)
+class ChainFinished(Event):
+    chain: Any
+
+
+@dataclass(frozen=True, kw_only=True)
+class MemberStarted(Event):
+    chain: Any
+    member: Any
+    input_pk: Any
+
+
+@dataclass(frozen=True, kw_only=True)
+class MemberFinished(Event):
+    chain: Any
+    member: Any
+    input_pk: Any
+    result: Any
+
+
+@dataclass(frozen=True, kw_only=True)
+class MemberFailed(Event):
+    chain: Any
+    member: Any
+    input_pk: Any
+    error: BaseException
+
+
+@dataclass(frozen=True, kw_only=True)
+class RowWritten(Event):
+    entity: type
+    pk: Any
+    columns: frozenset[str]
+
+
+@dataclass(frozen=True, kw_only=True)
+class RunStarted(Event):
+    pass
+
+
+@dataclass(frozen=True, kw_only=True)
+class PoolsRegistered(Event):
+    """Published once at the start of a run with the live pool gauges.
+
+    `capacity` is the per-pool token cap; `available` is the live remaining-token
+    mapping the scheduler mutates as tasks acquire and release tokens. Both are
+    references to the scheduler's own dicts so an observer can poll utilization
+    without per-task events. Pool keys are fixed at startup, so reading these
+    without the scheduler lock is safe for display purposes.
+    """
+
+    capacity: Mapping[str, int]
+    available: Mapping[str, int]
+
+
+@dataclass(frozen=True, kw_only=True)
+class RunFinished(Event):
+    error: BaseException | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class TaskEnqueued(Event):
+    """A (chain, input_pk) task was scheduled. Grows a chain's denominator."""
+
+    chain: Any
+
+
+@dataclass(frozen=True, kw_only=True)
+class TaskCompleted(Event):
+    """A task finished, either by running or by being skipped as fresh.
+
+    `skipped` is True when max_age suppressed the run (the work was done in a
+    prior run and is current in the store). Both cases grow a chain's numerator.
+    """
+
+    chain: Any
+    skipped: bool
+
+
+class Observer(Protocol):
+    def on_event(self, event: Event) -> None: ...