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/4] 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/4] 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/4] 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/4] 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"`), ); });