TS SDK: add conditional branching with dag.if

Author: jason810496

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

@@ -359,6 +359,30 @@ there. Set it once on the Dag and each task inherits it:
 ``queue`` on a task wins over the Dag's. See :ref:`typescript-sdk/coordinator-config` for the
 ``queue_to_coordinator`` entry that sends that queue to the coordinator.
 
+Conditional branching
+~~~~~~~~~~~~~~~~~~~~~
+
+``dag.if`` takes a task whose handler returns a boolean, and names the task each outcome runs:
+
+.. code-block:: typescript
+
+    const condition = dag.task("has_rows", async ({ rows }: { rows: number }) => rows > 0);
+    const gated = condition({ rows: extracted });
+
+    dag.if(gated).then(loaded).else(reportedEmpty);
+
+The condition is an ordinary task, so it is declared, typed and wired like any other, and the
+compiler checks that its handler really returns a boolean. ``else`` is optional: a one-sided
+condition skips its own branch when the condition fails and follows nothing.
+
+A guarded task takes no argument for the control edge, because a condition's boolean decides whether
+the task runs rather than what it runs on. Read a value from the condition with
+``getClient().getXCom``.
+
+The side not taken is skipped when the run reaches it, and stays skipped if you clear it later. Only
+the branches named here are skipped, so a task that several branches converge on still runs — unlike
+Python's ``@task.branch``, which skips every immediate downstream it did not follow.
+
 ``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
 ``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
 

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

