Skip to content

Commit fd127aa

Browse files
TS SDK: Add conditional branching with dag.if
Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com>
1 parent 5e476d7 commit fd127aa

10 files changed

Lines changed: 600 additions & 10 deletions

File tree

‎airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,34 @@ there. Set it once on the Dag and each task inherits it:
359359
``queue`` on a task wins over the Dag's. See :ref:`typescript-sdk/coordinator-config` for the
360360
``queue_to_coordinator`` entry that sends that queue to the coordinator.
361361

362+
Conditional branching
363+
~~~~~~~~~~~~~~~~~~~~~
364+
365+
``dag.if`` takes a handler that returns a boolean, and names the task each outcome runs:
366+
367+
.. code-block:: typescript
368+
369+
async function hasRows({ rows }: { rows: number }): Promise<boolean> {
370+
return rows > 0;
371+
}
372+
373+
const gate = dag.if(hasRows, { rows: extracted });
374+
gate.then(loaded).else(reportedEmpty);
375+
376+
``dag.if`` declares the condition as a task: its id is the function's name, and a trailing spec sets
377+
its id or other options, as in ``dag.if(hasRows, { rows }, { taskId: "has_rows" })``. The compiler
378+
checks that the handler returns a boolean. ``else`` is optional: a one-sided condition skips its own
379+
branch when the condition fails and follows nothing. The condition is also a node, so
380+
``notified.after(gate)`` orders a task after it.
381+
382+
A guarded task takes no argument for the control edge, because a condition's boolean decides whether
383+
the task runs rather than what it runs on. Read a value from the condition with
384+
``getClient().getXCom``.
385+
386+
The side not taken is skipped when the run reaches it, and stays skipped if you clear it later. Only
387+
the branches named here are skipped, so a task that several branches converge on still runs — unlike
388+
Python's ``@task.branch``, which skips every immediate downstream it did not follow.
389+
362390
``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
363391
``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
364392

