From 0490d681e544ef0952b3a0996d7ceaef43c4bc30 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:07 +0000 Subject: [PATCH 1/6] TS SDK: add task groups to the native Dag surface --- .../language-sdks/typescript.rst | 25 ++ ts-sdk/src/index.ts | 3 + ts-sdk/src/sdk/dag.ts | 330 ++++++++++++++-- ts-sdk/tests/public-api.test.ts | 27 +- ts-sdk/tests/sdk/dag.test.ts | 357 +++++++++++++++++- 5 files changed, 703 insertions(+), 39 deletions(-) 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 b89173b80d86a..fffde1335d7e0 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -311,6 +311,31 @@ already exists changes nothing. Each returns the reference it was called on, so Pass a value as an argument when the downstream task needs it, and use ``before`` or ``after`` when it only needs to run in order. +Task groups +~~~~~~~~~~~ + +``dag.taskGroup(groupId)`` opens a scope with the same ``task`` and ``taskGroup`` methods as the Dag, +prefixing the id of everything declared in it, as Python's ``prefix_group_id`` does: + +.. code-block:: typescript + + const staging = dag.taskGroup("staging"); + staging.task("stage_rows", stageRows)(); // task id "staging.stage_rows" + staging.taskGroup("checks").task("nulls", checkNulls)(); // "staging.checks.nulls" + + staging.before(loaded); // staging >> loaded + +A group is an edge endpoint in its own right, so ``before`` and ``after`` order a whole group against +a task or against another group. + +Tasks and groups share one id namespace, as they do in Python, so a Dag cannot hold both a task and a +group called ``staging``. A ``.`` is what separates a group from what it holds, so it cannot appear in +an id of either. + +Pass ``{ prefixGroupId: false }`` to keep the ids declared in a group as written, as ``prefix_group_id=False`` +does in Python; they then have to be unique across the Dag. A group id is made of letters, digits, dashes and +underscores, and is at most 200 characters. + ``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/src/index.ts b/ts-sdk/src/index.ts index ad45c3bad336c..4e7bee9a9669b 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -28,7 +28,10 @@ export type { ArgNameMap } from "./sdk/arg-names.js"; export type { Registerable } from "./sdk/bundle.js"; export type { DagSpec, + Node, TaskFactory, + TaskGroupOptions, + TaskGroupRef, TaskInput, TaskInputs, TaskOptions, diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index aead708fe6222..8f7a2042deadb 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -136,7 +136,26 @@ declare const RETURN_TYPE: unique symbol; * object is not one: an input also takes a plain JSON value, and without the * brand such an object could not be told apart from an upstream reference. */ -export interface TaskRef { +/** + * What an order-only edge can connect: a task, or a whole task group. + * + * The TypeScript counterpart of Python's `DAGNode` and the Go SDK's + * `airflow.Node`. A group carries edges as a task does, so `before` and + * `after` take either. + */ +export interface Node { + /** Identifier of the Dag this node belongs to. */ + readonly dagId: string; + /** + * Run this node before each of `downstream`, carrying no value — the + * TypeScript spelling of Python's `>>`. + */ + before(...downstream: readonly Node[]): Node; + /** Run this node after each of `upstream`, carrying no value — Python's `<<`. */ + after(...upstream: readonly Node[]): Node; +} + +export interface TaskRef extends Node { /** Identifier of the Dag this task belongs to. */ readonly dagId: string; /** Airflow task ID, including any TaskGroup prefix. */ @@ -155,7 +174,7 @@ export interface TaskRef { * than its arguments: a fan-out has no single "next" reference to hand back. * Declaring an edge that already exists changes nothing. */ - before(...downstream: readonly TaskRef[]): TaskRef; + before(...downstream: readonly Node[]): TaskRef; /** * Run this task after each of `upstream`, carrying no value — Python's `<<`. * @@ -168,11 +187,59 @@ export interface TaskRef { * direction has an answer: named keys when values flow, `after` when only * order does. */ - after(...upstream: readonly TaskRef[]): TaskRef; + after(...upstream: readonly Node[]): TaskRef; } /** - * An order-only edge of a Dag: upstream task ID, then downstream task ID. + * A scope that declares tasks under a shared id prefix, and carries edges as a + * whole — Python's `TaskGroup`, spelled to TypeScript convention. + * + * Offers the same `task` and `taskGroup` methods as the Dag, so nesting is the + * same call at every depth. Each id it declares is prefixed with the group id + * unless {@link TaskGroupOptions.prefixGroupId} turns that off, and the group + * itself stands at either end of an edge, so a whole group can be ordered + * against a task or against another group. + */ +export interface TaskGroupRef extends Node { + /** Identifier of the Dag this group belongs to. */ + readonly dagId: string; + /** Group ID, including any enclosing group's prefix. */ + readonly groupId: string; + /** Declare a task in this group; its id carries the group prefix unless it is turned off. */ + task( + taskId: string, + handler: (args: TArgs) => TReturn | Promise, + options?: TaskOptions, + ): TaskFactory; + /** Declare a task whose id is the handler's function name, prefixed the same way. */ + task( + handler: (args: TArgs) => TReturn | Promise, + options?: TaskOptions, + ): TaskFactory; + /** Nest a group inside this one. */ + taskGroup(groupId: string, options?: TaskGroupOptions): TaskGroupRef; + before(...downstream: readonly Node[]): TaskGroupRef; + after(...upstream: readonly Node[]): TaskGroupRef; +} + +/** Options of a task group, the trailing argument of `taskGroup()`. */ +export interface TaskGroupOptions { + /** + * Whether the ids declared in the group are prefixed with the group's, as + * in `staging.load`. On by default, as in Python; turning it off keeps them + * as written, as `prefix_group_id=False` does, so they have to be unique + * across the Dag. + */ + readonly prefixGroupId?: boolean; +} + +/** + * An order-only edge of a Dag, between two node IDs. + * + * An endpoint is a task ID or a group ID; the two share one namespace, so a + * bare ID names exactly one node. {@link TaskGroupRecord} is what tells a + * consumer which kind an endpoint is, and which tasks a group endpoint stands + * for. * * Kept apart from the wiring a factory call records, because an edge that * carries no value has no argument name to be recorded under. @@ -185,6 +252,41 @@ export interface OrderEdge { // A task id cannot hold a NUL, so a joined pair cannot collide with one. const EDGE_KEY_SEPARATOR = "\u0000"; +/** Internal: one group of a Dag, and the tree beneath it. */ +export interface TaskGroupRecord { + /** Group ID, including any enclosing group's prefix. */ + readonly groupId: string; + /** Enclosing group's ID, absent for a group declared on the Dag itself. */ + readonly parentGroupId?: string; + /** Whether the IDs declared in this group carry its ID as a prefix. */ + readonly prefixGroupId: boolean; + /** Task IDs declared directly in this group, in declaration order. */ + readonly taskIds: readonly string[]; + /** Group IDs nested directly in this group, in declaration order. */ + readonly childGroupIds: readonly string[]; +} + +/** Separates a group ID from what it contains, as Python's `prefix_group_id` does. */ +const GROUP_SEPARATOR = "."; + +/** The Dag's own view of a group, which it appends to as an author declares. */ +interface MutableTaskGroupRecord extends TaskGroupRecord { + readonly taskIds: string[]; + readonly childGroupIds: string[]; +} + +/** Whether `value` is a task group returned by any copy of this package. */ +function isTaskGroupRef(value: unknown): value is TaskGroupRef { + return hasBrand(value, "TaskGroupRef"); +} + +/** The ID an edge endpoint is recorded under: a task ID or a group ID. */ +function nodeId(node: Node): string | undefined { + if (isTaskRef(node)) return node.taskId; + if (isTaskGroupRef(node)) return node.groupId; + return undefined; +} + /** Whether `value` is a TaskRef returned by any copy of this package. */ function isTaskRef(value: unknown): value is TaskRef { return hasBrand(value, "TaskRef"); @@ -317,6 +419,7 @@ let taskRecordsOf: (dag: Dag) => ReadonlyMap; let inputsOf: (dag: Dag) => ReadonlyMap; let orderEdgesOf: (dag: Dag) => readonly OrderEdge[]; let definedInOf: (dag: Dag) => string | undefined; +let groupsOf: (dag: Dag) => ReadonlyMap; let finalizeOf: (dag: Dag) => void; /** Internal: whether `value` is a Dag built by any copy of this package. */ @@ -354,6 +457,12 @@ export class Dag { // insertion-ordered so the serialized Dag reads as written. readonly #orderEdges = new Map(); readonly #definedIn: string | undefined; + // 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(); + // The one reference each group was handed out as, so an edge endpoint can be + // checked by identity the way a task's is. + readonly #groupRefs = new Map(); #finalized = false; static { @@ -361,6 +470,7 @@ export class Dag { inputsOf = (dag) => dag.#inputs; orderEdgesOf = (dag) => [...dag.#orderEdges.values()]; definedInOf = (dag) => dag.#definedIn; + groupsOf = (dag) => dag.#groups; finalizeOf = (dag) => dag.#finalize(); } @@ -415,6 +525,31 @@ export class Dag { taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, maybeOptions?: TaskOptions, + ): TaskFactory { + return this.#addTask(undefined, taskIdOrHandler, handlerOrOptions, maybeOptions); + } + + /** + * Declare a task group of this Dag. + * + * The group prefixes the id of everything declared in it, unless + * {@link TaskGroupOptions.prefixGroupId} turns that off, and stands at + * either end of an order-only edge in its own right. + * + * ```ts + * const staging = dag.taskGroup("staging"); // tasks "staging." + * const checks = dag.taskGroup("checks", { prefixGroupId: false }); // ids as written + * ``` + */ + taskGroup(groupId: string, options?: TaskGroupOptions): TaskGroupRef { + return this.#addGroup(undefined, groupId, options); + } + + #addTask( + groupId: string | undefined, + taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), + handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + maybeOptions?: TaskOptions, ): TaskFactory { const idGiven = typeof taskIdOrHandler === "string"; const handler = (idGiven ? handlerOrOptions : taskIdOrHandler) as ( @@ -435,8 +570,8 @@ export class Dag { ); } const defaulted = idGiven ? undefined : (specTaskId ?? readFunctionName(handler)); - const taskId = idGiven ? taskIdOrHandler : defaulted; - if (taskId === undefined) { + const declared = idGiven ? taskIdOrHandler : defaulted; + if (declared === undefined) { throw new Error( `A task of Dag "${this.dagId}" has no id: its handler has no name to take one from. ` + 'Pass an id — dag.task("my_task", handler) — or give the handler a name. A bundler ' + @@ -453,6 +588,15 @@ export class Dag { 'underscores — dag.task("my_task", handler)', ); } + if (declared.includes(GROUP_SEPARATOR)) { + // The separator is what joins a group to what it holds, so a task id + // carrying one would name a group that does not exist. + throw new Error( + `Task ID "${declared}" of Dag "${this.dagId}" cannot contain "${GROUP_SEPARATOR}"; ` + + "declare a task group with taskGroup(...) and the prefix is added for you", + ); + } + const taskId = this.#childId(groupId, declared); if (typeof handler !== "function") { throw new Error(`handler for Dag "${this.dagId}" task "${taskId}" must be a function`); } @@ -464,9 +608,6 @@ export class Dag { "a handler takes one object of named arguments — async ({ rows, region }) => ...", ); } - if (this.#tasks.has(taskId)) { - throw new Error(`Task "${taskId}" is already registered for Dag "${this.dagId}"`); - } // A task added after the Dag was read could no longer be wired into it, and // would sit in the Dag unplaced and unreported. if (this.#finalized) { @@ -476,7 +617,9 @@ export class Dag { ); } const spec = this.#taskSpecOf(taskId, options); + this.#reserveNodeId(taskId, "Task"); const task = this.#createTaskRef(taskId); + if (groupId !== undefined) this.#groups.get(groupId)!.taskIds.push(taskId); this.#tasks.set(taskId, { task, // The runtime dispatches every handler through one instantiation, as it @@ -517,6 +660,122 @@ export class Dag { return spec as TaskSpec; } + #addGroup( + parentGroupId: string | undefined, + groupId: string, + options: TaskGroupOptions = {}, + ): TaskGroupRef { + if (typeof groupId !== "string" || groupId.length === 0) { + throw new Error(`A task group of Dag "${this.dagId}" must have a non-empty ID`); + } + if (groupId.includes(GROUP_SEPARATOR)) { + // The separator is what joins a group to what it holds, so one inside an + // ID would make the resulting task id ambiguous. + throw new Error( + `Task group ID "${groupId}" of Dag "${this.dagId}" cannot contain ` + + `"${GROUP_SEPARATOR}"; nest groups with taskGroup(...) instead`, + ); + } + // Python's validate_group_key. The length is in code points, as Python's + // len() counts, not in UTF-16 units. + const length = [...groupId].length; + if (length > GROUP_ID_MAX_LENGTH) { + throw new Error( + `Task group ID "${groupId}" of Dag "${this.dagId}" has ${length} characters; ` + + `at most ${GROUP_ID_MAX_LENGTH} are allowed`, + ); + } + if (!GROUP_ID_CHARACTERS.test(groupId)) { + throw new Error( + `Task group ID "${groupId}" of Dag "${this.dagId}" has to be made of letters, digits, ` + + "dashes and underscores", + ); + } + const prefixGroupId = this.#prefixGroupIdOf(groupId, options); + if (this.#finalized) { + throw new Error( + `Task group "${groupId}" cannot be added to Dag "${this.dagId}" after the Dag was read; ` + + "declare every group while the module is loading", + ); + } + const fullId = this.#childId(parentGroupId, groupId); + this.#reserveNodeId(fullId, "Task group"); + this.#groups.set(fullId, { + groupId: fullId, + ...(parentGroupId !== undefined && { parentGroupId }), + prefixGroupId, + taskIds: [], + childGroupIds: [], + }); + if (parentGroupId !== undefined) this.#groups.get(parentGroupId)!.childGroupIds.push(fullId); + const group = this.#createTaskGroupRef(fullId); + this.#groupRefs.set(fullId, group); + return group; + } + + // TypeScript is bypassable, as for #taskSpecOf, so a misspelled or mistyped + // option is rejected rather than read as the default. + #prefixGroupIdOf(groupId: string, options: TaskGroupOptions): boolean { + const value: unknown = options; + if (!isPlainRecord(value)) { + throw new Error(`options for Dag "${this.dagId}" task group "${groupId}" must be an object`); + } + for (const key of Reflect.ownKeys(value)) { + if (key !== "prefixGroupId") { + throw new Error( + `Unknown option "${String(key)}" for Dag "${this.dagId}" task group "${groupId}"`, + ); + } + } + const { prefixGroupId = true } = value; + if (typeof prefixGroupId !== "boolean") { + throw new Error( + `prefixGroupId for Dag "${this.dagId}" task group "${groupId}" must be a boolean`, + ); + } + return prefixGroupId; + } + + // Python's TaskGroup.child_id: a group adds its ID in front of what it holds + // only when prefixGroupId is on, and a nested group's ID is built the same way. + #childId(groupId: string | undefined, id: string): string { + if (groupId === undefined || !this.#groups.get(groupId)!.prefixGroupId) return id; + return `${groupId}${GROUP_SEPARATOR}${id}`; + } + + // Tasks and groups share one namespace, as they do in Python: a serialized + // Dag addresses both by a bare ID, so `dag.task("x")` and `dag.taskGroup("x")` + // cannot both exist. + #reserveNodeId(id: string, kind: "Task" | "Task group"): void { + if (this.#tasks.has(id) || this.#groups.has(id)) { + throw new Error(`${kind} "${id}" is already registered for Dag "${this.dagId}"`); + } + } + + #createTaskGroupRef(groupId: string): TaskGroupRef { + const group: TaskGroupRef = { + dagId: this.dagId, + groupId, + task: ( + taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), + handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + maybeOptions?: TaskOptions, + ) => this.#addTask(groupId, taskIdOrHandler, handlerOrOptions, maybeOptions), + taskGroup: (childId: string, options?: TaskGroupOptions) => + this.#addGroup(groupId, childId, options), + before: (...downstream) => { + for (const other of downstream) this.#addOrderEdge(group, other, "before"); + return group; + }, + after: (...upstream) => { + for (const other of upstream) this.#addOrderEdge(other, group, "after"); + return group; + }, + } as TaskGroupRef; + brand(group, "TaskGroupRef"); + return Object.freeze(group); + } + #createTaskRef(taskId: string): TaskRef { const task: TaskRef = { dagId: this.dagId, @@ -534,51 +793,56 @@ export class Dag { return Object.freeze(task); } - #addOrderEdge(upstream: TaskRef, downstream: TaskRef, verb: "before" | "after"): void { + #addOrderEdge(upstream: Node, downstream: Node, verb: "before" | "after"): void { if (this.#finalized) { throw new Error( `An edge was drawn on Dag "${this.dagId}" after the Dag was read; ` + "declare every edge while the module is loading", ); } - // The argument is the one that can be foreign: the receiver is a reference - // this Dag handed out, since it is what carries the method. + // 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.#validateOwnRef(other, verb); - if (upstream.taskId === downstream.taskId) { + this.#validateOwnNode(other, verb); + const upstreamId = nodeId(upstream); + const downstreamId = nodeId(downstream); + if (upstreamId === downstreamId) { throw new Error( - `${verb}() cannot draw an edge from task "${upstream.taskId}" of Dag "${this.dagId}" to ` + - "itself; an edge orders two different tasks", + `${verb}() cannot draw an edge from node "${upstreamId}" of Dag "${this.dagId}" to ` + + "itself; an edge orders two different nodes", ); } - const key = `${upstream.taskId}${EDGE_KEY_SEPARATOR}${downstream.taskId}`; + const key = `${upstreamId}${EDGE_KEY_SEPARATOR}${downstreamId}`; // Idempotent, so an edge drawn from both ends is one edge. if (!this.#orderEdges.has(key)) { this.#orderEdges.set( key, - Object.freeze({ upstream: upstream.taskId, downstream: downstream.taskId }), + Object.freeze({ upstream: upstreamId!, downstream: downstreamId! }), ); } } - #validateOwnRef(ref: TaskRef, verb: string): void { - if (!isTaskRef(ref)) { + #validateOwnNode(node: Node, verb: string): void { + const id = nodeId(node); + if (id === undefined) { throw new Error( - `${verb}() on Dag "${this.dagId}" takes task references returned by calling a task, ` + + `${verb}() on Dag "${this.dagId}" takes tasks and task groups this Dag handed out, ` + "not arbitrary values", ); } - if (ref.dagId !== this.dagId) { + if (node.dagId !== this.dagId) { throw new Error( - `${verb}() cannot draw an edge to Dag "${ref.dagId}" task "${ref.taskId}" from Dag ` + - `"${this.dagId}"; an edge joins two tasks of one Dag`, + `${verb}() cannot draw an edge to Dag "${node.dagId}" node "${id}" from Dag ` + + `"${this.dagId}"; an edge joins two nodes of one Dag`, ); } - // Identity, not the ID pair: two Dag objects can carry the same dagId, and - // a second resolved copy of this package brands its own references. - if (this.#tasks.get(ref.taskId)?.task !== ref) { + // Identity, not the ID: two Dag objects can carry the same dagId, and a + // second resolved copy of this package brands its own nodes. A group is + // checked the same way — matching on the ID alone would silently retarget + // 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 "${ref.taskId}" that this Dag did not hand out; ` + + `${verb}() 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}`, ); } @@ -681,6 +945,11 @@ export class Dag { // has to look like. const TASK_ID_CHARACTERS = /^[\p{L}\p{N}_.-]+$/u; +// Python's GROUP_KEY_REGEX and validate_group_key limit: a group ID has no dot, +// since the dot is what joins it to what it holds. +const GROUP_ID_CHARACTERS = /^[\p{L}\p{N}_-]+$/u; +const GROUP_ID_MAX_LENGTH = 200; + /** * A handler's own function name, or undefined when it has none. * @@ -716,6 +985,11 @@ export function getDagTaskRecords(dag: Dag): ReadonlyMap { return taskRecordsOf(dag); } +/** Internal: every task group of a Dag, keyed by full group ID. */ +export function getDagTaskGroups(dag: Dag): ReadonlyMap { + return groupsOf(dag); +} + /** Internal: the order-only edges of a Dag, in the order they were drawn. */ export function getDagOrderEdges(dag: Dag): readonly OrderEdge[] { return orderEdgesOf(dag); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index be48599841390..769a8342b0e54 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -24,12 +24,15 @@ import type { ConnectionResult, DagSpec, GetXComOpts, + Node, SetXComOpts, TaskClient, Registerable, TaskContext, TaskFactory, TaskFunction, + TaskGroupOptions, + TaskGroupRef, TaskInput, TaskInputs, TaskOptions, @@ -301,11 +304,31 @@ describe("public API", () => { // Variadic, and each returns its own receiver rather than its arguments, // return type included. expectTypeOf["before"]>().toEqualTypeOf< - (...downstream: readonly TaskRef[]) => TaskRef + (...downstream: readonly Node[]) => TaskRef >(); expectTypeOf["after"]>().toEqualTypeOf< - (...upstream: readonly TaskRef[]) => TaskRef + (...upstream: readonly Node[]) => TaskRef >(); + // A group is a scope and an edge endpoint, and nests the same way at + // every depth. + expectTypeOf>().toEqualTypeOf< + "dagId" | "groupId" | "task" | "taskGroup" | "before" | "after" + >(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf< + (groupId: string, options?: TaskGroupOptions) => TaskGroupRef + >(); + expectTypeOf().toEqualTypeOf< + (groupId: string, options?: TaskGroupOptions) => TaskGroupRef + >(); + // Python's prefix_group_id, and the only option a group takes so far. + expectTypeOf().toEqualTypeOf<{ readonly prefixGroupId?: boolean }>(); + expectTypeOf().toEqualTypeOf< + (...downstream: readonly Node[]) => TaskGroupRef + >(); + // Both satisfy Node, which is what lets an edge join either kind. + expectTypeOf().toExtend(); + expectTypeOf().toExtend(); // Wiring moved to the factory call, and the spec is the trailing argument // itself. expectTypeOf().toEqualTypeOf(); diff --git a/ts-sdk/tests/sdk/dag.test.ts b/ts-sdk/tests/sdk/dag.test.ts index 612cdc9cdb084..72c064a3df411 100644 --- a/ts-sdk/tests/sdk/dag.test.ts +++ b/ts-sdk/tests/sdk/dag.test.ts @@ -22,6 +22,7 @@ import { Dag, finalizeDag, getDagOrderEdges, + getDagTaskGroups, getDagTaskInputs, getDagTaskRecords, type DagSpec, @@ -585,10 +586,10 @@ describe("Dag", () => { const { refs } = placedDag("d", "a"); expect(() => refs.a!.before(refs.a!)).toThrowError( - /before\(\) cannot draw an edge from task "a" of Dag "d" to itself/, + /before\(\) cannot draw an edge from node "a" of Dag "d" to itself/, ); expect(() => refs.a!.after(refs.a!)).toThrowError( - /after\(\) cannot draw an edge from task "a" of Dag "d" to itself/, + /after\(\) cannot draw an edge from node "a" of Dag "d" to itself/, ); }); @@ -630,7 +631,7 @@ describe("Dag", () => { expect(() => draw(here.load!, there.cleanup!)).toThrowError( new RegExp( - `${verb}\\(\\) cannot draw an edge to Dag "there" task "cleanup" from Dag "here"`, + `${verb}\\(\\) cannot draw an edge to Dag "there" node "cleanup" from Dag "here"`, ), ); }); @@ -652,7 +653,7 @@ describe("Dag", () => { const { refs } = placedDag("d", "load"); expect(() => refs.load!.before(value as unknown as TaskRef)).toThrowError( - /before\(\) on Dag "d" takes task references returned by calling a task/, + /before\(\) on Dag "d" takes tasks and task groups this Dag handed out/, ); }); @@ -678,6 +679,332 @@ describe("Dag", () => { }); }); + describe("task groups", () => { + it("prefixes the id of every task declared in it", () => { + const dag = new Dag("grouped"); + const staging = dag.taskGroup("staging"); + staging.task("stage_rows", async () => undefined)(); + + expect(dag.taskIds).toEqual(["staging.stage_rows"]); + }); + + it("prefixes a task id defaulted from the handler name", () => { + const dag = new Dag("grouped"); + dag.taskGroup("staging").task(async function stageRows() {})(); + + expect(dag.taskIds).toEqual(["staging.stageRows"]); + }); + + it("nests, joining every enclosing group's id", () => { + const dag = new Dag("grouped"); + const outer = dag.taskGroup("outer"); + const inner = outer.taskGroup("inner"); + inner.task("deep", async () => undefined)(); + + expect(dag.taskIds).toEqual(["outer.inner.deep"]); + expect(inner.groupId).toBe("outer.inner"); + }); + + it("carries the Dag's own identity", () => { + const dag = new Dag("grouped"); + expect(dag.taskGroup("staging").dagId).toBe("grouped"); + }); + + it("records the tree, so nested groups serialize as one", () => { + const dag = new Dag("grouped"); + const outer = dag.taskGroup("outer"); + outer.task("first", async () => undefined)(); + const inner = outer.taskGroup("inner"); + inner.task("second", async () => undefined)(); + + expect([...getDagTaskGroups(dag).values()]).toEqual([ + { + groupId: "outer", + prefixGroupId: true, + taskIds: ["outer.first"], + childGroupIds: ["outer.inner"], + }, + { + groupId: "outer.inner", + parentGroupId: "outer", + prefixGroupId: true, + taskIds: ["outer.inner.second"], + childGroupIds: [], + }, + ]); + }); + + it("lets the same handler name be reused across groups", () => { + const dag = new Dag("grouped"); + dag.taskGroup("north").task(async function extract() {})(); + dag.taskGroup("south").task(async function extract() {})(); + + expect(dag.taskIds).toEqual(["north.extract", "south.extract"]); + }); + + it("lets a task in a group be wired to one outside it", () => { + const dag = new Dag("grouped"); + const extracted = dag.task("extract", async () => 1)(); + const staged = dag + .taskGroup("staging") + .task( + "stage", + async (_: { extracted: number }) => undefined, + )({ extracted }); + + expect(staged.taskId).toBe("staging.stage"); + expect(getDagTaskInputs(dag).get("staging.stage")).toEqual({ extracted }); + }); + + describe("as an edge endpoint", () => { + it("orders a whole group before a task", () => { + const dag = new Dag("grouped"); + const staging = dag.taskGroup("staging"); + staging.task("stage", async () => undefined)(); + const loaded = dag.task("load", async () => undefined)(); + + staging.before(loaded); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "staging", downstream: "load" }]); + }); + + it("orders a task before a whole group", () => { + const dag = new Dag("grouped"); + const extracted = dag.task("extract", async () => undefined)(); + const staging = dag.taskGroup("staging"); + staging.task("stage", async () => undefined)(); + + extracted.before(staging); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "extract", downstream: "staging" }]); + }); + + it("orders one group against another", () => { + const dag = new Dag("grouped"); + const first = dag.taskGroup("first"); + first.task("a", async () => undefined)(); + const second = dag.taskGroup("second"); + second.task("b", async () => undefined)(); + + second.after(first); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "first", downstream: "second" }]); + }); + + it("returns its own receiver, as a task does", () => { + const dag = new Dag("grouped"); + const staging = dag.taskGroup("staging"); + staging.task("stage", async () => undefined)(); + const loaded = dag.task("load", async () => undefined)(); + + expect(staging.before(loaded)).toBe(staging); + }); + + it("rejects a group from another Dag object with the same ID", () => { + const first = new Dag("same_id"); + const second = new Dag("same_id"); + const loaded = first.task("load", async () => undefined)(); + const foreign = second.taskGroup("staging"); + + expect(() => loaded.before(foreign)).toThrowError( + /before\(\) was given a reference to "staging" that this Dag did not hand out/, + ); + }); + + it("rejects a foreign group even when this Dag has one of the same ID", () => { + // Matching on the ID alone used to accept it, and the edge was then + // drawn at this Dag's own group of that name. + const first = new Dag("same_id"); + const second = new Dag("same_id"); + first.taskGroup("staging").task("stage", async () => undefined)(); + const loaded = first.task("load", async () => undefined)(); + const foreign = second.taskGroup("staging"); + + expect(() => loaded.before(foreign)).toThrowError( + /before\(\) was given a reference to "staging" that this Dag did not hand out/, + ); + expect(getDagOrderEdges(first)).toEqual([]); + }); + }); + + it.each([ + ["an empty ID", ""], + ["a non-string ID", 42], + ])("rejects %s", (_label, groupId) => { + const dag = new Dag("grouped"); + expect(() => dag.taskGroup(groupId as string)).toThrowError( + /A task group of Dag "grouped" must have a non-empty ID/, + ); + }); + + it("rejects a group ID holding the separator, which nesting is for", () => { + const dag = new Dag("grouped"); + expect(() => dag.taskGroup("outer.inner")).toThrowError( + /Task group ID "outer.inner" of Dag "grouped" cannot contain "\."; nest groups with taskGroup/, + ); + }); + + it.each([ + [ + "a task and a group", + (dag: Dag) => [() => dag.task("x", async () => undefined), () => dag.taskGroup("x")], + ], + [ + "a group and a task", + (dag: Dag) => [() => dag.taskGroup("x"), () => dag.task("x", async () => undefined)], + ], + ["two groups", (dag: Dag) => [() => dag.taskGroup("x"), () => dag.taskGroup("x")]], + ])( + "rejects %s sharing one ID, since a serialized Dag addresses both by it", + (_label, build) => { + const dag = new Dag("grouped"); + const [first, second] = build(dag); + + first!(); + expect(() => second!()).toThrowError(/"x" is already registered for Dag "grouped"/); + }, + ); + + it("allows the same group ID under different parents", () => { + const dag = new Dag("grouped"); + dag + .taskGroup("north") + .taskGroup("shared") + .task("t", async () => undefined)(); + dag + .taskGroup("south") + .taskGroup("shared") + .task("t", async () => undefined)(); + + expect(dag.taskIds).toEqual(["north.shared.t", "south.shared.t"]); + }); + + describe("with prefixGroupId off", () => { + it("keeps the ids declared in it as written", () => { + const dag = new Dag("grouped"); + const checks = dag.taskGroup("checks", { prefixGroupId: false }); + checks.task("nulls", async () => undefined)(); + const inner = checks.taskGroup("inner"); + + expect(dag.taskIds).toEqual(["nulls"]); + expect(checks.groupId).toBe("checks"); + expect(inner.groupId).toBe("inner"); + expect(getDagTaskGroups(dag).get("checks")?.prefixGroupId).toBe(false); + }); + + it("lets a nested group prefix again, as Python does", () => { + const dag = new Dag("grouped"); + const inner = dag.taskGroup("outer", { prefixGroupId: false }).taskGroup("inner"); + inner.task("t", async () => undefined)(); + + expect(inner.groupId).toBe("inner"); + expect(dag.taskIds).toEqual(["inner.t"]); + }); + + it("still takes its own place under a prefixing parent", () => { + const dag = new Dag("grouped"); + const inner = dag.taskGroup("outer").taskGroup("inner", { prefixGroupId: false }); + inner.task("t", async () => undefined)(); + + expect(inner.groupId).toBe("outer.inner"); + expect(dag.taskIds).toEqual(["t"]); + }); + + it("rejects an id that clashes with one outside the group", () => { + const dag = new Dag("grouped"); + dag.task("nulls", async () => undefined)(); + const checks = dag.taskGroup("checks", { prefixGroupId: false }); + + expect(() => checks.task("nulls", async () => undefined)).toThrowError( + /Task "nulls" is already registered for Dag "grouped"/, + ); + }); + }); + + it.each([ + [ + "an unknown option", + { prefix: false }, + /Unknown option "prefix" for Dag "grouped" task group "checks"/, + ], + [ + "a prefixGroupId that is not a boolean", + { prefixGroupId: "no" }, + /prefixGroupId for Dag "grouped" task group "checks" must be a boolean/, + ], + [ + "options that are not an object", + 7, + /options for Dag "grouped" task group "checks" must be an object/, + ], + ])("rejects %s", (_label, options, expected) => { + const dag = new Dag("grouped"); + + expect(() => dag.taskGroup("checks", options as never)).toThrowError(expected); + expect(getDagTaskGroups(dag).size).toBe(0); + }); + + describe("its ID, which follows Python's validate_group_key", () => { + it("may hold letters of any script, digits, dashes and underscores", () => { + const dag = new Dag("grouped"); + + expect(dag.taskGroup("región-1_組").groupId).toBe("región-1_組"); + }); + + it.each([ + ["a space", "two words"], + ["a slash", "a/b"], + ["a symbol", "cost$"], + ])("rejects %s", (_label, groupId) => { + const dag = new Dag("grouped"); + + expect(() => dag.taskGroup(groupId)).toThrowError( + `Task group ID "${groupId}" of Dag "grouped" has to be made of letters, digits, dashes and underscores`, + ); + }); + + it("allows 200 characters and rejects 201", () => { + const dag = new Dag("grouped"); + dag.taskGroup("g".repeat(200)); + + expect(() => dag.taskGroup("h".repeat(201))).toThrowError( + /has 201 characters; at most 200 are allowed/, + ); + }); + + it("counts characters as Python's len() does, not UTF-16 units", () => { + const dag = new Dag("grouped"); + // Each of these letters is two UTF-16 units, so the id is 400 units long. + const groupId = "\u{20000}".repeat(200); + + expect(dag.taskGroup(groupId).groupId).toBe(groupId); + }); + }); + + it("rejects a group declared after the Dag was read", () => { + const dag = new Dag("grouped"); + dag.task("extract", async () => undefined)(); + finalizeDag(dag); + + expect(() => dag.taskGroup("late")).toThrowError( + /Task group "late" cannot be added to Dag "grouped" after the Dag was read/, + ); + }); + + it("holds its tasks to the same every-task-is-called rule", () => { + const dag = new Dag("grouped"); + dag.taskGroup("staging").task("stage", async () => undefined); + + expect(() => finalizeDag(dag)).toThrowError( + /Task "staging.stage" of Dag "grouped" is never called/, + ); + }); + + it("has no groups before any are declared", () => { + expect(getDagTaskGroups(new Dag("plain")).size).toBe(0); + }); + }); + it("exposes its task IDs in attachment order", () => { const dag = new Dag("ordered_dag"); expect(dag.taskIds).toEqual([]); @@ -734,13 +1061,25 @@ describe("Dag", () => { ); }); - it("treats a dotted TaskGroup taskId as a single taskId (group.task)", () => { + it("rejects a dotted task id, which names a group that does not exist", () => { + // A dotted id used to be accepted verbatim, giving a task whose id carries + // a group prefix while the Dag holds no such group. `taskGroup(...)` is + // what puts a task under a prefix now. const dag = new Dag("example_dag"); - dag.task("transforms.normalize", async () => "ok")(); - const bundle = new Bundle(); - bundle.register(dag); + + expect(() => dag.task("transforms.normalize", async () => "ok")).toThrowError( + /Task ID "transforms.normalize" of Dag "example_dag" cannot contain "\."/, + ); + expect(dag.taskIds).toEqual([]); + }); + + it("names a task inside a group without the author writing the prefix", () => { + const dag = new Dag("example_dag"); + dag.taskGroup("transforms").task("normalize", async () => "ok")(); + const bundle = new Bundle(dag); + expect(bundle.getTaskHandler("example_dag", "transforms.normalize")).toBeDefined(); - // Should NOT accidentally match the prefix alone + // And not the prefix or the leaf on its own. expect(bundle.getTaskHandler("example_dag", "transforms")).toBeUndefined(); expect(bundle.getTaskHandler("example_dag", "normalize")).toBeUndefined(); }); From a2b0c391ebd9901ec4fd8883ce864fbdb11c084c Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:42 +0000 Subject: [PATCH 2/6] TS SDK: serialize a native Dag to Dag JSON --- .pre-commit-config.yaml | 15 + .../language-sdks/typescript.rst | 23 + scripts/ci/lang_sdk_serialization/__init__.py | 16 + scripts/ci/lang_sdk_serialization/compare.py | 230 +++++ .../serialize_python.py | 121 +++ .../ci/lang_sdk_serialization/test_dags.yaml | 166 ++++ .../ci/lang_sdk_serialization/__init__.py | 16 + .../ci/lang_sdk_serialization/test_compare.py | 234 ++++++ ts-sdk/.pre-commit-config.yaml | 14 + ts-sdk/package.json | 3 +- ts-sdk/pnpm-lock.yaml | 27 +- .../prek/check_serialization_conformance.py | 44 + ts-sdk/src/coordinator/serde.ts | 729 ++++++++++++++++ ts-sdk/src/sdk/dag.ts | 22 +- .../tests/conformance/serialize_typescript.ts | 163 ++++ ts-sdk/tests/coordinator/arg-binding.test.ts | 19 + ts-sdk/tests/coordinator/serde.test.ts | 785 ++++++++++++++++++ ts-sdk/tests/public-api.test.ts | 4 + 18 files changed, 2616 insertions(+), 15 deletions(-) create mode 100644 scripts/ci/lang_sdk_serialization/__init__.py create mode 100644 scripts/ci/lang_sdk_serialization/compare.py create mode 100644 scripts/ci/lang_sdk_serialization/serialize_python.py create mode 100644 scripts/ci/lang_sdk_serialization/test_dags.yaml create mode 100644 scripts/tests/ci/lang_sdk_serialization/__init__.py create mode 100644 scripts/tests/ci/lang_sdk_serialization/test_compare.py create mode 100755 ts-sdk/scripts/ci/prek/check_serialization_conformance.py create mode 100644 ts-sdk/src/coordinator/serde.ts create mode 100644 ts-sdk/tests/conformance/serialize_typescript.ts create mode 100644 ts-sdk/tests/coordinator/serde.test.ts diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a454338a4105e..dcc63d45a4e8c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -337,6 +337,21 @@ repos: (?x) ^airflow-core/src/airflow/serialization/schema\.json$| ^ts-sdk/schema/dag-schema\.json$ + # check-ts-sdk-serialization-conformance in ts-sdk/ covers the SDK's own files; this runs the same + # check when the shared harness or Airflow's serializer changes, which that project cannot see. + - id: check-ts-sdk-serialization-conformance-shared + name: Check the TS SDK serializes Dags the way Airflow does, after a shared change + description: "Serialize the shared test Dags with the TS SDK and with Airflow, and compare the two" + entry: ./ts-sdk/scripts/ci/prek/check_serialization_conformance.py + language: node + additional_dependencies: ['pnpm@10.28.1'] + pass_filenames: false + require_serial: true + files: > + (?x) + ^airflow-core/src/airflow/serialization/serialized_objects\.py$| + ^airflow-core/src/airflow/serialization/schema\.json$| + ^scripts/ci/lang_sdk_serialization/.*$ - id: check-go-version-in-sync name: Check Go toolchain version is consistent across build files entry: ./scripts/ci/prek/check_go_version_in_sync.py 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 fffde1335d7e0..c81b2b3dc5600 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -336,6 +336,29 @@ Pass ``{ prefixGroupId: false }`` to keep the ids declared in a group as written does in Python; they then have to be unique across the Dag. A group id is made of letters, digits, dashes and underscores, and is at most 200 characters. +Serialization +~~~~~~~~~~~~~ + +A native Dag serializes into the same Dag JSON a Python Dag produces, so the scheduler reads it +without knowing which language declared it. + +``schedule`` accepts what maps to a stock timetable: unset, ``@once``, ``@continuous``, or a cron +expression. A cron preset such as ``@daily`` is recorded as the expression it stands for. Anything +else names a Python object a TypeScript bundle cannot point at, and is rejected. + +Every task of a native Dag runs on the Node coordinator, so it needs the queue the deployment routes +there. Set it once on the Dag and each task inherits it: + +.. code-block:: typescript + + const dag = new Dag("ts_etl", { schedule: "@daily", queue: "typescript" }); + + // ...and one task that needs its own. + dag.task("heavy", heavyHandler, { queue: "typescript_large" })(); + +``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. + ``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/scripts/ci/lang_sdk_serialization/__init__.py b/scripts/ci/lang_sdk_serialization/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/scripts/ci/lang_sdk_serialization/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/scripts/ci/lang_sdk_serialization/compare.py b/scripts/ci/lang_sdk_serialization/compare.py new file mode 100644 index 0000000000000..596d0d2f0d8ad --- /dev/null +++ b/scripts/ci/lang_sdk_serialization/compare.py @@ -0,0 +1,230 @@ +# 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. +r""" +Check that a language SDK serializes Dags the way Airflow does. + +Runs the SDK's serializer and serialize_python.py over test_dags.yaml, each writing its serialization +to a JSON file in a temporary directory, and compares the two field by field. The Python side also +loads the SDK's output through Airflow's deserializer, so this fails both when the two serializers +drift apart and when Airflow cannot read what the SDK writes. + +The SDK's serializer is the command after ``--``. It runs from the repository root, with the paths of +test_dags.yaml and of the JSON file to write appended, and writes each Dag as +``DagSerialization.to_dict`` would, keyed by Dag id. The Python side needs ``uv``:: + + python3 scripts/ci/lang_sdk_serialization/compare.py --sdk typescript -- \ + pnpm --dir ts-sdk exec tsx tests/conformance/serialize_typescript.ts + +On a failure both files are kept, and their directory is printed. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +REPO_ROOT = HERE.parents[2] +TEST_DAGS = HERE / "test_dags.yaml" +SCHEMA = REPO_ROOT / "airflow-core" / "src" / "airflow" / "serialization" / "schema.json" + +# Name the file a Dag was declared in: Python's own Dag file, or the SDK's bundle. +DAG_KEYS_NOT_COMPARED = frozenset({"fileloc", "relative_fileloc", "_processor_dags_folder"}) + +TASK_KEYS_NOT_COMPARED = frozenset( + { + # Operator identity: Python names the Python class that ran, a language SDK a fixed pair. + # Neither is ever imported on the Airflow side. + "task_type", + "_task_module", + # The SDK's marker for a task in its language, which Python has no equivalent for. + "language", + # Python bookkeeping for mapped tasks and retry policies, neither of which a language SDK's + # Dag can declare. + "_needs_expansion", + "has_retry_policy", + # The stub contract every language SDK task carries, which the Python operator built here + # has no counterpart for. Loading the SDK's output through Airflow checks them. + "is_stub", + "_arg_bindings", + } +) + + +def run(command: list[str]) -> None: + if subprocess.run(command, cwd=REPO_ROOT, check=False).returncode: + raise SystemExit(f"`{' '.join(command)}` failed") + + +def get_task_defaults() -> dict[str, Any]: + """Map each task field the Dag schema gives a default to that default.""" + fields = json.loads(SCHEMA.read_text())["definitions"]["operator"]["properties"] + return {key: field["default"] for key, field in fields.items() if field.get("default") is not None} + + +def normalize_as_javascript(value: Any) -> Any: + """Read a JSON value as JavaScript does: one number type, and a bool that is not a number.""" + if isinstance(value, bool): + return ("bool", value) + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, list): + return [normalize_as_javascript(item) for item in value] + if isinstance(value, dict): + return {key: normalize_as_javascript(item) for key, item in value.items()} + return value + + +def is_same_json(python: Any, sdk: Any) -> bool: + return normalize_as_javascript(python) == normalize_as_javascript(sdk) + + +def find_differences(path: str, python: Any, sdk: Any) -> list[str]: + """List where two JSON values differ, down to the innermost key or index.""" + if isinstance(python, dict) and isinstance(sdk, dict): + problems = [] + for key in sorted(python.keys() | sdk.keys()): + if key not in sdk: + problems.append(f"{path}.{key} is missing, Python writes {python[key]!r}") + elif key not in python: + problems.append(f"{path}.{key} is {sdk[key]!r}, which Python does not write") + else: + problems.extend(find_differences(f"{path}.{key}", python[key], sdk[key])) + return problems + if isinstance(python, list) and isinstance(sdk, list) and len(python) == len(sdk): + return [ + problem + for index, (python_item, sdk_item) in enumerate(zip(python, sdk)) + for problem in find_differences(f"{path}[{index}]", python_item, sdk_item) + ] + if not is_same_json(python, sdk): + return [f"{path} is {sdk!r}, Python writes {python!r}"] + return [] + + +def compare_fields( + python: dict[str, Any], sdk: dict[str, Any], not_compared: frozenset[str], defaults: dict[str, Any] +) -> list[str]: + """ + Compare two serialized objects key by key. + + A key the SDK leaves out is fine when Python wrote its schema default. Python keeps such a value + when its ``client_defaults`` table disagrees with the schema, a table a language SDK is never sent, + and Airflow reads a missing field as its default anyway. + """ + problems = [] + for key in sorted((python.keys() | sdk.keys()) - not_compared): + if key in python and key in sdk: + problems.extend(find_differences(key, python[key], sdk[key])) + elif key in sdk: + problems.append(f"{key} is {sdk[key]!r}, which Python does not write") + elif key not in defaults or not is_same_json(python[key], defaults[key]): + problems.append(f"{key} is missing, Python writes {python[key]!r}") + return problems + + +def compare_dag(python: dict[str, Any], sdk: dict[str, Any], defaults: dict[str, Any]) -> list[str]: + problems = compare_fields(python, sdk, DAG_KEYS_NOT_COMPARED | {"tasks"}, {}) + python_ids = [task["__var"]["task_id"] for task in python["tasks"]] + sdk_ids = [task["__var"]["task_id"] for task in sdk["tasks"]] + if sdk_ids != python_ids: + return [*problems, f"tasks are {sdk_ids}, Python writes {python_ids}"] + for python_task, sdk_task in zip(python["tasks"], sdk["tasks"]): + task_id = python_task["__var"]["task_id"] + if sdk_task["__type"] != python_task["__type"]: + problems.append(f"task {task_id} is a {sdk_task['__type']!r}, not a {python_task['__type']!r}") + problems.extend( + f"task {task_id}: {problem}" + for problem in compare_fields( + python_task["__var"], sdk_task["__var"], TASK_KEYS_NOT_COMPARED, defaults + ) + ) + return problems + + +def compare(python: dict[str, Any], sdk: dict[str, Any], task_defaults: dict[str, Any]) -> list[str]: + """List every way the SDK's serialization differs from Python's.""" + if sdk.keys() != python.keys(): + return [f"the Dags are {sorted(sdk)}, Python writes {sorted(python)}"] + problems = [] + for dag_id, python_dag in python.items(): + sdk_dag = sdk[dag_id] + if sdk_dag["__version"] != python_dag["__version"]: + problems.append( + f"{dag_id}: __version is {sdk_dag['__version']}, Python writes {python_dag['__version']}" + ) + problems.extend( + f"{dag_id}: {problem}" + for problem in compare_dag(python_dag["dag"], sdk_dag["dag"], task_defaults) + ) + return problems + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--sdk", required=True, help="the SDK's name, as in serialized_.json") + parser.add_argument("command", nargs="+", help="the SDK's serializer, after --") + args = parser.parse_args(argv) + + directory = Path(tempfile.mkdtemp(prefix=f"{args.sdk}-serialization-")) + python_output = directory / "serialized_python.json" + sdk_output = directory / f"serialized_{args.sdk}.json" + run([*args.command, str(TEST_DAGS), str(sdk_output)]) + run( + [ + "uv", + "run", + "--project", + "airflow-core", + # airflow-core's dev group pulls in providers and extras that build native code; the + # serializer only needs airflow-core itself. + "--no-dev", + "python", + str(HERE / "serialize_python.py"), + str(TEST_DAGS), + str(python_output), + "--deserialize", + str(sdk_output), + ] + ) + + python = json.loads(python_output.read_text()) + problems = compare(python, json.loads(sdk_output.read_text()), get_task_defaults()) + if problems: + print( + f"The {args.sdk} serialization differs from Python's in {len(problems)} place(s):", + file=sys.stderr, + ) + for problem in problems: + print(f" {problem}", file=sys.stderr) + print(f"Both serializations are kept in {directory}", file=sys.stderr) + return 1 + shutil.rmtree(directory) + print(f"The {args.sdk} SDK serializes all {len(python)} Dags of {TEST_DAGS.name} as Python does") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/ci/lang_sdk_serialization/serialize_python.py b/scripts/ci/lang_sdk_serialization/serialize_python.py new file mode 100644 index 0000000000000..39752347ff665 --- /dev/null +++ b/scripts/ci/lang_sdk_serialization/serialize_python.py @@ -0,0 +1,121 @@ +# 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. +r""" +Serialize the Dags of test_dags.yaml with Airflow's own serializer. + +Each Dag is built with the Python authoring API and written as ``DagSerialization.to_dict`` returns it, +keyed by Dag id. ``--deserialize`` also loads every Dag a language SDK wrote through +``DagSerialization.validate_schema`` and ``from_dict``, as Airflow does with the Dags that SDK sends. +compare.py runs it as:: + + uv run --project airflow-core --no-dev python scripts/ci/lang_sdk_serialization/serialize_python.py \ + scripts/ci/lang_sdk_serialization/test_dags.yaml serialized_python.json \ + --deserialize serialized_typescript.json +""" + +from __future__ import annotations + +import argparse +import datetime +import json +import sys +from pathlib import Path +from typing import Any + +import yaml + +from airflow.sdk import DAG, BaseOperator, TaskGroup +from airflow.serialization.serialized_objects import DagSerialization + + +class NoopOperator(BaseOperator): + """Stands in for a language SDK's task: a task with no Python behaviour.""" + + def execute(self, context): + return None + + +def construct_datetime(loader: yaml.SafeLoader, node: yaml.ScalarNode) -> datetime.datetime: + return datetime.datetime.fromisoformat(loader.construct_scalar(node)) + + +def construct_timedelta(loader: yaml.SafeLoader, node: yaml.ScalarNode) -> datetime.timedelta: + return datetime.timedelta(seconds=float(loader.construct_scalar(node))) + + +yaml.SafeLoader.add_constructor("!datetime", construct_datetime) +yaml.SafeLoader.add_constructor("!timedelta", construct_timedelta) + + +def build_dag(case: dict[str, Any]) -> DAG: + dag = DAG(case["dag_id"], **case.get("spec", {})) + groups: dict[str, TaskGroup] = {} + for group_id in case.get("groups", []): + # A group id is fully qualified, so its parent is whatever comes before the last dot. + parent_id, _, local_id = group_id.rpartition(".") + groups[group_id] = TaskGroup(local_id, dag=dag, parent_group=groups[parent_id] if parent_id else None) + for task in case["tasks"]: + NoopOperator( + task_id=task["task_id"], dag=dag, task_group=groups.get(task.get("group")), **task.get("spec", {}) + ) + for task in case["tasks"]: + task_id = f"{task['group']}.{task['task_id']}" if "group" in task else task["task_id"] + for upstream in task.get("upstream", []): + dag.get_task(upstream) >> dag.get_task(task_id) + for upstream, downstream in case.get("order_edges", []): + get_node(dag, groups, upstream) >> get_node(dag, groups, downstream) + return dag + + +def get_node(dag: DAG, groups: dict[str, TaskGroup], node_id: str): + """Resolve an edge endpoint: the task group with that id if there is one, else the task.""" + return groups[node_id] if node_id in groups else dag.get_task(node_id) + + +def deserialize(path: Path) -> None: + for dag_id, data in json.loads(path.read_text()).items(): + try: + DagSerialization.validate_schema(data) + DagSerialization.from_dict(data) + except Exception: + print(f"Airflow cannot load Dag {dag_id!r} as the SDK wrote it", file=sys.stderr) + raise + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("test_dags", type=Path, help="the test cases, test_dags.yaml") + parser.add_argument("output", type=Path, help="the JSON file to write") + parser.add_argument( + "--deserialize", + type=Path, + metavar="SDK_OUTPUT", + help="a JSON file a language SDK wrote", + ) + args = parser.parse_args() + + cases = yaml.safe_load(args.test_dags.read_text())["dags"] + serialized = {case["dag_id"]: DagSerialization.to_dict(build_dag(case)) for case in cases} + args.output.write_text(json.dumps(serialized, indent=2) + "\n") + if args.deserialize: + deserialize(args.deserialize) + + +if __name__ == "__main__": + main() diff --git a/scripts/ci/lang_sdk_serialization/test_dags.yaml b/scripts/ci/lang_sdk_serialization/test_dags.yaml new file mode 100644 index 0000000000000..d80c90bbb403f --- /dev/null +++ b/scripts/ci/lang_sdk_serialization/test_dags.yaml @@ -0,0 +1,166 @@ +# 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. +--- +# The Dags compare.py checks a language SDK's serializer on. serialize_python.py and the SDK's own +# serializer each build every Dag here with their SDK's authoring API, and the two serializations +# have to agree. +# +# `spec` keys are the serialized Dag schema's snake_case names, which each side maps onto +# its authoring API. A task is called with the tasks named in its `upstream`, which come +# before it. Task and group ids are fully qualified, and a group comes after its parent. +# `order_edges` are `[upstream, downstream]` pairs of task or group ids. +# +# `!datetime` takes an ISO 8601 timestamp, and `!timedelta` a number of seconds. +dags: + # No schedule and no options: NullTimetable plus the fields Python always emits. + - dag_id: conformance_minimal + tasks: + - task_id: solo + + # A cron schedule with every Dag option set, and one fully configured task followed by + # one left entirely at its schema defaults. + - dag_id: conformance_cron + spec: + schedule: "0 3 * * *" + description: a Dag that sets everything + dag_display_name: Conformance Cron + doc_md: "# notes" + tags: [gamma, alpha, gamma] + start_date: !datetime 2026-01-01T00:00:00+00:00 + end_date: !datetime 2026-12-31T23:30:15+00:00 + dagrun_timeout: !timedelta 300 + catchup: true + render_template_as_native_obj: true + disable_bundle_versioning: true + is_paused_upon_creation: true + max_active_tasks: 8 + max_active_runs: 3 + max_consecutive_failed_dag_runs: 2 + tasks: + - task_id: extract + spec: + owner: data-team + start_date: !datetime 2026-02-01T00:00:00+00:00 + end_date: !datetime 2026-11-30T00:00:00+00:00 + trigger_rule: all_done + depends_on_past: true + ignore_first_depends_on_past: true + wait_for_past_depends_before_skipping: true + wait_for_downstream: true + retries: 2 + queue: typescript + pool: tiny + pool_slots: 2 + execution_timeout: !timedelta 120 + retry_delay: !timedelta 600 + retry_exponential_backoff: 2 + max_retry_delay: !timedelta 900 + priority_weight: 5 + weight_rule: upstream + executor: LocalExecutor + do_xcom_push: false + email_on_failure: false + email_on_retry: false + doc_md: extracts things + map_index_template: "{{ task.task_id }}" + max_active_tis_per_dag: 3 + max_active_tis_per_dagrun: 4 + # Every value explicitly at its schema default, so all of them are omitted. + - task_id: transform + upstream: [extract] + spec: + owner: airflow + trigger_rule: all_success + depends_on_past: false + ignore_first_depends_on_past: false + wait_for_past_depends_before_skipping: false + wait_for_downstream: false + retries: 0 + queue: default + pool: default_pool + pool_slots: 1 + retry_delay: !timedelta 300 + retry_exponential_backoff: 0 + priority_weight: 1 + weight_rule: downstream + do_xcom_push: true + email_on_failure: true + email_on_retry: true + + # The @once preset, fail_fast, and a task feeding two downstreams. fail_fast is set here + # rather than on conformance_cron, since Airflow only allows it when every task uses the + # all_success trigger rule. + - dag_id: conformance_once + spec: + schedule: "@once" + fail_fast: true + tasks: + - task_id: seed + - task_id: beta + upstream: [seed] + - task_id: alpha + upstream: [seed] + + # The @continuous preset. Airflow caps a continuous Dag at one active run, so + # max_active_runs is set rather than left at the default of 16. + - dag_id: conformance_continuous + spec: + schedule: "@continuous" + max_active_runs: 1 + tasks: + - task_id: watch + + # Fan-out and fan-in, so downstream_task_ids has to be sorted rather than kept in wiring + # order. The schedule is not written "@daily" because Python expands a preset before + # serializing it, which serde.test.ts covers separately. + - dag_id: conformance_diamond + spec: + schedule: "0 0 * * *" + tasks: + - task_id: root + - task_id: right + upstream: [root] + - task_id: left + upstream: [root] + - task_id: join + upstream: [right, left] + + # Nested task groups and order-only edges. A group edge is recorded on the group, and is + # also expanded onto the task graph, where it reaches the group's roots and leaves + # rather than every task it holds: staging.stage feeds staging.checks.nulls, so only + # the latter is a leaf. Both sides have to agree on that expansion, and on the local + # _group_id a nested group carries. + - dag_id: conformance_groups + spec: + schedule: "@once" + groups: [staging, staging.checks, publish] + tasks: + - task_id: extract + - task_id: stage + group: staging + - task_id: nulls + group: staging.checks + upstream: [staging.stage] + - task_id: push + group: publish + - task_id: load + - task_id: notify + order_edges: + - [extract, staging] + - [staging, publish] + - [publish, load] + - [load, notify] diff --git a/scripts/tests/ci/lang_sdk_serialization/__init__.py b/scripts/tests/ci/lang_sdk_serialization/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/scripts/tests/ci/lang_sdk_serialization/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/scripts/tests/ci/lang_sdk_serialization/test_compare.py b/scripts/tests/ci/lang_sdk_serialization/test_compare.py new file mode 100644 index 0000000000000..13101587d3654 --- /dev/null +++ b/scripts/tests/ci/lang_sdk_serialization/test_compare.py @@ -0,0 +1,234 @@ +# 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. +from __future__ import annotations + +import copy +import json +import subprocess +from pathlib import Path +from unittest import mock + +import pytest +from ci.lang_sdk_serialization.compare import TEST_DAGS, compare, get_task_defaults, main + +DEFAULTS = {"retry_delay": 300.0} + + +def build_serialized(fileloc: str, tasks: list[dict]) -> dict: + return { + "d": { + "__version": 3, + "dag": { + "dag_id": "d", + "fileloc": fileloc, + "timezone": "UTC", + "tags": ["a"], + "task_group": {"prefix_group_id": True, "children": {"extract": ["operator", "extract"]}}, + "tasks": [{"__type": "operator", "__var": task} for task in tasks], + }, + } + } + + +PYTHON = build_serialized( + "/dags/d.py", + [ + { + "task_id": "extract", + "retries": 2, + "pool_slots": 1, + "retry_delay": 300.0, + "task_type": "NoopOperator", + }, + {"task_id": "load", "retry_delay": 300.0, "task_type": "NoopOperator"}, + ], +) + +# Differs from PYTHON only where an SDK may: where it was declared, what stands for its +# tasks, a number JavaScript writes without a fraction, and a field left at its default. +SDK = build_serialized( + "/bundles/app/bundle.mjs", + [ + {"task_id": "extract", "retries": 2, "pool_slots": 1.0, "task_type": "Task", "is_stub": True}, + {"task_id": "load", "task_type": "Task", "language": "typescript", "_arg_bindings": []}, + ], +) + + +def get_task(serialized: dict, index: int = 0) -> dict: + return serialized["d"]["dag"]["tasks"][index]["__var"] + + +def test_accepts_the_differences_an_sdk_is_allowed(): + assert compare(PYTHON, SDK, DEFAULTS) == [] + + +@pytest.mark.parametrize( + ("change", "expected"), + [ + pytest.param( + lambda sdk: get_task(sdk).update(retries=3), + ["d: task extract: retries is 3, Python writes 2"], + id="task-value", + ), + pytest.param( + lambda sdk: get_task(sdk).update(pool_slots=True), + ["d: task extract: pool_slots is True, Python writes 1"], + id="bool-is-not-a-number", + ), + pytest.param( + lambda sdk: get_task(sdk).pop("retries"), + ["d: task extract: retries is missing, Python writes 2"], + id="task-key-missing", + ), + pytest.param( + lambda sdk: get_task(sdk).update(owner="me"), + ["d: task extract: owner is 'me', which Python does not write"], + id="task-key-extra", + ), + pytest.param( + lambda sdk: sdk["d"]["dag"].pop("timezone"), + ["d: timezone is missing, Python writes 'UTC'"], + id="dag-key-missing", + ), + pytest.param( + lambda sdk: sdk["d"]["dag"].update(tags=["b"]), + ["d: tags[0] is 'b', Python writes 'a'"], + id="nested-list", + ), + pytest.param( + lambda sdk: sdk["d"]["dag"]["task_group"].update(prefix_group_id=False), + ["d: task_group.prefix_group_id is False, Python writes True"], + id="nested-dict", + ), + pytest.param( + lambda sdk: sdk["d"]["dag"]["tasks"].reverse(), + ["d: tasks are ['load', 'extract'], Python writes ['extract', 'load']"], + id="task-order", + ), + pytest.param( + lambda sdk: sdk["d"]["dag"]["tasks"][0].update(__type="taskgroup"), + ["d: task extract is a 'taskgroup', not a 'operator'"], + id="task-encoding", + ), + pytest.param( + lambda sdk: sdk["d"].update(__version=4), + ["d: __version is 4, Python writes 3"], + id="version", + ), + pytest.param( + lambda sdk: sdk.update(e=sdk["d"]), + ["the Dags are ['d', 'e'], Python writes ['d']"], + id="dag-ids", + ), + ], +) +def test_reports_each_difference(change, expected): + sdk = copy.deepcopy(SDK) + change(sdk) + + assert compare(PYTHON, sdk, DEFAULTS) == expected + + +def test_accepts_a_left_out_task_key_only_at_its_schema_default(): + python = copy.deepcopy(PYTHON) + get_task(python)["retry_delay"] = 600.0 + + assert compare(python, SDK, DEFAULTS) == ["d: task extract: retry_delay is missing, Python writes 600.0"] + + +def test_reads_the_task_defaults_from_the_dag_schema(): + defaults = get_task_defaults() + + assert defaults["retry_delay"] == 300.0 + # A null default and no default both mean the field has none. + assert "render_template_as_native_obj" not in defaults + assert "task_id" not in defaults + + +def write_outputs(sdk: dict): + """Stand in for both serializers, each writing its output to the path compare.py hands it.""" + + def run(command, **kwargs): + if command[0] == "uv": + Path(command[-3]).write_text(json.dumps(PYTHON)) + else: + Path(command[-1]).write_text(json.dumps(sdk)) + return subprocess.CompletedProcess(command, 0) + + return run + + +@mock.patch("ci.lang_sdk_serialization.compare.get_task_defaults", autospec=True, return_value=DEFAULTS) +@mock.patch("ci.lang_sdk_serialization.compare.tempfile.mkdtemp", autospec=True) +@mock.patch("ci.lang_sdk_serialization.compare.subprocess.run", autospec=True) +def test_main_runs_both_serializers_and_removes_their_output_when_they_agree( + mock_run, mock_mkdtemp, mock_get_task_defaults, tmp_path, capsys +): + mock_mkdtemp.return_value = str(tmp_path) + mock_run.side_effect = write_outputs(SDK) + + assert main(["--sdk", "typescript", "--", "pnpm", "exec", "tsx", "serialize.ts"]) == 0 + + sdk_output = str(tmp_path / "serialized_typescript.json") + assert [call.args[0] for call in mock_run.call_args_list] == [ + ["pnpm", "exec", "tsx", "serialize.ts", str(TEST_DAGS), sdk_output], + [ + "uv", + "run", + "--project", + "airflow-core", + "--no-dev", + "python", + str(TEST_DAGS.parent / "serialize_python.py"), + str(TEST_DAGS), + str(tmp_path / "serialized_python.json"), + "--deserialize", + sdk_output, + ], + ] + assert not tmp_path.exists() + assert "serializes all 1 Dags" in capsys.readouterr().out + + +@mock.patch("ci.lang_sdk_serialization.compare.get_task_defaults", autospec=True, return_value=DEFAULTS) +@mock.patch("ci.lang_sdk_serialization.compare.tempfile.mkdtemp", autospec=True) +@mock.patch("ci.lang_sdk_serialization.compare.subprocess.run", autospec=True) +def test_main_reports_the_differences_and_keeps_the_output( + mock_run, mock_mkdtemp, mock_get_task_defaults, tmp_path, capsys +): + mock_mkdtemp.return_value = str(tmp_path) + sdk = copy.deepcopy(SDK) + get_task(sdk)["retries"] = 3 + mock_run.side_effect = write_outputs(sdk) + + assert main(["--sdk", "typescript", "--", "serialize"]) == 1 + + assert (tmp_path / "serialized_typescript.json").exists() + assert "d: task extract: retries is 3, Python writes 2" in capsys.readouterr().err + + +@mock.patch("ci.lang_sdk_serialization.compare.tempfile.mkdtemp", autospec=True) +@mock.patch("ci.lang_sdk_serialization.compare.subprocess.run", autospec=True) +def test_main_stops_when_a_serializer_fails(mock_run, mock_mkdtemp, tmp_path): + mock_mkdtemp.return_value = str(tmp_path) + mock_run.return_value = subprocess.CompletedProcess([], 1) + + with pytest.raises(SystemExit, match="`serialize .*` failed"): + main(["--sdk", "typescript", "--", "serialize"]) + + assert mock_run.call_count == 1 diff --git a/ts-sdk/.pre-commit-config.yaml b/ts-sdk/.pre-commit-config.yaml index ad2bcb99c1710..4346cce9f4a6c 100644 --- a/ts-sdk/.pre-commit-config.yaml +++ b/ts-sdk/.pre-commit-config.yaml @@ -61,6 +61,20 @@ repos: additional_dependencies: ['pnpm@10.28.1'] pass_filenames: false require_serial: true + - id: check-ts-sdk-serialization-conformance + name: Check the TS SDK serializes Dags the way Airflow does + description: "Serialize the shared test Dags with the TS SDK and with Airflow, and compare the two" + entry: ./scripts/ci/prek/check_serialization_conformance.py + language: node + files: > + (?x) + ^schema/dag-schema\.json$| + ^src/coordinator/serde\.ts$| + ^src/sdk/dag\.ts$| + ^tests/conformance/.*$ + additional_dependencies: ['pnpm@10.28.1'] + pass_filenames: false + require_serial: true - id: check-ts-sdk-docs-package-version-in-sync name: Check ts-sdk/docs package.json versions match ts-sdk/package.json entry: ../scripts/ci/prek/check_ts_sdk_docs_package_version_in_sync.py diff --git a/ts-sdk/package.json b/ts-sdk/package.json index bd91883d55ece..ab205567beeab 100644 --- a/ts-sdk/package.json +++ b/ts-sdk/package.json @@ -86,6 +86,7 @@ "tsx": "^4.21.0", "typescript": "^6.0.2", "typescript-eslint": "^8.60.0", - "vitest": "^4.1.7" + "vitest": "^4.1.7", + "yaml": "^2.9.1" } } diff --git a/ts-sdk/pnpm-lock.yaml b/ts-sdk/pnpm-lock.yaml index 543171cdcff53..f1cd0251f6fe7 100644 --- a/ts-sdk/pnpm-lock.yaml +++ b/ts-sdk/pnpm-lock.yaml @@ -41,7 +41,10 @@ importers: version: 8.67.0(eslint@10.4.0)(typescript@6.0.3) vitest: specifier: ^4.1.7 - version: 4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)) + version: 4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1)) + yaml: + specifier: ^2.9.1 + version: 2.9.1 example: dependencies: @@ -1183,6 +1186,11 @@ packages: resolution: {integrity: sha512-BN22B5eaMMI9UMtjrGd5g5eCYPpCPDUy0FJXbYsaT5zYxjFOckS53SQDE3pWkVoWpHXVb3BrYcEN4Twa55B5cA==} engines: {node: '>=0.10.0'} + yaml@2.9.1: + resolution: {integrity: sha512-3NxN8+78OdzbT7C/WjGsyfPAtJaN3FNDsWxv7Y7mcDsT/oOmgW8BpyQQFFBnvZE3j9Y2Sdz1ULFLezL7Eb2yFw==} + engines: {node: '>= 14.6'} + hasBin: true + yocto-queue@0.1.0: resolution: {integrity: sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==} engines: {node: '>=10'} @@ -1571,7 +1579,7 @@ snapshots: obug: 2.1.4 std-env: 4.2.0 tinyrainbow: 3.1.1 - vitest: 4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)) + vitest: 4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1)) optional: true '@vitest/expect@4.1.7': @@ -1583,13 +1591,13 @@ snapshots: chai: 6.2.2 tinyrainbow: 3.1.0 - '@vitest/mocker@4.1.7(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12))': + '@vitest/mocker@4.1.7(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1))': dependencies: '@vitest/spy': 4.1.7 estree-walker: 3.0.3 magic-string: 0.30.21 optionalDependencies: - vite: 8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12) + vite: 8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1) '@vitest/pretty-format@4.1.7': dependencies: @@ -2115,7 +2123,7 @@ snapshots: dependencies: punycode: 2.3.1 - vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12): + vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1): dependencies: lightningcss: 1.33.0 picomatch: 4.0.7 @@ -2127,11 +2135,12 @@ snapshots: esbuild: 0.28.2 fsevents: 2.3.3 tsx: 4.23.12 + yaml: 2.9.1 - vitest@4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)): + vitest@4.1.7(@types/node@26.2.0)(@vitest/coverage-v8@4.1.7)(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1)): dependencies: '@vitest/expect': 4.1.7 - '@vitest/mocker': 4.1.7(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)) + '@vitest/mocker': 4.1.7(vite@8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1)) '@vitest/pretty-format': 4.1.7 '@vitest/runner': 4.1.7 '@vitest/snapshot': 4.1.7 @@ -2148,7 +2157,7 @@ snapshots: tinyexec: 1.1.1 tinyglobby: 0.2.16 tinyrainbow: 3.1.0 - vite: 8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12) + vite: 8.0.10(@types/node@26.2.0)(esbuild@0.28.2)(tsx@4.23.12)(yaml@2.9.1) why-is-node-running: 2.3.0 optionalDependencies: '@types/node': 26.2.0 @@ -2167,4 +2176,6 @@ snapshots: word-wrap@1.2.5: {} + yaml@2.9.1: {} + yocto-queue@0.1.0: {} diff --git a/ts-sdk/scripts/ci/prek/check_serialization_conformance.py b/ts-sdk/scripts/ci/prek/check_serialization_conformance.py new file mode 100755 index 0000000000000..c51dfffe1834b --- /dev/null +++ b/ts-sdk/scripts/ci/prek/check_serialization_conformance.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +# 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. +"""Check the TS SDK serializes Dags as Airflow does, with scripts/ci/lang_sdk_serialization/compare.py.""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[4] / "scripts" / "ci" / "prek")) + +from common_prek_utils import AIRFLOW_ROOT_PATH, run_command + +if __name__ not in ("__main__", "__mp_main__"): + raise SystemExit( + "This file is intended to be executed as an executable program. You cannot use it as a module." + f"To run this script, run the ./{__file__} command" + ) + +if __name__ == "__main__": + run_command( + ["pnpm", "install", "--frozen-lockfile", "--config.confirmModulesPurge=false"], + cwd=AIRFLOW_ROOT_PATH / "ts-sdk", + ) + compare = AIRFLOW_ROOT_PATH / "scripts" / "ci" / "lang_sdk_serialization" / "compare.py" + serializer = ["pnpm", "--dir", "ts-sdk", "exec", "tsx", "tests/conformance/serialize_typescript.ts"] + command = [sys.executable, str(compare), "--sdk", "typescript", "--", *serializer] + sys.exit(subprocess.run(command, check=False).returncode) diff --git a/ts-sdk/src/coordinator/serde.ts b/ts-sdk/src/coordinator/serde.ts new file mode 100644 index 0000000000000..26a1d844a3dc3 --- /dev/null +++ b/ts-sdk/src/coordinator/serde.ts @@ -0,0 +1,729 @@ +/*! + * 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. + */ + +// Turns a Dag declared in TypeScript into Airflow's DagSerialization v3 JSON — +// what the Dag processor stores and the scheduler reads. The format is +// Airflow-internal rather than an SDK schema, so it is reimplemented per +// language against `airflow-core/src/airflow/serialization/schema.json`; see +// `airflow-core/adr/lang-sdk/0004-dag-parsing.md` for the field table this +// follows, and the Java SDK's `Serde.kt` for the same job in another language. +// +// Byte-parity with Python's serializer is not the goal — Python omits fields +// against a `client_defaults` table this SDK does not receive. What has to hold +// is that `DagSerialization.from_dict` rebuilds the same Dag, which +// scripts/ci/lang_sdk_serialization/compare.py checks against Python's own serialization. +// +// This module only produces the payload. Answering a Dag-parsing request with +// it is the bundle's job, once the coordinator has a parse request to answer. + +import { relative as relativePath } from "node:path"; + +import { + DAG_SCHEMA_FIELDS, + TASK_SCHEMA_FIELDS, + type SchemaField, +} from "../generated/dag-schema-fields.js"; +import type { JsonValue } from "../sdk/client-types.js"; +import { + getDagOrderEdges, + getDagTaskGroups, + getDagTaskInputs, + getDagTaskRecords, + isTaskRef, + type Dag, + type RecordedInputs, + type TaskGroupRecord, +} from "../sdk/dag.js"; + +/** A serialized Dag: JSON, by the time it reaches the supervisor as msgpack. */ +type SerializedValue = JsonValue; + +/** Airflow's type/var encoding, as `BaseSerialization.serialize()` emits it. */ +interface TypeEncoded { + readonly __type: string; + readonly __var: SerializedValue; +} + +/** + * Identity every TypeScript task carries, in place of the Python operator class + * a Python Dag would name. + * + * Fixed rather than derived: nothing on the Airflow side imports `_task_module` + * — `SerializedBaseOperator.populate_operator` only compares the pair as + * strings when matching plugin extra links — so the pair is free to name the + * coordinator that actually runs the task, which makes every TypeScript task + * greppable in the UI and the metadata DB. + */ +const TASK_TYPE = "TypeScriptOperator"; +const TASK_MODULE = "airflow.sdk.coordinators.node"; + +/** + * Marks the tasks this SDK serialized, as the Java SDK marks its own. + * + * Nothing in airflow-core reads it today; the `operator` schema definition + * allows additional properties, so it rides along as a marker for tooling that + * wants to tell language-native tasks apart without parsing `_task_module`. + */ +const TASK_LANGUAGE = "typescript"; + +// Python resolves these from [core]/[scheduler] config when the Dag leaves them +// unset, and its serializer always writes the resolved value — there is no +// schema default to omit against. A bundle cannot read airflow.cfg, so the +// stock defaults stand in. +const DAG_CONFIG_FALLBACKS: Readonly> = { + max_active_tasks: 16, // [core] max_active_tasks_per_dag + max_active_runs: 16, // [core] max_active_runs_per_dag + max_consecutive_failed_dag_runs: 0, + catchup: false, // [scheduler] catchup_by_default + disable_bundle_versioning: false, +}; + +/** How one set of authoring fields is written into a serialized object. */ +interface FieldRules { + readonly fields: Readonly>; + /** + * Fields whose {__type, __var} wrapper survives; every other field is + * unwrapped to the bare __var, as Python's `serialize_to_json` does. + * + * Neither set overlaps the generated authoring fields today, so in practice + * everything is unwrapped. They are named so that a decorated field added to + * the schema later takes the right path rather than silently losing its + * wrapper. Source: `DagSerialization._decorated_fields` and + * `OperatorSerialization._decorated_fields`. + */ + readonly decorated: ReadonlySet; + /** Fields never written, whatever the author set. */ + readonly omitted: ReadonlySet; +} + +const DAG_FIELD_RULES: FieldRules = { + fields: DAG_SCHEMA_FIELDS, + decorated: new Set(["default_args", "access_control"]), + omitted: new Set(), +}; + +const TASK_FIELD_RULES: FieldRules = { + fields: TASK_SCHEMA_FIELDS, + decorated: new Set(["executor_config"]), + // Python drops both unless the operator names an email recipient + // (`OperatorSerialization._serialize_node`). A TypeScript task has no `email` + // field to name one, so writing them would describe a notification that can + // never be sent. + omitted: new Set(["email_on_failure", "email_on_retry"]), +}; + +const NULL_TIMETABLE = "airflow.timetables.simple.NullTimetable"; +const ONCE_TIMETABLE = "airflow.timetables.simple.OnceTimetable"; +const CONTINUOUS_TIMETABLE = "airflow.timetables.simple.ContinuousTimetable"; +const CRON_TIMETABLE = "airflow.timetables.trigger.CronTriggerTimetable"; + +/** Serialize one Dag to the `dag` object of a DagSerialization v3 payload. */ +export function serializeDag( + dag: Dag, + fileloc: string, + relativeFileloc: string, +): Record { + const graph = buildDagGraph(dag); + const inputs = getDagTaskInputs(dag); + const data: Record = { + dag_id: dag.dagId, + fileloc, + relative_fileloc: relativeFileloc, + timezone: "UTC", + timetable: serializeTimetable(dag.spec.schedule, dag.dagId), + tasks: [...getDagTaskRecords(dag)].map(([taskId, record]) => + serializeTask( + dag.dagId, + taskId, + withDagQueue(record.spec, dag.spec.queue), + graph.downstreamTaskIds.get(taskId), + inputs.get(taskId), + ), + ), + dag_dependencies: [], + task_group: serializeTaskGroups(dag, graph), + edge_info: {}, + params: [], + // Always written by Python's serializer, so a Dag without either still + // round-trips to the same object. + deadline: null, + allowed_run_types: null, + }; + applySchemaFields(data, dag.spec, DAG_FIELD_RULES, `Dag "${dag.dagId}"`); + for (const [key, fallback] of Object.entries(DAG_CONFIG_FALLBACKS)) { + data[key] ??= fallback; + } + return data; +} + +/** + * The task's spec with the Dag's queue filled in, when the task named none. + * + * Merged before the fields are written rather than after, so a queue that + * happens to equal the schema default is omitted the way any other defaulted + * field is. + */ +function withDagQueue(spec: object, dagQueue: string | undefined): object { + if (dagQueue === undefined || "queue" in spec) return spec; + return { ...spec, queue: dagQueue }; +} + +/** Serialize one task, with its downstream edges sorted for a stable payload. */ +function serializeTask( + dagId: string, + taskId: string, + spec: object, + downstream: ReadonlySet | undefined, + inputs: RecordedInputs | undefined, +): SerializedValue { + const data: Record = { + task_id: taskId, + task_type: TASK_TYPE, + _task_module: TASK_MODULE, + language: TASK_LANGUAGE, + // Python's operator serializer always emits this — its list value never + // matches the tuple default it is compared against. A TypeScript task has + // no Jinja templating, so the list is empty rather than absent. + template_fields: [], + // What marks a task whose arguments the API server resolves per instance + // and sends to a foreign runtime, as `@task.stub` does on the Python side. + // `get_arg_bindings` reads nothing without it. + is_stub: true, + }; + const bindings = serializeArgBindings(inputs); + if (bindings) data["_arg_bindings"] = bindings; + applySchemaFields(data, spec, TASK_FIELD_RULES, `task "${taskId}" of Dag "${dagId}"`); + if (downstream?.size) { + data["downstream_task_ids"] = [...downstream].sort(); + } + return { __type: "operator", __var: data }; +} + +/** + * The task's arguments as the binding spec the API server hands back at run + * time, one entry per argument in the order the call named them. + * + * A reference becomes an `xcom` binding naming the upstream task, and anything + * else a `literal` carrying the value, matching the `TaskArgBinding` union the + * execution API declares. `value_schema` is left out: it constrains the decode + * side, and nothing records a TypeScript argument's type at pack time, so + * omitting it says "unconstrained" rather than asserting a wrong type. + * + * `undefined` for a task called with no arguments, which needs no spec. + */ +function serializeArgBindings(inputs: RecordedInputs | undefined): SerializedValue | undefined { + if (inputs === undefined) return undefined; + const entries = Object.entries(inputs); + if (entries.length === 0) return undefined; + return entries.map(([name, value]): SerializedValue => + isTaskRef(value) + ? { name, kind: "xcom", task_id: value.taskId } + : { name, kind: "literal", value: serializeValue(value) }, + ); +} + +/** + * The Dag's task-group tree, rooted at the group Python builds for every Dag. + * + * `children[label] = [kind, value]`, as `serialize_for_task_group()` writes it: + * a task is `["operator", task_id]` and a nested group is `["taskgroup", ...]` + * holding that group's own object, so the tree nests by embedding rather than + * by reference. `_group_id` is the group's *local* segment, not its qualified + * id, matching what Python records. + */ +function serializeTaskGroups(dag: Dag, graph: DagGraph): SerializedValue { + const groups = getDagTaskGroups(dag); + const grouped = new Set(); + for (const group of groups.values()) { + for (const taskId of group.taskIds) grouped.add(taskId); + } + const rootTaskIds = [...getDagTaskRecords(dag).keys()].filter((id) => !grouped.has(id)); + const rootGroupIds = [...groups.values()] + .filter((group) => group.parentGroupId === undefined) + .map((group) => group.groupId); + + return taskGroupObject(null, rootTaskIds, rootGroupIds, groups, graph); +} + +/** One group object: the root when `groupId` is null, otherwise a nested one. */ +function taskGroupObject( + groupId: string | null, + taskIds: readonly string[], + childGroupIds: readonly string[], + groups: ReadonlyMap, + graph: DagGraph, +): SerializedValue { + const children: Record = {}; + for (const taskId of taskIds) { + children[taskId] = ["operator", taskId]; + } + for (const childId of childGroupIds) { + const child = groups.get(childId)!; + children[childId] = [ + "taskgroup", + taskGroupObject(childId, child.taskIds, child.childGroupIds, groups, graph), + ]; + } + const own = groupId === null ? undefined : graph.groupEdges.get(groupId); + return { + // The local segment: Python's TaskGroup stores the id it was given, and + // rebuilds the qualified one from where the group sits in the tree. + _group_id: groupId === null ? null : localGroupId(groupId), + group_display_name: "", + prefix_group_id: groupId === null || groups.get(groupId)!.prefixGroupId, + tooltip: "", + ui_color: "CornflowerBlue", + ui_fgcolor: "#000", + children, + upstream_group_ids: sorted(own?.upstreamGroups), + downstream_group_ids: sorted(own?.downstreamGroups), + upstream_task_ids: sorted(own?.upstreamTasks), + downstream_task_ids: sorted(own?.downstreamTasks), + }; +} + +function localGroupId(groupId: string): string { + const cut = groupId.lastIndexOf(GROUP_SEPARATOR); + return cut === -1 ? groupId : groupId.slice(cut + 1); +} + +function sorted(values: ReadonlySet | undefined): string[] { + return [...(values ?? [])].sort(); +} + +const GROUP_SEPARATOR = "."; + +interface GroupEdgeSets { + readonly upstreamGroups: Set; + readonly downstreamGroups: Set; + readonly upstreamTasks: Set; + readonly downstreamTasks: Set; +} + +/** Both views of a Dag's edges: the task graph, and what each group records. */ +interface DagGraph { + /** Each task's downstream task IDs, with every group endpoint expanded. */ + readonly downstreamTaskIds: Map>; + readonly groupEdges: Map; +} + +/** + * Resolve a Dag's two kinds of edge into the two views a serialized Dag holds. + * + * An order-only edge with a group at either end lands in both: the group object + * records it for the UI, and the task graph records it expanded, because the + * scheduler only ever reads task-to-task edges. A group expands to its *roots* + * when it is downstream and its *leaves* when it is upstream — an edge into a + * group reaches the tasks that start it, and one out of a group leaves from the + * tasks that finish it — which is what `TaskGroup.set_upstream` does in Python. + */ +function buildDagGraph(dag: Dag): DagGraph { + const groups = getDagTaskGroups(dag); + const downstreamTaskIds = new Map>(); + const groupEdges = new Map(); + const link = (upstream: string, downstream: string): void => { + // Two arguments fed by the same upstream are one edge, as is an order-only + // edge redeclaring one the wiring already drew. + const edges = downstreamTaskIds.get(upstream) ?? new Set(); + edges.add(downstream); + downstreamTaskIds.set(upstream, edges); + }; + + for (const [taskId, inputs] of getDagTaskInputs(dag)) { + for (const value of Object.values(inputs)) { + if (isTaskRef(value)) link(value.taskId, taskId); + } + } + const orderEdges = getDagOrderEdges(dag); + for (const { upstream, downstream } of orderEdges) { + if (!groups.has(upstream) && !groups.has(downstream)) link(upstream, downstream); + } + + // Roots and leaves are read off the task-to-task graph, which holds every + // edge that can sit inside a group by now: wiring, and any order-only edge + // between two tasks. Python resolves them at `>>` time and so is equally + // order-sensitive, which is what keeps the two in step. + const ends = new GroupEnds(groups, downstreamTaskIds); + + // Which group edges each endpoint has, for stepping over a group that holds + // no tasks. + const upstreamsOf = new Map>(); + const downstreamsOf = new Map>(); + for (const { upstream, downstream } of orderEdges) { + addTo(downstreamsOf, upstream, downstream); + addTo(upstreamsOf, downstream, upstream); + } + + /** + * The tasks an edge endpoint stands for: the task itself, or a group's leaves + * when it is upstream and its roots when it is downstream. + * + * A group holding no tasks has neither, so the edge steps over it and + * continues along the group edges beyond — `x >> empty >> y` still runs `y` + * after `x`, as Python's `find_leaves` walk does. + */ + const tasksAt = (id: string, side: "upstream" | "downstream"): string[] => { + if (!groups.has(id)) return [id]; + const own = side === "upstream" ? ends.leaves(id) : ends.roots(id); + return own.length > 0 ? own : tasksBeyond(id, side, new Set()); + }; + const tasksBeyond = ( + id: string, + side: "upstream" | "downstream", + seen: Set, + ): string[] => { + if (seen.has(id)) return []; + seen.add(id); + const next = side === "upstream" ? upstreamsOf.get(id) : downstreamsOf.get(id); + return [...(next ?? [])].flatMap((other) => { + if (!groups.has(other)) return [other]; + const own = side === "upstream" ? ends.leaves(other) : ends.roots(other); + return own.length > 0 ? own : tasksBeyond(other, side, seen); + }); + }; + const setsFor = (groupId: string): GroupEdgeSets => { + let sets = groupEdges.get(groupId); + if (!sets) { + sets = { + upstreamGroups: new Set(), + downstreamGroups: new Set(), + upstreamTasks: new Set(), + downstreamTasks: new Set(), + }; + groupEdges.set(groupId, sets); + } + return sets; + }; + + for (const { upstream, downstream } of orderEdges) { + const upstreamIsGroup = groups.has(upstream); + const downstreamIsGroup = groups.has(downstream); + if (!upstreamIsGroup && !downstreamIsGroup) continue; + + const from = tasksAt(upstream, "upstream"); + const to = tasksAt(downstream, "downstream"); + for (const tail of from) { + for (const head of to) link(tail, head); + } + + if (downstreamIsGroup) { + const sets = setsFor(downstream); + for (const tail of from) sets.upstreamTasks.add(tail); + if (upstreamIsGroup) sets.upstreamGroups.add(upstream); + } + // Only a group whose downstream is a plain task records it as a task; when + // both ends are groups the pair is recorded as a group edge on this side + // and as the expanded tasks on the other, which is how Python leaves it. + if (upstreamIsGroup) { + const sets = setsFor(upstream); + if (downstreamIsGroup) sets.downstreamGroups.add(downstream); + else sets.downstreamTasks.add(downstream); + } + } + return { downstreamTaskIds, groupEdges }; +} + +function addTo(index: Map>, key: string, value: string): void { + const existing = index.get(key) ?? new Set(); + existing.add(value); + index.set(key, existing); +} + +/** The tasks an edge reaches when it points at a group, cached per group. */ +class GroupEnds { + readonly #groups: ReadonlyMap; + readonly #downstream: ReadonlyMap>; + readonly #members = new Map>(); + + constructor( + groups: ReadonlyMap, + downstream: ReadonlyMap>, + ) { + this.#groups = groups; + this.#downstream = downstream; + } + + /** Tasks in the group with no upstream inside it: where an edge in arrives. */ + roots(groupId: string): string[] { + const members = this.#membersOf(groupId); + const hasInternalUpstream = new Set(); + for (const [upstream, downstream] of this.#downstream) { + if (!members.has(upstream)) continue; + for (const task of downstream) { + if (members.has(task)) hasInternalUpstream.add(task); + } + } + return [...members].filter((task) => !hasInternalUpstream.has(task)); + } + + /** Tasks in the group with no downstream inside it: where an edge out leaves. */ + leaves(groupId: string): string[] { + const members = this.#membersOf(groupId); + return [...members].filter( + (task) => ![...(this.#downstream.get(task) ?? [])].some((other) => members.has(other)), + ); + } + + /** Every task the group holds, nested groups included. */ + #membersOf(groupId: string): Set { + const cached = this.#members.get(groupId); + if (cached) return cached; + const members = new Set(); + const pending = [groupId]; + for (let i = 0; i < pending.length; i += 1) { + const group = this.#groups.get(pending[i]!); + if (!group) continue; + for (const taskId of group.taskIds) members.add(taskId); + pending.push(...group.childGroupIds); + } + this.#members.set(groupId, members); + return members; + } +} + +/** + * Lower a `schedule` onto the timetable the scheduler reconstructs. + * + * Only the four schedules that map to a stock timetable are accepted. Anything + * else — an asset expression, a custom timetable — is a Python object the + * scheduler has to import, which a TypeScript bundle cannot name, so it is + * rejected here rather than serialized into a Dag that fails to deserialize. + */ +function serializeTimetable(schedule: unknown, dagId: string): SerializedValue { + if (schedule === undefined || schedule === null) { + return simpleTimetable(NULL_TIMETABLE); + } + if (typeof schedule !== "string") { + throw new Error( + `schedule for Dag "${dagId}" must be "@once", "@continuous", or a cron expression; ` + + `${describeType(schedule)} schedule names a Python object this SDK cannot serialize`, + ); + } + if (schedule.trim() === "") { + throw new Error( + `schedule for Dag "${dagId}" is empty; leave it unset for a Dag with no schedule`, + ); + } + if (schedule === "@once") return simpleTimetable(ONCE_TIMETABLE); + if (schedule === "@continuous") return simpleTimetable(CONTINUOUS_TIMETABLE); + const expression = CRON_PRESETS[schedule] ?? schedule; + if (!isCronExpression(expression)) { + throw new Error( + `schedule ${JSON.stringify(schedule)} for Dag "${dagId}" is not a cron expression or a ` + + `preset (${Object.keys(CRON_PRESETS).join(", ")}, @once, @continuous); a schedule the ` + + "scheduler cannot parse would leave the Dag unschedulable", + ); + } + // TODO: honour [scheduler] create_cron_data_intervals, which switches Python + // to CronDataIntervalTimetable. A bundle cannot read airflow.cfg, so the + // supervisor has to send the flag first; tracked at + // https://github.com/apache/airflow/issues/67938 + return { + __type: CRON_TIMETABLE, + __var: { expression, timezone: "UTC", interval: 0, run_immediately: false }, + }; +} + +/** + * Presets expanded the way `CronMixin.__init__` expands them, so the serialized + * expression is the one Python records — which the Dag's summary and its hash + * are both taken from. Mirrors `airflow.utils.dates.cron_presets`. + */ +const CRON_PRESETS: Readonly> = { + "@hourly": "0 * * * *", + "@daily": "0 0 * * *", + "@weekly": "0 0 * * 0", + "@monthly": "0 0 1 * *", + "@quarterly": "0 0 1 */3 *", + "@yearly": "0 0 1 1 *", +}; + +/** + * Whether `expression` has the shape croniter accepts: five or six + * space-separated fields of cron characters. + * + * A shape check, not a parse: croniter validates the ranges, and repeating that + * here would be a second implementation to keep in step. What it does catch is + * prose — `"every tuesday"` — which would otherwise be written into a Dag that + * the scheduler then fails to build a timetable for. + */ +function isCronExpression(expression: string): boolean { + const fields = expression.trim().split(/\s+/); + if (fields.length !== 5 && fields.length !== 6) return false; + return fields.every( + (field) => /^[\d*,\-/?LW#]+$/i.test(field) || /^[A-Z]{3}(-[A-Z]{3})?$/i.test(field), + ); +} + +function simpleTimetable(type: string): SerializedValue { + return { __type: type, __var: {} }; +} + +/** + * Write the fields a spec set onto `data`, skipping any left at its schema + * default — Python's serializer omits what the scheduler re-derives. + */ +function applySchemaFields( + data: Record, + spec: object, + rules: FieldRules, + label: string, +): void { + const values = spec as Record; + for (const [name, field] of Object.entries(rules.fields)) { + // A virtual field names a key the serializer derives rather than writes; + // `schedule` becomes `timetable`. + if (field.virtual || rules.omitted.has(field.key)) continue; + const value = values[name]; + if (value === undefined || value === field.default) continue; + const encoded = encodeField(field, value, `${name} for ${label}`); + data[field.key] = rules.decorated.has(field.key) ? encoded : unwrapTypeEncoding(encoded); + } +} + +/** Encode one authoring value as the schema's type for that field. Rejects a + * value of the wrong type: this is the last point before the scheduler, and a + * mistyped field would otherwise surface as an unreadable Dag. */ +function encodeField(field: SchemaField, value: unknown, label: string): SerializedValue { + switch (field.type) { + case "string": + if (typeof value !== "string") throw typeError(label, "a string", value); + return value; + case "boolean": + if (typeof value !== "boolean") throw typeError(label, "a boolean", value); + return value; + case "number": + if (typeof value !== "number" || !Number.isFinite(value)) { + throw typeError(label, "a finite number", value); + } + return value; + case "timedelta": + if (typeof value !== "number" || !Number.isFinite(value)) { + throw typeError(label, "a duration in seconds", value); + } + return { __type: "timedelta", __var: value }; + case "datetime": + if (!(value instanceof Date) || Number.isNaN(value.getTime())) { + throw typeError(label, "a valid Date", value); + } + return serializeValue(value); + case "string[]": { + if (!Array.isArray(value) || value.some((item) => typeof item !== "string")) { + throw typeError(label, "an array of strings", value); + } + // Python holds these in a set, so duplicates collapse and the order is + // the sorted one that keeps a Dag's hash stable across runs. + return serializeValue(new Set(value as string[])); + } + } +} + +function typeError(label: string, expected: string, value: unknown): Error { + return new Error(`${label} must be ${expected}, not ${describeType(value)}`); +} + +function describeType(value: unknown): string { + if (value === null) return "null"; + if (Array.isArray(value)) return "an array"; + if (value instanceof Date) return "a Date"; + if (value instanceof Set) return "a Set"; + const noun = typeof value; + return `${/^[aeiou]/.test(noun) ? "an" : "a"} ${noun}`; +} + +/** + * Encode a value the way `BaseSerialization.serialize()` does. + * + * A duration has no distinct runtime type in TypeScript — it is a number of + * seconds — so `timedelta` is applied by {@link encodeField} from the schema + * rather than inferred here. + */ +export function serializeValue(value: unknown): SerializedValue { + if (value === null || value === undefined) return null; + if (typeof value === "string" || typeof value === "boolean") return value; + if (typeof value === "number") { + if (!Number.isFinite(value)) { + throw new Error(`Cannot serialize the non-finite number ${String(value)}`); + } + return value; + } + if (value instanceof Date) { + if (Number.isNaN(value.getTime())) throw new Error("Cannot serialize an invalid Date"); + return { __type: "datetime", __var: value.getTime() / 1000 }; + } + if (value instanceof Set) { + return { + __type: "set", + __var: [...value].map(serializeValue).sort(compareSerialized), + }; + } + if (Array.isArray(value)) return value.map(serializeValue); + if (value instanceof Map) { + return { __type: "dict", __var: serializeEntries(value.entries()) }; + } + if (typeof value === "object") { + return { __type: "dict", __var: serializeEntries(Object.entries(value)) }; + } + throw new Error(`Cannot serialize a ${typeof value}`); +} + +function serializeEntries(entries: Iterable<[unknown, unknown]>): Record { + const encoded: Record = {}; + for (const [key, item] of entries) { + encoded[String(key)] = serializeValue(item); + } + return encoded; +} + +// Python sorts a set's members before writing them; JSON's default sort is +// lexicographic on the string form, which matches for the string sets this +// SDK produces and stays total for anything else. +function compareSerialized(left: SerializedValue, right: SerializedValue): number { + const a = typeof left === "string" ? left : JSON.stringify(left); + const b = typeof right === "string" ? right : JSON.stringify(right); + return a < b ? -1 : a > b ? 1 : 0; +} + +/** + * Strip the type encoding from a non-decorated field, as Python's + * `serialize_to_json` does: it serializes every field, then keeps only the + * `__var` of the ones outside its decorated set. + */ +export function unwrapTypeEncoding(value: SerializedValue): SerializedValue { + if (!isTypeEncoded(value)) return value; + return value.__var; +} + +function isTypeEncoded(value: SerializedValue): value is TypeEncoded & SerializedValue { + return ( + typeof value === "object" && + value !== null && + !Array.isArray(value) && + "__type" in value && + "__var" in value + ); +} + +/** Where the Dag file sits inside its bundle, as Airflow records it. */ +export function computeRelativeFileloc(fileloc: string, bundlePath: string): string { + if (!fileloc) return ""; + if (!bundlePath) return "."; + const result = relativePath(bundlePath, fileloc); + return result === "" ? "." : result; +} diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index 8f7a2042deadb..d4c18c198c3b4 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -84,7 +84,7 @@ function kindOf(value: object): string { return prototype?.constructor?.name ?? "value"; } -const DAG_SPEC_KEYS: ReadonlySet = new Set(Object.keys(DAG_SCHEMA_FIELDS)); +const DAG_SPEC_KEYS: ReadonlySet = new Set([...Object.keys(DAG_SCHEMA_FIELDS), "queue"]); // `taskId` is hand-written rather than generated: the schema's task_id is // serializer-owned, and this is the authoring surface's own way to set it. const TASK_SPEC_KEYS: ReadonlySet = new Set([...Object.keys(TASK_SCHEMA_FIELDS), "taskId"]); @@ -97,10 +97,20 @@ const TASK_SPEC_KEYS: ReadonlySet = new Set([...Object.keys(TASK_SCHEMA_ * break a call site. An unknown key is rejected, so a misspelled field is an * error rather than a Dag that quietly ignores it. * - * Setting a field records it. A Dag declared in TypeScript is not served to - * Airflow yet, so nothing reads it. + * `queue` is the one hand-written field: Airflow's schema has no Dag-level + * queue, but every task of a native Dag runs on the same coordinator, so the + * queue that routes them there belongs on the Dag rather than on each task. */ -export type DagSpec = GeneratedDagFields; +export interface DagSpec extends GeneratedDagFields { + /** + * Queue the Dag's tasks run on, unless a task names its own. + * + * A native Dag's tasks are executed by the Node coordinator, which the + * deployment's `queue_to_coordinator` maps a queue to, so this is what + * routes them there. `queue` on a {@link TaskSpec} wins for that task. + */ + readonly queue?: string; +} /** * Task-level options: the retries, the pool, the trigger rule, and the rest of @@ -287,8 +297,8 @@ function nodeId(node: Node): string | undefined { return undefined; } -/** Whether `value` is a TaskRef returned by any copy of this package. */ -function isTaskRef(value: unknown): value is TaskRef { +/** Internal: whether `value` is a TaskRef returned by any copy of this package. */ +export function isTaskRef(value: unknown): value is TaskRef { return hasBrand(value, "TaskRef"); } diff --git a/ts-sdk/tests/conformance/serialize_typescript.ts b/ts-sdk/tests/conformance/serialize_typescript.ts new file mode 100644 index 0000000000000..948d12a1688b5 --- /dev/null +++ b/ts-sdk/tests/conformance/serialize_typescript.ts @@ -0,0 +1,163 @@ +/*! + * 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. + */ + +// Serializes the Dags of scripts/ci/lang_sdk_serialization/test_dags.yaml with this SDK, as +// serialize_python.py there does with Airflow's serializer, and writes them keyed by Dag id. The +// check-ts-sdk-serialization-conformance hook has that directory's compare.py run it as: +// +// pnpm --dir ts-sdk exec tsx tests/conformance/serialize_typescript.ts + +import { readFileSync, writeFileSync } from "node:fs"; +import { argv } from "node:process"; + +import { parse, type ScalarTag } from "yaml"; + +import { serializeDag } from "../../src/coordinator/serde.js"; +import { + DAG_SCHEMA_FIELDS, + SERIALIZATION_VERSION, + TASK_SCHEMA_FIELDS, + type SchemaField, +} from "../../src/generated/dag-schema-fields.js"; +import { + Dag, + type DagSpec, + type TaskGroupRef, + type TaskRef, + type TaskSpec, +} from "../../src/index.js"; + +interface TaskCase { + readonly task_id: string; + readonly group?: string; + readonly upstream?: readonly string[]; + readonly spec?: Readonly>; +} + +interface DagCase { + readonly dag_id: string; + readonly spec?: Readonly>; + readonly groups?: readonly string[]; + readonly tasks: readonly TaskCase[]; + readonly order_edges?: readonly (readonly [string, string])[]; +} + +// A duration is already a number of seconds in this SDK's authoring API. +const TAGS: ScalarTag[] = [ + { tag: "!datetime", resolve: (value) => new Date(value) }, + { tag: "!timedelta", resolve: (value) => Number(value) }, +]; + +/** + * Map each schema key of a generated table to the authoring name it is set by. + * + * A virtual field is keyed by its authoring name instead, since its schema key names what the + * serializer derives: `schedule` is spelled that way in both Python's `DAG()` and this SDK, but + * its schema key is `timetable`. + */ +function authoringNames(table: Readonly>): Map { + return new Map( + Object.entries(table).map(([name, field]) => [field.virtual ? name : field.key, name]), + ); +} + +const DAG_NAMES = authoringNames(DAG_SCHEMA_FIELDS); +const TASK_NAMES = authoringNames(TASK_SCHEMA_FIELDS); + +function toAuthoringSpec( + spec: Readonly>, + names: Map, + label: string, +): Record { + const authoring: Record = {}; + for (const [key, value] of Object.entries(spec)) { + const name = names.get(key); + if (name === undefined) throw new Error(`${label}: "${key}" is not an authoring field`); + authoring[name] = value; + } + return authoring; +} + +function qualifiedTaskId(task: TaskCase): string { + return task.group === undefined ? task.task_id : `${task.group}.${task.task_id}`; +} + +function buildDag(dagCase: DagCase): Dag { + const dag = new Dag( + dagCase.dag_id, + toAuthoringSpec(dagCase.spec ?? {}, DAG_NAMES, dagCase.dag_id) as DagSpec, + ); + // A group id is fully qualified, so its parent is whatever comes before the last dot. + const groups = new Map(); + for (const groupId of dagCase.groups ?? []) { + const cut = groupId.lastIndexOf("."); + const scope = cut === -1 ? dag : groups.get(groupId.slice(0, cut))!; + groups.set(groupId, scope.taskGroup(groupId.slice(cut + 1))); + } + + const factories = new Map) => TaskRef>(); + for (const task of dagCase.tasks) { + const label = `${dagCase.dag_id}.${qualifiedTaskId(task)}`; + const spec = toAuthoringSpec(task.spec ?? {}, TASK_NAMES, label) as TaskSpec; + const scope = task.group === undefined ? dag : groups.get(task.group)!; + factories.set( + qualifiedTaskId(task), + scope.task(task.task_id, async (_inputs: Record) => undefined, spec), + ); + } + const refs = new Map(); + for (const task of dagCase.tasks) { + const inputs: Record = {}; + for (const upstream of task.upstream ?? []) { + const ref = refs.get(upstream); + if (!ref) + throw new Error(`${dagCase.dag_id}: "${upstream}" has to come before its downstream`); + inputs[`from_${upstream}`] = ref; + } + refs.set(qualifiedTaskId(task), factories.get(qualifiedTaskId(task))!(inputs)); + } + for (const [upstream, downstream] of dagCase.order_edges ?? []) { + const from = groups.get(upstream) ?? refs.get(upstream); + const to = groups.get(downstream) ?? refs.get(downstream); + if (!from || !to) + throw new Error(`${dagCase.dag_id}: no node "${from ? downstream : upstream}"`); + from.before(to); + } + return dag; +} + +const [testDags, output] = argv.slice(2); +if (testDags === undefined || output === undefined) { + throw new Error("Usage: serialize_typescript.ts "); +} +const { dags } = parse(readFileSync(testDags, "utf-8"), { customTags: TAGS }) as { + dags: DagCase[]; +}; +// compare.py leaves fileloc out, as it names the file a Dag was declared in, so any bundle path +// does. Airflow still needs one to load the Dag. +const serialized = Object.fromEntries( + dags.map((dagCase) => [ + dagCase.dag_id, + { + __version: SERIALIZATION_VERSION, + dag: serializeDag(buildDag(dagCase), "/bundles/app/bundle.mjs", "bundle.mjs"), + }, + ]), +); +writeFileSync(output, `${JSON.stringify(serialized, null, 2)}\n`); diff --git a/ts-sdk/tests/coordinator/arg-binding.test.ts b/ts-sdk/tests/coordinator/arg-binding.test.ts index 1e3f1161d34cb..663462613e880 100644 --- a/ts-sdk/tests/coordinator/arg-binding.test.ts +++ b/ts-sdk/tests/coordinator/arg-binding.test.ts @@ -23,7 +23,9 @@ import { foldArgName, resolveArgs, type BoundArgs } from "../../src/coordinator/ import type { CoordinatorClient, XComEntry } from "../../src/coordinator/client.js"; import type { LogChannel } from "../../src/coordinator/log-channel.js"; import type { ArgBindings } from "../../src/generated/supervisor.js"; +import { Bundle } from "../../src/sdk/bundle.js"; import type { GetXComOpts } from "../../src/sdk/client-types.js"; +import { Dag } from "../../src/sdk/dag.js"; function literal(name: string, value: unknown, extra: Record = {}) { return { name, kind: "literal" as const, value, ...extra }; @@ -76,6 +78,23 @@ async function bind( return { ...bound, warning, pulls }; } +describe("a native task's bound arguments", () => { + it("reach a named handler as the object it destructures", async () => { + const dag = new Dag("d"); + const seen: unknown[] = []; + const store = dag.task("store", async ({ rows }: { rows: number }) => { + seen.push(rows); + }); + store({ rows: 1 }); + const handler = new Bundle(dag).getTaskHandler("d", "store")!; + + const { args } = await bind([literal("rows", 7)]); + await handler(args as never); + + expect(seen).toEqual([7]); + }); +}); + describe("foldArgName", () => { it.each([ ["region_code", "regioncode"], diff --git a/ts-sdk/tests/coordinator/serde.test.ts b/ts-sdk/tests/coordinator/serde.test.ts new file mode 100644 index 0000000000000..84893ebdd6c89 --- /dev/null +++ b/ts-sdk/tests/coordinator/serde.test.ts @@ -0,0 +1,785 @@ +/*! + * 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 } from "vitest"; + +import { + computeRelativeFileloc, + serializeDag, + serializeValue, + unwrapTypeEncoding, +} from "../../src/coordinator/serde.js"; +import { + Dag, + type DagSpec, + type TaskGroupRef, + type TaskRef, + type TaskSpec, +} from "../../src/sdk/dag.js"; + +type Json = Record; +type WiredArgs = Record; + +/** Declare a task and place it, optionally behind some upstreams. */ +function place( + scope: Dag | TaskGroupRef, + taskId: string, + upstream: readonly TaskRef[] = [], + spec?: TaskSpec, +) { + const factory = scope.task(taskId, async (_args: WiredArgs) => undefined, spec); + const inputs: Record = {}; + upstream.forEach((ref, index) => { + inputs[`in_${index}`] = ref; + }); + return factory(inputs); +} + +/** A one-task Dag, serialized. */ +function serializeWith(spec: DagSpec, taskSpec?: TaskSpec): Json { + const dag = new Dag("d", spec); + place(dag, "t", [], taskSpec); + return serializeDag(dag, "/bundles/app/bundle.mjs", "bundle.mjs") as Json; +} + +function taskVar(serialized: Json, index = 0): Json { + const tasks = serialized["tasks"] as { __type: string; __var: Json }[]; + expect(tasks[index]!.__type).toBe("operator"); + return tasks[index]!.__var; +} + +/** Serialized tasks keyed by task id. */ +function taskMap(serialized: Json): Map { + const tasks = serialized["tasks"] as { __var: Json }[]; + return new Map(tasks.map(({ __var }) => [__var["task_id"] as string, __var])); +} + +describe("serializeDag", () => { + it("writes the fields Airflow always expects", () => { + const serialized = serializeWith({}); + + expect(serialized).toMatchObject({ + dag_id: "d", + fileloc: "/bundles/app/bundle.mjs", + relative_fileloc: "bundle.mjs", + timezone: "UTC", + dag_dependencies: [], + edge_info: {}, + params: [], + deadline: null, + allowed_run_types: null, + max_active_tasks: 16, + max_active_runs: 16, + max_consecutive_failed_dag_runs: 0, + catchup: false, + disable_bundle_versioning: false, + }); + }); + + it("puts every task in the flat root group", () => { + const dag = new Dag("d"); + const first = place(dag, "first"); + place(dag, "second", [first]); + + expect(serializeDag(dag, "", ".")["task_group"]).toEqual({ + _group_id: null, + group_display_name: "", + prefix_group_id: true, + tooltip: "", + ui_color: "CornflowerBlue", + ui_fgcolor: "#000", + children: { first: ["operator", "first"], second: ["operator", "second"] }, + upstream_group_ids: [], + downstream_group_ids: [], + upstream_task_ids: [], + downstream_task_ids: [], + }); + }); + + it("identifies a task as TypeScript rather than as a Python operator", () => { + expect(taskVar(serializeWith({}))).toEqual({ + task_id: "t", + task_type: "TypeScriptOperator", + _task_module: "airflow.sdk.coordinators.node", + language: "typescript", + template_fields: [], + is_stub: true, + }); + }); + + describe("queue", () => { + it("gives every task the Dag's queue, which is what routes them here", () => { + const dag = new Dag("d", { queue: "typescript" }); + place(dag, "one"); + place(dag, "two"); + + const tasks = taskMap(serializeDag(dag, "", ".") as Json); + expect([...tasks.values()].map((task) => task["queue"])).toEqual([ + "typescript", + "typescript", + ]); + }); + + it("lets a task name its own queue instead", () => { + const dag = new Dag("d", { queue: "typescript" }); + place(dag, "light"); + place(dag, "heavy", [], { queue: "typescript_large" }); + + const tasks = taskMap(serializeDag(dag, "", ".") as Json); + expect(tasks.get("light")?.["queue"]).toBe("typescript"); + expect(tasks.get("heavy")?.["queue"]).toBe("typescript_large"); + }); + + it("writes no queue when neither the Dag nor the task names one", () => { + expect(taskVar(serializeWith({}))).not.toHaveProperty("queue"); + }); + + it("omits a queue that is already the schema default", () => { + // The scheduler re-derives it, as it does any other defaulted field. + const dag = new Dag("d", { queue: "default" }); + place(dag, "one"); + + expect(taskMap(serializeDag(dag, "", ".") as Json).get("one")).not.toHaveProperty("queue"); + }); + + it("is not written onto the Dag itself, which has no queue field", () => { + expect(serializeWith({ queue: "typescript" })).not.toHaveProperty("queue"); + }); + }); + + describe("arg bindings", () => { + it("binds an upstream reference as the xcom the API server resolves", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + const transform = dag.task("transform", async (_: { extracted: unknown }) => undefined); + transform({ extracted }); + + expect(taskMap(serializeDag(dag, "", ".") as Json).get("transform")).toMatchObject({ + is_stub: true, + _arg_bindings: [{ name: "extracted", kind: "xcom", task_id: "extract" }], + }); + }); + + it("binds a literal as the value itself", () => { + const dag = new Dag("d"); + const transform = dag.task( + "transform", + async (_: { regionCode: string; limits: number[] }) => undefined, + ); + transform({ regionCode: "us", limits: [1, 2] }); + + expect(taskMap(serializeDag(dag, "", ".") as Json).get("transform")).toMatchObject({ + _arg_bindings: [ + { name: "regionCode", kind: "literal", value: "us" }, + { name: "limits", kind: "literal", value: [1, 2] }, + ], + }); + }); + + it("keeps the order the call named the arguments in", () => { + const dag = new Dag("d"); + const north = place(dag, "north"); + const south = place(dag, "south"); + const summarize = dag.task( + "summarize", + async (_: { south: unknown; north: unknown; label: string }) => undefined, + ); + summarize({ south, north, label: "both" }); + + const bindings = taskMap(serializeDag(dag, "", ".") as Json).get("summarize")![ + "_arg_bindings" + ] as { name: string }[]; + expect(bindings.map(({ name }) => name)).toEqual(["south", "north", "label"]); + }); + + it("leaves the key out for a task called with no arguments", () => { + expect(taskVar(serializeWith({}))).not.toHaveProperty("_arg_bindings"); + }); + }); + + describe("timetable", () => { + it.each([ + ["unset", undefined, { __type: "airflow.timetables.simple.NullTimetable", __var: {} }], + ["@once", "@once", { __type: "airflow.timetables.simple.OnceTimetable", __var: {} }], + [ + "@continuous", + "@continuous", + { __type: "airflow.timetables.simple.ContinuousTimetable", __var: {} }, + ], + [ + "a cron expression", + "0 3 * * *", + { + __type: "airflow.timetables.trigger.CronTriggerTimetable", + __var: { + expression: "0 3 * * *", + timezone: "UTC", + interval: 0, + run_immediately: false, + }, + }, + ], + ])("maps %s onto the matching timetable", (_name, schedule, expected) => { + expect(serializeWith({ schedule })["timetable"]).toEqual(expected); + }); + + it.each([ + ["@hourly", "0 * * * *"], + ["@daily", "0 0 * * *"], + ["@weekly", "0 0 * * 0"], + ["@monthly", "0 0 1 * *"], + ["@quarterly", "0 0 1 */3 *"], + ["@yearly", "0 0 1 1 *"], + ])("expands the preset %s the way Python records it", (schedule, expression) => { + // Both spellings rebuild the same timetable, but the expression is what + // the Dag's summary and its hash are taken from, so writing the preset + // verbatim would disagree with the same Dag written in Python. + const timetable = serializeWith({ schedule })["timetable"] as Json; + expect(timetable["__type"]).toBe("airflow.timetables.trigger.CronTriggerTimetable"); + expect((timetable["__var"] as Json)["expression"]).toBe(expression); + }); + + it.each([ + ["prose", "every tuesday"], + ["too few fields", "0 0 * *"], + ["too many fields", "0 0 * * * * *"], + ["an unknown preset", "@fortnightly"], + ])("rejects %s rather than writing a schedule nothing can parse", (_label, schedule) => { + expect(() => serializeWith({ schedule })).toThrowError( + /is not a cron expression or a preset/, + ); + }); + + it.each([ + ["a six-field expression", "0 0 0 * * *"], + ["named weekdays", "0 0 * * MON-FRI"], + ["a step", "*/15 * * * *"], + ])("accepts %s", (_label, schedule) => { + expect(() => serializeWith({ schedule })).not.toThrow(); + }); + + it.each([ + ["an asset expression", { assets: ["s3://bucket/key"] }, /an object schedule names a Python/], + ["a number", 86400, /a number schedule names a Python/], + ["an empty string", "", /schedule for Dag "d" is empty/], + ["a blank string", " ", /schedule for Dag "d" is empty/], + ])("rejects %s", (_name, schedule, expected) => { + expect(() => serializeWith({ schedule } as DagSpec)).toThrowError(expected); + }); + }); + + describe("non-decorated fields", () => { + it("writes them as bare values, without the type encoding", () => { + const serialized = serializeWith({ + startDate: new Date("2026-01-01T00:00:00Z"), + endDate: new Date("2026-12-31T23:30:15Z"), + dagrunTimeout: 300, + tags: ["gamma", "alpha"], + description: "demo", + }); + + expect(serialized).toMatchObject({ + start_date: 1767225600, + end_date: 1798759815, + dagrun_timeout: 300, + // Python holds tags in a set and writes them sorted. + tags: ["alpha", "gamma"], + description: "demo", + }); + }); + + it("keeps only what a decorated field would keep", () => { + // The two halves of the split: everything above went through + // serializeValue and then lost its wrapper, because no authoring field is + // in Python's decorated set. A decorated field would stop at the first. + const wrapped = serializeValue(new Date("2026-01-01T00:00:00Z")); + expect(wrapped).toEqual({ __type: "datetime", __var: 1767225600 }); + expect(unwrapTypeEncoding(wrapped)).toBe(1767225600); + }); + + it("collapses duplicate tags, as the set Python holds them in does", () => { + expect(serializeWith({ tags: ["b", "a", "b"] })["tags"]).toEqual(["a", "b"]); + }); + }); + + describe("omit-if-default", () => { + it("omits a Dag field left at its schema default", () => { + const serialized = serializeWith({ failFast: false, renderTemplateAsNativeObj: false }); + expect(serialized).not.toHaveProperty("fail_fast"); + expect(serialized).not.toHaveProperty("render_template_as_native_obj"); + }); + + it("writes a Dag field that differs from its schema default", () => { + expect(serializeWith({ failFast: true })).toMatchObject({ fail_fast: true }); + }); + + it("omits task fields left at their schema defaults", () => { + const task = taskVar( + serializeWith( + {}, + { retries: 0, queue: "default", pool: "default_pool", retryDelay: 300, owner: "airflow" }, + ), + ); + for (const key of ["retries", "queue", "pool", "retry_delay", "owner"]) { + expect(task).not.toHaveProperty(key); + } + }); + + it("writes task fields that differ from their schema defaults", () => { + const task = taskVar( + serializeWith( + {}, + { retries: 2, queue: "typescript", retryDelay: 600, executionTimeout: 5 }, + ), + ); + expect(task).toMatchObject({ + retries: 2, + queue: "typescript", + retry_delay: 600, + execution_timeout: 5, + }); + }); + + it("never writes the email flags, which have no recipient to reach", () => { + const task = taskVar(serializeWith({}, { emailOnFailure: false, emailOnRetry: false })); + expect(task).not.toHaveProperty("email_on_failure"); + expect(task).not.toHaveProperty("email_on_retry"); + }); + }); + + describe("downstream_task_ids", () => { + it("inverts the recorded wiring, sorted", () => { + const dag = new Dag("d"); + const root = place(dag, "root"); + const right = place(dag, "right", [root]); + const left = place(dag, "left", [root]); + place(dag, "join", [right, left]); + + const serialized = serializeDag(dag, "", ".") as Json; + expect(taskVar(serialized, 0)["downstream_task_ids"]).toEqual(["left", "right"]); + expect(taskVar(serialized, 1)["downstream_task_ids"]).toEqual(["join"]); + expect(taskVar(serialized, 3)).not.toHaveProperty("downstream_task_ids"); + }); + + it("counts one edge when two arguments come from the same upstream", () => { + const dag = new Dag("d"); + const upstream = place(dag, "up"); + place(dag, "down", [upstream, upstream]); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "down", + ]); + }); + + it("ignores a literal argument that looks like a reference", () => { + const dag = new Dag("d"); + place(dag, "up"); + const factory = dag.task("down", async (_args: WiredArgs) => undefined); + factory({ config: { dagId: "d", taskId: "up" } as unknown as TaskRef }); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)).not.toHaveProperty( + "downstream_task_ids", + ); + }); + }); + + describe("order-only edges", () => { + it("writes an edge between two tasks onto the upstream task", () => { + const dag = new Dag("d"); + const loaded = place(dag, "load"); + const cleaned = place(dag, "cleanup"); + loaded.before(cleaned); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "cleanup", + ]); + }); + + it("joins the wiring on one graph, since the serialized Dag has only one", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + place(dag, "transform", [extracted]); + const cleaned = place(dag, "cleanup"); + extracted.before(cleaned); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "cleanup", + "transform", + ]); + }); + + it("counts one edge when the wiring already drew it", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + const transformed = place(dag, "transform", [extracted]); + extracted.before(transformed); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "transform", + ]); + }); + }); + + describe("task_group", () => { + /** The Dag's root task group, as the serialized payload carries it. */ + function rootGroup(dag: Dag): Json { + return (serializeDag(dag, "", ".") as Json)["task_group"] as Json; + } + + it("puts a Dag's top-level tasks in the root group", () => { + const dag = new Dag("d"); + place(dag, "alpha"); + place(dag, "beta"); + + expect(rootGroup(dag)).toMatchObject({ + _group_id: null, + prefix_group_id: true, + children: { alpha: ["operator", "alpha"], beta: ["operator", "beta"] }, + }); + }); + + it("nests a group by embedding its own object, as Python does", () => { + const dag = new Dag("d"); + place(dag, "top"); + const staging = dag.taskGroup("staging"); + place(staging, "stage"); + + const root = rootGroup(dag); + expect(root["children"]).toMatchObject({ top: ["operator", "top"] }); + const [kind, nested] = (root["children"] as Json)["staging"] as [string, Json]; + expect(kind).toBe("taskgroup"); + expect(nested).toMatchObject({ + _group_id: "staging", + children: { "staging.stage": ["operator", "staging.stage"] }, + }); + }); + + it("carries prefixGroupId off, with the ids it holds as written", () => { + const dag = new Dag("d"); + place(dag.taskGroup("checks", { prefixGroupId: false }), "nulls"); + + const [, checks] = (rootGroup(dag)["children"] as Json)["checks"] as [string, Json]; + expect(checks).toMatchObject({ + _group_id: "checks", + prefix_group_id: false, + children: { nulls: ["operator", "nulls"] }, + }); + }); + + it("keeps a grouped task out of the root group's children", () => { + const dag = new Dag("d"); + place(dag.taskGroup("staging"), "stage"); + + expect(Object.keys(rootGroup(dag)["children"] as Json)).toEqual(["staging"]); + }); + + it("nests to any depth", () => { + const dag = new Dag("d"); + const outer = dag.taskGroup("outer"); + place(outer.taskGroup("inner"), "deep"); + + const root = rootGroup(dag); + const [, outerGroup] = (root["children"] as Json)["outer"] as [string, Json]; + const [, innerGroup] = (outerGroup["children"] as Json)["outer.inner"] as [string, Json]; + // The local segment, as Python records it: the qualified id is rebuilt + // from where the group sits in the tree. + expect(innerGroup).toMatchObject({ + _group_id: "inner", + children: { "outer.inner.deep": ["operator", "outer.inner.deep"] }, + }); + }); + + it("records an edge between a group and a task on the group", () => { + const dag = new Dag("d"); + const staging = dag.taskGroup("staging"); + place(staging, "stage"); + const loaded = place(dag, "load"); + staging.before(loaded); + + const [, group] = (rootGroup(dag)["children"] as Json)["staging"] as [string, Json]; + expect(group).toMatchObject({ + downstream_task_ids: ["load"], + upstream_task_ids: [], + downstream_group_ids: [], + upstream_group_ids: [], + }); + // And expanded onto the task graph, which is the only one the scheduler + // reads: the group's leaves carry the edge to the task it points at. + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "load", + ]); + }); + + it("records an edge between two groups on both of them", () => { + const dag = new Dag("d"); + const first = dag.taskGroup("first"); + place(first, "a"); + const second = dag.taskGroup("second"); + place(second, "b"); + first.before(second); + + const children = rootGroup(dag)["children"] as Json; + const [, firstGroup] = children["first"] as [string, Json]; + const [, secondGroup] = children["second"] as [string, Json]; + expect(firstGroup).toMatchObject({ + downstream_group_ids: ["second"], + upstream_group_ids: [], + }); + expect(secondGroup).toMatchObject({ + upstream_group_ids: ["first"], + downstream_group_ids: [], + }); + }); + + it("records an upstream task on the group it points at", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + const staging = dag.taskGroup("staging"); + place(staging, "stage"); + extracted.before(staging); + + const [, group] = (rootGroup(dag)["children"] as Json)["staging"] as [string, Json]; + expect(group).toMatchObject({ upstream_task_ids: ["extract"], downstream_task_ids: [] }); + }); + + it("expands an edge into a group onto its roots, not every task it holds", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + const staging = dag.taskGroup("staging"); + const staged = place(staging, "stage"); + place(staging, "check", [staged]); + extracted.before(staging); + + // "staging.check" already runs after "staging.stage", so the edge in + // reaches only the task that starts the group. + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "staging.stage", + ]); + }); + + it("expands an edge out of a group from its leaves", () => { + const dag = new Dag("d"); + const staging = dag.taskGroup("staging"); + const staged = place(staging, "stage"); + place(staging, "check", [staged]); + const loaded = place(dag, "load"); + staging.before(loaded); + + const serialized = serializeDag(dag, "", ".") as Json; + // "staging.stage" keeps its own intra-group edge and nothing more: the + // edge out of the group leaves from the task that finishes it. + expect(taskVar(serialized, 0)["downstream_task_ids"]).toEqual(["staging.check"]); + expect(taskVar(serialized, 1)["downstream_task_ids"]).toEqual(["load"]); + }); + + it("joins one group's leaves to the next group's roots", () => { + const dag = new Dag("d"); + const first = dag.taskGroup("first"); + const firstHead = place(first, "head"); + place(first, "tail", [firstHead]); + const second = dag.taskGroup("second"); + const secondHead = place(second, "head"); + place(second, "tail", [secondHead]); + first.before(second); + + const serialized = serializeDag(dag, "", ".") as Json; + expect(taskVar(serialized, 1)["downstream_task_ids"]).toEqual(["second.head"]); + const children = rootGroup(dag)["children"] as Json; + const [, secondGroup] = children["second"] as [string, Json]; + // The downstream group records the upstream's leaves as well as the + // group itself; the upstream group records only the group edge. + expect(secondGroup).toMatchObject({ + upstream_group_ids: ["first"], + upstream_task_ids: ["first.tail"], + }); + const [, firstGroup] = children["first"] as [string, Json]; + expect(firstGroup).toMatchObject({ + downstream_group_ids: ["second"], + downstream_task_ids: [], + }); + }); + + it("reaches a task held by a nested group", () => { + const dag = new Dag("d"); + const extracted = place(dag, "extract"); + const outer = dag.taskGroup("outer"); + place(outer.taskGroup("inner"), "deep"); + extracted.before(outer); + + expect(taskVar(serializeDag(dag, "", ".") as Json, 0)["downstream_task_ids"]).toEqual([ + "outer.inner.deep", + ]); + }); + + it("carries an empty group, which holds nothing and constrains nothing", () => { + const dag = new Dag("d"); + place(dag, "solo"); + dag.taskGroup("empty"); + + const [, group] = (rootGroup(dag)["children"] as Json)["empty"] as [string, Json]; + expect(group).toMatchObject({ _group_id: "empty", children: {} }); + }); + + it("steps over an empty group so a chain through it still orders its ends", () => { + // An empty group has no roots and no leaves, so the two edges used to + // expand into nothing at all and `after` ran beside `before`. Python + // bridges the gap by walking up through the group's own upstreams. + const dag = new Dag("d"); + const before = place(dag, "before"); + const after = place(dag, "after"); + const empty = dag.taskGroup("empty"); + before.before(empty); + empty.before(after); + + const tasks = taskMap(serializeDag(dag, "", ".") as Json); + expect(tasks.get("before")?.["downstream_task_ids"]).toEqual(["after"]); + }); + + it("steps over a chain of empty groups", () => { + const dag = new Dag("d"); + const before = place(dag, "before"); + const after = place(dag, "after"); + const first = dag.taskGroup("first"); + const second = dag.taskGroup("second"); + before.before(first); + first.before(second); + second.before(after); + + const tasks = taskMap(serializeDag(dag, "", ".") as Json); + expect(tasks.get("before")?.["downstream_task_ids"]).toEqual(["after"]); + }); + + it("draws nothing when an empty group has no other side to reach", () => { + const dag = new Dag("d"); + const before = place(dag, "before"); + before.before(dag.taskGroup("empty")); + + const tasks = taskMap(serializeDag(dag, "", ".") as Json); + expect(tasks.get("before")).not.toHaveProperty("downstream_task_ids"); + }); + }); + + describe("rejects a value the schema cannot carry", () => { + it.each([ + ["startDate", { startDate: "2026-01-01" }, /startDate for Dag "d" must be a valid Date/], + ["an invalid Date", { startDate: new Date("nope") }, /must be a valid Date/], + ["description", { description: 7 }, /description for Dag "d" must be a string/], + ["catchup", { catchup: "yes" }, /catchup for Dag "d" must be a boolean/], + [ + "maxActiveRuns", + { maxActiveRuns: "3" }, + /maxActiveRuns for Dag "d" must be a finite number/, + ], + ["dagrunTimeout", { dagrunTimeout: Infinity }, /must be a duration in seconds/], + ["tags", { tags: ["a", 2] }, /tags for Dag "d" must be an array of strings/], + ])("on %s", (_name, spec, expected) => { + expect(() => serializeWith(spec as DagSpec)).toThrowError(expected); + }); + + it("names the task a bad task field belongs to", () => { + expect(() => serializeWith({}, { retries: "two" } as unknown as TaskSpec)).toThrowError( + /retries for task "t" of Dag "d" must be a finite number/, + ); + }); + }); +}); + +describe("serializeValue", () => { + it.each([ + ["a string", "x", "x"], + ["a boolean", true, true], + ["a number", 1.5, 1.5], + ["null", null, null], + ["undefined", undefined, null], + ["a list, without a wrapper", [1, "a"], [1, "a"]], + ])("passes %s through", (_name, value, expected) => { + expect(serializeValue(value)).toEqual(expected); + }); + + it("encodes a Date as fractional epoch seconds", () => { + expect(serializeValue(new Date("2026-01-01T00:00:00.500Z"))).toEqual({ + __type: "datetime", + __var: 1767225600.5, + }); + }); + + it("encodes a Set as a sorted list", () => { + expect(serializeValue(new Set(["gamma", "alpha", "beta"]))).toEqual({ + __type: "set", + __var: ["alpha", "beta", "gamma"], + }); + }); + + it.each([ + ["an object", { b: 1, a: "x" }], + [ + "a Map", + new Map([ + ["b", 1], + ["a", "x"], + ]), + ], + ])("encodes %s as a dict", (_name, value) => { + expect(serializeValue(value)).toEqual({ __type: "dict", __var: { b: 1, a: "x" } }); + }); + + it("recurses into nested values", () => { + expect(serializeValue({ when: new Date("2026-01-01T00:00:00Z"), items: [{ n: 1 }] })).toEqual({ + __type: "dict", + __var: { + when: { __type: "datetime", __var: 1767225600 }, + items: [{ __type: "dict", __var: { n: 1 } }], + }, + }); + }); + + it.each([ + ["a non-finite number", Number.NaN, /non-finite number/], + ["an invalid Date", new Date("nope"), /invalid Date/], + ["a function", () => undefined, /Cannot serialize a function/], + ])("rejects %s", (_name, value, expected) => { + expect(() => serializeValue(value)).toThrowError(expected); + }); +}); + +describe("unwrapTypeEncoding", () => { + it("takes the __var of an encoded value", () => { + expect(unwrapTypeEncoding({ __type: "timedelta", __var: 300 })).toBe(300); + }); + + it.each([ + ["a primitive", 5], + ["a list", [1, 2]], + ["an object that is not encoded", { __var: 1 }], + ])("leaves %s alone", (_name, value) => { + expect(unwrapTypeEncoding(value)).toEqual(value); + }); +}); + +describe("computeRelativeFileloc", () => { + it.each([ + ["a file inside the bundle", "/bundles/app/dags/bundle.mjs", "/bundles/app", "dags/bundle.mjs"], + ["a file at the bundle root", "/bundles/app/bundle.mjs", "/bundles/app", "bundle.mjs"], + ["a file that is the bundle", "/bundles/app", "/bundles/app", "."], + ["an unknown bundle path", "/bundles/app/bundle.mjs", "", "."], + ["an unknown file", "", "/bundles/app", ""], + ])("resolves %s", (_name, fileloc, bundlePath, expected) => { + expect(computeRelativeFileloc(fileloc, bundlePath)).toBe(expected); + }); +}); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 769a8342b0e54..b6cb9541627a9 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -360,6 +360,10 @@ describe("public API", () => { expectTypeOf().toEqualTypeOf(); // Both specs are all-optional, so `{}` stays assignable and a field the // schema gains later cannot break a call site. + // `queue` is the one Dag field Airflow's schema does not have: a native + // Dag's tasks all run on the same coordinator, so the queue that routes + // them there belongs on the Dag. + expectTypeOf().toEqualTypeOf(); const emptyDagSpec: DagSpec = {}; const emptyTaskSpec: TaskSpec = {}; expect([emptyDagSpec, emptyTaskSpec]).toEqual([{}, {}]); From 5214a35320a0577a4b1ddbfb0c5f1ae901fdf2a9 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:51 +0000 Subject: [PATCH 3/6] TS SDK: answer the Dag-parsing request from bundle.serve --- ts-sdk/src/coordinator/runtime.ts | 78 ++++++++++++-- ts-sdk/src/sdk/bundle.ts | 18 ++++ ts-sdk/tests/coordinator/integration.test.ts | 104 ++++++++++++++++++- 3 files changed, 185 insertions(+), 15 deletions(-) diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index bf9764f3e6618..8029e8422b73e 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -55,7 +55,10 @@ import { type StartupDetails, } from "./protocol.js"; import { getArgNames } from "../sdk/arg-names.js"; -import { bundleDagTaskIds, type Bundle } from "../sdk/bundle.js"; +import { bundleDagTaskIds, listBundleNativeDags, type Bundle } from "../sdk/bundle.js"; +import type { Dag } from "../sdk/dag.js"; +import { SERIALIZATION_VERSION } from "../generated/dag-schema-fields.js"; +import { computeRelativeFileloc, serializeDag } from "./serde.js"; import { runInTaskScope, type TaskContext } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; @@ -262,22 +265,77 @@ export function createRuntimeAbort( }; } +/** + * Answer a parse request with the Dags this bundle declared in TypeScript. + * + * No handler body runs: a `TaskRef` is inert, so reading a Dag only walks what + * its module already built. Reading it is also what enforces that every task + * was called exactly once, which is why a Dag that is not fully laid out + * surfaces here. + * + * A Dag that cannot be serialized becomes an import error against this file, + * as a Python Dag file that raises does, rather than failing the whole parse: + * one broken Dag must not take out the others a bundle serves. + */ function handleParse( request: { file: string; bundle_path: string }, bundle: Bundle, logs: LogChannel, ): RuntimeDagFileParsingResult { - // TypeScript-native Dag parsing is not yet supported. - // Respond with an empty result so the Python-stub-Dag workflow works. - logs.info("Parse-mode response (TS Dag parsing not yet supported)", { - registered_tasks: Object.fromEntries(bundleDagTaskIds(bundle)), + const fileloc = request.file; + const relativeFileloc = computeRelativeFileloc(fileloc, request.bundle_path); + const serializedDags: { data: Record }[] = []; + // Airflow keys an import error by the bundle-relative path and holds one row + // per file (`DagFileProcessorManager.update_import_errors`), so every failure + // in this bundle is reported under that one key, naming its Dag in the + // message. An absolute path, or one with a Dag id appended, would give a row + // the UI cannot tie back to the file, and would leave the file itself looking + // healthy while its Dags had vanished. + const failures: string[] = []; + + let dags: Dag[]; + try { + dags = listBundleNativeDags(bundle); + } catch (err) { + // Reading the bundle closes every native Dag in it, so a fault in one is + // reported against the file rather than leaving the request unanswered. + const detail = err instanceof Error ? err.message : String(err); + logs.error("Bundle could not be read for parsing", { fileloc, detail }); + return { + type: "DagFileParsingResult", + fileloc, + serialized_dags: [], + import_errors: { [relativeFileloc]: detail }, + }; + } + + for (const dag of dags) { + try { + serializedDags.push({ + data: { + __version: SERIALIZATION_VERSION, + dag: serializeDag(dag, fileloc, relativeFileloc), + }, + }); + } catch (err) { + const detail = err instanceof Error ? err.message : String(err); + logs.error("Dag could not be serialized", { dag_id: dag.dagId, detail }); + failures.push(`Dag "${dag.dagId}": ${detail}`); + } + } + + logs.info("Parse-mode response", { + fileloc, + dag_ids: dags.map((dag) => dag.dagId), + serialized: serializedDags.length, + import_errors: failures.length, }); - const response: RuntimeDagFileParsingResult = { + return { type: "DagFileParsingResult", - fileloc: request.file, - serialized_dags: [], - }; - return response; + fileloc, + serialized_dags: serializedDags, + ...(failures.length > 0 && { import_errors: { [relativeFileloc]: failures.join("\n") } }), + } as RuntimeDagFileParsingResult; } async function handleTask( diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index a3e54f1a23172..34b6351032ce9 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -239,3 +239,21 @@ export function bundleDagTaskIds(bundle: Bundle): Map { } return byDag; } + +/** + * Internal: the Dags this bundle declared in TypeScript, in registration order. + * + * What a parse request answers with, and what closes each of them: a Dag known + * only through task handlers is not one of them, because its graph belongs to + * the Python Dag file that declares it, and serializing it here would register + * a second Dag with the same `dag_id` from a different `fileloc`. + * + * Supersedes {@link finalizeBundleDags} for a caller that also wants the Dags. + */ +export function listBundleNativeDags(bundle: Bundle): Dag[] { + const dags = [...dagsOf(bundle).values()]; + for (const dag of dags) { + finalizeDag(dag); + } + return dags; +} diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index d0b6d53f83fff..e20a780e8f0f9 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -798,18 +798,112 @@ describe("coordinator runtime integration", () => { expect(calledSecondDag).toBe(false); }); - it("returns empty serialized_dags for DagFileParseRequest", async () => { + describe("DagFileParseRequest", () => { const parseRequest = { type: "DagFileParseRequest", file: "/dags/test.mjs", bundle_path: "/dags", }; - const result = await driveSupervisor(parseRequest); + async function parse(): Promise> { + const result = await driveSupervisor(parseRequest); + return result.firstResponse!.body as Record; + } + + it("answers with the Dags the bundle declared in TypeScript", async () => { + testDag.task("extract", async () => undefined)(); + otherDag.task("stage", async () => undefined)(); + + const body = await parse(); + + expect(body.type).toBe("DagFileParsingResult"); + expect(body.fileloc).toBe("/dags/test.mjs"); + const dags = body.serialized_dags as { data: Record }[]; + expect(dags.map((entry) => (entry.data.dag as Record).dag_id)).toEqual([ + "test_dag", + "other_dag", + ]); + expect(dags[0]!.data.__version).toBe(3); + expect((dags[0]!.data.dag as Record).relative_fileloc).toBe("test.mjs"); + expect(body.import_errors).toBeUndefined(); + }); + + it("runs no handler body while parsing", async () => { + let ran = false; + testDag.task("extract", async () => { + ran = true; + })(); + otherDag.task("stage", async () => undefined)(); + + await parse(); + + expect(ran).toBe(false); + }); + + it("answers with nothing when the bundle only binds handlers to Python Dags", async () => { + // A Python Dag's graph belongs to the Python file that declares it, so a + // bundle of task handlers has no Dag of its own to serialize. + bundle = new Bundle(new TaskHandler("py_dag", "transform", async () => undefined)); + + const body = await parse(); + + expect(body.serialized_dags).toEqual([]); + expect(body.import_errors).toBeUndefined(); + }); + + it("reports an uncalled task as an import error rather than failing the parse", async () => { + testDag.task("extract", async () => undefined)(); + testDag.task("orphan", async () => undefined); - const body = result.firstResponse!.body as Record; - expect(body.type).toBe("DagFileParsingResult"); - expect(body.serialized_dags).toEqual([]); + const body = await parse(); + + expect(body.serialized_dags).toEqual([]); + // Keyed by the bundle-relative path, which is how Airflow ties the row + // to the file it came from. + expect(body.import_errors).toEqual({ + "test.mjs": expect.stringContaining('Task "orphan" of Dag "test_dag" is never called'), + }); + }); + + it("merges several failing Dags into the one row Airflow keeps per file", async () => { + testDag.task("extract", async () => undefined)(); + for (const dagId of ["broken_a", "broken_b"]) { + const broken = new Dag(dagId, { schedule: "" }); + broken.task("t", async () => undefined)(); + bundle.register(broken); + } + + const body = await parse(); + + const errors = body.import_errors as Record; + expect(Object.keys(errors)).toEqual(["test.mjs"]); + expect(errors["test.mjs"]).toContain('Dag "broken_a"'); + expect(errors["test.mjs"]).toContain('Dag "broken_b"'); + }); + + it("keeps the Dags it can serialize when one of them cannot be", async () => { + testDag.task("extract", async () => undefined)(); + // An empty schedule is rejected by the serializer, and only that Dag is + // lost: the rest of the bundle still parses. + const broken = new Dag("broken_dag", { schedule: "" }); + broken.task("t", async () => undefined)(); + bundle.register(broken); + + const body = await parse(); + + const dags = body.serialized_dags as { data: Record }[]; + expect(dags.map((entry) => (entry.data.dag as Record).dag_id)).toEqual([ + "test_dag", + "other_dag", + ]); + // One row per file, so the failing Dag is named in the message rather + // than appended to the key. + expect(body.import_errors).toEqual({ + "test.mjs": expect.stringContaining( + 'Dag "broken_dag": schedule for Dag "broken_dag" is empty', + ), + }); + }); }); it("auto-pushes return_value XCom when handler returns a value", async () => { From f394f5ef181652a126cd0cc79c20f3d891be4a6b Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:52 +0000 Subject: [PATCH 4/6] TS SDK: add conditional branching with dag.if --- .../language-sdks/typescript.rst | 24 ++ ts-sdk/adr/0002-native-dag-interface.md | 37 +++ ts-sdk/src/coordinator/client.ts | 11 + ts-sdk/src/coordinator/protocol.ts | 1 + ts-sdk/src/coordinator/serde.ts | 6 + ts-sdk/src/index.ts | 2 + ts-sdk/src/sdk/client.ts | 13 + ts-sdk/src/sdk/dag.ts | 207 +++++++++++- ts-sdk/tests/sdk/conditional.test.ts | 296 ++++++++++++++++++ ts-sdk/tests/sdk/dag.test.ts | 4 +- 10 files changed, 592 insertions(+), 9 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 c81b2b3dc5600..ace8a627c0fb9 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,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. diff --git a/ts-sdk/adr/0002-native-dag-interface.md b/ts-sdk/adr/0002-native-dag-interface.md index 6614ee62091bb..7d9f63e8987bb 100644 --- a/ts-sdk/adr/0002-native-dag-interface.md +++ b/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, diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index ba0be84faa240..c8c908186f6a4 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"; @@ -202,6 +203,16 @@ export function createCoordinatorClient( await rpc("SetXCom", null, msg, () => undefined, "throw"); }, + // ---- Control flow ---- + + async skipDownstreamTasks(taskIds: readonly string[]): Promise { + // 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 { 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 26a1d844a3dc3..3227b374a157b 100644 --- a/ts-sdk/src/coordinator/serde.ts +++ b/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 | undefined, inputs: RecordedInputs | undefined, + canSkipDownstream: boolean, ): SerializedValue { const data: Record = { 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(); diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 4e7bee9a9669b..0c8027b55a6e2 100644 --- a/ts-sdk/src/index.ts +++ b/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, diff --git a/ts-sdk/src/sdk/client.ts b/ts-sdk/src/sdk/client.ts index 0facdd9695c5d..448664474c0e8 100644 --- a/ts-sdk/src/sdk/client.ts +++ b/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; + + /** + * 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; } /** Error thrown by {@link TaskClient.getVariableOrThrow}. */ diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index d4c18c198c3b4..fc5ffb8617a09 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -31,7 +31,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"; function isPlainRecord(value: unknown): value is Record { if (typeof value !== "object" || value === null || Array.isArray(value)) return false; @@ -400,13 +400,58 @@ export type TaskFactory = */ export type TaskOptions = TaskSpec; +/** + * What `dag.if(condition)` returns: name the task the condition runs. + * + * ```ts + * dag.if(gated).then(loaded).else(reportedEmpty); + * ``` + * + * `then` is required, `else` optional. A one-sided condition follows nothing + * when it is false, skipping only its own branch rather than the whole + * downstream closure the way `ShortCircuitOperator` does — decision 3 of + * `airflow-core/adr/lang-sdk/0008-control-flow-constructs.md`. + */ +export interface Condition { + then(taskRef: TaskRef): ConditionElse; +} + +/** What `.then(...)` returns: the other side, which a one-sided condition omits. */ +export interface ConditionElse { + else(taskRef: TaskRef): void; +} + /** 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. + * + * Serialized as `_can_skip_downstream`, which is what makes Airflow consult + * the task's `skipmixin_key` XCom when a skipped downstream is cleared. + */ + readonly canSkipDownstream?: boolean; +} + +/** The branches a condition chooses between, filled in as the chain is written. */ +interface ConditionRecord { + whenTrue?: TaskRef; + whenFalse?: TaskRef; } +/** + * The XCom a task deciding skips leaves behind, so clearing a skipped task + * re-skips it rather than running it. + * + * `NotPreviouslySkippedDep` reads this key off each upstream that declares + * `_can_skip_downstream`, and skips a task named under `skipped`. 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. * @@ -467,6 +512,10 @@ export class Dag { // insertion-ordered so the serialized Dag reads as written. readonly #orderEdges = new Map(); readonly #definedIn: string | undefined; + // Which tasks decide a branch, so one task cannot decide twice and a + // condition with no branch named is reported when the Dag is read. + readonly #conditions = new Map(); + readonly #branches = 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(); @@ -539,6 +588,136 @@ export class Dag { return this.#addTask(undefined, taskIdOrHandler, handlerOrOptions, maybeOptions); } + /** + * Make an existing task a condition, and name the task each outcome runs — + * TypeScript's `if`/`else`, spelled as the Dag's own control flow rather + * than after an operator class. + * + * ```ts + * const gated = dag.task("has_rows", async ({ rows }: { rows: number }) => rows > 0)({ + * rows: extracted, + * }); + * + * dag.if(gated).then(loaded).else(reportedEmpty); + * ``` + * + * The condition is an ordinary task whose handler returns a boolean, so the + * compiler checks the type and nothing depends on its id or its function + * name. It becomes a branch: the side not taken is skipped when the run + * reaches it, and stays skipped if it is cleared later. + * + * A guarded task takes no argument for the control edge — a condition's + * boolean is a signal, not data. + */ + if(condition: TaskRef): Condition { + const taskId = this.#beginBranch(condition, "dag.if"); + 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", + ); + } + // The side not taken. A one-sided condition that holds skips nothing. + 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", taskRef: TaskRef): void => { + // `await` on a thenable calls `then(resolve, reject)`. This object only + // ever takes a task reference, so saying why beats "not a task". + if (typeof taskRef === "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", + ); + } + this.#validateOwnNode(taskRef, `the "${side}" branch of "${taskId}"`); + 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; + // The control edge carries no value, so it is an ordinary order-only + // edge; what makes it a branch is the skip, which only exists at run time. + condition.before(taskRef); + }; + + const elseStep: ConditionElse = { + else: (taskRef) => name("else", taskRef), + }; + return { + then: (taskRef) => { + name("then", taskRef); + return elseStep; + }, + }; + } + + /** + * Check that `condition` is a task of this Dag that no other construct has + * already claimed, and return its id. + */ + #beginBranch(condition: TaskRef, verb: string): string { + if (this.#finalized) { + throw new Error( + `${verb}(...) cannot be used on Dag "${this.dagId}" after the Dag was read; ` + + "declare every branch while the module is loading", + ); + } + this.#validateOwnNode(condition, `the condition given to ${verb}`); + const taskId = condition.taskId; + if (this.#conditions.has(taskId) || this.#branches.has(taskId)) { + throw new Error( + `Task "${taskId}" of Dag "${this.dagId}" already decides a branch; ` + + "one task decides one way", + ); + } + return taskId; + } + + /** + * Replace a task's handler with one that also sends the skip its decision + * implies. + * + * The decision is read when the task runs, so the branches it chooses + * between can still be named after this call — which is what lets the chain + * read `dag.if(gate).then(a).else(b)`. + */ + #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(); + // Written before the skip so a downstream cleared later is re-skipped + // by `NotPreviouslySkippedDep` rather than run, which is what Python's + // SkipMixin does. Keyed on the task, so its own id is not needed here. + 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. * @@ -813,7 +992,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) { @@ -832,17 +1011,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`, ); } @@ -852,7 +1031,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}`, ); } @@ -937,6 +1116,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(condition).then(task)", + ); + } + } for (const taskId of this.#tasks.keys()) { if (!this.#inputs.has(taskId)) { throw new Error( @@ -967,6 +1154,14 @@ const GROUP_ID_MAX_LENGTH = 200; * `keepNames`, so the name survives minification; a bundler that drops names * leaves the empty string, which is why the empty string is not an id. */ +/** 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} ${JSON.stringify(value)}`; +} + function readFunctionName(handler: unknown): string | undefined { if (typeof handler !== "function") return undefined; const { name } = handler as { name?: unknown }; diff --git a/ts-sdk/tests/sdk/conditional.test.ts b/ts-sdk/tests/sdk/conditional.test.ts new file mode 100644 index 0000000000000..b873f197acde8 --- /dev/null +++ b/ts-sdk/tests/sdk/conditional.test.ts @@ -0,0 +1,296 @@ +/*! + * 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, + 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"; + +/** A Dag with two placed tasks and a placed boolean condition between them. */ +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.task("has_rows", async () => holds as boolean)(); + return { dag, ifReady, fallback, gated }; +} + +/** Run a registered handler with a client in scope, as the runtime does. */ +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("leaves the condition an ordinary task of the Dag", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + dag.if(gated).then(ifReady).else(fallback); + + expect(dag.taskIds).toEqual(["load_if_ready", "load_fallback", "has_rows"]); + }); + + it("marks the condition as deciding skips, which Airflow reads on a clear", () => { + const { dag, ifReady, gated } = gatedDag(); + dag.if(gated).then(ifReady); + + expect(getDagTaskRecords(dag).get("has_rows")?.canSkipDownstream).toBe(true); + expect(getDagTaskRecords(dag).get("load_if_ready")?.canSkipDownstream).toBeUndefined(); + }); + + it("serializes the condition as a task that decides skips", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + dag.if(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("draws an order-only edge to each branch, since no value flows", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + dag.if(gated).then(ifReady).else(fallback); + + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "has_rows", downstream: "load_if_ready" }, + { upstream: "has_rows", downstream: "load_fallback" }, + ]); + }); + + it("draws one edge for a one-sided condition", () => { + const { dag, ifReady, gated } = gatedDag(); + dag.if(gated).then(ifReady); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "has_rows", downstream: "load_if_ready" }]); + }); + + 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 extract = dag.task("extract", async () => 3); + const condition = dag.task("has_rows", async ({ rows }: { rows: number }) => rows > 0); + const gated = condition({ rows: extract() }); + dag.if(gated).then(loaded); + + expect((await runHandler(dag, "has_rows", { rows: 2 })).returned).toBe(true); + expect((await runHandler(dag, "has_rows", { rows: 0 })).returned).toBe(false); + }); + + describe("at run time", () => { + it("skips the else branch when the condition holds", async () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + dag.if(gated).then(ifReady).else(fallback); + + const { skipDownstreamTasks, setXCom } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_fallback"]); + // Recorded so a cleared "load_fallback" is skipped again rather than run. + 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); + dag.if(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(); + dag.if(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 () => { + // Not the whole downstream closure: this is a branch with one candidate, + // so a task several branches converge on keeps running. + const { dag, ifReady, gated } = gatedDag("d", false); + dag.if(gated).then(ifReady); + + const { skipDownstreamTasks } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_if_ready"]); + }); + + it("reads the branches at run time, so the chain can be written after", async () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + const chain = dag.if(gated); + chain.then(ifReady).else(fallback); + + const { skipDownstreamTasks } = await runHandler(dag, "has_rows"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_fallback"]); + }); + + it("awaits a condition that resolves later", async () => { + const dag = new Dag("d"); + const ifReady = dag.task("load_if_ready", async () => undefined)(); + const fallback = dag.task("load_fallback", async () => undefined)(); + const gated = dag.task("has_rows", async () => { + await Promise.resolve(); + return false; + })(); + dag.if(gated).then(ifReady).else(fallback); + + const { returned, skipDownstreamTasks } = await runHandler(dag, "has_rows"); + + expect(returned).toBe(false); + expect(skipDownstreamTasks).toHaveBeenCalledWith(["load_if_ready"]); + }); + + it.each([ + ["a string", "yes"], + ["null", null], + ["a number", 1], + ])("fails the task when the condition returns %s", async (_label, value) => { + const { dag, ifReady, gated } = gatedDag("d", value); + dag.if(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)(); + dag.if(gated).then(ifReady); + + gated.after(extracted); + + expect(getDagOrderEdges(dag)).toContainEqual({ + upstream: "extract", + downstream: "has_rows", + }); + }); + + describe("rejects", () => { + it("a condition taken from another Dag", () => { + const { dag } = gatedDag("here"); + const { gated: foreign } = gatedDag("there"); + + expect(() => dag.if(foreign)).toThrowError( + /the condition given to dag.if cannot reach Dag "there" node "has_rows"/, + ); + }); + + it("a branch taken from another Dag", () => { + const { dag, gated } = gatedDag("here"); + const { ifReady: foreign } = gatedDag("there"); + + expect(() => dag.if(gated).then(foreign)).toThrowError( + /the "then" branch of "has_rows" cannot reach Dag "there" node "load_if_ready"/, + ); + }); + + it("an else branch taken from another Dag", () => { + const { dag, ifReady, gated } = gatedDag("here"); + const { fallback: foreign } = gatedDag("there"); + + expect(() => dag.if(gated).then(ifReady).else(foreign)).toThrowError( + /the "else" branch of "has_rows" cannot reach Dag "there"/, + ); + }); + + 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 { dag, gated } = gatedDag(); + + expect(() => dag.if(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 { dag, ifReady, gated } = gatedDag(); + + expect(() => dag.if(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 { dag, ifReady, fallback, gated } = gatedDag(); + const chain = dag.if(gated).then(ifReady); + chain.else(fallback); + + expect(() => chain.else(fallback)).toThrowError( + /Condition "has_rows" of Dag "d" already has an? "else" branch/, + ); + }); + + it("a task made a condition twice", () => { + const { dag, ifReady, fallback, gated } = gatedDag(); + dag.if(gated).then(ifReady); + + expect(() => dag.if(gated).then(fallback)).toThrowError( + /Task "has_rows" of Dag "d" already decides a branch/, + ); + }); + + it("a condition that names no branch, when the Dag is read", () => { + const { dag, gated } = gatedDag(); + dag.if(gated); + + 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 () => { + // An object with a callable `then` is a thenable, so `await` hands it a + // resolve function where a task reference belongs. + const { dag, ifReady, gated } = gatedDag(); + dag.task("other", async () => undefined)(); + + await expect(Promise.resolve(dag.if(gated) as unknown as Promise)).rejects.toThrow( + /dag.if\(\.\.\.\) of Dag "d" was awaited/, + ); + void ifReady; + }); + }); +}); 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"`), ); }); From 7918d19dc4b33dcc402ffa8053703cbc9460797d Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:54 +0000 Subject: [PATCH 5/6] TS SDK: add multi-way branching with dag.switch --- .../language-sdks/typescript.rst | 25 ++ ts-sdk/adr/0002-native-dag-interface.md | 29 ++ ts-sdk/src/index.ts | 1 + ts-sdk/src/sdk/dag.ts | 89 +++++++ ts-sdk/tests/sdk/branching.test.ts | 247 ++++++++++++++++++ 5 files changed, 391 insertions(+) create mode 100644 ts-sdk/tests/sdk/branching.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 ace8a627c0fb9..56fd2881e3df5 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -383,6 +383,31 @@ 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 task whose handler returns one of the cases it is given. + +.. code-block:: typescript + + const decider = dag.task("pick_path", async ({ rows }: { rows: number }) => + rows > 1000 ? handleLong : handleShort, + ); + const picked = decider({ rows: extracted }); + + dag.switch(picked).case(handleLong).case(handleShort); + +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. diff --git a/ts-sdk/adr/0002-native-dag-interface.md b/ts-sdk/adr/0002-native-dag-interface.md index 7d9f63e8987bb..7c08f4f600916 100644 --- a/ts-sdk/adr/0002-native-dag-interface.md +++ b/ts-sdk/adr/0002-native-dag-interface.md @@ -172,6 +172,35 @@ 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(decider)` 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 one too. + +```ts +const decider = dag.task("pick_path", async ({ rows }: { rows: number }) => + rows > 1000 ? handleLong : handleShort, +); +const picked = decider({ rows: extracted }); + +dag.switch(picked).case(handleLong).case(handleShort); +``` + +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, diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 0c8027b55a6e2..e5beb163c48ae 100644 --- a/ts-sdk/src/index.ts +++ b/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, DagSpec, diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index fc5ffb8617a09..ab732ce7ebaf4 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -421,6 +421,21 @@ export interface ConditionElse { else(taskRef: TaskRef): void; } +/** + * What `dag.switch(decider)` returns: name each task the decider chooses + * between. + * + * ```ts + * dag.switch(picked).case(publishDaily).case(publishWeekly); + * ``` + * + * Chainable, so the cases read as a list. There is no default case, and a + * branch with no case at all is rejected when the Dag is read. + */ +export interface Branch { + case(taskRef: TaskRef): Branch; +} + /** Per-task record a Dag retains: the reference, the handler, and its spec. */ export interface TaskRecord { readonly task: TaskRef; @@ -718,6 +733,72 @@ export class Dag { this.#tasks.set(taskId, { ...record, canSkipDownstream: true, fn: wrapped }); } + /** + * Make an existing task a multi-way branch, and name the tasks it chooses + * between — TypeScript's `switch`/`case`. + * + * ```ts + * const decider = dag.task("pick_path", async ({ rows }: { rows: number }) => + * rows > 1000 ? handleLong : handleShort, + * ); + * const picked = decider({ rows: extracted }); + * + * dag.switch(picked).case(handleLong).case(handleShort); + * ``` + * + * The decider is an ordinary task whose handler returns one of the case + * references the Dag already handed back, so the compiler checks the + * candidate exists and renaming a handler cannot silently rewire a Dag. The + * task's own value on the wire is the chosen task's id, which is what a + * downstream task reads back from its XCom. + * + * **There is no default case.** A decider that returns a reference outside + * the declared cases fails the task, the way Python's `skip_all_except` + * raises on a task id the Dag does not hold. + * + * **Exactly one case is selected.** Python's branch callable may return a + * list of task ids; a single-reference return cannot express that, and no + * Lang SDK offers it for now. A Dag that needs several paths together puts + * them behind one task, or gates each with its own condition. + */ + switch(decider: TaskRef): Branch { + const taskId = this.#beginBranch(decider, "dag.switch"); + const candidates: TaskRef[] = []; + this.#branches.set(taskId, candidates); + this.#wrapDecider(taskId, async (chosen: unknown) => { + const known = candidates.map((ref) => ref.taskId); + if (!candidates.includes(chosen as TaskRef)) { + 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(", ")}`, + ); + } + const picked = (chosen as TaskRef).taskId; + // The task_id is what a branch puts on the wire, so it is also what a + // downstream task reads back from this one's XCom. + return { skip: known.filter((id) => id !== picked), result: picked }; + }); + + const branch: Branch = { + case: (taskRef) => { + this.#validateOwnNode(taskRef, `a case of "${taskId}"`); + // Two cases naming one task cannot be told apart by a decision, and one + // of them would always be skipped alongside itself. + 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; + }, + }; + return branch; + } + /** * Declare a task group of this Dag. * @@ -1124,6 +1205,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(decider).case(task)", + ); + } + } for (const taskId of this.#tasks.keys()) { if (!this.#inputs.has(taskId)) { throw new Error( diff --git a/ts-sdk/tests/sdk/branching.test.ts b/ts-sdk/tests/sdk/branching.test.ts new file mode 100644 index 0000000000000..ab5481f3545fa --- /dev/null +++ b/ts-sdk/tests/sdk/branching.test.ts @@ -0,0 +1,247 @@ +/*! + * 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, + 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"; + +/** A Dag with three placed tasks and a placed decider that returns one of them. */ +function branchedDag(dagId = "d", choose?: (refs: Record) => unknown) { + const dag = new Dag(dagId); + const long = dag.task("handle_long", async () => undefined)(); + const short = dag.task("handle_short", async () => undefined)(); + const other = dag.task("handle_other", async () => undefined)(); + const refs = { long, short, other }; + const decider = dag.task("pick_path", async () => (choose ? (choose(refs) as TaskRef) : long))(); + return { dag, long, short, other, decider }; +} + +/** Run a registered handler with a client in scope, as the runtime does. */ +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.switch", () => { + it("leaves the decider an ordinary task of the Dag", () => { + const { dag, long, short, decider } = branchedDag(); + dag.switch(decider).case(long).case(short); + + expect(dag.taskIds).toEqual(["handle_long", "handle_short", "handle_other", "pick_path"]); + }); + + it("marks the decider as deciding skips, which Airflow reads on a clear", () => { + const { dag, long, short, decider } = branchedDag(); + dag.switch(decider).case(long).case(short); + + 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("pick_path")?.["_can_skip_downstream"]).toBe(true); + expect(byId.get("handle_long")).not.toHaveProperty("_can_skip_downstream"); + }); + + it("draws an order-only edge to each case, since no value flows", () => { + const { dag, long, short, decider } = branchedDag(); + dag.switch(decider).case(long).case(short); + + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "pick_path", downstream: "handle_long" }, + { upstream: "pick_path", downstream: "handle_short" }, + ]); + }); + + it("takes a single case, which is then the only side that can run", () => { + const { dag, long, decider } = branchedDag(); + dag.switch(decider).case(long); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "pick_path", downstream: "handle_long" }]); + }); + + describe("at run time", () => { + it("returns the chosen task's id, which is what a branch puts on the wire", async () => { + const { dag, long, short, decider } = branchedDag("d", ({ short: s }) => s); + dag.switch(decider).case(long).case(short); + + expect((await runHandler(dag, "pick_path")).returned).toBe("handle_short"); + }); + + it("skips every case it did not choose", async () => { + const { dag, long, short, other, decider } = branchedDag("d", ({ long: l }) => l); + dag.switch(decider).case(long).case(short).case(other); + + const { skipDownstreamTasks, setXCom } = await runHandler(dag, "pick_path"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["handle_short", "handle_other"]); + expect(setXCom).toHaveBeenCalledWith({ + key: "skipmixin_key", + value: { skipped: ["handle_short", "handle_other"] }, + }); + }); + + it("skips nothing when its one case is the one chosen", async () => { + const { dag, long, decider } = branchedDag("d", ({ long: l }) => l); + dag.switch(decider).case(long); + + const { skipDownstreamTasks } = await runHandler(dag, "pick_path"); + + expect(skipDownstreamTasks).not.toHaveBeenCalled(); + }); + + it("reads the cases at run time, so the chain can be written after", async () => { + const { dag, long, short, decider } = branchedDag("d", ({ long: l }) => l); + const chain = dag.switch(decider); + chain.case(long).case(short); + + const { skipDownstreamTasks } = await runHandler(dag, "pick_path"); + + expect(skipDownstreamTasks).toHaveBeenCalledWith(["handle_short"]); + }); + + it("fails the task when the decider chooses a task that is not a case", async () => { + const { dag, long, short, decider } = branchedDag("d", ({ other }) => other); + dag.switch(decider).case(long).case(short); + + await expect(runHandler(dag, "pick_path")).rejects.toThrow( + /Task "pick_path" of Dag "d" chose "handle_other", which is not one of its cases: handle_long, handle_short/, + ); + }); + + it.each([ + ["a string", "handle_long"], + ["a look-alike object", { dagId: "d", taskId: "handle_long" }], + ["null", null], + ])("fails the task when the decider returns %s", async (_label, value) => { + const { dag, long, short, decider } = branchedDag("d", () => value); + dag.switch(decider).case(long).case(short); + + await expect(runHandler(dag, "pick_path")).rejects.toThrow(/which is not one of its cases/); + }); + + it("skips nothing and fails when the decider chose wrongly", async () => { + const { dag, long, short, decider } = branchedDag("d", ({ other }) => other); + dag.switch(decider).case(long).case(short); + + const skipDownstreamTasks = vi.fn(async () => undefined); + const client = { skipDownstreamTasks, setXCom: vi.fn() } as unknown as TaskClient; + const ctx = { dagId: "d", taskId: "pick_path" } as unknown as TaskContext; + const handler = getDagTaskRecords(dag).get("pick_path")!.fn; + + await expect( + runInTaskScope({ ctx, client }, () => (handler as (args: unknown) => Promise)({})), + ).rejects.toThrow(); + expect(skipDownstreamTasks).not.toHaveBeenCalled(); + }); + }); + + it("carries edges of its own, like any task", () => { + const { dag, long, decider } = branchedDag(); + const extracted = dag.task("extract", async () => undefined)(); + dag.switch(decider).case(long); + + decider.after(extracted); + + expect(getDagOrderEdges(dag)).toContainEqual({ + upstream: "extract", + downstream: "pick_path", + }); + }); + + describe("rejects", () => { + it("a decider taken from another Dag", () => { + const { dag } = branchedDag("here"); + const { decider: foreign } = branchedDag("there"); + + expect(() => dag.switch(foreign)).toThrowError( + /the condition given to dag.switch cannot reach Dag "there" node "pick_path"/, + ); + }); + + it("a case taken from another Dag", () => { + const { dag, decider } = branchedDag("here"); + const { long: foreign } = branchedDag("there"); + + expect(() => dag.switch(decider).case(foreign)).toThrowError( + /a case of "pick_path" cannot reach Dag "there" node "handle_long"/, + ); + }); + + it.each([ + ["a plain object", { dagId: "d", taskId: "handle_long" }], + ["a string", "handle_long"], + ["null", null], + ])("%s where a case belongs", (_label, value) => { + const { dag, decider } = branchedDag(); + + expect(() => dag.switch(decider).case(value as TaskRef)).toThrowError( + /a case of "pick_path" on Dag "d" takes tasks and task groups/, + ); + }); + + it("the same task listed twice", () => { + const { dag, long, decider } = branchedDag(); + + expect(() => dag.switch(decider).case(long).case(long)).toThrowError( + /Dag "d" branch "pick_path" lists "handle_long" twice; each case names a different task/, + ); + }); + + it("a task made a branch twice", () => { + const { dag, long, short, decider } = branchedDag(); + dag.switch(decider).case(long); + + expect(() => dag.switch(decider).case(short)).toThrowError( + /Task "pick_path" of Dag "d" already decides a branch/, + ); + }); + + it("a task that is both a condition and a branch", () => { + const { dag, long, short, decider } = branchedDag(); + dag.switch(decider).case(long); + + expect(() => dag.if(decider as unknown as TaskRef).then(short)).toThrowError( + /Task "pick_path" of Dag "d" already decides a branch/, + ); + }); + + it("a branch with no case, when the Dag is read", () => { + const { dag, decider } = branchedDag(); + dag.switch(decider); + + expect(() => finalizeDag(dag)).toThrowError( + /Branch "pick_path" of Dag "d" has no cases, so it decides nothing/, + ); + }); + }); +}); From 172b5802d7d1358e2f6771e1b0185b8a60de7078 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 12:11:54 +0000 Subject: [PATCH 6/6] TS SDK: add triggerDagRun to the native Dag surface --- .../lang-sdk/0008-control-flow-constructs.md | 10 +- .../language-sdks/typescript.rst | 44 +++ ts-sdk/src/coordinator/client.ts | 53 +++ ts-sdk/src/coordinator/protocol.ts | 8 + ts-sdk/src/coordinator/runtime.ts | 27 +- ts-sdk/src/coordinator/serde.ts | 89 ++++- ts-sdk/src/coordinator/trigger-runner.ts | 239 ++++++++++++ ts-sdk/src/index.ts | 5 + ts-sdk/src/sdk/bundle.ts | 12 + ts-sdk/src/sdk/dag.ts | 100 ++++- ts-sdk/src/sdk/trigger-dag-run.ts | 215 +++++++++++ ts-sdk/tests/coordinator/integration.test.ts | 341 ++++++++++++++++++ ts-sdk/tests/public-api.test.ts | 2 +- ts-sdk/tests/sdk/trigger-dag-run.test.ts | 293 +++++++++++++++ 14 files changed, 1414 insertions(+), 24 deletions(-) create mode 100644 ts-sdk/src/coordinator/trigger-runner.ts create mode 100644 ts-sdk/src/sdk/trigger-dag-run.ts create mode 100644 ts-sdk/tests/sdk/trigger-dag-run.test.ts diff --git a/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md b/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md index 121020491900d..75c4e78b94c7a 100644 --- a/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md +++ b/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md @@ -30,7 +30,8 @@ Proposed. 2. **A branch selects a task, not a string.** A case is the task reference the SDK already handed back, and the value on the wire is that task's task_id, so the compiler checks the candidate exists and no label has to be kept in step with it. 3. **No default case.** `BranchPythonOperator` has none to serialize, and a one-sided `If` whose condition is false follows nothing at all. A decider returning a ref that is not one of the declared cases is a run-time error the SDK raises, a narrower check than `skip_all_except`, which only rejects a task_id missing from the whole Dag. -4. **Triggering a Dag run is an ordinary DSL task**, it is pure DSL purpose instead of a new runtime. +4. **Triggering a Dag run is an ordinary task that the SDK's own runtime runs.** It has no handler. + The runtime sends the `TriggerDagRun` request itself and follows `TriggerDagRunOperator` for every option it offers, including `wait_for_completion` and `deferrable`. 5. **Grouping keeps Python's semantics**: a scope offering the same task and nesting methods as the Dag, prefixing each task_id with the group id (`prefix_group_id`), and can be ordered against a task or another group, as `group1 >> group2` does in Python. @@ -89,7 +90,12 @@ That operator defaults to skipping every task in its downstream closure and igno a single-ref return cannot express that, and no Lang SDK offers it for now. An author needing several paths together puts them behind one task, or gates each with its own condition. This limitation is accepted rather than open. -- **No Lang SDK needs a deferral mechanism for now** to offer `deferrable` or `wait_for_completion`, because the trigger task runs in Python. +- **Each Lang SDK runtime implements the trigger task**, so a native Dag needs no Python worker. + `deferrable` needs no trigger written in that language: the runtime defers to the standard provider's `DagStateTrigger`, which runs in the Python triggerer, and the task keeps its queue, so it resumes in the same runtime with `next_method` set. +- **The runtime reads the config the trigger task needs from its environment**, since it cannot read `airflow.cfg`. + It reads `AIRFLOW__API__BASE_URL`, `AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE` and `AIRFLOW__TRIGGERER__QUEUES_ENABLED`, and an unset one takes Python's fallback. + The coordinator exports them from Airflow's config (#73842). +- **The trigger's arguments are not Jinja-rendered**, because no Python code runs the task. A value is sent as the author wrote it. - **Nothing extra reaches the Dag JSON**, which carries no branch-candidate field at all. A ref is a task_id by the time the decision is sent, so each SDK stores its cases and nothing else. - **A group edge needs one base type per SDK** that both a task and a group satisfy, since either can sit at the end of an edge. Python already has it: `TaskGroup(TaskGroupMixin, DAGNode)` (`task-sdk/src/airflow/sdk/definitions/taskgroup.py:96`) and every operator inherit `DependencyMixin` (`.../definitions/_internal/mixins.py:35`), where `set_upstream` and `set_downstream` live. 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 56fd2881e3df5..cb2846b0f52eb 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -408,6 +408,50 @@ Exactly one case is selected. Python's branch callable may return a list of task SDK offers that yet: put the paths that run together behind one task, or gate each with its own condition. +Triggering another Dag +~~~~~~~~~~~~~~~~~~~~~~ + +``dag.triggerDagRun`` starts another Dag's run, as a task of this one: + +.. code-block:: typescript + + const trigger = dag.triggerDagRun({ + taskId: "trigger_downstream", + dagId: "downstream_etl", + waitForCompletion: true, + }); + + trigger.after(loaded); + +There is nothing to call: the task takes no TypeScript arguments, so it hands back the reference +directly. A second options object carries the task's own spec, as ``dag.task`` takes one. + +The task runs in the TypeScript runtime, like the Dag's other tasks, and inherits the Dag's queue. It +behaves as ``TriggerDagRunOperator`` does. ``waitForCompletion`` polls the run every ``pokeInterval`` +seconds until it reaches one of ``allowedStates`` or ``failedStates``. With ``deferrable`` as well, the +task defers to ``DagStateTrigger`` instead of holding a worker slot, and resumes in this runtime when +the run finishes. ``DagStateTrigger`` runs in the Python triggerer, so the triggerer needs the standard +provider installed. + +The task pushes the ``trigger_run_id`` XCom, and the task's "Triggered DAG" link opens the run it +started. + +The task follows these config options, with Python's fallback when one is unset: + +- ``[api] base_url`` is the base of the "Triggered DAG" link. It falls back to ``/``. +- ``[operators] default_deferrable`` is the default of ``deferrable``. It falls back to ``false``. +- ``[triggerer] queues_enabled`` decides whether the deferred trigger gets the task's queue. It falls + back to ``false``, so the trigger gets none. + +Some of what ``TriggerDagRunOperator`` does is not offered: + +- ``logical_date`` and ``run_after``. The triggered run's logical date is the time the task triggers + it, as in Python when neither is set. +- Jinja. Values are sent as written, so ``{{ ds }}`` in ``conf`` reaches the triggered run as that + literal string. +- OpenLineage parent injection (``openlineage_inject_parent_info``). The runtime does not add the + parent task's OpenLineage details to ``conf``. + ``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/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index c8c908186f6a4..4debf7dbb4d94 100644 --- a/ts-sdk/src/coordinator/client.ts +++ b/ts-sdk/src/coordinator/client.ts @@ -31,6 +31,9 @@ import type { SetXCom, GetConnection, SkipDownstreamTasks, + TriggerDagRun, + GetDagRunState, + GetDag, ConnectionResult as WireConnectionResult, } from "./protocol.js"; @@ -83,6 +86,15 @@ export interface CoordinatorClient extends TaskClient { * output fails the task, while one that pushed null binds null. */ getXComEntry(opts: GetXComOpts): Promise; + /** + * Trigger a Dag run. `"already_exists"` is the one refusal a trigger task + * handles itself; any other error throws. + */ + triggerDagRun(msg: Omit): Promise<"triggered" | "already_exists">; + /** The state of a Dag run, as `GetDagRunState` reports it. */ + getDagRunState(dagId: string, runId: string): Promise; + /** Whether a Dag is paused. */ + isDagPaused(dagId: string): Promise; } export function createCoordinatorClient( @@ -213,6 +225,47 @@ export function createCoordinatorClient( await rpc("SkipDownstreamTasks", null, msg, () => undefined, "throw"); }, + // ---- Dag runs ---- + + async triggerDagRun(msg: Omit) { + const request: TriggerDagRun = { type: "TriggerDagRun", ...msg }; + logs?.debug("TriggerDagRun request"); + const frame = await comm.request(request); + const err = parseFrameError(frame); + // Sent as an ErrorResponse body, not a frame error; see + // `ActivitySubprocess._handle_trigger_dag_run`. + if (err?.code === "DAGRUN_ALREADY_EXISTS") return "already_exists"; + if (err) { + logs?.warning("TriggerDagRun failed", { error: err.code }); + throw new Error(`TriggerDagRun failed: ${err.code}`); + } + return "triggered"; + }, + + async getDagRunState(dagId: string, runId: string): Promise { + const msg: GetDagRunState = { type: "GetDagRunState", dag_id: dagId, run_id: runId }; + const state = await rpc( + "GetDagRunState", + "DagRunStateResult", + msg, + (body) => body!.state as string, + "throw", + ); + return state!; + }, + + async isDagPaused(dagId: string): Promise { + const msg: GetDag = { type: "GetDag", dag_id: dagId }; + const paused = await rpc( + "GetDag", + "DagResult", + msg, + (body) => body!.is_paused === true, + "throw", + ); + return paused!; + }, + // ---- Connections ---- async getConnection(connId: string): Promise { diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts index 1c95fee6f907d..9f2f97b949c3d 100644 --- a/ts-sdk/src/coordinator/protocol.ts +++ b/ts-sdk/src/coordinator/protocol.ts @@ -28,6 +28,7 @@ import type { DagFileParseRequest, ErrorResponse, DagFileParsingResult, + DeferTask, RetryTask, SucceedTask, TaskState, @@ -47,6 +48,7 @@ export type { export type { DagFileParsingResult, + DeferTask, RetryTask, SucceedTask, TaskState, @@ -65,6 +67,11 @@ export type { SetXCom, GetConnection, SkipDownstreamTasks, + TriggerDagRun, + GetDagRunState, + DagRunStateResult, + GetDag, + DagResult, } from "../generated/supervisor.js"; // -------- Frames from supervisor -------- @@ -82,6 +89,7 @@ export type MsgFromSupervisor = export type RuntimeTaskState = WithRequiredType; export type RuntimeRetryTask = WithRequiredType; +export type RuntimeDeferTask = WithRequiredType; export type RuntimeSucceedTask = WithRequiredType & { task_outlets: NonNullable; outlet_events: NonNullable; diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index 8029e8422b73e..b3b6686c12ef8 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -49,16 +49,23 @@ import { asMsgFromSupervisor, SUPERVISOR_API_VERSION, type RuntimeDagFileParsingResult, + type RuntimeDeferTask, type RuntimeRetryTask, type RuntimeSucceedTask, type RuntimeTaskState, type StartupDetails, } from "./protocol.js"; import { getArgNames } from "../sdk/arg-names.js"; -import { bundleDagTaskIds, listBundleNativeDags, type Bundle } from "../sdk/bundle.js"; +import { + bundleDagTaskIds, + getBundleTrigger, + listBundleNativeDags, + type Bundle, +} from "../sdk/bundle.js"; import type { Dag } from "../sdk/dag.js"; import { SERIALIZATION_VERSION } from "../generated/dag-schema-fields.js"; import { computeRelativeFileloc, serializeDag } from "./serde.js"; +import { runTriggerDagRun } from "./trigger-runner.js"; import { runInTaskScope, type TaskContext } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; @@ -195,6 +202,8 @@ export async function startCoordinator( await sendSupervisorResponse(firstFrame.id, response, comm, runtimeLogs); if (response.type === "SucceedTask") { runtimeLogs.info("Task succeeded", { task_id: body.ti.task_id }); + } else if (response.type === "DeferTask") { + runtimeLogs.info("Task deferred", { task_id: body.ti.task_id }); } } else { const errMsg = `First frame must be DagFileParseRequest or StartupDetails, got ${body.type}`; @@ -345,8 +354,22 @@ async function handleTask( logs: LogChannel, clientLogs: LogChannel, signal: AbortSignal, -): Promise { +): Promise { const ti = details.ti; + const trigger = getBundleTrigger(bundle, ti.dag_id, ti.task_id); + if (trigger) { + const ctx = buildContext(details, signal); + const client = createCoordinatorClient(comm, ctx, clientLogs); + const fail = (message: string) => { + logs.error("Task failed", { task_id: ctx.taskId, error: message }); + return buildFailureResponse(details, message); + }; + try { + return await runTriggerDagRun(details, trigger, client, logs, ctx.signal, fail); + } catch (err) { + return fail((err as Error).message ?? String(err)); + } + } const handler = bundle.getTaskHandler(ti.dag_id, ti.task_id); if (!handler) { diff --git a/ts-sdk/src/coordinator/serde.ts b/ts-sdk/src/coordinator/serde.ts index 3227b374a157b..cec2cc083142f 100644 --- a/ts-sdk/src/coordinator/serde.ts +++ b/ts-sdk/src/coordinator/serde.ts @@ -49,6 +49,7 @@ import { type Dag, type RecordedInputs, type TaskGroupRecord, + type TaskRecord, } from "../sdk/dag.js"; /** A serialized Dag: JSON, by the time it reaches the supervisor as msgpack. */ @@ -82,6 +83,39 @@ const TASK_MODULE = "airflow.sdk.coordinators.node"; */ const TASK_LANGUAGE = "typescript"; +/** + * The Dags this one reaches, as the UI's dependency graph and + * `airflow dags show-dependencies` read them. + * + * Python derives these from the live operator with `detect_dependencies`. Only + * a trigger task creates one here. + */ +function serializeDagDependencies(dag: Dag): SerializedValue { + const dependencies: SerializedValue[] = []; + for (const [taskId, record] of getDagTaskRecords(dag)) { + if (!record.trigger) continue; + dependencies.push({ + source: dag.dagId, + target: record.trigger.dagId, + label: taskId, + dependency_type: "trigger", + dependency_id: taskId, + }); + } + return dependencies; +} + +/** + * What a trigger task adds to a TypeScript task, so the UI draws it as + * `TriggerDagRunOperator` and shows its "Triggered DAG" link. The runtime + * pushes the link's XCom under the key named here. + */ +const TRIGGER_DAG_RUN_FIELDS: Readonly> = { + _operator_name: "TriggerDagRunOperator", + ui_color: "#ffefeb", + _operator_extra_links: { "Triggered DAG": "_link_TriggerDagRunLink" }, +}; + // Python resolves these from [core]/[scheduler] config when the Dag leaves them // unset, and its serializer always writes the resolved value — there is no // schema default to omit against. A bundle cannot read airflow.cfg, so the @@ -151,13 +185,13 @@ export function serializeDag( serializeTask( dag.dagId, taskId, - withDagQueue(record.spec, dag.spec.queue), + record, graph.downstreamTaskIds.get(taskId), inputs.get(taskId), - record.canSkipDownstream === true, + dag.spec.queue, ), ), - dag_dependencies: [], + dag_dependencies: serializeDagDependencies(dag), task_group: serializeTaskGroups(dag, graph), edge_info: {}, params: [], @@ -189,10 +223,10 @@ function withDagQueue(spec: object, dagQueue: string | undefined): object { function serializeTask( dagId: string, taskId: string, - spec: object, + record: TaskRecord, downstream: ReadonlySet | undefined, inputs: RecordedInputs | undefined, - canSkipDownstream: boolean, + dagQueue: string | undefined, ): SerializedValue { const data: Record = { task_id: taskId, @@ -208,12 +242,21 @@ function serializeTask( // `get_arg_bindings` reads nothing without it. is_stub: true, }; + if (record.trigger) { + // `conf` reaches Airflow at run time rather than in the Dag JSON, so a + // value it cannot store is caught here, as this Dag's import error. + if (record.trigger.conf !== undefined) { + plainJsonArg(record.trigger.conf, `conf of task "${taskId}" of Dag "${dagId}"`); + } + Object.assign(data, structuredClone(TRIGGER_DAG_RUN_FIELDS)); + } 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; + if (record.canSkipDownstream === true) data["_can_skip_downstream"] = true; + const spec = withDagQueue(record.spec, dagQueue); applySchemaFields(data, spec, TASK_FIELD_RULES, `task "${taskId}" of Dag "${dagId}"`); if (downstream?.size) { data["downstream_task_ids"] = [...downstream].sort(); @@ -221,6 +264,40 @@ function serializeTask( return { __type: "operator", __var: data }; } +/** + * A trigger task's `conf`, checked to be plain JSON. + * + * A value JSON cannot carry, such as a `Date`, a `Map` or a function, is + * rejected, because msgpack would otherwise hand Airflow something it cannot + * store as a Dag run's conf. + */ +function plainJsonArg(value: unknown, label: string): SerializedValue { + if (value === null || value === undefined) return null; + if (typeof value === "string" || typeof value === "boolean") return value; + if (typeof value === "number") { + if (!Number.isFinite(value)) throw new Error(`${label} is not a finite number`); + return value; + } + if (Array.isArray(value)) + return value.map((item, index) => plainJsonArg(item, `${label}[${index}]`)); + if (isPlainObject(value)) { + const copy: Record = {}; + for (const [key, item] of Object.entries(value)) + copy[key] = plainJsonArg(item, `${label}.${key}`); + return copy; + } + throw new Error( + `${label} is ${describeType(value)}, which a Dag run's conf cannot carry; ` + + "pass a string, a number, a boolean, null, an array or a plain object", + ); +} + +function isPlainObject(value: unknown): value is Record { + if (typeof value !== "object" || value === null) return false; + const prototype = Object.getPrototypeOf(value); + return prototype === Object.prototype || prototype === null; +} + /** * The task's arguments as the binding spec the API server hands back at run * time, one entry per argument in the order the call named them. diff --git a/ts-sdk/src/coordinator/trigger-runner.ts b/ts-sdk/src/coordinator/trigger-runner.ts new file mode 100644 index 0000000000000..156812e67e129 --- /dev/null +++ b/ts-sdk/src/coordinator/trigger-runner.ts @@ -0,0 +1,239 @@ +/*! + * 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. + */ + +// Runs a `dag.triggerDagRun(...)` task. +// +// Mirrors `TriggerDagRunOperator.execute` and `execute_complete` in +// `providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py` +// and `_handle_trigger_dag_run` in +// `task-sdk/src/airflow/sdk/execution_time/task_runner.py`. The config Python +// reads here, `[api] base_url` and `[triggerer] queues_enabled`, comes from the +// environment instead, with Python's fallbacks when it is unset. + +import { setTimeout as sleep } from "node:timers/promises"; +import type { CoordinatorClient } from "./client.js"; +import type { LogChannel } from "./log-channel.js"; +import type { + RuntimeDeferTask, + RuntimeRetryTask, + RuntimeSucceedTask, + RuntimeTaskState, + StartupDetails, +} from "./protocol.js"; +import { getBooleanEnv, type TriggerDagRunTask } from "../sdk/trigger-dag-run.js"; + +export type TriggerOutcome = + RuntimeSucceedTask | RuntimeRetryTask | RuntimeTaskState | RuntimeDeferTask; + +/** A failure the task's retries apply to, as `_handle_current_task_failed` decides. */ +export type FailTask = (message: string) => RuntimeRetryTask | RuntimeTaskState; + +const DAG_STATE_TRIGGER = "airflow.providers.standard.triggers.external_task.DagStateTrigger"; +/** `TriggerDagRunLink().xcom_key`: the XCom the "Triggered DAG" extra link reads. */ +const LINK_XCOM_KEY = "_link_TriggerDagRunLink"; +const RUN_ID_XCOM_KEY = "trigger_run_id"; +/** `TRIGGER_FAIL_REPR`: the `next_method` a failed or timed-out trigger resumes with. */ +const TRIGGER_FAIL = "__fail__"; +const EXECUTE_COMPLETE = "execute_complete"; + +export async function runTriggerDagRun( + details: StartupDetails, + trigger: TriggerDagRunTask, + client: CoordinatorClient, + logs: LogChannel, + signal: AbortSignal, + fail: FailTask, +): Promise { + const nextMethod = details.ti_context.next_method; + if (nextMethod) return resume(nextMethod, details.ti_context.next_kwargs, trigger, logs, fail); + + // No logicalDate option, so Python's default applies: the trigger time. + const logicalDate = new Date(); + const runId = trigger.runId ?? `manual__${pythonIsoformat(logicalDate)}`; + + if (trigger.failWhenDagIsPaused && (await client.isDagPaused(trigger.dagId))) { + return fail(`Dag ${trigger.dagId} is paused`); + } + + await client.setXCom({ key: LINK_XCOM_KEY, value: dagRunUrl(trigger.dagId, runId) }); + + logs.info("Triggering Dag Run.", { trigger_dag_id: trigger.dagId }); + const triggered = await client.triggerDagRun({ + dag_id: trigger.dagId, + run_id: runId, + logical_date: logicalDate.toISOString(), + run_after: null, + conf: (trigger.conf as Record | undefined) ?? null, + reset_dag_run: trigger.resetDagRun, + note: trigger.note ?? null, + }); + if (triggered === "already_exists") { + if (trigger.skipWhenAlreadyExists) { + logs.info( + "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", + { dag_id: trigger.dagId }, + ); + return { type: "TaskState", state: "skipped", end_date: new Date().toISOString() }; + } + logs.error("Dag Run already exists, marking task as failed.", { dag_id: trigger.dagId }); + return { type: "TaskState", state: "failed", end_date: new Date().toISOString() }; + } + logs.info("Dag Run triggered successfully.", { trigger_dag_id: trigger.dagId }); + await client.setXCom({ key: RUN_ID_XCOM_KEY, value: runId }); + + if (!trigger.waitForCompletion) { + if (trigger.deferrable) { + logs.info( + "Ignoring deferrable=True because wait_for_completion=False. " + + "Task will complete immediately without waiting for the triggered DAG run.", + { trigger_dag_id: trigger.dagId }, + ); + } + return succeeded(); + } + + if (trigger.deferrable) { + logs.info("Pausing task as DEFERRED.", { trigger_dag_id: trigger.dagId, run_id: runId }); + return { + type: "DeferTask", + state: "deferred", + classpath: DAG_STATE_TRIGGER, + // `DagStateTrigger.serialize()`, key for key. + trigger_kwargs: { + dag_id: trigger.dagId, + states: [...trigger.allowedStates, ...trigger.failedStates], + poll_interval: trigger.pokeInterval, + run_ids: [runId], + execution_dates: null, + }, + trigger_timeout: null, + // `_defer_task` hands the trigger the task's queue only when triggerer queues are enabled. + queue: getBooleanEnv("AIRFLOW__TRIGGERER__QUEUES_ENABLED", false) + ? (details.ti.queue ?? null) + : null, + next_method: EXECUTE_COMPLETE, + next_kwargs: {}, + }; + } + + while (true) { + logs.info("Waiting for dag run to complete execution in allowed state.", { + dag_id: trigger.dagId, + run_id: runId, + allowed_state: trigger.allowedStates, + }); + await sleep(trigger.pokeInterval * 1000, undefined, { signal }); + const state = await client.getDagRunState(trigger.dagId, runId); + if (includes(trigger.failedStates, state)) { + logs.error("DagRun finished with failed state.", { dag_id: trigger.dagId, state }); + return fail(`${trigger.dagId} failed with failed state ${state}`); + } + if (includes(trigger.allowedStates, state)) { + logs.info("DagRun finished with allowed state.", { dag_id: trigger.dagId, state }); + return succeeded(); + } + logs.debug("DagRun not yet in allowed or failed state.", { dag_id: trigger.dagId, state }); + } +} + +/** `BaseOperator.resume_execution` for this task: `__fail__` or `execute_complete`. */ +function resume( + nextMethod: string, + nextKwargs: unknown, + trigger: TriggerDagRunTask, + logs: LogChannel, + fail: FailTask, +): TriggerOutcome { + const kwargs = isRecord(nextKwargs) ? nextKwargs : {}; + if (nextMethod === TRIGGER_FAIL) { + const traceback = kwargs["traceback"]; + if (Array.isArray(traceback)) logs.error(`Trigger failed:\n${traceback.join("\n")}`); + return fail(String(kwargs["error"] ?? "Unknown")); + } + if (nextMethod !== EXECUTE_COMPLETE) { + return fail(`Task cannot resume with next_method "${nextMethod}"`); + } + const eventData = decodeEvent(kwargs["event"]); + const runIds = Array.isArray(eventData?.["run_ids"]) ? (eventData["run_ids"] as unknown[]) : []; + const failedRunIds: string[] = []; + for (const runId of runIds) { + const state = eventData?.[String(runId)]; + if (includes(trigger.failedStates, state)) { + failedRunIds.push(String(runId)); + continue; + } + if (includes(trigger.allowedStates, state)) { + logs.info("Triggered Dag run finished with allowed state.", { + dag_id: trigger.dagId, + state, + run_id: runId, + }); + } + } + if (failedRunIds.length > 0) { + return fail( + `${trigger.dagId} failed with failed states ${JSON.stringify(trigger.failedStates)} ` + + `for run_ids ${JSON.stringify(failedRunIds)}`, + ); + } + return succeeded(); +} + +/** + * The payload `DagStateTrigger` fired with: `(classpath, data)`. The triggerer + * stores it with serde, which encodes a tuple as `{__classname__, __data__}`. + */ +function decodeEvent(event: unknown): Record | undefined { + const pair = + isRecord(event) && event["__classname__"] === "builtins.tuple" ? event["__data__"] : event; + if (!Array.isArray(pair) || pair.length !== 2 || !isRecord(pair[1])) return undefined; + return pair[1]; +} + +function succeeded(): RuntimeSucceedTask { + return { + type: "SucceedTask", + end_date: new Date().toISOString(), + task_outlets: [], + outlet_events: [], + }; +} + +/** `datetime.isoformat()` of a UTC instant, which is how Python spells a run ID's date. */ +export function pythonIsoformat(date: Date): string { + const iso = date.toISOString(); + const seconds = iso.slice(0, 19); + const millis = date.getUTCMilliseconds(); + const fraction = millis === 0 ? "" : `.${String(millis * 1000).padStart(6, "0")}`; + return `${seconds}${fraction}+00:00`; +} + +/** `build_airflow_dagrun_url`, on `[api] base_url` from the environment, or "/" when unset. */ +export function dagRunUrl(dagId: string, runId: string): string { + const base = process.env["AIRFLOW__API__BASE_URL"] || "/"; + return `${base.replace(/\/+$/, "")}/dags/${dagId}/runs/${runId}`; +} + +function includes(states: readonly string[], state: unknown): boolean { + return typeof state === "string" && states.includes(state); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index e5beb163c48ae..52344fefaf777 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -26,6 +26,11 @@ export { ConnectionNotFoundError, VariableNotFoundError } from "./sdk/client.js" 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 { + DagRunState, + TriggerDagRunOptions, + TriggerDagRunSpec, +} from "./sdk/trigger-dag-run.js"; export type { Branch, Condition, diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index 34b6351032ce9..57cb0b2fe92c3 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -23,6 +23,7 @@ import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; import { Dag, finalizeDag, getDagTaskRecords, isDag } from "./dag.js"; import { getTaskHandlerFunction, isTaskHandler, TaskHandler } from "./task-handler.js"; import type { TaskFunction } from "./task.js"; +import type { TriggerDagRunTask } from "./trigger-dag-run.js"; // Assigned inside Bundle's static block, as Dag does for its tasks. let dagsOf: (bundle: Bundle) => ReadonlyMap; @@ -211,6 +212,17 @@ export function validateOwnBundle(value: unknown, accessor: string): asserts val ); } +/** Internal: the `triggerDagRun(...)` task at `taskId` of a Dag this bundle + * declared in TypeScript, which the runtime runs without a handler. */ +export function getBundleTrigger( + bundle: Bundle, + dagId: string, + taskId: string, +): TriggerDagRunTask | undefined { + const dag = dagsOf(bundle).get(dagId); + return dag === undefined ? undefined : getDagTaskRecords(dag).get(taskId)?.trigger; +} + /** Internal: finalize every Dag this bundle declared in TypeScript, so no task * can be added or wired afterwards. */ export function finalizeBundleDags(bundle: Bundle): void { diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index ab732ce7ebaf4..6e18bbe306435 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -31,6 +31,13 @@ import { import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; import type { JsonValue } from "./client-types.js"; import { getCurrentModuleSource } from "./module-source.js"; +import { + isTriggerDagRunTask, + triggerDagRun, + type TriggerDagRunOptions, + type TriggerDagRunSpec, + type TriggerDagRunTask, +} from "./trigger-dag-run.js"; import { getClient, type TaskFunction } from "./task.js"; function isPlainRecord(value: unknown): value is Record { @@ -226,6 +233,8 @@ export interface TaskGroupRef extends Node { handler: (args: TArgs) => TReturn | Promise, options?: TaskOptions, ): TaskFactory; + /** Trigger another Dag's run from this group; its id carries the group prefix. */ + triggerDagRun(spec: TriggerDagRunSpec, taskSpec?: TaskSpec): TaskRef; /** Nest a group inside this one. */ taskGroup(groupId: string, options?: TaskGroupOptions): TaskGroupRef; before(...downstream: readonly Node[]): TaskGroupRef; @@ -436,10 +445,15 @@ export interface Branch { case(taskRef: TaskRef): Branch; } -/** Per-task record a Dag retains: the reference, the handler, and its spec. */ +/** Per-task record a Dag retains: the reference, what runs it, and its spec. + * + * Exactly one of `fn` and `trigger` is set: a task has either a handler, or + * the options of a `triggerDagRun(...)` the runtime carries out itself. */ export interface TaskRecord { readonly task: TaskRef; - readonly fn: TaskFunction; + readonly fn?: TaskFunction; + /** What a `triggerDagRun(...)` task triggers, when it has no handler. */ + readonly trigger?: TriggerDagRunTask; readonly spec: TaskSpec; /** * Whether this task decides which of its downstream tasks to skip. @@ -695,6 +709,13 @@ export class Dag { } this.#validateOwnNode(condition, `the condition given to ${verb}`); const taskId = condition.taskId; + // A trigger task returns nothing, so there is no decision to read. + if (this.#tasks.get(taskId)!.fn === undefined) { + throw new Error( + `Task "${taskId}" of Dag "${this.dagId}" is a triggerDagRun task, so it cannot decide a ` + + `branch; give ${verb}(...) a task with a TypeScript handler`, + ); + } if (this.#conditions.has(taskId) || this.#branches.has(taskId)) { throw new Error( `Task "${taskId}" of Dag "${this.dagId}" already decides a branch; ` + @@ -717,7 +738,7 @@ export class Dag { decide: (returned: unknown) => Promise<{ skip: string[]; result: unknown }>, ): void { const record = this.#tasks.get(taskId)!; - const inner = record.fn; + const inner = record.fn!; const wrapped: TaskFunction = async (args) => { const { skip, result } = await decide(await inner(args as never)); if (skip.length > 0) { @@ -799,6 +820,56 @@ export class Dag { return branch; } + /** + * Trigger another Dag's run, as a task of this Dag. + * + * ```ts + * const trigger = dag.triggerDagRun({ + * taskId: "trigger_downstream", + * dagId: "downstream_etl", + * waitForCompletion: true, + * }); + * trigger.after(loaded); + * ``` + * + * The task wraps no TypeScript function. This runtime runs it the way + * Python runs `TriggerDagRunOperator`: it triggers the run, and with + * `waitForCompletion` polls the run's state, or with `deferrable` defers to + * `DagStateTrigger`, which the Python triggerer runs, and resumes when the + * run finishes. Like every task of this Dag it inherits the Dag's queue. + * + * There is nothing to call: a task with no TypeScript arguments has nothing + * for a factory call to supply, so this returns the reference directly. + * + * Values are sent as written. There is no Jinja rendering, so `{{ ds }}` in + * `conf` reaches the triggered run as that literal string. + */ + triggerDagRun(spec: TriggerDagRunSpec, taskSpec: TaskSpec = {}): TaskRef { + return this.#addTriggerDagRun(undefined, spec, taskSpec); + } + + #addTriggerDagRun( + groupId: string | undefined, + spec: TriggerDagRunSpec, + taskSpec: TaskSpec, + ): TaskRef { + const { taskId, ...options } = isPlainRecord(spec) + ? (spec as TriggerDagRunSpec & { taskId?: unknown }) + : ({} as TriggerDagRunSpec & { taskId?: unknown }); + if (typeof taskId !== "string" || taskId.length === 0) { + throw new Error( + `A triggerDagRun task of Dag "${this.dagId}" has no taskId; name it with ` + + 'dag.triggerDagRun({ taskId: "trigger_downstream", dagId: "..." })', + ); + } + return this.#addTask( + groupId, + taskId, + triggerDagRun(options as TriggerDagRunOptions), + taskSpec, + )(); + } + /** * Declare a task group of this Dag. * @@ -818,13 +889,14 @@ export class Dag { #addTask( groupId: string | undefined, taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), - handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + handlerOrOptions?: + ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask | TaskOptions, maybeOptions?: TaskOptions, ): TaskFactory { const idGiven = typeof taskIdOrHandler === "string"; - const handler = (idGiven ? handlerOrOptions : taskIdOrHandler) as ( - args: TArgs, - ) => TReturn | Promise; + const body = idGiven ? handlerOrOptions : taskIdOrHandler; + const trigger = isTriggerDagRunTask(body) ? body : undefined; + const handler = body as (args: TArgs) => TReturn | Promise; const given = idGiven ? maybeOptions : (handlerOrOptions as TaskOptions | undefined); // Defaulted only when absent: an explicit `null` is a bad spec, not an // omitted one, and #taskSpecOf is what reports it. @@ -835,12 +907,12 @@ export class Dag { // wins and the spec's is dropped without a word. if (idGiven && specTaskId !== undefined) { throw new Error( - `Task "${taskIdOrHandler}" of Dag "${this.dagId}" also carries taskId "${specTaskId}" in ` + - "its spec; give the id once, either positionally or in the spec", + `Task "${taskIdOrHandler as string}" of Dag "${this.dagId}" also carries taskId ` + + `"${specTaskId}" in its spec; give the id once, either positionally or in the spec`, ); } const defaulted = idGiven ? undefined : (specTaskId ?? readFunctionName(handler)); - const declared = idGiven ? taskIdOrHandler : defaulted; + const declared = idGiven ? (taskIdOrHandler as string) : defaulted; if (declared === undefined) { throw new Error( `A task of Dag "${this.dagId}" has no id: its handler has no name to take one from. ` + @@ -867,12 +939,12 @@ export class Dag { ); } const taskId = this.#childId(groupId, declared); - if (typeof handler !== "function") { + if (trigger === undefined && typeof handler !== "function") { throw new Error(`handler for Dag "${this.dagId}" task "${taskId}" must be a function`); } // TypeScript already says so, but a plain-JavaScript author lands here // with the argument list Python would take. - if (handler.length > 1) { + if (trigger === undefined && handler.length > 1) { throw new Error( `Handler for Dag "${this.dagId}" task "${taskId}" declares ${handler.length} parameters; ` + "a handler takes one object of named arguments — async ({ rows, region }) => ...", @@ -894,7 +966,7 @@ export class Dag { task, // The runtime dispatches every handler through one instantiation, as it // does a registered TaskHandler. - fn: handler as unknown as TaskFunction, + ...(trigger ? { trigger } : { fn: handler as unknown as TaskFunction }), spec: freezeSpec(spec, () => `The spec for Dag "${this.dagId}" task "${taskId}"`), }); return ((...inputs: unknown[]) => { @@ -1031,6 +1103,8 @@ export class Dag { handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, maybeOptions?: TaskOptions, ) => this.#addTask(groupId, taskIdOrHandler, handlerOrOptions, maybeOptions), + triggerDagRun: (spec: TriggerDagRunSpec, taskSpec: TaskSpec = {}) => + this.#addTriggerDagRun(groupId, spec, taskSpec), taskGroup: (childId: string, options?: TaskGroupOptions) => this.#addGroup(groupId, childId, options), before: (...downstream) => { diff --git a/ts-sdk/src/sdk/trigger-dag-run.ts b/ts-sdk/src/sdk/trigger-dag-run.ts new file mode 100644 index 0000000000000..6f8bb300b48cb --- /dev/null +++ b/ts-sdk/src/sdk/trigger-dag-run.ts @@ -0,0 +1,215 @@ +/*! + * 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. + */ + +// `dag.triggerDagRun(...)`: a task that starts another Dag's run. +// +// It carries no handler. The runtime runs it the way Python's task runner runs +// `TriggerDagRunOperator`: it sends the `TriggerDagRun` request itself, then +// waits, defers or returns. See `runTriggerDagRun` in +// `../coordinator/trigger-runner.ts`. +// +// Hand-written rather than generated, because triggering a Dag run is control +// flow (`airflow-core/adr/lang-sdk/0008-control-flow-constructs.md`, decision 4). + +import { brand, hasBrand } from "./brand.js"; +import type { JsonValue } from "./client-types.js"; + +/** A Dag run state, as `allowedStates` and `failedStates` name one. */ +export type DagRunState = "queued" | "running" | "success" | "failed"; + +const DAG_RUN_STATES: ReadonlySet = new Set([ + "queued", + "running", + "success", + "failed", +]); + +/** The `TriggerDagRunOperator` options, named as TypeScript spells them. */ +export interface TriggerDagRunOptions { + /** Identifier of the Dag to trigger. */ + readonly dagId: string; + /** Run ID for the triggered run; generated from the trigger time when unset. */ + readonly runId?: string; + /** Configuration the triggered run is started with. */ + readonly conf?: Readonly>; + /** Clear an existing run with the same ID instead of failing. */ + readonly resetDagRun?: boolean; + /** Hold this task open until the triggered run finishes. */ + readonly waitForCompletion?: boolean; + /** Seconds between checks while waiting. Defaults to 60. */ + readonly pokeInterval?: number; + /** Run states that count as success when waiting. Defaults to `["success"]`. */ + readonly allowedStates?: readonly DagRunState[]; + /** Run states that count as failure when waiting. Defaults to `["failed"]`. */ + readonly failedStates?: readonly DagRunState[]; + /** Skip rather than fail when the run already exists. */ + readonly skipWhenAlreadyExists?: boolean; + /** Fail rather than trigger when the target Dag is paused. */ + readonly failWhenDagIsPaused?: boolean; + /** Note recorded against the triggered run. */ + readonly note?: string; + /** + * While waiting, free the worker slot and defer to `DagStateTrigger`, which + * the Python triggerer runs. Defaults to `[operators] default_deferrable`, or + * to false when that is unset. + */ + readonly deferrable?: boolean; +} + +/** + * What `dag.triggerDagRun(...)` takes: the operator's options, plus the task ID + * the trigger takes in the Dag that declares it. + */ +export interface TriggerDagRunSpec extends TriggerDagRunOptions { + /** Airflow task ID of the trigger task itself. */ + readonly taskId: string; +} + +/** Internal: a trigger task's options with `TriggerDagRunOperator`'s defaults applied. */ +export interface TriggerDagRunTask { + readonly dagId: string; + readonly runId: string | undefined; + readonly conf: Readonly> | undefined; + readonly resetDagRun: boolean; + readonly waitForCompletion: boolean; + readonly pokeInterval: number; + readonly allowedStates: readonly DagRunState[]; + readonly failedStates: readonly DagRunState[]; + readonly skipWhenAlreadyExists: boolean; + readonly failWhenDagIsPaused: boolean; + readonly note: string | undefined; + readonly deferrable: boolean; +} + +/** Internal: whether `value` is a trigger task built by any copy of this package. */ +export function isTriggerDagRunTask(value: unknown): value is TriggerDagRunTask { + return hasBrand(value, "TriggerDagRunTask"); +} + +const OPTION_NAMES: ReadonlySet = new Set([ + "dagId", + "runId", + "conf", + "resetDagRun", + "waitForCompletion", + "pokeInterval", + "allowedStates", + "failedStates", + "skipWhenAlreadyExists", + "failWhenDagIsPaused", + "note", + "deferrable", +]); + +function checkStates(name: string, states: unknown): DagRunState[] | undefined { + if (states === undefined) return undefined; + if (!Array.isArray(states)) { + throw new Error(`triggerDagRun(...) option "${name}" must be an array of Dag run states`); + } + for (const state of states) { + if (typeof state !== "string" || !DAG_RUN_STATES.has(state)) { + throw new Error( + `triggerDagRun(...) option "${name}" holds ${JSON.stringify(state)}, which is not a Dag ` + + `run state; use one of ${[...DAG_RUN_STATES].join(", ")}`, + ); + } + } + return [...(states as DagRunState[])]; +} + +function checkType(name: string, value: unknown, type: "string" | "boolean"): void { + if (value !== undefined && typeof value !== type) { + throw new Error(`triggerDagRun(...) option "${name}" must be a ${type}`); + } +} + +/** + * Internal: a boolean Airflow option from the environment, such as + * `AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE`, or `fallback` when it is unset. The + * runtime cannot read airflow.cfg, so the coordinator passes these on. It + * accepts the values `conf.getboolean` accepts. + */ +export function getBooleanEnv(name: string, fallback: boolean): boolean { + const raw = process.env[name]; + if (raw === undefined) return fallback; + const value = raw.trim().toLowerCase(); + if (value === "t" || value === "true" || value === "1") return true; + if (value === "f" || value === "false" || value === "0") return false; + throw new Error(`${name} is ${JSON.stringify(raw)}, which is not a boolean; use true or false`); +} + +/** + * Internal: `Dag.triggerDagRun` is the authoring surface, and it strips the + * `taskId` before calling this. Defaults are those of `TriggerDagRunOperator`. + */ +export function triggerDagRun(spec: TriggerDagRunOptions): TriggerDagRunTask { + const value: unknown = spec; + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("triggerDagRun(...) takes an options object"); + } + const options = value as Record; + for (const name of Object.keys(options)) { + if (!OPTION_NAMES.has(name)) throw new Error(`Unknown option "${name}" for triggerDagRun(...)`); + } + if (typeof options["dagId"] !== "string" || options["dagId"].length === 0) { + throw new Error("triggerDagRun(...) needs the dagId of the Dag to trigger"); + } + const conf = options["conf"]; + if (conf !== undefined && (typeof conf !== "object" || conf === null || Array.isArray(conf))) { + throw new Error('triggerDagRun(...) option "conf" must be an object'); + } + checkType("runId", options["runId"], "string"); + checkType("note", options["note"], "string"); + for (const flag of [ + "resetDagRun", + "waitForCompletion", + "skipWhenAlreadyExists", + "failWhenDagIsPaused", + "deferrable", + ]) { + checkType(flag, options[flag], "boolean"); + } + const pokeInterval = options["pokeInterval"] ?? 60; + if (typeof pokeInterval !== "number" || !Number.isFinite(pokeInterval) || pokeInterval < 0) { + throw new Error('triggerDagRun(...) option "pokeInterval" must be a non-negative number'); + } + const allowedStates = checkStates("allowedStates", options["allowedStates"]); + const failedStates = checkStates("failedStates", options["failedStates"]); + const task: TriggerDagRunTask = { + dagId: options["dagId"], + runId: (options["runId"] as string | undefined) || undefined, + conf: conf as TriggerDagRunTask["conf"], + resetDagRun: options["resetDagRun"] === true, + waitForCompletion: options["waitForCompletion"] === true, + pokeInterval, + // Python falls back on an empty `allowed_states` too, but keeps an empty + // `failed_states` as given. + allowedStates: allowedStates?.length ? allowedStates : ["success"], + failedStates: failedStates ?? ["failed"], + skipWhenAlreadyExists: options["skipWhenAlreadyExists"] === true, + failWhenDagIsPaused: options["failWhenDagIsPaused"] === true, + note: (options["note"] as string | undefined) || undefined, + // Read where the Dag is built, as Python reads it for the operator's default. + deferrable: + (options["deferrable"] as boolean | undefined) ?? + getBooleanEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", false), + }; + brand(task, "TriggerDagRunTask"); + return Object.freeze(task); +} diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index e20a780e8f0f9..1200a178fe7b2 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -938,4 +938,345 @@ describe("coordinator runtime integration", () => { const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); expect(setXComReqs).toHaveLength(0); }); + + describe("triggerDagRun", () => { + const DAG_STATE_TRIGGER = "airflow.providers.standard.triggers.external_task.DagStateTrigger"; + const OK = { body: { type: "OKResponse", ok: true } }; + + /** Answers each request type with the next queued reply, acking the rest. */ + function replies(queued: Record): Responder { + return (msgType) => queued[msgType]?.shift() ?? { body: null }; + } + + function requestsOf(result: MockResult, type: string) { + return result.runtimeRequests.filter((r) => r.type === type).map((r) => r.body); + } + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("triggers the run, pushes the link and run id, and succeeds", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + conf: { source: "ts" }, + note: "from ts", + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(result.runtimeRequests.map((r) => r.type)).toEqual([ + "SetXCom", + "TriggerDagRun", + "SetXCom", + ]); + const [sent] = requestsOf(result, "TriggerDagRun"); + const runId = sent!["run_id"] as string; + // `DagRun.generate_run_id` for a set logical date: manual__. + expect(runId).toMatch(/^manual__\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d(\.\d{6})?\+00:00$/); + expect(new Date(sent!["logical_date"] as string).getTime()).toBe( + new Date(runId.slice("manual__".length)).getTime(), + ); + expect(sent).toMatchObject({ + dag_id: "downstream", + conf: { source: "ts" }, + reset_dag_run: false, + note: "from ts", + run_after: null, + }); + expect(requestsOf(result, "SetXCom")).toEqual([ + expect.objectContaining({ + key: "_link_TriggerDagRunLink", + value: `/dags/downstream/runs/${runId}`, + task_id: "trigger", + }), + expect.objectContaining({ key: "trigger_run_id", value: runId, task_id: "trigger" }), + ]); + }); + + it("builds the link on AIRFLOW__API__BASE_URL when the environment sets it", async () => { + vi.stubEnv("AIRFLOW__API__BASE_URL", "https://airflow.example.com/sub/"); + testDag.triggerDagRun({ taskId: "trigger", dagId: "downstream", runId: "given" }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(requestsOf(result, "TriggerDagRun")[0]).toMatchObject({ run_id: "given" }); + expect(requestsOf(result, "SetXCom")[0]).toMatchObject({ + value: "https://airflow.example.com/sub/dags/downstream/runs/given", + }); + }); + + it.each([ + [true, "skipped"], + [false, "failed"], + ])( + "with skipWhenAlreadyExists=%s, ends %s when the run exists, retries or not", + async (skip, state) => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + skipWhenAlreadyExists: skip, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [{ body: { type: "ErrorResponse", error: "DAGRUN_ALREADY_EXISTS" } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state }); + expect(requestsOf(result, "SetXCom").map((b) => b["key"])).toEqual([ + "_link_TriggerDagRunLink", + ]); + }, + ); + + it("fails like any task error when the trigger request itself fails", async () => { + testDag.triggerDagRun({ taskId: "trigger", dagId: "missing" }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [ + { + body: { + type: "ErrorResponse", + error: "API_SERVER_ERROR", + detail: { status_code: 404 }, + }, + }, + ], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "TriggerDagRun failed: API_SERVER_ERROR", + }); + }); + + it("refuses to trigger a paused Dag when failWhenDagIsPaused is set", async () => { + testDag.triggerDagRun({ taskId: "trigger", dagId: "downstream", failWhenDagIsPaused: true }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ + GetDag: [{ body: { type: "DagResult", dag_id: "downstream", is_paused: true } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "failed" }); + expect(requestsOf(result, "GetDag")).toEqual([{ type: "GetDag", dag_id: "downstream" }]); + expect(requestsOf(result, "TriggerDagRun")).toEqual([]); + }); + + it("polls the run's state until it reaches an allowed state", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + pokeInterval: 0, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ + TriggerDagRun: [OK], + GetDagRunState: [ + { body: { type: "DagRunStateResult", state: "queued" } }, + { body: { type: "DagRunStateResult", state: "running" } }, + { body: { type: "DagRunStateResult", state: "success" } }, + ], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + const runId = requestsOf(result, "TriggerDagRun")[0]!["run_id"]; + expect(requestsOf(result, "GetDagRunState")).toEqual( + Array(3).fill({ type: "GetDagRunState", dag_id: "downstream", run_id: runId }), + ); + }); + + it("fails, with retries honoured, when the polled run reaches a failed state", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + pokeInterval: 0, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [OK], + GetDagRunState: [{ body: { type: "DagRunStateResult", state: "failed" } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "downstream failed with failed state failed", + }); + }); + + it("defers to DagStateTrigger with the kwargs Python's serializer writes", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + pokeInterval: 5, + allowedStates: ["success"], + failedStates: ["failed", "queued"], + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + const runId = requestsOf(result, "TriggerDagRun")[0]!["run_id"]; + expect(result.firstResponse!.body).toEqual({ + type: "DeferTask", + state: "deferred", + classpath: DAG_STATE_TRIGGER, + trigger_kwargs: { + dag_id: "downstream", + states: ["success", "failed", "queued"], + poll_interval: 5, + run_ids: [runId], + execution_dates: null, + }, + trigger_timeout: null, + queue: null, + next_method: "execute_complete", + next_kwargs: {}, + }); + expect(requestsOf(result, "GetDagRunState")).toEqual([]); + }); + + it.each([ + ["False", null], + ["True", "default"], + ])( + "with AIRFLOW__TRIGGERER__QUEUES_ENABLED=%s, hands the trigger the queue %s", + async (enabled, queue) => { + vi.stubEnv("AIRFLOW__TRIGGERER__QUEUES_ENABLED", enabled); + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "DeferTask", queue }); + }, + ); + + it("ignores deferrable without waitForCompletion, as Python does", async () => { + testDag.triggerDagRun({ taskId: "trigger", dagId: "downstream", deferrable: true }); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + }); + + /** `next_kwargs` as `handle_event_submit` stores DagStateTrigger's event. */ + function resumedWith(state: string, runId = "manual__2026-09-29T12:00:00+00:00") { + return { + next_method: "execute_complete", + next_kwargs: { + event: { + __classname__: "builtins.tuple", + __version__: 1, + __data__: [ + DAG_STATE_TRIGGER, + { + dag_id: "downstream", + states: ["success", "failed"], + poll_interval: 60, + run_ids: [runId], + execution_dates: null, + [runId]: state, + }, + ], + }, + }, + }; + } + + it("succeeds on resume when the triggered run finished in an allowed state", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", resumedWith("success")), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + // Resuming runs execute_complete only: nothing is triggered or pushed again. + expect(result.runtimeRequests).toEqual([]); + }); + + it("fails on resume when the triggered run finished in a failed state", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", resumedWith("failed")), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "failed" }); + }); + + it("fails on resume through __fail__, when the trigger failed or timed out", async () => { + testDag.triggerDagRun({ + taskId: "trigger", + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + }); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { + should_retry: true, + next_method: "__fail__", + next_kwargs: { error: "Trigger timeout", traceback: ["Traceback", "TimeoutError"] }, + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "Trigger timeout", + }); + expect( + result.logRecords.some((r) => String(r["event"]).includes("Trigger failed:\nTraceback")), + ).toBe(true); + }); + }); }); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index b6cb9541627a9..fc05e1d9f5fc6 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -312,7 +312,7 @@ describe("public API", () => { // A group is a scope and an edge endpoint, and nests the same way at // every depth. expectTypeOf>().toEqualTypeOf< - "dagId" | "groupId" | "task" | "taskGroup" | "before" | "after" + "dagId" | "groupId" | "task" | "taskGroup" | "triggerDagRun" | "before" | "after" >(); expectTypeOf().toEqualTypeOf(); expectTypeOf().toEqualTypeOf< diff --git a/ts-sdk/tests/sdk/trigger-dag-run.test.ts b/ts-sdk/tests/sdk/trigger-dag-run.test.ts new file mode 100644 index 0000000000000..2db4cba9b9c99 --- /dev/null +++ b/ts-sdk/tests/sdk/trigger-dag-run.test.ts @@ -0,0 +1,293 @@ +/*! + * 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 { afterEach, describe, expect, it, vi } from "vitest"; +import { serializeDag } from "../../src/coordinator/serde.js"; +import { Bundle, finalizeBundleDags, getBundleTrigger } from "../../src/sdk/bundle.js"; +import { Dag, getDagOrderEdges, getDagTaskRecords } from "../../src/sdk/dag.js"; +import type { TriggerDagRunSpec } from "../../src/sdk/trigger-dag-run.js"; + +type Json = Record; + +/** The serialized Dag, and its tasks keyed by task id. */ +function serialize(dag: Dag) { + const serialized = serializeDag(dag, "", ".") as Json; + const tasks = serialized["tasks"] as { __var: Json }[]; + return { + serialized, + tasks: new Map(tasks.map(({ __var }) => [__var["task_id"] as string, __var])), + }; +} + +/** A Dag holding one trigger task, built from `spec`. */ +function triggered(spec: Partial = {}, taskSpec = {}) { + const dag = new Dag("d"); + const trigger = dag.triggerDagRun( + { taskId: "trigger_downstream", dagId: "downstream_etl", ...spec } as TriggerDagRunSpec, + taskSpec, + ); + return { dag, trigger }; +} + +describe("dag.triggerDagRun", () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("returns the task's reference, since there is nothing to call", () => { + const { dag, trigger } = triggered(); + + expect(trigger).toMatchObject({ dagId: "d", taskId: "trigger_downstream" }); + expect(dag.taskIds).toEqual(["trigger_downstream"]); + }); + + it("carries no handler, and its options with the operator's defaults", () => { + const { dag } = triggered(); + + const record = getDagTaskRecords(dag).get("trigger_downstream")!; + expect(record.fn).toBeUndefined(); + expect(record.trigger).toEqual({ + dagId: "downstream_etl", + runId: undefined, + conf: undefined, + resetDagRun: false, + waitForCompletion: false, + pokeInterval: 60, + allowedStates: ["success"], + failedStates: ["failed"], + skipWhenAlreadyExists: false, + failWhenDagIsPaused: false, + note: undefined, + deferrable: false, + }); + const bundle = new Bundle(dag); + expect(bundle.getTaskHandler("d", "trigger_downstream")).toBeUndefined(); + expect(getBundleTrigger(bundle, "d", "trigger_downstream")).toBe(record.trigger); + }); + + it("falls back on the default for empty allowedStates, not for empty failedStates", () => { + const { dag } = triggered({ allowedStates: [], failedStates: [] }); + + expect(getDagTaskRecords(dag).get("trigger_downstream")!.trigger).toMatchObject({ + allowedStates: ["success"], + failedStates: [], + }); + }); + + it.each([ + ["True", {}, true], + ["false", {}, false], + ["1", { deferrable: false }, false], + ["f", { deferrable: true }, true], + ])( + "with AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE=%s and %o, defers: %s", + (value, spec, deferrable) => { + vi.stubEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", value); + + const { dag } = triggered(spec); + + expect(getDagTaskRecords(dag).get("trigger_downstream")!.trigger).toMatchObject({ + deferrable, + }); + }, + ); + + it("counts as called, so the Dag is complete without a factory call", () => { + const { dag } = triggered(); + + expect(() => finalizeBundleDags(new Bundle(dag))).not.toThrow(); + }); + + it("serializes as a TypeScript task, drawn as TriggerDagRunOperator", () => { + const { dag } = triggered({ waitForCompletion: true, pokeInterval: 30 }); + + const task = serialize(dag).tasks.get("trigger_downstream")!; + expect(task).toMatchObject({ + task_id: "trigger_downstream", + task_type: "TypeScriptOperator", + _task_module: "airflow.sdk.coordinators.node", + language: "typescript", + is_stub: true, + template_fields: [], + _operator_name: "TriggerDagRunOperator", + ui_color: "#ffefeb", + _operator_extra_links: { "Triggered DAG": "_link_TriggerDagRunLink" }, + }); + // The runtime reads its options from the Dag it built, not from the Dag JSON. + expect(task).not.toHaveProperty("trigger_dag_id"); + expect(task).not.toHaveProperty("wait_for_completion"); + expect(task).not.toHaveProperty("_arg_bindings"); + }); + + it("records the triggered Dag as a dependency of this one", () => { + const { dag } = triggered(); + + expect(serialize(dag).serialized["dag_dependencies"]).toEqual([ + { + source: "d", + target: "downstream_etl", + label: "trigger_downstream", + dependency_type: "trigger", + dependency_id: "trigger_downstream", + }, + ]); + }); + + it("inherits the Dag's queue, so it runs where the Dag's other tasks run", () => { + const dag = new Dag("d", { queue: "typescript" }); + dag.task("extract", async () => undefined)(); + dag.triggerDagRun({ taskId: "trigger_downstream", dagId: "downstream_etl" }); + + const { tasks } = serialize(dag); + expect(tasks.get("extract")?.["queue"]).toBe("typescript"); + expect(tasks.get("trigger_downstream")?.["queue"]).toBe("typescript"); + }); + + it("still takes a task spec", () => { + const { dag } = triggered({}, { retries: 2, queue: "typescript_heavy" }); + + expect(serialize(dag).tasks.get("trigger_downstream")).toMatchObject({ + retries: 2, + queue: "typescript_heavy", + }); + }); + + it("participates in before and after like any task", () => { + const dag = new Dag("d"); + const loaded = dag.task("load", async () => undefined)(); + const trigger = dag.triggerDagRun({ taskId: "trigger_downstream", dagId: "downstream_etl" }); + + trigger.after(loaded); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "load", downstream: "trigger_downstream" }]); + expect(serialize(dag).tasks.get("load")?.["downstream_task_ids"]).toEqual([ + "trigger_downstream", + ]); + }); + + it("sits in a task group like any task, prefix included", () => { + const dag = new Dag("d"); + dag.taskGroup("publish").triggerDagRun({ taskId: "trigger", dagId: "downstream_etl" }); + + expect(dag.taskIds).toEqual(["publish.trigger"]); + }); + + describe("rejects", () => { + it.each([ + ["no taskId", {}], + ["an empty taskId", { taskId: "" }], + ["a taskId that is not a string", { taskId: 7 }], + ])("%s", (_label, spec) => { + const dag = new Dag("d"); + + expect(() => + dag.triggerDagRun({ dagId: "downstream_etl", ...spec } as TriggerDagRunSpec), + ).toThrowError(/A triggerDagRun task of Dag "d" has no taskId/); + }); + + it.each([ + ["no dagId", {}], + ["an empty dagId", { dagId: "" }], + ["a dagId that is not a string", { dagId: 7 }], + ])("%s", (_label, spec) => { + const dag = new Dag("d"); + + expect(() => + dag.triggerDagRun({ taskId: "t", ...spec } as unknown as TriggerDagRunSpec), + ).toThrowError(/triggerDagRun\(\.\.\.\) needs the dagId of the Dag to trigger/); + }); + + it.each([ + ["null", null], + ["a string", "downstream_etl"], + ["an array", []], + ])("%s where the options belong", (_label, spec) => { + const dag = new Dag("d"); + + expect(() => dag.triggerDagRun(spec as unknown as TriggerDagRunSpec)).toThrowError( + /A triggerDagRun task of Dag "d" has no taskId/, + ); + }); + + it("an option the operator does not declare", () => { + const dag = new Dag("d"); + + expect(() => + dag.triggerDagRun({ + taskId: "t", + dagId: "d2", + waitFor: true, + } as unknown as TriggerDagRunSpec), + ).toThrowError(/Unknown option "waitFor" for triggerDagRun\(\.\.\.\)/); + }); + + it.each([ + ["a Date", new Date("2026-01-01")], + ["a Map", new Map()], + ["a function", () => 1], + ])("%s inside conf, which the Dag JSON cannot carry", (_label, value) => { + const { dag } = triggered({ conf: { bad: value } as never }); + + expect(() => serialize(dag)).toThrowError( + /conf of task "trigger_downstream" of Dag "d"\.bad is .*, which a Dag run's conf cannot carry/, + ); + }); + + it.each([ + ["allowedStates", { allowedStates: ["done"] }], + ["failedStates", { failedStates: ["error"] }], + ])("a %s entry that is not a Dag run state", (name, spec) => { + const dag = new Dag("d"); + + expect(() => + dag.triggerDagRun({ taskId: "t", dagId: "d2", ...spec } as unknown as TriggerDagRunSpec), + ).toThrowError(new RegExp(`option "${name}" holds "\\w+", which is not a Dag run state`)); + }); + + it.each([ + ["a negative pokeInterval", { pokeInterval: -1 }, /"pokeInterval" must be a non-negative/], + ["a non-boolean deferrable", { deferrable: "yes" }, /"deferrable" must be a boolean/], + ["a non-string runId", { runId: 7 }, /"runId" must be a string/], + ["a conf that is an array", { conf: [1] }, /"conf" must be an object/], + ])("%s", (_label, spec, error) => { + const dag = new Dag("d"); + + expect(() => + dag.triggerDagRun({ taskId: "t", dagId: "d2", ...spec } as unknown as TriggerDagRunSpec), + ).toThrowError(error); + }); + + it("an AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE that is not a boolean", () => { + vi.stubEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", "yes"); + + expect(() => triggered()).toThrowError( + 'AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE is "yes", which is not a boolean; use true or false', + ); + }); + + it("a condition built on a trigger task, which has no handler to decide with", () => { + const { dag, trigger } = triggered(); + const other = dag.task("other", async () => undefined)(); + + expect(() => dag.if(trigger as never).then(other)).toThrowError( + /Task "trigger_downstream" of Dag "d" is a triggerDagRun task, so it cannot decide a branch/, + ); + }); + }); +});