Implement Cluster 1 Lakeflow control-flow fidelity

README.md

@@ -106,6 +106,29 @@ Schema flags:
 
 `partitionBy` columns must exist on the DataFrame. `option("replaceWhere", "<predicate>")` with `mode("overwrite")` deletes matching stored rows then appends the new frame. `bucketBy` / `sortBy` are accepted no-ops (bucketing is not simulated). `df.writeTo(table).using(...).create()` / `.replace()` / `.append()` maps onto the same table writer; `createOrReplace` and `overwritePartitions` raise `NotImplementedError` with a `saveAsTable` hint.
 
+## Repair-and-rerun
+
+`LocalWorkflowRunner.run_workflow` can re-run a subgraph without executing the full DAG:
+
+```python
+runner.run_workflow(extra_globals={"spark": spark}, only=["build_summary"])
+runner.run_workflow(extra_globals={"spark": spark}, from_task="enrich_customers")
+```
+
+- `only=["task_a", ...]` runs those task keys in topological order.
+- `from_task="task_a"` runs that task and every downstream dependent.
+- Tasks not selected are treated as already `SUCCESS` so `run_if` / `depends_on` still resolve.
+
+Overwrite table writes (`mode("overwrite")` / default) are naturally idempotent on re-run. `append` writes are not: a repair will insert another copy of the rows. For append tables, use a repair policy in the notebook (skip if the target already has the batch, dedupe on a key, or pass a flag that switches the write to overwrite for that run).
+
+Lakeflow JSON extras the runner understands:
+
+- `dbutils.jobs.taskValues.set` / `.get` across tasks (shared immediately; isolated `notebook.run` commits on return)
+- `run_if` and `depends_on[].outcome` (`ALL_SUCCESS`, `ALL_FAILED`, `AT_LEAST_ONE_SUCCESS`, `ALL_DONE`, `NONE_FAILED`, `AT_LEAST_ONE_FAILED`); ineligible tasks are `SKIPPED`
+- `condition_task` (`EQUAL` / `EQUAL_TO`, `NOT_EQUAL`, `GREATER_THAN`, …) with `{{tasks.<key>.values.<name>}}` operands
+- `for_each_task` sequential expansion of `inputs` with `{{input}}` parameters
+- `max_retries` / `min_retry_interval_millis` (retry notebook and for_each tasks); `timeout_seconds` is accepted and logged, not enforced
+
 ## Key Modules
 1. `SparkProxy` - A Spark proxy that manipulates incoming Delta table reads and writes and redirects them to interactions with CSV files stored locally
 2. `LocalWorkflowRunner` - A notebook orchestrator that takes the notebook .py files as defined in a Databricks Workflow JSON file and executes them as per the DAG definition. Databricks comment magics `%run`, `%sh` (`# %sh` / `# MAGIC %sh`, including `%sh -e`), and `%fs` (`# %fs` / `# MAGIC %fs`) work in those notebooks.

docs/superpowers/specs/2026-09-03-spark-write-lakeflow-roadmap-design.md

@@ -33,12 +33,12 @@ Storage stays `{base_path}/{schema}/{table}.csv` + temp views `{schema}_{table}`
 
 ### Cluster 1 — Lakeflow control flow (top priority, in order)
 
-- [ ] **T1. `taskValues` propagation between tasks** — add `dbutils.jobs.taskValues.set(key, value)` / `.get(task_key, key)` backed by a runner-scoped dict, seeded from and visible to `base_parameters`; define semantics for shared-namespace `run_workflow` (visible immediately) vs isolated `dbutils.notebook.run()` (visible after return). Small dict + widget/env bridge, no persistence.
-- [ ] **T2. `run_if` / `depends_on` outcome conditions** — support `depends_on[].outcome` / task-level `run_if` (`ALL_SUCCESS`, `ALL_FAILED`, `AT_LEAST_ONE_SUCCESS`, `ALL_DONE`, `NONE_FAILED`); skip non-eligible tasks with a clear `SKIPPED` status instead of running them; keep topological order. Pure scheduler logic in `_build_graphs` + `run_workflow` loop.
-- [ ] **T3. `if/else` condition tasks** — support `condition_task` (`op`: `EQUAL`, `NOT_EQUAL`, `GREATER_THAN`, etc. over left/right operands referencing `taskValues`/params); route to the matching branch and mark the other branch skipped. Reuses T1 values + T2 skip machinery.
-- [ ] **T4. `for_each` parameter expansion** — support `for_each_task` with `inputs` (literal list or `taskValues` reference) + `concurrency`; execute the nested notebook once per input with per-iteration params/env isolation; iterate sequentially first, parallel as a later stretch. Builds on executor `run_isolated` semantics.
-- [ ] **T5. Repair-and-rerun + idempotent re-execution** — allow `run_workflow(only=["task_a"])` / `from_task` to re-run a subgraph; document that `overwrite` saves are naturally idempotent while `append` saves need a `repair` policy (skip, dedupe-key, or explicit overwrite flag). Runner-only change + docs, no catalog change.
-- [ ] **T6. Per-task retries / timeouts** — parse `max_retries` / `min_retry_interval_millis` / `timeout_seconds` from workflow JSON; retry failed notebook tasks N times, accept-but-log timeout (matching current `dbutils.notebook.run(timeout)` behavior of accept-not-enforce, or enforce via signals as a stretch). Loop wrapper in `run_workflow`.
+- [x] **T1. `taskValues` propagation between tasks** — add `dbutils.jobs.taskValues.set(key, value)` / `.get(task_key, key)` backed by a runner-scoped dict, seeded from and visible to `base_parameters`; define semantics for shared-namespace `run_workflow` (visible immediately) vs isolated `dbutils.notebook.run()` (visible after return). Small dict + widget/env bridge, no persistence.
+- [x] **T2. `run_if` / `depends_on` outcome conditions** — support `depends_on[].outcome` / task-level `run_if` (`ALL_SUCCESS`, `ALL_FAILED`, `AT_LEAST_ONE_SUCCESS`, `ALL_DONE`, `NONE_FAILED`); skip non-eligible tasks with a clear `SKIPPED` status instead of running them; keep topological order. Pure scheduler logic in `_build_graphs` + `run_workflow` loop.
+- [x] **T3. `if/else` condition tasks** — support `condition_task` (`op`: `EQUAL`, `NOT_EQUAL`, `GREATER_THAN`, etc. over left/right operands referencing `taskValues`/params); route to the matching branch and mark the other branch skipped. Reuses T1 values + T2 skip machinery.
+- [x] **T4. `for_each` parameter expansion** — support `for_each_task` with `inputs` (literal list or `taskValues` reference) + `concurrency`; execute the nested notebook once per input with per-iteration params/env isolation; iterate sequentially first, parallel as a later stretch. Builds on executor `run_isolated` semantics.
+- [x] **T5. Repair-and-rerun + idempotent re-execution** — allow `run_workflow(only=["task_a"])` / `from_task` to re-run a subgraph; document that `overwrite` saves are naturally idempotent while `append` saves need a `repair` policy (skip, dedupe-key, or explicit overwrite flag). Runner-only change + docs, no catalog change.
+- [x] **T6. Per-task retries / timeouts** — parse `max_retries` / `min_retry_interval_millis` / `timeout_seconds` from workflow JSON; retry failed notebook tasks N times, accept-but-log timeout (matching current `dbutils.notebook.run(timeout)` behavior of accept-not-enforce, or enforce via signals as a stretch). Loop wrapper in `run_workflow`.
 
 ### Cluster 2 — `spark.write` fidelity (in order)
 