‎ts-sdk/adr/0002-native-dag-interface.md‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,49 @@ data is the wiring object itself — `summarize({ north: extractNorth(), south:
135135
does. This matches `Before`/`After` in the Go SDK's native Dag interface, spelled to TypeScript
136136
convention.
137137

138+
### Conditional branching: `if` and `else`
139+
140+
`dag.if(handler)` is TypeScript's spelling of the construct
141+
[`airflow-core/adr/lang-sdk/0008`](../../airflow-core/adr/lang-sdk/0008-control-flow-constructs.md)
142+
names after the host language's control flow:
143+
144+
```ts
145+
async function hasRows({ rows }: { rows: number }): Promise<boolean> {
146+
return rows > 0;
147+
}
148+
149+
const gate = dag.if(hasRows, { rows: validated });
150+
gate.then(loadIfReady).else(loadFallback);
151+
```
152+
153+
**The condition is a handler, declared and wired in one call.** That ADR writes
154+
`dag.If(hasRows, airflow.Inputs(validated), airflow.TaskSpec{...})` in Go, and
155+
`dag.if(hasRows, { rows: validated }, { ... })` reads the same way: inputs, then an optional
156+
`TaskSpec`. The deciding task takes its id from the function's name unless the spec's `taskId` sets
157+
it, and the spec carries its other options, such as `queue` and `retries`. The compiler checks that
158+
the handler returns a `boolean`, and that the inputs match its argument.
159+
160+
**A condition is a node.** What `dag.if` returns carries `before` and `after`, and stands at the
161+
other end of an edge as the deciding task, so `notified.after(gate)` reads as Go's `.After(gate)`. It
162+
is also a valid branch of another condition.
163+
164+
**A `then` chain is a thenable, and is guarded rather than avoided.** An object with a callable
165+
`then` is a *thenable*: were the condition `dag.if` returns to reach an `await`, the runtime would hand
166+
its `then` a resolve function where a task reference belongs. Two things contain that. `.then(...)`
167+
returns an object carrying only `.else`, so nothing past the first step is awaitable at all; and
168+
`.then` rejects a function argument by naming the cause, so an author who does await it reads "this
169+
builds a branch, drop the await" rather than a type error about references.
170+
171+
**A branch is a real branch to Airflow.** The control edges serialize as ordinary order-only edges
172+
and carry no branch-candidate field, as that ADR's consequences require. The condition task is
173+
serialized with `_can_skip_downstream`, and it writes the `skipmixin_key` XCom alongside the skip, so
174+
clearing a skipped branch re-skips it the way a Python `@task.branch` does rather than running the
175+
side the condition rejected.
176+
177+
A one-sided `if` is a branch with one candidate — it skips `then` and follows nothing — rather than a
178+
`ShortCircuitOperator`, which would also skip the whole downstream closure and ignore trigger rules.
179+
A guarded task takes no argument for the control edge: a condition's boolean is a signal, not data.
180+
138181
## Consequences
139182

140183
- One authoring surface (`dag.task()` plus its factory) covers the graph and each task's arguments,

‎ts-sdk/src/coordinator/client.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import type {
3030
GetXCom,
3131
SetXCom,
3232
GetConnection,
33+
SkipDownstreamTasks,
3334
ConnectionResult as WireConnectionResult,
3435
} from "./protocol.js";
3536

@@ -82,6 +83,9 @@ export interface CoordinatorClient extends TaskClient {
8283
* output fails the task, while one that pushed null binds null.
8384
*/
8485
getXComEntry(opts: GetXComOpts): Promise<XComEntry>;
86+
87+
/** Mark direct downstream tasks of the running task as skipped; none is a no-op. */
88+
skipDownstreamTasks(taskIds: readonly string[]): Promise<void>;
8589
}
8690

8791
export function createCoordinatorClient(
@@ -202,6 +206,14 @@ export function createCoordinatorClient(
202206
await rpc("SetXCom", null, msg, () => undefined, "throw");
203207
},
204208

209+
// ---- Control flow ----
210+
211+
async skipDownstreamTasks(taskIds: readonly string[]): Promise<void> {
212+
if (taskIds.length === 0) return;
213+
const msg: SkipDownstreamTasks = { type: "SkipDownstreamTasks", tasks: [...taskIds] };
214+
await rpc("SkipDownstreamTasks", null, msg, () => undefined, "throw");
215+
},
216+
205217
// ---- Connections ----
206218

207219
async getConnection(connId: string): Promise<ConnectionResult | null> {

‎ts-sdk/src/coordinator/protocol.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ export type {
6464
GetXCom,
6565
SetXCom,
6666
GetConnection,
67+
SkipDownstreamTasks,
6768
} from "../generated/supervisor.js";
6869

6970
// -------- Frames from supervisor --------

‎ts-sdk/src/coordinator/serde.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ export function serializeDag(
143143
withDagQueue(record.spec, dag.spec.queue),
144144
graph.downstreamTaskIds.get(taskId),
145145
inputs.get(taskId),
146+
record.canSkipDownstream === true,
146147
),
147148
),
148149
dag_dependencies: [],
@@ -180,6 +181,7 @@ function serializeTask(
180181
spec: object,
181182
downstream: ReadonlySet<string> | undefined,
182183
inputs: RecordedInputs | undefined,
184+
canSkipDownstream: boolean,
183185
): SerializedValue {
184186
const data: Record<string, SerializedValue> = {
185187
task_id: taskId,
@@ -198,6 +200,8 @@ function serializeTask(
198200
const label = `task "${taskId}" of Dag "${dagId}"`;
199201
const bindings = serializeArgBindings(inputs, label);
200202
if (bindings) data["_arg_bindings"] = bindings;
203+
// Lets `NotPreviouslySkippedDep` re-skip a cleared downstream, as Python's SkipMixin does.
204+
if (canSkipDownstream) data["_can_skip_downstream"] = true;
201205
applySchemaFields(data, spec, TASK_FIELD_RULES, label);
202206
if (downstream?.size) {
203207
data["downstream_task_ids"] = [...downstream].sort();

‎ts-sdk/src/index.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ export { SUPERVISOR_API_VERSION } from "./coordinator/index.js";
2727
export type { ArgNameMap } from "./sdk/arg-names.js";
2828
export type { Registerable } from "./sdk/bundle.js";
2929
export type {
30+
Condition,
31+
ConditionElse,
32+
DeciderArgs,
3033
DagSpec,
3134
Node,
3235
TaskFactory,

0 commit comments

Comments
 (0)