@@ -135,6 +135,43 @@ data is the wiring object itself — `summarize({ north: extractNorth(), south:
 does. This matches `Before`/`After` in the Go SDK's native Dag interface, spelled to TypeScript
 convention.
 
+### Conditional branching: `if` and `else`
+
+`dag.if(condition)` is TypeScript's spelling of the construct
+[`airflow-core/adr/lang-sdk/0008`](../../airflow-core/adr/lang-sdk/0008-control-flow-constructs.md)
+names after the host language's control flow:
+
+```ts
+const gated = dag.task("has_rows", async ({ rows }: { rows: number }) => rows > 0)({
+  rows: extracted,
+});
+
+dag.if(gated).then(loadIfReady).else(loadFallback);
+```
+
+**The condition is a task reference, not a task id and a function.** `dag.if` takes a `TaskRef` the
+Dag already handed back, whose handler's return type the compiler checks is `boolean`. So nothing
+depends on the task's id or on its function name, and a condition is declared, named and typed the
+same way every other task is. The same reasoning that makes a *case* a reference in that ADR's
+decision 2 applies to the condition itself.
+
+**A `then` chain is a thenable, and is guarded rather than avoided.** An object with a callable
+`then` is a *thenable*: were the object `dag.if` returns to reach an `await`, the runtime would hand
+its `then` a resolve function where a task reference belongs. Two things contain that. `.then(...)`
+returns an object carrying only `.else`, so nothing past the first step is awaitable at all; and
+`.then` rejects a function argument by naming the cause, so an author who does await it reads "this
+builds a branch, drop the await" rather than a type error about references.
+
+**A branch is a real branch to Airflow.** The control edges serialize as ordinary order-only edges
+and carry no branch-candidate field, as that ADR's consequences require. The condition task is
+serialized with `_can_skip_downstream`, and it writes the `skipmixin_key` XCom alongside the skip, so
+clearing a skipped branch re-skips it the way a Python `@task.branch` does rather than running the
+side the condition rejected.
+
+A one-sided `if` is a branch with one candidate — it skips `then` and follows nothing — rather than a
+`ShortCircuitOperator`, which would also skip the whole downstream closure and ignore trigger rules.
+A guarded task takes no argument for the control edge: a condition's boolean is a signal, not data.
+
 ## Consequences
 
 - One authoring surface (`dag.task()` plus its factory) covers the graph and each task's arguments,

ts-sdk/src/coordinator/client.ts

@@ -30,6 +30,7 @@ import type {
   GetXCom,
   SetXCom,
   GetConnection,
+  SkipDownstreamTasks,
   ConnectionResult as WireConnectionResult,
 } from "./protocol.js";
 
@@ -202,6 +203,16 @@ export function createCoordinatorClient(
       await rpc("SetXCom", null, msg, () => undefined, "throw");
     },
 
+    // ---- Control flow ----
+
+    async skipDownstreamTasks(taskIds: readonly string[]): Promise<void> {
+      // The supervisor treats an empty list as "skip nothing", and sending it
+      // would cost a round trip to say so.
+      if (taskIds.length === 0) return;
+      const msg: SkipDownstreamTasks = { type: "SkipDownstreamTasks", tasks: [...taskIds] };
+      await rpc("SkipDownstreamTasks", null, msg, () => undefined, "throw");
+    },
+
     // ---- Connections ----
 
     async getConnection(connId: string): Promise<ConnectionResult | null> {

ts-sdk/src/coordinator/protocol.ts

@@ -64,6 +64,7 @@ export type {
   GetXCom,
   SetXCom,
   GetConnection,
+  SkipDownstreamTasks,
 } from "../generated/supervisor.js";
 
 // -------- Frames from supervisor --------

ts-sdk/src/coordinator/serde.ts

@@ -154,6 +154,7 @@ export function serializeDag(
         withDagQueue(record.spec, dag.spec.queue),
         graph.downstreamTaskIds.get(taskId),
         inputs.get(taskId),
+        record.canSkipDownstream === true,
       ),
     ),
     dag_dependencies: [],
@@ -191,6 +192,7 @@ function serializeTask(
   spec: object,
   downstream: ReadonlySet<string> | undefined,
   inputs: RecordedInputs | undefined,
+  canSkipDownstream: boolean,
 ): SerializedValue {
   const data: Record<string, SerializedValue> = {
     task_id: taskId,
@@ -208,6 +210,10 @@ function serializeTask(
   };
   const bindings = serializeArgBindings(inputs);
   if (bindings) data["_arg_bindings"] = bindings;
+  // What Python writes for a SkipMixin operator, and what makes
+  // `NotPreviouslySkippedDep` consult this task's `skipmixin_key` XCom when one
+  // of its skipped downstream tasks is cleared.
+  if (canSkipDownstream) data["_can_skip_downstream"] = true;
   applySchemaFields(data, spec, TASK_FIELD_RULES, `task "${taskId}" of Dag "${dagId}"`);
   if (downstream?.size) {
     data["downstream_task_ids"] = [...downstream].sort();

ts-sdk/src/index.ts

@@ -27,6 +27,8 @@ export { SUPERVISOR_API_VERSION } from "./coordinator/index.js";
 export type { ArgNameMap } from "./sdk/arg-names.js";
 export type { Registerable } from "./sdk/bundle.js";
 export type {
+  Condition,
+  ConditionElse,
   DagSpec,
   Node,
   TaskFactory,

ts-sdk/src/sdk/client.ts

@@ -107,6 +107,19 @@ export interface TaskClient {
    * @throws {@link Exceptions!ConnectionNotFoundError | ConnectionNotFoundError} when the connection does not exist.
    */
   getConnectionOrThrow(connId: string): Promise<ConnectionResult>;
+
+  /**
+   * Mark downstream tasks of the running task as skipped.
+   *
+   * The mechanism behind `dag.if(...)`: a condition's boolean is a run-time
+   * signal, so which branch a Dag run takes is decided here rather than
+   * recorded in the Dag. Python reaches the same supervisor message through
+   * `skip()` and `skip_all_except()`.
+   *
+   * Each id must name a direct downstream of the running task. Passing none is
+   * a no-op, so a branch that skips nothing needs no guard at the call site.
+   */
+  skipDownstreamTasks(taskIds: readonly string[]): Promise<void>;
 }
 
 /** Error thrown by {@link TaskClient.getVariableOrThrow}. */