diff --git a/README.md b/README.md index f5ff75a..6131f15 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,29 @@ Schema flags: `partitionBy` columns must exist on the DataFrame. `option("replaceWhere", "")` 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..values.}}` 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. diff --git a/docs/superpowers/specs/2026-09-03-spark-write-lakeflow-roadmap-design.md b/docs/superpowers/specs/2026-09-03-spark-write-lakeflow-roadmap-design.md index 94d8bb4..7e0d040 100644 --- a/docs/superpowers/specs/2026-09-03-spark-write-lakeflow-roadmap-design.md +++ b/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) diff --git a/src/testbricks/dbutils/dbutils_mock.py b/src/testbricks/dbutils/dbutils_mock.py index c9e27fa..548b385 100644 --- a/src/testbricks/dbutils/dbutils_mock.py +++ b/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() diff --git a/src/testbricks/dbutils/jobs.py b/src/testbricks/dbutils/jobs.py new file mode 100644 index 0000000..11a242c --- /dev/null +++ b/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() diff --git a/src/testbricks/local_workflow_runner.py b/src/testbricks/local_workflow_runner.py index 18a0f41..3ebe769 100644 --- a/src/testbricks/local_workflow_runner.py +++ b/src/testbricks/local_workflow_runner.py @@ -1,5 +1,7 @@ import json import os +import re +import time from contextlib import contextmanager from graphlib import CycleError, TopologicalSorter @@ -7,12 +9,87 @@ __all__ = ["LocalWorkflowRunner", "transform_run_commands"] +RUN_IF_VALUES = { + "ALL_SUCCESS", + "ALL_FAILED", + "AT_LEAST_ONE_SUCCESS", + "ALL_DONE", + "NONE_FAILED", + "AT_LEAST_ONE_FAILED", +} + +CONDITION_OPS = { + "EQUAL", + "EQUAL_TO", + "NOT_EQUAL", + "NOT_EQUAL_TO", + "GREATER_THAN", + "GREATER_THAN_OR_EQUAL", + "LESS_THAN", + "LESS_THAN_OR_EQUAL", +} + +TASK_VALUE_RE = re.compile(r"^\{\{\s*tasks\.([^.]+)\.values\.([^}]+?)\s*\}\}$") +JOB_PARAM_RE = re.compile(r"^\{\{\s*job\.parameters\.([^}]+?)\s*\}\}$") + def _require(condition, message): if not condition: raise ValueError(message) +def _matches_run_if(run_if, statuses): + if not statuses: + return True + if run_if == "ALL_SUCCESS": + return all(status == "SUCCESS" for status in statuses) + if run_if == "ALL_FAILED": + return all(status == "FAILED" for status in statuses) + if run_if == "AT_LEAST_ONE_SUCCESS": + return any(status == "SUCCESS" for status in statuses) + if run_if == "ALL_DONE": + return True + if run_if == "NONE_FAILED": + return all(status != "FAILED" for status in statuses) + if run_if == "AT_LEAST_ONE_FAILED": + return any(status == "FAILED" for status in statuses) + return False + + +def _as_numbers(left, right): + try: + return float(left), float(right) + except (TypeError, ValueError): + return None + + +def _compare_condition(op, left, right): + normalized = op.upper() + nums = _as_numbers(left, right) + if normalized in {"EQUAL", "EQUAL_TO"}: + if nums is not None: + return nums[0] == nums[1] + return str(left) == str(right) + if normalized in {"NOT_EQUAL", "NOT_EQUAL_TO"}: + if nums is not None: + return nums[0] != nums[1] + return str(left) != str(right) + if nums is None: + raise ValueError( + f"Condition op '{op}' requires numeric operands, got {left!r} and {right!r}" + ) + left_n, right_n = nums + if normalized == "GREATER_THAN": + return left_n > right_n + if normalized == "GREATER_THAN_OR_EQUAL": + return left_n >= right_n + if normalized == "LESS_THAN": + return left_n < right_n + if normalized == "LESS_THAN_OR_EQUAL": + return left_n <= right_n + raise ValueError(f"Unsupported condition op '{op}'") + + class LocalWorkflowRunner: def __init__(self, source_dir, workflow_json_path, base_path): self.source_dir = source_dir @@ -25,9 +102,22 @@ def __init__(self, source_dir, workflow_json_path, base_path): _require(isinstance(tasks, list), "Workflow JSON must contain a 'tasks' list") self._task_to_notebook = {} + self._notebook_to_task = {} self._task_dependencies = {} + self._task_dep_specs = {} + self._task_run_if = {} + self._task_kind = {} + self._condition_task = {} + self._for_each_task = {} + self._task_max_retries = {} + self._task_retry_interval_ms = {} + self._task_timeout_seconds = {} + self._task_insertion_order = [] self._notebook_insertion_order = [] + self._task_base_params = {} self._notebook_base_params = {} + self.task_statuses = {} + self.task_results = {} self._parse_tasks(tasks) self.dag, predecessors = self._build_graphs() @@ -42,28 +132,78 @@ def _parse_tasks(self, tasks): task_key = task.get("task_key") _require(task_key, "Each task must include a non-empty 'task_key'") _require( - task_key not in self._task_to_notebook, + task_key not in self._task_kind, f"Duplicate task_key found: {task_key}", ) notebook_task = task.get("notebook_task") + condition_task = task.get("condition_task") + for_each_task = task.get("for_each_task") + has_notebook = isinstance(notebook_task, dict) + has_condition = isinstance(condition_task, dict) + has_for_each = isinstance(for_each_task, dict) _require( - isinstance(notebook_task, dict), + has_notebook or has_condition or has_for_each, f"Task '{task_key}' is missing 'notebook_task'", ) - notebook_name = self._extract_notebook_name( - notebook_task.get("notebook_path"), task_key - ) _require( - notebook_name not in self._notebook_insertion_order, - f"Duplicate notebook name found: {notebook_name}", + [has_notebook, has_condition, has_for_each].count(True) == 1, + f"Task '{task_key}' cannot mix notebook, condition, and for_each tasks", ) - base_parameters = notebook_task.get("base_parameters", {}) - _require( - isinstance(base_parameters, dict), - f"Task '{task_key}' has invalid 'base_parameters' format", - ) + notebook_name = None + base_parameters = {} + if has_notebook: + notebook_name = self._extract_notebook_name( + notebook_task.get("notebook_path"), task_key + ) + _require( + notebook_name not in self._notebook_insertion_order, + f"Duplicate notebook name found: {notebook_name}", + ) + base_parameters = notebook_task.get("base_parameters", {}) + _require( + isinstance(base_parameters, dict), + f"Task '{task_key}' has invalid 'base_parameters' format", + ) + elif has_condition: + _require( + condition_task.get("op"), + f"Task '{task_key}' has invalid 'condition_task'", + ) + op = str(condition_task.get("op")).upper() + _require( + op in CONDITION_OPS, + f"Unsupported condition op '{condition_task.get('op')}' on task '{task_key}'", + ) + _require( + "left" in condition_task and "right" in condition_task, + f"Task '{task_key}' condition_task requires 'left' and 'right'", + ) + condition_task = {**condition_task, "op": op} + else: + nested = for_each_task.get("task") + _require( + isinstance(nested, dict), + f"Task '{task_key}' for_each_task requires a nested 'task'", + ) + nested_notebook = nested.get("notebook_task") + _require( + isinstance(nested_notebook, dict), + f"Task '{task_key}' for_each nested task is missing 'notebook_task'", + ) + self._extract_notebook_name( + nested_notebook.get("notebook_path"), task_key + ) + nested_params = nested_notebook.get("base_parameters", {}) + _require( + isinstance(nested_params, dict), + f"Task '{task_key}' has invalid nested 'base_parameters' format", + ) + _require( + "inputs" in for_each_task, + f"Task '{task_key}' for_each_task requires 'inputs'", + ) depends_on = task.get("depends_on", []) _require( @@ -71,17 +211,52 @@ def _parse_tasks(self, tasks): f"Task '{task_key}' has invalid 'depends_on' format", ) dependency_keys = [] + dep_specs = [] for dependency in depends_on: _require( isinstance(dependency, dict) and dependency.get("task_key"), f"Task '{task_key}' has malformed dependency entry", ) dependency_keys.append(dependency["task_key"]) + dep_specs.append((dependency["task_key"], dependency.get("outcome"))) + + run_if = task.get("run_if") or "ALL_SUCCESS" + _require( + isinstance(run_if, str) and run_if.upper() in RUN_IF_VALUES, + f"Unsupported run_if '{run_if}' on task '{task_key}'", + ) + if has_condition: + kind = "condition" + elif has_for_each: + kind = "for_each" + else: + kind = "notebook" + self._task_kind[task_key] = kind + self._condition_task[task_key] = condition_task if has_condition else None + self._for_each_task[task_key] = for_each_task if has_for_each else None self._task_to_notebook[task_key] = notebook_name + if notebook_name is not None: + self._notebook_to_task[notebook_name] = task_key + self._notebook_insertion_order.append(notebook_name) + self._notebook_base_params[notebook_name] = base_parameters self._task_dependencies[task_key] = dependency_keys - self._notebook_insertion_order.append(notebook_name) - self._notebook_base_params[notebook_name] = base_parameters + self._task_dep_specs[task_key] = dep_specs + self._task_run_if[task_key] = run_if.upper() + self._task_max_retries[task_key] = self._parse_non_negative_int( + task.get("max_retries"), 0, "max_retries", task_key + ) + self._task_retry_interval_ms[task_key] = self._parse_non_negative_int( + task.get("min_retry_interval_millis"), + 0, + "min_retry_interval_millis", + task_key, + ) + self._task_timeout_seconds[task_key] = self._parse_non_negative_int( + task.get("timeout_seconds"), 0, "timeout_seconds", task_key + ) + self._task_insertion_order.append(task_key) + self._task_base_params[task_key] = base_parameters def _extract_notebook_name(self, notebook_path, task_key): _require( @@ -92,30 +267,148 @@ def _extract_notebook_name(self, notebook_path, task_key): _require(notebook_name, f"Task '{task_key}' has an invalid notebook_path") return notebook_name + def _parse_non_negative_int(self, raw, default, field, task_key): + if raw is None: + return default + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Task '{task_key}' has invalid '{field}': {raw!r}" + ) from exc + _require(value >= 0, f"Task '{task_key}' has invalid '{field}': {raw!r}") + return value + def _build_graphs(self): - successors = {name: set() for name in self._notebook_insertion_order} - predecessors = {name: set() for name in self._notebook_insertion_order} + successors = {key: set() for key in self._task_insertion_order} + predecessors = {key: set() for key in self._task_insertion_order} for task_key, dependency_keys in self._task_dependencies.items(): - current = self._task_to_notebook[task_key] for dependency_key in dependency_keys: _require( - dependency_key in self._task_to_notebook, + dependency_key in self._task_kind, f"Task '{task_key}' depends on unknown task '{dependency_key}'", ) - dependency = self._task_to_notebook[dependency_key] - successors[dependency].add(current) - predecessors[current].add(dependency) + successors[dependency_key].add(task_key) + predecessors[task_key].add(dependency_key) return successors, predecessors + def _outcome_matches(self, dep_key, status, outcome): + wanted = str(outcome).upper() + if wanted in {"SUCCESS", "FAILED", "SKIPPED"}: + return status == wanted + result = self.task_results.get(dep_key) + if result is not None: + return str(result).lower() == str(outcome).lower() + return status == wanted + + def _is_eligible(self, task_key): + specs = self._task_dep_specs.get(task_key, []) + if not specs: + return True + for dep_key, outcome in specs: + status = self.task_statuses.get(dep_key) + if status is None: + return False + if outcome is not None and not self._outcome_matches( + dep_key, status, outcome + ): + return False + statuses = [self.task_statuses[dep_key] for dep_key, _ in specs] + return _matches_run_if(self._task_run_if[task_key], statuses) + + def _resolve_operand(self, raw, store): + if raw is None: + return "" + if not isinstance(raw, str): + return str(raw) + text = raw.strip() + match = TASK_VALUE_RE.match(text) + if match: + return store.get(taskKey=match.group(1), key=match.group(2).strip()) + match = JOB_PARAM_RE.match(text) + if match: + key = match.group(1).strip() + if key not in os.environ: + raise ValueError(f"Job parameter '{key}' not found") + return os.environ[key] + return raw + + def _evaluate_condition(self, task_key, store): + spec = self._condition_task[task_key] + left = self._resolve_operand(spec.get("left"), store) + right = self._resolve_operand(spec.get("right"), store) + matched = _compare_condition(spec["op"], left, right) + return "true" if matched else "false" + def format_dag(self): lines = [] - for notebook_name in self._notebook_insertion_order: - outgoing = sorted(self.dag.get(notebook_name, set())) - lines.append(f"- {notebook_name} -> [{', '.join(outgoing)}]") + for task_key in self._task_insertion_order: + outgoing = sorted(self.dag.get(task_key, set())) + lines.append(f"- {task_key} -> [{', '.join(outgoing)}]") lines.append("Execution order:") lines.append(" -> ".join(self.execution_order)) return "\n".join(lines) + def _selected_tasks(self, only, from_task): + if only is not None and from_task is not None: + raise ValueError("Pass only one of 'only' or 'from_task'") + if only is None and from_task is None: + return None + if only is not None: + selected = set(only) + unknown = selected - set(self._task_kind) + _require(not unknown, f"unknown task in only: {sorted(unknown)}") + return selected + _require( + from_task in self._task_kind, + f"unknown task in from_task: '{from_task}'", + ) + selected = {from_task} + stack = [from_task] + while stack: + current = stack.pop() + for successor in self.dag.get(current, ()): + if successor not in selected: + selected.add(successor) + stack.append(successor) + return selected + + def _run_task_with_retries( + self, task_key, executor, store, execution_globals, extra_globals + ): + timeout = self._task_timeout_seconds.get(task_key) or 0 + if timeout: + print( + f"Task '{task_key}' timeout_seconds={timeout} accepted but not enforced" + ) + + def action(): + kind = self._task_kind[task_key] + if kind == "condition": + self.task_results[task_key] = self._evaluate_condition(task_key, store) + elif kind == "for_each": + self._run_for_each_task(task_key, executor, store, extra_globals) + else: + self._run_notebook_task(task_key, executor, store, execution_globals) + + retries = self._task_max_retries.get(task_key, 0) + interval_ms = self._task_retry_interval_ms.get(task_key, 0) + attempt = 0 + while True: + try: + action() + return + except Exception as exc: + attempt += 1 + if attempt > retries: + raise + print( + f"Task '{task_key}' failed " + f"(attempt {attempt}/{retries + 1}), retrying: {exc}" + ) + if interval_ms: + time.sleep(interval_ms / 1000.0) + def _executor(self): from testbricks.dbutils import dbutils @@ -141,20 +434,128 @@ def _seeded_env(self, values): else: os.environ[key] = original - def run_workflow(self, extra_globals=None): - from testbricks.dbutils import configure, dbutils + def _run_notebook_task(self, task_key, executor, store, execution_globals): + notebook_name = self._task_to_notebook[task_key] + notebook_path = os.path.join(self.source_dir, f"{notebook_name}.py") + if not os.path.exists(notebook_path): + raise FileNotFoundError(f"Notebook file not found: {notebook_path}") from testbricks.dbutils.widgets import argument_override_context + base_params = self._task_base_params.get(task_key, {}) + with ( + self._seeded_env(base_params), + argument_override_context(base_params.keys()), + store.current_task(task_key), + ): + for param_key, param_value in base_params.items(): + store.set(key=param_key, value=param_value, update_env=False) + executor.exec_file(notebook_path, execution_globals, top_level=True) + + def _input_as_str(self, item): + if isinstance(item, (dict, list)): + return json.dumps(item) + return str(item) + + def _render_input_template(self, value, item, index): + if not isinstance(value, str): + return str(value) + text = value.strip() + if re.match(r"^\{\{\s*input\s*\}\}$", text): + return self._input_as_str(item) + field = re.match(r"^\{\{\s*input\.([^}]+?)\s*\}\}$", text) + if field: + key = field.group(1).strip() + if isinstance(item, dict): + return self._input_as_str(item.get(key)) + raise ValueError(f"for_each input is not an object; cannot read '{key}'") + if "{{input}}" in value: + return value.replace("{{input}}", self._input_as_str(item)) + return value + + def _resolve_for_each_inputs(self, raw, store): + if isinstance(raw, list): + return raw + if not isinstance(raw, str): + raise ValueError(f"for_each inputs must be a JSON list, got {raw!r}") + text = raw.strip() + match = TASK_VALUE_RE.match(text) + if match: + text = store.get(taskKey=match.group(1), key=match.group(2).strip()) + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError(f"for_each inputs must be a JSON list: {raw!r}") from exc + if not isinstance(parsed, list): + raise ValueError(f"for_each inputs must be a JSON list: {raw!r}") + return parsed + + def _run_for_each_task(self, task_key, executor, store, extra_globals): + spec = self._for_each_task[task_key] + inputs = self._resolve_for_each_inputs(spec.get("inputs"), store) + nested = spec["task"] + notebook_task = nested["notebook_task"] + notebook_name = self._extract_notebook_name( + notebook_task.get("notebook_path"), task_key + ) + nested_key = nested.get("task_key") or task_key + base_parameters = notebook_task.get("base_parameters", {}) or {} + caller = os.path.join(self.source_dir, "_workflow.py") + for index, item in enumerate(inputs): + params = { + key: self._render_input_template(value, item, index) + for key, value in base_parameters.items() + } + saved_env = {key: os.environ.get(key) for key in params} + try: + with ( + store.current_task(nested_key), + executor.caller_context(caller), + ): + executor.run_isolated( + f"/Workspace/{notebook_name}", + arguments=params, + extra=extra_globals, + ) + finally: + for key, original in saved_env.items(): + if original is None: + os.environ.pop(key, None) + else: + os.environ[key] = original + + def run_workflow(self, extra_globals=None, only=None, from_task=None): + from testbricks.dbutils import configure, dbutils + configure(self.base_path, source_dir=self.source_dir) executor = dbutils.executor + store = dbutils.jobs.taskValues + store.clear() execution_globals = executor.namespace("", extra=extra_globals) print(f"\nExecuting workflow: {self.workflow_json_path}\n") print("==========================================") print(self.format_dag()) - for notebook_name in self.execution_order: - notebook_path = os.path.join(self.source_dir, f"{notebook_name}.py") - if not os.path.exists(notebook_path): - raise FileNotFoundError(f"Notebook file not found: {notebook_path}") - base_params = self._notebook_base_params.get(notebook_name, {}) - with self._seeded_env(base_params), argument_override_context(base_params.keys()): - executor.exec_file(notebook_path, execution_globals, top_level=True) + self.task_statuses = {} + self.task_results = {} + selected = self._selected_tasks(only, from_task) + first_error = None + for task_key in self.execution_order: + if selected is not None and task_key not in selected: + self.task_statuses[task_key] = "SUCCESS" + continue + if not self._is_eligible(task_key): + self.task_statuses[task_key] = "SKIPPED" + print(f"Skipping task '{task_key}': run_if not met") + continue + try: + self._run_task_with_retries( + task_key, executor, store, execution_globals, extra_globals + ) + except Exception as exc: + self.task_statuses[task_key] = "FAILED" + if first_error is None: + first_error = exc + print(f"Task '{task_key}' failed: {exc}") + continue + self.task_statuses[task_key] = "SUCCESS" + if first_error is not None: + raise first_error diff --git a/src/testbricks/notebook_executor.py b/src/testbricks/notebook_executor.py index 0914fdb..aebea49 100644 --- a/src/testbricks/notebook_executor.py +++ b/src/testbricks/notebook_executor.py @@ -244,7 +244,7 @@ def run_shared(self, path, namespace): notebook_path = self.resolve_path(path, caller_file=namespace.get("__file__")) self.exec_file(notebook_path, namespace, top_level=False) - def run_isolated(self, path, arguments=None): + def run_isolated(self, path, arguments=None, extra=None): from testbricks.dbutils.widgets import argument_override_context arguments = arguments or {} @@ -252,8 +252,11 @@ def run_isolated(self, path, arguments=None): for key, value in arguments.items(): os.environ[key] = str(value) - namespace = self.namespace(notebook_path) - with argument_override_context(arguments.keys()): + namespace = self.namespace(notebook_path, extra=extra) + with ( + argument_override_context(arguments.keys()), + self._dbutils.jobs.taskValues.isolated_context(), + ): try: self.exec_file(notebook_path, namespace, top_level=False) except NotebookExit as exc: diff --git a/tests/test_dbutils_fs.py b/tests/test_dbutils_fs.py index cb9902d..8f8b612 100644 --- a/tests/test_dbutils_fs.py +++ b/tests/test_dbutils_fs.py @@ -212,7 +212,7 @@ def test_unimplemented_fs_method_returns_true(self, tmp_path): def test_unimplemented_submodule_returns_true(self, tmp_path): configure(str(tmp_path)) - assert dbutils.jobs.taskValues.get(key="k") is True + assert dbutils.credentials.get(key="k") is True class TestWorkflowRunnerIntegration: diff --git a/tests/test_dbutils_jobs.py b/tests/test_dbutils_jobs.py new file mode 100644 index 0000000..393d104 --- /dev/null +++ b/tests/test_dbutils_jobs.py @@ -0,0 +1,104 @@ +import sys +import os + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +import pytest + +from testbricks.dbutils import DbutilsError, configure, dbutils + + +@pytest.fixture(autouse=True) +def reset_dbutils(tmp_path): + configure(str(tmp_path), source_dir=str(tmp_path)) + dbutils.jobs.taskValues.clear() + yield + dbutils.jobs.taskValues.clear() + dbutils.widgets.removeAll() + + +class TestTaskValuesSetGet: + def test_set_and_get_within_current_task(self): + with dbutils.jobs.taskValues.current_task("producer"): + dbutils.jobs.taskValues.set(key="region", value="eu") + assert ( + dbutils.jobs.taskValues.get(taskKey="producer", key="region") == "eu" + ) + + def test_get_task_key_alias(self): + with dbutils.jobs.taskValues.current_task("producer"): + dbutils.jobs.taskValues.set(key="region", value="eu") + assert ( + dbutils.jobs.taskValues.get(task_key="producer", key="region") == "eu" + ) + + def test_non_string_value_is_stringified(self): + with dbutils.jobs.taskValues.current_task("producer"): + dbutils.jobs.taskValues.set(key="count", value=3) + assert dbutils.jobs.taskValues.get(taskKey="producer", key="count") == "3" + + def test_missing_key_raises_in_job(self): + with dbutils.jobs.taskValues.current_task("producer"): + with pytest.raises(DbutilsError, match="not found"): + dbutils.jobs.taskValues.get(taskKey="producer", key="missing") + + def test_default_used_when_key_missing_in_job(self): + with dbutils.jobs.taskValues.current_task("producer"): + assert ( + dbutils.jobs.taskValues.get( + taskKey="producer", key="missing", default="fallback" + ) + == "fallback" + ) + + def test_debug_value_used_outside_job(self): + assert ( + dbutils.jobs.taskValues.get( + taskKey="producer", key="region", debugValue="local" + ) + == "local" + ) + + def test_set_updates_environ(self): + with dbutils.jobs.taskValues.current_task("producer"): + dbutils.jobs.taskValues.set(key="bridge_key", value="from_task") + assert os.environ["bridge_key"] == "from_task" + os.environ.pop("bridge_key", None) + + def test_get_requires_task_key(self): + with pytest.raises(DbutilsError, match="taskKey"): + dbutils.jobs.taskValues.get(key="k") + + +class TestTaskValuesIsolation: + def test_isolated_sets_visible_after_context_exits(self): + store = dbutils.jobs.taskValues + with store.current_task("parent"): + with store.isolated_context(): + store.set(key="from_child", value="secret") + assert store.get(taskKey="parent", key="from_child") == "secret" + assert store.get(taskKey="parent", key="from_child") == "secret" + + def test_isolated_sets_not_committed_if_context_raises(self): + store = dbutils.jobs.taskValues + with store.current_task("parent"): + with pytest.raises(RuntimeError): + with store.isolated_context(): + store.set(key="from_child", value="secret") + raise RuntimeError("boom") + with pytest.raises(DbutilsError, match="not found"): + store.get(taskKey="parent", key="from_child") + + def test_isolated_notebook_run_commits_on_return(self, tmp_path): + child = tmp_path / "child.py" + child.write_text( + 'dbutils.jobs.taskValues.set(key="from_child", value="ok")\n', + encoding="utf-8", + ) + parent = tmp_path / "parent.py" + store = dbutils.jobs.taskValues + with store.current_task("parent"), dbutils.executor.caller_context( + str(parent) + ): + dbutils.notebook.run("./child", 60) + assert store.get(taskKey="parent", key="from_child") == "ok" diff --git a/tests/test_local_workflow_runner.py b/tests/test_local_workflow_runner.py index ef04aae..044b7c5 100644 --- a/tests/test_local_workflow_runner.py +++ b/tests/test_local_workflow_runner.py @@ -19,26 +19,26 @@ class TestSampleWorkflowDAG: def test_builds_dag_using_notebook_names(self): runner = LocalWorkflowRunner("src", WORKFLOW_SAMPLE_PATH, DEFAULT_BASE_PATH) - assert runner.dag["auxillary_dims"] == {"data_quality"} - assert runner.dag["reviews_fact"] == {"data_quality"} - assert runner.dag["data_quality"] == {"semantic_layer"} + assert runner.dag["dimensions"] == {"quality_checks"} + assert runner.dag["reviews_fact"] == {"quality_checks"} + assert runner.dag["quality_checks"] == {"semantic_layer"} assert runner.dag["semantic_layer"] == set() def test_execution_order_respects_dependencies(self): runner = LocalWorkflowRunner("src", WORKFLOW_SAMPLE_PATH, DEFAULT_BASE_PATH) order = runner.execution_order - assert order.index("data_quality") > order.index("auxillary_dims") - assert order.index("data_quality") > order.index("reviews_fact") - assert order.index("semantic_layer") > order.index("data_quality") + assert order.index("quality_checks") > order.index("dimensions") + assert order.index("quality_checks") > order.index("reviews_fact") + assert order.index("semantic_layer") > order.index("quality_checks") def test_format_dag_includes_nodes_and_execution_order(self): runner = LocalWorkflowRunner("src", WORKFLOW_SAMPLE_PATH, DEFAULT_BASE_PATH) formatted = runner.format_dag() - assert "- auxillary_dims -> [data_quality]" in formatted - assert "- reviews_fact -> [data_quality]" in formatted - assert "- data_quality -> [semantic_layer]" in formatted + assert "- dimensions -> [quality_checks]" in formatted + assert "- reviews_fact -> [quality_checks]" in formatted + assert "- quality_checks -> [semantic_layer]" in formatted assert "- semantic_layer -> []" in formatted assert "Execution order:" in formatted assert " -> ".join(runner.execution_order) in formatted @@ -68,7 +68,7 @@ def test_extract_notebook_name(self, tmp_path, notebook_path, expected): runner = LocalWorkflowRunner(str(tmp_path), str(workflow_path), str(tmp_path)) assert runner._task_to_notebook["task_1"] == expected - assert runner.execution_order == [expected] + assert runner.execution_order == ["task_1"] class TestWorkflowExecution: @@ -196,6 +196,822 @@ def test_base_parameters_do_not_clobber_existing_env(self, tmp_path): assert marker.read_text(encoding="utf-8") == "from_test" +class TestTaskValuesPropagation: + def test_downstream_task_reads_upstream_task_value(self, tmp_path): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + marker = tmp_path / "value.txt" + (source_dir / "producer.py").write_text( + 'dbutils.jobs.taskValues.set(key="region", value="eu")\n', + encoding="utf-8", + ) + (source_dir / "consumer.py").write_text( + 'value = dbutils.jobs.taskValues.get(taskKey="producer", key="region")\n' + f'with open(r"{marker}", "w") as f: f.write(value)\n', + encoding="utf-8", + ) + workflow = { + "tasks": [ + { + "task_key": "producer", + "notebook_task": {"notebook_path": "/Workspace/any/producer"}, + }, + { + "task_key": "consumer", + "depends_on": [{"task_key": "producer"}], + "notebook_task": {"notebook_path": "/Workspace/any/consumer"}, + }, + ] + } + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps(workflow), encoding="utf-8") + + runner = LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + runner.run_workflow() + + assert marker.read_text(encoding="utf-8") == "eu" + + def test_base_parameters_seed_task_values(self, tmp_path): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + marker = tmp_path / "seeded_tv.txt" + (source_dir / "main.py").write_text( + 'value = dbutils.jobs.taskValues.get(taskKey="task_1", key="mode")\n' + f'with open(r"{marker}", "w") as f: f.write(value)\n', + encoding="utf-8", + ) + workflow = { + "tasks": [ + { + "task_key": "task_1", + "notebook_task": { + "notebook_path": "/Workspace/any/main", + "base_parameters": {"mode": "default"}, + }, + } + ] + } + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps(workflow), encoding="utf-8") + + runner = LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + runner.run_workflow() + + assert marker.read_text(encoding="utf-8") == "default" + + def test_shared_namespace_sees_set_immediately_in_later_task(self, tmp_path): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + marker = tmp_path / "shared.txt" + (source_dir / "first.py").write_text( + 'dbutils.jobs.taskValues.set(key="flag", value="on")\n', + encoding="utf-8", + ) + (source_dir / "second.py").write_text( + 'value = dbutils.jobs.taskValues.get(taskKey="first_task", key="flag")\n' + f'with open(r"{marker}", "w") as f: f.write(value)\n', + encoding="utf-8", + ) + workflow = { + "tasks": [ + { + "task_key": "first_task", + "notebook_task": {"notebook_path": "/Workspace/any/first"}, + }, + { + "task_key": "second_task", + "depends_on": [{"task_key": "first_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/second"}, + }, + ] + } + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps(workflow), encoding="utf-8") + + runner = LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + runner.run_workflow() + + assert marker.read_text(encoding="utf-8") == "on" + + +def _write_notebooks(source_dir, mapping): + for name, body in mapping.items(): + (source_dir / f"{name}.py").write_text(body, encoding="utf-8") + + +class TestRunIfConditions: + def _runner(self, tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + + def test_all_success_skips_when_dependency_fails(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "downstream", + "run_if": "ALL_SUCCESS", + "depends_on": [{"task_key": "fail_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/down"}, + }, + ], + { + "fail": "raise RuntimeError('boom')\n", + "down": f'with open(r"{log_file}", "a") as f: f.write("down\\n")\n', + }, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["fail_task"] == "FAILED" + assert runner.task_statuses["downstream"] == "SKIPPED" + assert not log_file.exists() + + def test_all_failed_runs_when_dependency_fails(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "cleanup", + "run_if": "ALL_FAILED", + "depends_on": [{"task_key": "fail_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/cleanup"}, + }, + ], + { + "fail": "raise RuntimeError('boom')\n", + "cleanup": f'with open(r"{log_file}", "a") as f: f.write("cleanup\\n")\n', + }, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["fail_task"] == "FAILED" + assert runner.task_statuses["cleanup"] == "SUCCESS" + assert log_file.read_text(encoding="utf-8").splitlines() == ["cleanup"] + + def test_all_failed_skips_when_dependency_succeeds(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "ok_task", + "notebook_task": {"notebook_path": "/Workspace/any/ok"}, + }, + { + "task_key": "cleanup", + "run_if": "ALL_FAILED", + "depends_on": [{"task_key": "ok_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/cleanup"}, + }, + ], + { + "ok": "pass\n", + "cleanup": f'with open(r"{log_file}", "a") as f: f.write("cleanup\\n")\n', + }, + ) + runner.run_workflow() + assert runner.task_statuses["ok_task"] == "SUCCESS" + assert runner.task_statuses["cleanup"] == "SKIPPED" + assert not log_file.exists() + + def test_all_done_runs_after_failed_dependency(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "always", + "run_if": "ALL_DONE", + "depends_on": [{"task_key": "fail_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/always"}, + }, + ], + { + "fail": "raise RuntimeError('boom')\n", + "always": f'with open(r"{log_file}", "a") as f: f.write("always\\n")\n', + }, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["always"] == "SUCCESS" + assert log_file.read_text(encoding="utf-8").splitlines() == ["always"] + + def test_none_failed_skips_when_dependency_fails(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "next", + "run_if": "NONE_FAILED", + "depends_on": [{"task_key": "fail_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/next"}, + }, + ], + {"fail": "raise RuntimeError('boom')\n", "next": "pass\n"}, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["next"] == "SKIPPED" + + def test_none_failed_runs_when_dependency_skipped(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "mid", + "run_if": "ALL_SUCCESS", + "depends_on": [{"task_key": "fail_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/mid"}, + }, + { + "task_key": "tail", + "run_if": "NONE_FAILED", + "depends_on": [{"task_key": "mid"}], + "notebook_task": {"notebook_path": "/Workspace/any/tail"}, + }, + ], + { + "fail": "raise RuntimeError('boom')\n", + "mid": "pass\n", + "tail": "pass\n", + }, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["mid"] == "SKIPPED" + assert runner.task_statuses["tail"] == "SUCCESS" + + def test_at_least_one_success_runs_if_any_dep_succeeded(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "ok_task", + "notebook_task": {"notebook_path": "/Workspace/any/ok"}, + }, + { + "task_key": "fail_task", + "notebook_task": {"notebook_path": "/Workspace/any/fail"}, + }, + { + "task_key": "join", + "run_if": "AT_LEAST_ONE_SUCCESS", + "depends_on": [ + {"task_key": "ok_task"}, + {"task_key": "fail_task"}, + ], + "notebook_task": {"notebook_path": "/Workspace/any/join"}, + }, + ], + { + "ok": "pass\n", + "fail": "raise RuntimeError('boom')\n", + "join": f'with open(r"{log_file}", "a") as f: f.write("join\\n")\n', + }, + ) + with pytest.raises(RuntimeError, match="boom"): + runner.run_workflow() + assert runner.task_statuses["join"] == "SUCCESS" + assert log_file.read_text(encoding="utf-8").splitlines() == ["join"] + + def test_depends_on_outcome_skips_when_status_does_not_match(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "ok_task", + "notebook_task": {"notebook_path": "/Workspace/any/ok"}, + }, + { + "task_key": "only_on_fail", + "depends_on": [{"task_key": "ok_task", "outcome": "FAILED"}], + "notebook_task": {"notebook_path": "/Workspace/any/only_on_fail"}, + }, + ], + {"ok": "pass\n", "only_on_fail": "pass\n"}, + ) + runner.run_workflow() + assert runner.task_statuses["only_on_fail"] == "SKIPPED" + + def test_unknown_run_if_raises(self, tmp_path): + workflow = { + "tasks": [ + { + "task_key": "t1", + "run_if": "SOMETIMES", + "notebook_task": {"notebook_path": "/a/b"}, + } + ] + } + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps(workflow), encoding="utf-8") + with pytest.raises(ValueError, match="Unsupported run_if"): + LocalWorkflowRunner(str(tmp_path), str(workflow_path), str(tmp_path)) + + +class TestConditionTasks: + def _runner(self, tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + + def test_true_branch_runs_and_false_branch_skipped(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "producer", + "notebook_task": {"notebook_path": "/Workspace/any/producer"}, + }, + { + "task_key": "check", + "depends_on": [{"task_key": "producer"}], + "condition_task": { + "op": "EQUAL_TO", + "left": "{{tasks.producer.values.region}}", + "right": "eu", + }, + }, + { + "task_key": "on_true", + "depends_on": [{"task_key": "check", "outcome": "true"}], + "notebook_task": {"notebook_path": "/Workspace/any/on_true"}, + }, + { + "task_key": "on_false", + "depends_on": [{"task_key": "check", "outcome": "false"}], + "notebook_task": {"notebook_path": "/Workspace/any/on_false"}, + }, + ], + { + "producer": 'dbutils.jobs.taskValues.set(key="region", value="eu")\n', + "on_true": f'with open(r"{log_file}", "a") as f: f.write("true\\n")\n', + "on_false": f'with open(r"{log_file}", "a") as f: f.write("false\\n")\n', + }, + ) + runner.run_workflow() + assert runner.task_statuses["check"] == "SUCCESS" + assert runner.task_results["check"] == "true" + assert runner.task_statuses["on_true"] == "SUCCESS" + assert runner.task_statuses["on_false"] == "SKIPPED" + assert log_file.read_text(encoding="utf-8").splitlines() == ["true"] + + def test_false_branch_runs_when_condition_fails(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "producer", + "notebook_task": {"notebook_path": "/Workspace/any/producer"}, + }, + { + "task_key": "check", + "depends_on": [{"task_key": "producer"}], + "condition_task": { + "op": "EQUAL", + "left": "{{tasks.producer.values.region}}", + "right": "eu", + }, + }, + { + "task_key": "on_true", + "depends_on": [{"task_key": "check", "outcome": "true"}], + "notebook_task": {"notebook_path": "/Workspace/any/on_true"}, + }, + { + "task_key": "on_false", + "depends_on": [{"task_key": "check", "outcome": "false"}], + "notebook_task": {"notebook_path": "/Workspace/any/on_false"}, + }, + ], + { + "producer": 'dbutils.jobs.taskValues.set(key="region", value="us")\n', + "on_true": f'with open(r"{log_file}", "a") as f: f.write("true\\n")\n', + "on_false": f'with open(r"{log_file}", "a") as f: f.write("false\\n")\n', + }, + ) + runner.run_workflow() + assert runner.task_results["check"] == "false" + assert runner.task_statuses["on_true"] == "SKIPPED" + assert runner.task_statuses["on_false"] == "SUCCESS" + assert log_file.read_text(encoding="utf-8").splitlines() == ["false"] + + def test_greater_than_compares_numerically(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "producer", + "notebook_task": {"notebook_path": "/Workspace/any/producer"}, + }, + { + "task_key": "check", + "depends_on": [{"task_key": "producer"}], + "condition_task": { + "op": "GREATER_THAN", + "left": "{{tasks.producer.values.count}}", + "right": "10", + }, + }, + { + "task_key": "on_true", + "depends_on": [{"task_key": "check", "outcome": "true"}], + "notebook_task": {"notebook_path": "/Workspace/any/on_true"}, + }, + ], + { + "producer": 'dbutils.jobs.taskValues.set(key="count", value="12")\n', + "on_true": f'with open(r"{log_file}", "a") as f: f.write("true\\n")\n', + }, + ) + runner.run_workflow() + assert runner.task_results["check"] == "true" + assert log_file.read_text(encoding="utf-8").splitlines() == ["true"] + + def test_condition_task_without_notebook_file(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "check", + "condition_task": { + "op": "EQUAL_TO", + "left": "a", + "right": "a", + }, + } + ], + {}, + ) + runner.run_workflow() + assert runner.task_statuses["check"] == "SUCCESS" + assert runner.task_results["check"] == "true" + + def test_missing_task_type_still_raises(self, tmp_path): + workflow = {"tasks": [{"task_key": "t1"}]} + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps(workflow), encoding="utf-8") + with pytest.raises(ValueError, match="is missing 'notebook_task'"): + LocalWorkflowRunner(str(tmp_path), str(workflow_path), str(tmp_path)) + + +class TestForEachTasks: + def _runner(self, tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + + def test_runs_nested_notebook_once_per_input(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "loop", + "for_each_task": { + "inputs": '["a", "b", "c"]', + "concurrency": 1, + "task": { + "task_key": "process", + "notebook_task": { + "notebook_path": "/Workspace/any/process", + "base_parameters": {"item": "{{input}}"}, + }, + }, + }, + } + ], + { + "process": ( + "import os\n" + f'with open(r"{log_file}", "a") as f: f.write(os.environ["item"] + "\\n")\n' + ) + }, + ) + runner.run_workflow() + assert runner.task_statuses["loop"] == "SUCCESS" + assert log_file.read_text(encoding="utf-8").splitlines() == ["a", "b", "c"] + + def test_inputs_from_task_values_json_list(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "producer", + "notebook_task": {"notebook_path": "/Workspace/any/producer"}, + }, + { + "task_key": "loop", + "depends_on": [{"task_key": "producer"}], + "for_each_task": { + "inputs": "{{tasks.producer.values.items}}", + "task": { + "task_key": "process", + "notebook_task": { + "notebook_path": "/Workspace/any/process", + "base_parameters": {"item": "{{input}}"}, + }, + }, + }, + }, + ], + { + "producer": 'dbutils.jobs.taskValues.set(key="items", value=\'["x", "y"]\')\n', + "process": ( + "import os\n" + f'with open(r"{log_file}", "a") as f: f.write(os.environ["item"] + "\\n")\n' + ), + }, + ) + runner.run_workflow() + assert log_file.read_text(encoding="utf-8").splitlines() == ["x", "y"] + + def test_child_failure_fails_for_each_task(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "loop", + "for_each_task": { + "inputs": '["ok", "bad", "later"]', + "task": { + "task_key": "process", + "notebook_task": { + "notebook_path": "/Workspace/any/process", + "base_parameters": {"item": "{{input}}"}, + }, + }, + }, + } + ], + { + "process": ( + "import os\n" + "item = os.environ['item']\n" + f'with open(r"{log_file}", "a") as f: f.write(item + "\\n")\n' + "if item == 'bad':\n" + " raise RuntimeError('bad item')\n" + ) + }, + ) + with pytest.raises(RuntimeError, match="bad item"): + runner.run_workflow() + assert runner.task_statuses["loop"] == "FAILED" + assert log_file.read_text(encoding="utf-8").splitlines() == ["ok", "bad"] + + def test_literal_list_inputs(self, tmp_path): + log_file = tmp_path / "execution.log" + runner = self._runner( + tmp_path, + [ + { + "task_key": "loop", + "for_each_task": { + "inputs": [1, 2], + "task": { + "task_key": "process", + "notebook_task": { + "notebook_path": "/Workspace/any/process", + "base_parameters": {"item": "{{input}}"}, + }, + }, + }, + } + ], + { + "process": ( + "import os\n" + f'with open(r"{log_file}", "a") as f: f.write(os.environ["item"] + "\\n")\n' + ) + }, + ) + runner.run_workflow() + assert log_file.read_text(encoding="utf-8").splitlines() == ["1", "2"] + + +class TestRepairAndRerun: + def _runner(self, tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + + def test_only_runs_selected_tasks(self, tmp_path): + log_file = tmp_path / "execution.log" + tasks = [ + { + "task_key": "first_task", + "notebook_task": {"notebook_path": "/Workspace/any/first"}, + }, + { + "task_key": "second_task", + "depends_on": [{"task_key": "first_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/second"}, + }, + { + "task_key": "third_task", + "depends_on": [{"task_key": "second_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/third"}, + }, + ] + notebooks = { + name: f'with open(r"{log_file}", "a") as f: f.write("{name}\\n")\n' + for name in ["first", "second", "third"] + } + runner = self._runner(tmp_path, tasks, notebooks) + runner.run_workflow(only=["second_task"]) + assert log_file.read_text(encoding="utf-8").splitlines() == ["second"] + assert runner.task_statuses["second_task"] == "SUCCESS" + + def test_from_task_runs_subgraph(self, tmp_path): + log_file = tmp_path / "execution.log" + tasks = [ + { + "task_key": "first_task", + "notebook_task": {"notebook_path": "/Workspace/any/first"}, + }, + { + "task_key": "second_task", + "depends_on": [{"task_key": "first_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/second"}, + }, + { + "task_key": "third_task", + "depends_on": [{"task_key": "second_task"}], + "notebook_task": {"notebook_path": "/Workspace/any/third"}, + }, + ] + notebooks = { + name: f'with open(r"{log_file}", "a") as f: f.write("{name}\\n")\n' + for name in ["first", "second", "third"] + } + runner = self._runner(tmp_path, tasks, notebooks) + runner.run_workflow(from_task="second_task") + assert log_file.read_text(encoding="utf-8").splitlines() == ["second", "third"] + + def test_unknown_only_task_raises(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "task_1", + "notebook_task": {"notebook_path": "/Workspace/any/main"}, + } + ], + {"main": "pass\n"}, + ) + with pytest.raises(ValueError, match="unknown task"): + runner.run_workflow(only=["missing"]) + + def test_only_and_from_task_are_exclusive(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "task_1", + "notebook_task": {"notebook_path": "/Workspace/any/main"}, + } + ], + {"main": "pass\n"}, + ) + with pytest.raises(ValueError, match="only one of"): + runner.run_workflow(only=["task_1"], from_task="task_1") + + +class TestRetriesAndTimeouts: + def _runner(self, tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + + def test_retries_until_success(self, tmp_path): + counter = tmp_path / "counter.txt" + counter.write_text("0", encoding="utf-8") + runner = self._runner( + tmp_path, + [ + { + "task_key": "flaky", + "max_retries": 2, + "min_retry_interval_millis": 0, + "notebook_task": {"notebook_path": "/Workspace/any/flaky"}, + } + ], + { + "flaky": ( + f"from pathlib import Path\n" + f"p = Path(r'{counter}')\n" + "n = int(p.read_text())\n" + "p.write_text(str(n + 1))\n" + "if n < 2:\n" + " raise RuntimeError('not yet')\n" + ) + }, + ) + runner.run_workflow() + assert runner.task_statuses["flaky"] == "SUCCESS" + assert counter.read_text(encoding="utf-8") == "3" + + def test_retry_exhaustion_reraises(self, tmp_path): + runner = self._runner( + tmp_path, + [ + { + "task_key": "always_fail", + "max_retries": 1, + "notebook_task": {"notebook_path": "/Workspace/any/always_fail"}, + } + ], + {"always_fail": "raise RuntimeError('still bad')\n"}, + ) + with pytest.raises(RuntimeError, match="still bad"): + runner.run_workflow() + assert runner.task_statuses["always_fail"] == "FAILED" + + def test_retry_interval_sleeps(self, tmp_path, monkeypatch): + slept = [] + monkeypatch.setattr( + "testbricks.local_workflow_runner.time.sleep", + lambda seconds: slept.append(seconds), + ) + runner = self._runner( + tmp_path, + [ + { + "task_key": "flaky", + "max_retries": 1, + "min_retry_interval_millis": 50, + "notebook_task": {"notebook_path": "/Workspace/any/flaky"}, + } + ], + {"flaky": "raise RuntimeError('nope')\n"}, + ) + with pytest.raises(RuntimeError, match="nope"): + runner.run_workflow() + assert slept == [0.05] + + def test_timeout_seconds_accepted_not_enforced(self, tmp_path, capsys): + runner = self._runner( + tmp_path, + [ + { + "task_key": "task_1", + "timeout_seconds": 1, + "notebook_task": {"notebook_path": "/Workspace/any/main"}, + } + ], + {"main": "pass\n"}, + ) + runner.run_workflow() + captured = capsys.readouterr() + assert "timeout_seconds=1" in captured.out + assert "not enforced" in captured.out + assert runner.task_statuses["task_1"] == "SUCCESS" + + class TestWorkflowValidation: def test_missing_tasks_raises(self, tmp_path): workflow_path = tmp_path / "workflow.json"