src/testbricks/dbutils/dbutils_mock.py

@@ -1,5 +1,6 @@
 from .data import DataMock
 from .fs import FsMock
+from .jobs import JobsMock
 from .library import LibraryMock
 from .noop import NoOpModule
 from .notebook import NotebookMock
@@ -19,7 +20,7 @@ def __init__(self):
         self.secrets = SecretsMock()
         self.widgets = WidgetsMock()
         self.notebook = NotebookMock(self._executor)
-        self.jobs = NoOpModule()
+        self.jobs = JobsMock()
         self.library = LibraryMock()
         self.data = DataMock()
 

src/testbricks/dbutils/jobs.py

@@ -0,0 +1,107 @@
+import os
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+from .errors import DbutilsError
+
+_MISSING = object()
+_current_task_key: ContextVar[str | None] = ContextVar(
+    "task_values_current_task", default=None
+)
+
+
+def _stringify(value):
+    return value if isinstance(value, str) else str(value)
+
+
+class TaskValuesStore:
+    """In-memory `dbutils.jobs.taskValues` store for a single workflow run.
+
+    Shared-namespace workflow tasks write through immediately. Isolated
+    `dbutils.notebook.run()` calls buffer `set` until the child returns.
+    """
+
+    def __init__(self):
+        self.clear()
+
+    def clear(self):
+        self._committed: dict[str, dict[str, str]] = {}
+        self._buffers: list[dict[str, dict[str, str]]] = []
+
+    @contextmanager
+    def current_task(self, task_key):
+        token = _current_task_key.set(task_key)
+        try:
+            yield
+        finally:
+            _current_task_key.reset(token)
+
+    @contextmanager
+    def isolated_context(self):
+        self._buffers.append({})
+        try:
+            yield
+        except BaseException:
+            self._buffers.pop()
+            raise
+        inner = self._buffers.pop()
+        parent = self._buffers[-1] if self._buffers else self._committed
+        for task_key, values in inner.items():
+            parent.setdefault(task_key, {}).update(values)
+
+    def _write_target(self):
+        return self._buffers[-1] if self._buffers else self._committed
+
+    def _read_values(self, task_key):
+        merged = dict(self._committed.get(task_key, {}))
+        current = _current_task_key.get()
+        if task_key == current:
+            for buffer in self._buffers:
+                merged.update(buffer.get(task_key, {}))
+        return merged
+
+    def set(self, key=None, value=None, *, update_env=True):
+        if key is None:
+            raise DbutilsError("taskValues.set requires 'key'")
+        if value is None:
+            raise DbutilsError("taskValues.set requires 'value'")
+        task_key = _current_task_key.get()
+        if not task_key:
+            task_key = ""
+        rendered = _stringify(value)
+        self._write_target().setdefault(task_key, {})[str(key)] = rendered
+        if update_env:
+            os.environ[str(key)] = rendered
+        return None
+
+    def get(
+        self,
+        taskKey=None,
+        key=None,
+        default=_MISSING,
+        debugValue=_MISSING,
+        task_key=None,
+    ):
+        resolved_task = taskKey if taskKey is not None else task_key
+        if resolved_task is None:
+            raise DbutilsError("taskValues.get requires 'taskKey'")
+        if key is None:
+            raise DbutilsError("taskValues.get requires 'key'")
+
+        values = self._read_values(resolved_task)
+        if key in values:
+            return values[key]
+
+        in_job = _current_task_key.get() is not None
+        if not in_job and debugValue is not _MISSING:
+            return _stringify(debugValue)
+        if default is not _MISSING:
+            return _stringify(default)
+        raise DbutilsError(
+            f"Task value '{key}' not found for task '{resolved_task}'"
+        )
+
+
+class JobsMock:
+    def __init__(self):
+        self.taskValues = TaskValuesStore()