TS SDK: Add multi-way branching with dag.switch

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

@@ -419,6 +419,30 @@ The side not taken is skipped when the run reaches it, and stays skipped if you
 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.
 
+Multi-way branching
+~~~~~~~~~~~~~~~~~~~
+
+``dag.switch`` is the multi-way form: a handler that returns one of the cases it is given.
+
+.. code-block:: typescript
+
+    async function pickPath({ rows }: { rows: number }): Promise<TaskRef> {
+      return rows > 1000 ? handleLong : handleShort;
+    }
+
+    dag.switch(pickPath, { rows: extracted }).case(handleLong).case(handleShort);
+
+``dag.switch`` declares the decider the way ``dag.if`` does. A case is the task reference itself, so
+the compiler checks the candidate exists and renaming a handler cannot silently rewire a Dag. The
+task's own value is the chosen task's id, which a downstream task can read from its XCom.
+
+There is no default case. A decider that returns anything outside its cases fails the task, naming
+what it chose and what it could have chosen.
+
+Exactly one case is selected. Python's branch callable may return a list of task ids, and no language
+SDK offers that yet: put the paths that run together behind one task, or gate each with its own
+condition.
+
 ``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

@@ -178,6 +178,38 @@ A one-sided `if` is a branch with one candidate — it skips `then` and follows
 `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.
 
+### Multi-way branching: `switch` and `case`
+
+`dag.switch(pickPath)` follows
+[`airflow-core/adr/lang-sdk/0008`](../../airflow-core/adr/lang-sdk/0008-control-flow-constructs.md)
+decision 2 without divergence: a case **is** the reference the SDK handed back, not a label kept in
+step with one. The decider is a handler, declared and wired in one call as `dag.if` declares a condition, as
+Go's `dag.Switch(pickPath)` does.
+
+```ts
+async function pickPath({ rows }: { rows: number }): Promise<TaskRef> {
+  return rows > 1000 ? handleLong : handleShort;
+}
+
+dag.switch(pickPath, { rows: extracted }).case(handleLong).case(handleShort);
+```
+
+A case is a task, never a condition or a branch: the decider returns the case from an async handler,
+and a condition carries `then`, so returning one would be awaited as a thenable instead.
+
+An earlier draft selected a case by a string label the author writes, on the grounds that a handler's
+function name does not survive bundling. That concern does not apply: a `TaskRef` carries the task's
+own id, which the SDK fixed when the task was declared and esbuild never touches. Selecting by
+reference keeps the compiler checking that a candidate exists, which a label cannot.
+
+The cases chain, as they do in Go. `case` reads the candidate list when the task runs rather than
+when it is declared, which is what lets the chain follow the `dag.switch` call; and unlike a
+condition's `then`, `case` is not a thenable trap, so nothing has to be guarded here.
+
+**No default case**, per decision 3, and **exactly one case is selected**, the limitation that ADR
+records for every Lang SDK. A branch with no case at all decides nothing, and is rejected when the
+Dag is read.
+
 ## Consequences
 
 - One authoring surface (`dag.task()` plus its factory) covers the graph and each task's arguments,

ts-sdk/src/index.ts

@@ -27,6 +27,7 @@ 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 {
+  Branch,
   Condition,
   ConditionElse,
   DeciderArgs,

ts-sdk/src/sdk/dag.ts

@@ -418,21 +418,36 @@ export type TaskOptions = TaskSpec;
 export interface Condition extends Node {
   /** Airflow task ID of the deciding task. */
   readonly taskId: string;
-  then(taskRef: TaskRef | Condition): ConditionElse;
+  then(taskRef: TaskRef | Condition | Branch): ConditionElse;
   before(...downstream: readonly Node[]): Condition;
   after(...upstream: readonly Node[]): Condition;
 }
 
 /** What `.then(...)` returns: the other side, which a one-sided condition omits. */
 export interface ConditionElse {
-  else(taskRef: TaskRef | Condition): void;
+  else(taskRef: TaskRef | Condition | Branch): void;
 }
 
 /** The arguments after a decider's handler: its inputs, then its {@link TaskSpec}. */
 export type DeciderArgs<TArgs extends object | void> = [TArgs] extends [void]
   ? [inputs?: undefined, spec?: TaskOptions]
   : [inputs: TaskInputs<TArgs>, spec?: TaskOptions];
 
+/**
+ * A placed multi-way branch: name each task the decider chooses between.
+ *
+ * ```ts
+ * dag.switch(pickPath, { rows: extracted }).case(handleLong).case(handleShort);
+ * ```
+ */
+export interface Branch extends Node {
+  /** Airflow task ID of the deciding task. */
+  readonly taskId: string;
+  case(taskRef: TaskRef): Branch;
+  before(...downstream: readonly Node[]): Branch;
+  after(...upstream: readonly Node[]): Branch;
+}
+
 /** Per-task record a Dag retains: the reference, the handler, and its spec. */
 export interface TaskRecord {
   readonly task: TaskRef;
@@ -512,6 +527,7 @@ export class Dag {
   readonly #orderEdges = new Map<string, OrderEdge>();
   readonly #definedIn: string | undefined;
   readonly #conditions = new Map<string, ConditionRecord>();
+  readonly #branches = new Map<string, readonly TaskRef[]>();
   // Keyed by full group ID; a group's own record holds what it declares, so
   // the tree is reconstructed by walking from the roots.
   readonly #groups = new Map<string, MutableTaskGroupRecord>();
@@ -621,7 +637,7 @@ export class Dag {
     });
 
     const named = new Set<"then" | "else">();
-    const name = (side: "then" | "else", target: TaskRef | Condition): void => {
+    const name = (side: "then" | "else", target: TaskRef | Condition | Branch): void => {
       // `await` calls `then(resolve, reject)`, so a function here means the condition was awaited.
       if (typeof target === "function") {
         throw new Error(
@@ -665,14 +681,7 @@ export class Dag {
         name("then", taskRef);
         return elseStep;
       },
-      before: (...downstream) => {
-        condition.before(...downstream);
-        return placed;
-      },
-      after: (...upstream) => {
-        condition.after(...upstream);
-        return placed;
-      },
+      ...this.#forwardEdges(condition, () => placed),
     };
     conditionTasks.set(placed, condition);
     return Object.freeze(placed);
@@ -697,6 +706,76 @@ export class Dag {
     this.#tasks.set(taskId, { ...record, canSkipDownstream: true, fn: wrapped });
   }
 
+  /**
+   * Declare a task that returns one of its cases; every other case is skipped.
+   *
+   * ```ts
+   * dag.switch(pickPath, { rows: extracted }).case(handleLong).case(handleShort);
+   * ```
+   */
+  switch<TArgs extends object | void = void>(
+    handler: (args: TArgs) => TaskRef | Promise<TaskRef>,
+    ...args: DeciderArgs<NoInfer<TArgs>>
+  ): Branch {
+    return this.#placeBranch(this.#placeDecider(handler, args) as TaskRef<TaskRef>);
+  }
+
+  #forwardEdges<T>(ref: TaskRef, self: () => T) {
+    return {
+      before: (...downstream: readonly Node[]) => {
+        ref.before(...downstream);
+        return self();
+      },
+      after: (...upstream: readonly Node[]) => {
+        ref.after(...upstream);
+        return self();
+      },
+    };
+  }
+
+  #placeBranch(decider: TaskRef<TaskRef>): Branch {
+    const taskId = decider.taskId;
+    const candidates: TaskRef[] = [];
+    this.#branches.set(taskId, candidates);
+    this.#wrapDecider(taskId, async (chosen: unknown) => {
+      const known = candidates.map((ref) => ref.taskId);
+      const picked = isTaskRef(chosen) && chosen.dagId === this.dagId ? chosen.taskId : undefined;
+      if (picked === undefined || !known.includes(picked)) {
+        throw new Error(
+          `Task "${taskId}" of Dag "${this.dagId}" chose ` +
+            `${isTaskRef(chosen) ? `"${chosen.taskId}"` : describeValue(chosen)}, ` +
+            `which is not one of its cases: ${known.join(", ")}`,
+        );
+      }
+      return { skip: known.filter((id) => id !== picked), result: picked };
+    });
+
+    const branch: Branch = {
+      dagId: this.dagId,
+      taskId,
+      case: (taskRef) => {
+        this.#validateOwnNode(taskRef, `a case of "${taskId}"`);
+        if (!isTaskRef(taskRef)) {
+          throw new Error(
+            `A case of Dag "${this.dagId}" branch "${taskId}" has to be a task, not a task group`,
+          );
+        }
+        if (candidates.some((candidate) => candidate.taskId === taskRef.taskId)) {
+          throw new Error(
+            `Dag "${this.dagId}" branch "${taskId}" lists "${taskRef.taskId}" twice; ` +
+              "each case names a different task",
+          );
+        }
+        candidates.push(taskRef);
+        decider.before(taskRef);
+        return branch;
+      },
+      ...this.#forwardEdges(decider, () => branch),
+    };
+    conditionTasks.set(branch, decider);
+    return Object.freeze(branch);
+  }
+
   /**
    * Declare a task group of this Dag.
    *
@@ -1105,6 +1184,14 @@ export class Dag {
         );
       }
     }
+    for (const [taskId, cases] of this.#branches) {
+      if (cases.length === 0) {
+        throw new Error(
+          `Branch "${taskId}" of Dag "${this.dagId}" has no cases, so it decides nothing; ` +
+            "give it the tasks to choose between with dag.switch(handler, inputs).case(task)",
+        );
+      }
+    }
     for (const taskId of this.#tasks.keys()) {
       if (!this.#inputs.has(taskId)) {
         throw new Error(