From fd127aa9714ff5cbf04c89052707a39ccca26baa Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:46:22 +0800 Subject: [PATCH] TS SDK: Add conditional branching with dag.if Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com> --- .../language-sdks/typescript.rst | 28 ++ ts-sdk/adr/0002-native-dag-interface.md | 43 +++ ts-sdk/src/coordinator/client.ts | 12 + ts-sdk/src/coordinator/protocol.ts | 1 + ts-sdk/src/coordinator/serde.ts | 4 + ts-sdk/src/index.ts | 3 + ts-sdk/src/sdk/dag.ts | 189 ++++++++++- ts-sdk/tests/coordinator/client.test.ts | 18 + ts-sdk/tests/sdk/conditional.test.ts | 308 ++++++++++++++++++ ts-sdk/tests/sdk/dag.test.ts | 4 +- 10 files changed, 600 insertions(+), 10 deletions(-) create mode 100644 ts-sdk/tests/sdk/conditional.test.ts diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index 47070c581da15..864f5f4591325 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -359,6 +359,34 @@ 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 handler that returns a boolean, and names the task each outcome runs: + +.. code-block:: typescript + + async function hasRows({ rows }: { rows: number }): Promise { + return rows > 0; + } + + const gate = dag.if(hasRows, { rows: extracted }); + gate.then(loaded).else(reportedEmpty); + +``dag.if`` declares the condition as a task: its id is the function's name, and a trailing spec sets +its id or other options, as in ``dag.if(hasRows, { rows }, { taskId: "has_rows" })``. The compiler +checks that the handler returns a boolean. ``else`` is optional: a one-sided condition skips its own +branch when the condition fails and follows nothing. The condition is also a node, so +``notified.after(gate)`` orders a task after it. + +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. diff --git a/ts-sdk/adr/0002-native-dag-interface.md b/ts-sdk/adr/0002-native-dag-interface.md index 6614ee62091bb..f1c4faa8dcfb1 100644 --- a/ts-sdk/adr/0002-native-dag-interface.md +++ b/ts-sdk/adr/0002-native-dag-interface.md @@ -135,6 +135,49 @@ 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(handler)` 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 +async function hasRows({ rows }: { rows: number }): Promise { + return rows > 0; +} + +const gate = dag.if(hasRows, { rows: validated }); +gate.then(loadIfReady).else(loadFallback); +``` + +**The condition is a handler, declared and wired in one call.** That ADR writes +`dag.If(hasRows, airflow.Inputs(validated), airflow.TaskSpec{...})` in Go, and +`dag.if(hasRows, { rows: validated }, { ... })` reads the same way: inputs, then an optional +`TaskSpec`. The deciding task takes its id from the function's name unless the spec's `taskId` sets +it, and the spec carries its other options, such as `queue` and `retries`. The compiler checks that +the handler returns a `boolean`, and that the inputs match its argument. + +**A condition is a node.** What `dag.if` returns carries `before` and `after`, and stands at the +other end of an edge as the deciding task, so `notified.after(gate)` reads as Go's `.After(gate)`. It +is also a valid branch of another condition. + +**A `then` chain is a thenable, and is guarded rather than avoided.** An object with a callable +`then` is a *thenable*: were the condition `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, diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index ba0be84faa240..9103d824d58de 100644 --- a/ts-sdk/src/coordinator/client.ts +++ b/ts-sdk/src/coordinator/client.ts @@ -30,6 +30,7 @@ import type { GetXCom, SetXCom, GetConnection, + SkipDownstreamTasks, ConnectionResult as WireConnectionResult, } from "./protocol.js"; @@ -82,6 +83,9 @@ export interface CoordinatorClient extends TaskClient { * output fails the task, while one that pushed null binds null. */ getXComEntry(opts: GetXComOpts): Promise; + + /** Mark direct downstream tasks of the running task as skipped; none is a no-op. */ + skipDownstreamTasks(taskIds: readonly string[]): Promise; } export function createCoordinatorClient( @@ -202,6 +206,14 @@ export function createCoordinatorClient( await rpc("SetXCom", null, msg, () => undefined, "throw"); }, + // ---- Control flow ---- + + async skipDownstreamTasks(taskIds: readonly string[]): Promise { + 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 { diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts index 323765036ea19..1c95fee6f907d 100644 --- a/ts-sdk/src/coordinator/protocol.ts +++ b/ts-sdk/src/coordinator/protocol.ts @@ -64,6 +64,7 @@ export type { GetXCom, SetXCom, GetConnection, + SkipDownstreamTasks, } from "../generated/supervisor.js"; // -------- Frames from supervisor -------- diff --git a/ts-sdk/src/coordinator/serde.ts b/ts-sdk/src/coordinator/serde.ts index 900b41b6887b6..bf7969a318642 100644 --- a/ts-sdk/src/coordinator/serde.ts +++ b/ts-sdk/src/coordinator/serde.ts @@ -143,6 +143,7 @@ export function serializeDag( withDagQueue(record.spec, dag.spec.queue), graph.downstreamTaskIds.get(taskId), inputs.get(taskId), + record.canSkipDownstream === true, ), ), dag_dependencies: [], @@ -180,6 +181,7 @@ function serializeTask( spec: object, downstream: ReadonlySet | undefined, inputs: RecordedInputs | undefined, + canSkipDownstream: boolean, ): SerializedValue { const data: Record = { task_id: taskId, @@ -198,6 +200,8 @@ function serializeTask( const label = `task "${taskId}" of Dag "${dagId}"`; const bindings = serializeArgBindings(inputs, label); if (bindings) data["_arg_bindings"] = bindings; + // Lets `NotPreviouslySkippedDep` re-skip a cleared downstream, as Python's SkipMixin does. + if (canSkipDownstream) data["_can_skip_downstream"] = true; applySchemaFields(data, spec, TASK_FIELD_RULES, label); if (downstream?.size) { data["downstream_task_ids"] = [...downstream].sort(); diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 4e7bee9a9669b..09eb66af399b4 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -27,6 +27,9 @@ 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, + DeciderArgs, DagSpec, Node, TaskFactory, diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index f643e290f45ec..f2d1e66f3b383 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -22,6 +22,7 @@ // Dag and supplies its arguments, the way calling a TaskFlow function does in // Python. +import type { CoordinatorClient } from "../coordinator/client.js"; import { DAG_SCHEMA_FIELDS, TASK_SCHEMA_FIELDS, @@ -31,7 +32,7 @@ import { import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; import type { JsonValue } from "./client-types.js"; import { getCurrentModuleSource } from "./module-source.js"; -import type { TaskFunction } from "./task.js"; +import { getClient, type TaskFunction } from "./task.js"; /** Internal: whether `value` is an object literal, not an array or a class instance. */ export function isPlainRecord(value: unknown): value is Record { @@ -303,6 +304,12 @@ export function isTaskRef(value: unknown): value is TaskRef { return hasBrand(value, "TaskRef"); } +const conditionTasks = new WeakMap(); + +function resolveNode(node: Node): Node { + return (typeof node === "object" && node !== null && conditionTasks.get(node)) || node; +} + /** * The first reference reachable inside `value`, or `undefined` when there is * none. @@ -401,13 +408,49 @@ export type TaskFactory = */ export type TaskOptions = TaskSpec; +/** + * A placed condition: name the task each outcome runs. `else` is optional. + * + * ```ts + * dag.if(hasRows, { rows: validated }).then(loadIfReady).else(loadFallback); + * ``` + */ +export interface Condition extends Node { + /** Airflow task ID of the deciding task. */ + readonly taskId: string; + then(taskRef: TaskRef | Condition): 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; +} + +/** The arguments after a decider's handler: its inputs, then its {@link TaskSpec}. */ +export type DeciderArgs = [TArgs] extends [void] + ? [inputs?: undefined, spec?: TaskOptions] + : [inputs: TaskInputs, spec?: TaskOptions]; + /** Per-task record a Dag retains: the reference, the handler, and its spec. */ export interface TaskRecord { readonly task: TaskRef; readonly fn: TaskFunction; readonly spec: TaskSpec; + /** Whether this task decides which of its downstream tasks to skip. */ + readonly canSkipDownstream?: boolean; +} + +interface ConditionRecord { + whenTrue?: TaskRef; + whenFalse?: TaskRef; } +// Mirrors `SkipMixin.skip` in `task-sdk/src/airflow/sdk/bases/skipmixin.py`. +const SKIPMIXIN_XCOM_KEY = "skipmixin_key"; +const SKIPMIXIN_SKIPPED = "skipped"; + /** * Internal: what one call to a task factory recorded, by argument name. * @@ -468,6 +511,7 @@ export class Dag { // insertion-ordered so the serialized Dag reads as written. readonly #orderEdges = new Map(); readonly #definedIn: string | undefined; + readonly #conditions = new Map(); // 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(); @@ -540,6 +584,119 @@ export class Dag { return this.#addTask(undefined, taskIdOrHandler, handlerOrOptions, maybeOptions); } + /** + * Declare a task whose boolean picks a branch; the side not taken is skipped. + * + * ```ts + * dag.if(hasRows, { rows: validated }).then(loadIfReady).else(loadFallback); + * dag.if(hasRows, { rows: validated }, { taskId: "has_rows", retries: 2 }).then(loadIfReady); + * ``` + */ + if( + handler: (args: TArgs) => boolean | Promise, + ...args: DeciderArgs> + ): Condition { + return this.#placeCondition(this.#placeDecider(handler, args) as TaskRef); + } + + #placeDecider(handler: (args: never) => unknown, args: readonly unknown[]): TaskRef { + const [inputs, spec] = args as [unknown, TaskOptions | undefined]; + const factory = this.#addTask(undefined, handler, spec) as (...inputs: unknown[]) => TaskRef; + return inputs === undefined ? factory() : factory(inputs); + } + + #placeCondition(condition: TaskRef): Condition { + const taskId = condition.taskId; + const branches: ConditionRecord = {}; + this.#conditions.set(taskId, branches); + this.#wrapDecider(taskId, async (held: unknown) => { + if (typeof held !== "boolean") { + throw new Error( + `Condition "${taskId}" of Dag "${this.dagId}" returned ${describeValue(held)} ` + + "rather than a boolean, so there is no branch to take", + ); + } + const skipped = held ? branches.whenFalse : branches.whenTrue; + return { skip: skipped ? [skipped.taskId] : [], result: held }; + }); + + const named = new Set<"then" | "else">(); + const name = (side: "then" | "else", target: TaskRef | Condition): void => { + // `await` calls `then(resolve, reject)`, so a function here means the condition was awaited. + if (typeof target === "function") { + throw new Error( + `dag.if(...) of Dag "${this.dagId}" was awaited. It builds a branch rather than ` + + "doing work, so there is nothing to wait for; drop the await", + ); + } + if (named.has(side)) { + throw new Error( + `Condition "${taskId}" of Dag "${this.dagId}" already has a "${side}" branch; ` + + "a condition names each side once", + ); + } + const taskRef = resolveNode(target) as TaskRef; + this.#validateOwnNode(taskRef, `the "${side}" branch of "${taskId}"`); + if (!isTaskRef(taskRef)) { + throw new Error( + `The "${side}" branch of Dag "${this.dagId}" condition "${taskId}" has to be a task, ` + + "not a task group", + ); + } + if (side === "else" && taskRef === branches.whenTrue) { + throw new Error( + `Both branches of Dag "${this.dagId}" condition "${taskId}" are ` + + `"${taskRef.taskId}", so the condition decides nothing; drop the else branch`, + ); + } + named.add(side); + if (side === "then") branches.whenTrue = taskRef; + else branches.whenFalse = taskRef; + condition.before(taskRef); + }; + + const elseStep: ConditionElse = { + else: (taskRef) => name("else", taskRef), + }; + const placed: Condition = { + dagId: this.dagId, + taskId, + then: (taskRef) => { + name("then", taskRef); + return elseStep; + }, + before: (...downstream) => { + condition.before(...downstream); + return placed; + }, + after: (...upstream) => { + condition.after(...upstream); + return placed; + }, + }; + conditionTasks.set(placed, condition); + return Object.freeze(placed); + } + + #wrapDecider( + taskId: string, + decide: (returned: unknown) => Promise<{ skip: string[]; result: unknown }>, + ): void { + const record = this.#tasks.get(taskId)!; + const inner = record.fn; + const wrapped: TaskFunction = async (args) => { + const { skip, result } = await decide(await inner(args as never)); + if (skip.length > 0) { + const client = getClient() as CoordinatorClient; + // Written before the skip so a cleared downstream is re-skipped, as SkipMixin does. + await client.setXCom({ key: SKIPMIXIN_XCOM_KEY, value: { [SKIPMIXIN_SKIPPED]: skip } }); + await client.skipDownstreamTasks(skip); + } + return result; + }; + this.#tasks.set(taskId, { ...record, canSkipDownstream: true, fn: wrapped }); + } + /** * Declare a task group of this Dag. * @@ -804,7 +961,9 @@ export class Dag { return Object.freeze(task); } - #addOrderEdge(upstream: Node, downstream: Node, verb: "before" | "after"): void { + #addOrderEdge(upstreamNode: Node, downstreamNode: Node, verb: "before" | "after"): void { + const upstream = resolveNode(upstreamNode); + const downstream = resolveNode(downstreamNode); if (this.#finalized) { throw new Error( `An edge was drawn on Dag "${this.dagId}" after the Dag was read; ` + @@ -814,7 +973,7 @@ export class Dag { // The argument is the one that can be foreign: the receiver is a node this // Dag handed out, since it is what carries the method. const other = verb === "before" ? downstream : upstream; - this.#validateOwnNode(other, verb); + this.#validateOwnNode(other, `${verb}()`); const upstreamId = nodeId(upstream); const downstreamId = nodeId(downstream); if (upstreamId === downstreamId) { @@ -833,17 +992,17 @@ export class Dag { } } - #validateOwnNode(node: Node, verb: string): void { + #validateOwnNode(node: Node, label: string): void { const id = nodeId(node); if (id === undefined) { throw new Error( - `${verb}() on Dag "${this.dagId}" takes tasks and task groups this Dag handed out, ` + + `${label} on Dag "${this.dagId}" takes tasks and task groups this Dag handed out, ` + "not arbitrary values", ); } if (node.dagId !== this.dagId) { throw new Error( - `${verb}() cannot draw an edge to Dag "${node.dagId}" node "${id}" from Dag ` + + `${label} cannot reach Dag "${node.dagId}" node "${id}" from Dag ` + `"${this.dagId}"; an edge joins two nodes of one Dag`, ); } @@ -853,7 +1012,7 @@ export class Dag { // the edge at this Dag's own group of that name. if (isTaskRef(node) ? this.#tasks.get(id)?.task !== node : this.#groupRefs.get(id) !== node) { throw new Error( - `${verb}() was given a reference to "${id}" that this Dag did not hand out; ` + + `${label} was given a reference to "${id}" that this Dag did not hand out; ` + `it comes from another Dag object with the same ID, or ${DUPLICATE_COPY_HINT}`, ); } @@ -938,6 +1097,14 @@ export class Dag { #finalize(): void { if (this.#finalized) return; + for (const [taskId, branches] of this.#conditions) { + if (branches.whenTrue === undefined) { + throw new Error( + `Condition "${taskId}" of Dag "${this.dagId}" names no branch, so it decides nothing; ` + + "give it one with dag.if(handler, inputs).then(task)", + ); + } + } for (const taskId of this.#tasks.keys()) { if (!this.#inputs.has(taskId)) { throw new Error( @@ -961,6 +1128,14 @@ const TASK_ID_CHARACTERS = /^[\p{L}\p{N}_.-]+$/u; const GROUP_ID_CHARACTERS = /^[\p{L}\p{N}_-]+$/u; const GROUP_ID_MAX_LENGTH = 200; +/** A value as an error message names it: its type, or the literal when short. */ +function describeValue(value: unknown): string { + if (value === null) return "null"; + if (value === undefined) return "undefined"; + if (typeof value === "object") return Array.isArray(value) ? "an array" : "an object"; + return `the ${typeof value} ${typeof value === "string" ? JSON.stringify(value) : String(value)}`; +} + /** * A handler's own function name, or undefined when it has none. * diff --git a/ts-sdk/tests/coordinator/client.test.ts b/ts-sdk/tests/coordinator/client.test.ts index 2934eec32f732..cba7154b51675 100644 --- a/ts-sdk/tests/coordinator/client.test.ts +++ b/ts-sdk/tests/coordinator/client.test.ts @@ -134,6 +134,24 @@ describe("deleteVariable", () => { }); }); +describe("skipDownstreamTasks", () => { + it("sends SkipDownstreamTasks with the task ids", async () => { + const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true }); + + await c.skipDownstreamTasks(["load_fallback"]); + + expect(sent).toEqual([{ type: "SkipDownstreamTasks", tasks: ["load_fallback"] }]); + }); + + it("sends nothing for an empty list", async () => { + const { client: c, sent } = recordingClient(); + + await c.skipDownstreamTasks([]); + + expect(sent).toEqual([]); + }); +}); + describe("writes do not read a supervisor 404 as absence", () => { it.each([ ["setVariable", "PutVariable", (c: TaskClient) => c.setVariable("k", "v")], diff --git a/ts-sdk/tests/sdk/conditional.test.ts b/ts-sdk/tests/sdk/conditional.test.ts new file mode 100644 index 0000000000000..f2487f2af6a80 --- /dev/null +++ b/ts-sdk/tests/sdk/conditional.test.ts @@ -0,0 +1,308 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, expect, it, vi } from "vitest"; +import { + Dag, + finalizeDag, + getDagOrderEdges, + getDagTaskInputs, + getDagTaskRecords, + type TaskRef, +} from "../../src/sdk/dag.js"; +import { serializeDag } from "../../src/coordinator/serde.js"; +import type { TaskClient } from "../../src/sdk/client.js"; +import { runInTaskScope, type TaskContext } from "../../src/sdk/task.js"; + +function gatedDag(dagId = "d", holds: boolean | unknown = true) { + const dag = new Dag(dagId); + const ifReady = dag.task("load_if_ready", async () => undefined)(); + const fallback = dag.task("load_fallback", async () => undefined)(); + const gated = dag.if(async () => holds as boolean, undefined, { taskId: "has_rows" }); + return { dag, ifReady, fallback, gated }; +} + +async function runHandler(dag: Dag, taskId: string, args: unknown = {}) { + const skipDownstreamTasks = vi.fn(async () => undefined); + const setXCom = vi.fn(async () => undefined); + const client = { skipDownstreamTasks, setXCom } as unknown as TaskClient; + const ctx = { dagId: dag.dagId, taskId } as unknown as TaskContext; + const handler = getDagTaskRecords(dag).get(taskId)!.fn; + const returned = await runInTaskScope({ ctx, client }, () => + (handler as (args: unknown) => Promise)(args), + ); + return { returned, skipDownstreamTasks, setXCom }; +} + +describe("dag.if", () => { + it("serializes the condition as a task that decides skips", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + gated.then(ifReady).else(fallback); + + const tasks = (serializeDag(dag, "", ".") as { tasks: { __var: Record }[] }) + .tasks; + const byId = new Map(tasks.map(({ __var }) => [__var["task_id"] as string, __var])); + expect(byId.get("has_rows")?.["_can_skip_downstream"]).toBe(true); + expect(byId.get("load_if_ready")).not.toHaveProperty("_can_skip_downstream"); + }); + + it("serializes the spec given to the condition", () => { + const { dag, ifReady } = gatedDag(); + dag + .if(async () => true, undefined, { taskId: "is_weekday", retries: 2, queue: "q" }) + .then(ifReady); + + const tasks = (serializeDag(dag, "", ".") as { tasks: { __var: Record }[] }) + .tasks; + expect(tasks.find(({ __var }) => __var["task_id"] === "is_weekday")?.__var).toMatchObject({ + retries: 2, + queue: "q", + }); + }); + + it("draws an order-only edge to each branch, since no value flows", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + gated.then(ifReady).else(fallback); + + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "has_rows", downstream: "load_if_ready" }, + { upstream: "has_rows", downstream: "load_fallback" }, + ]); + }); + + it("keeps the condition's own arguments and its return value", async () => { + const dag = new Dag("d"); + const loaded = dag.task("load", async () => undefined)(); + const extracted = dag.task("extract", async () => 3)(); + dag + .if( + async ({ rows }: { rows: number }) => rows > 0, + { rows: extracted }, + { taskId: "has_rows" }, + ) + .then(loaded); + + expect(getDagTaskInputs(dag).get("has_rows")).toEqual({ rows: extracted }); + expect((await runHandler(dag, "has_rows", { rows: 2 })).returned).toBe(true); + expect((await runHandler(dag, "has_rows", { rows: 0 })).returned).toBe(false); + }); + + it("takes its inputs without a task id", () => { + const dag = new Dag("d"); + const loaded = dag.task("load", async () => undefined)(); + const extracted = dag.task("extract", async () => 3)(); + async function hasRows({ rows }: { rows: number }): Promise { + return rows > 0; + } + dag.if(hasRows, { rows: extracted }).then(loaded); + + expect(getDagTaskInputs(dag).get("hasRows")).toEqual({ rows: extracted }); + }); + + it("requires the handler's inputs at compile time", () => { + const dag = new Dag("d"); + async function hasRows({ rows }: { rows: number }): Promise { + return rows > 0; + } + + // @ts-expect-error -- hasRows needs { rows }. + expect(() => dag.if(hasRows, undefined, { taskId: "has_rows" })).not.toThrow(); + }); + + it("takes the condition's id from the handler's name", () => { + const dag = new Dag("d"); + const loaded = dag.task("load", async () => undefined)(); + async function hasRows(): Promise { + return true; + } + const gate = dag.if(hasRows); + gate.then(loaded); + + expect(gate.taskId).toBe("hasRows"); + expect(dag.taskIds).toEqual(["load", "hasRows"]); + }); + + describe("at run time", () => { + it("skips the else branch when the condition holds", async () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + gated.then(ifReady).else(fallback); + + const { skipDownstreamTasks, setXCom } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_fallback"]); + expect(setXCom).toHaveBeenCalledWith({ + key: "skipmixin_key", + value: { skipped: ["load_fallback"] }, + }); + }); + + it("skips the then branch when it does not", async () => { + const { dag, ifReady, fallback, gated } = gatedDag("d", false); + gated.then(ifReady).else(fallback); + + const { skipDownstreamTasks } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_if_ready"]); + }); + + it("skips nothing, and records nothing, when a one-sided condition holds", async () => { + const { dag, ifReady, gated } = gatedDag(); + gated.then(ifReady); + + const { skipDownstreamTasks, setXCom } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).not.toHaveBeenCalled(); + expect(setXCom).not.toHaveBeenCalled(); + }); + + it("skips only its own branch when a one-sided condition fails", async () => { + const { dag, ifReady, gated } = gatedDag("d", false); + gated.then(ifReady); + + const { skipDownstreamTasks } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_if_ready"]); + }); + + it.each([ + ["a string", "yes"], + ["null", null], + ["a number", 1], + ["a bigint", 1n], + ])("fails the task when the condition returns %s", async (_label, value) => { + const { dag, ifReady, gated } = gatedDag("d", value); + gated.then(ifReady); + + await expect(runHandler(dag, "has_rows")).rejects.toThrow( + /Condition "has_rows" of Dag "d" returned .* rather than a boolean/, + ); + }); + }); + + it("carries edges of its own, like any task", () => { + const { dag, ifReady, gated } = gatedDag(); + const extracted = dag.task("extract", async () => undefined)(); + gated.then(ifReady); + + gated.after(extracted); + + expect(getDagOrderEdges(dag)).toContainEqual({ + upstream: "extract", + downstream: "has_rows", + }); + }); + + it("stands at the other end of an edge, as Go's .After(gate) does", () => { + const { dag, ifReady, gated } = gatedDag(); + const notified = dag.task("notify", async () => undefined)(); + gated.then(ifReady); + + notified.after(gated); + + expect(getDagOrderEdges(dag)).toContainEqual({ upstream: "has_rows", downstream: "notify" }); + }); + + it("takes another condition as a branch", () => { + const { dag, ifReady, gated } = gatedDag(); + const inner = dag.if(async () => true, undefined, { taskId: "is_weekday" }); + inner.then(ifReady); + + gated.then(inner); + + expect(getDagOrderEdges(dag)).toContainEqual({ + upstream: "has_rows", + downstream: "is_weekday", + }); + }); + + describe("rejects", () => { + it("a task group as a branch", () => { + const { dag, gated } = gatedDag(); + const group = dag.taskGroup("staging"); + + expect(() => gated.then(group as never)).toThrowError( + /The "then" branch of Dag "d" condition "has_rows" has to be a task, not a task group/, + ); + }); + + it("a branch taken from another Dag", () => { + const { gated } = gatedDag("here"); + const { ifReady: foreign } = gatedDag("there"); + + expect(() => gated.then(foreign)).toThrowError( + /the "then" branch of "has_rows" cannot reach Dag "there" node "load_if_ready"/, + ); + }); + + it.each([ + ["a plain object", { dagId: "d", taskId: "load_if_ready" }], + ["a string", "load_if_ready"], + ["null", null], + ])("%s where a branch belongs", (_label, value) => { + const { gated } = gatedDag(); + + expect(() => gated.then(value as TaskRef)).toThrowError( + /the "then" branch of "has_rows" on Dag "d" takes tasks and task groups/, + ); + }); + + it("both branches naming the same task, which decides nothing", () => { + const { ifReady, gated } = gatedDag(); + + expect(() => gated.then(ifReady).else(ifReady)).toThrowError( + /Both branches of Dag "d" condition "has_rows" are "load_if_ready", so the condition decides nothing/, + ); + }); + + it("a second else branch", () => { + const { ifReady, fallback, gated } = gatedDag(); + const chain = gated.then(ifReady); + chain.else(fallback); + + expect(() => chain.else(fallback)).toThrowError( + /Condition "has_rows" of Dag "d" already has an? "else" branch/, + ); + }); + + it("a second then branch", () => { + const { ifReady, fallback, gated } = gatedDag(); + gated.then(ifReady); + + expect(() => gated.then(fallback)).toThrowError( + /Condition "has_rows" of Dag "d" already has a "then" branch/, + ); + }); + + it("a condition that names no branch, when the Dag is read", () => { + const { dag } = gatedDag(); + + expect(() => finalizeDag(dag)).toThrowError( + /Condition "has_rows" of Dag "d" names no branch, so it decides nothing/, + ); + }); + + it("being awaited, which would resolve the chain rather than build it", async () => { + const { gated } = gatedDag(); + + await expect(Promise.resolve(gated as unknown as Promise)).rejects.toThrow( + /dag.if\(\.\.\.\) of Dag "d" was awaited/, + ); + }); + }); +}); diff --git a/ts-sdk/tests/sdk/dag.test.ts b/ts-sdk/tests/sdk/dag.test.ts index 72c064a3df411..2a5061041c3e2 100644 --- a/ts-sdk/tests/sdk/dag.test.ts +++ b/ts-sdk/tests/sdk/dag.test.ts @@ -630,9 +630,7 @@ describe("Dag", () => { const { refs: there } = placedDag("there", "cleanup"); expect(() => draw(here.load!, there.cleanup!)).toThrowError( - new RegExp( - `${verb}\\(\\) cannot draw an edge to Dag "there" node "cleanup" from Dag "here"`, - ), + new RegExp(`${verb}\\(\\) cannot reach Dag "there" node "cleanup" from Dag "here"`), ); });