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 e40b1363c0ad7..0900ba3a40a74 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -289,6 +289,28 @@ The task id may be omitted, in which case it is the handler's function name: name of its own, such as an arrow function passed inline, has nothing to take an id from and needs one: either positionally or as ``taskId`` in its spec. Give it in one place only, not both. +Order-only edges +~~~~~~~~~~~~~~~~ + +An edge that carries no value has no argument name to travel under, so it is drawn between the +references themselves with ``before`` and ``after``, the TypeScript pair for Python's ``>>`` and +``<<``: + +.. code-block:: typescript + + const loaded = load({ transformed }); + const cleaned = cleanup(); + + loaded.before(cleaned); // loaded >> cleaned + cleaned.after(loaded, transformed); // [loaded, transformed] >> cleaned + +Both take any number of references, so one call draws several edges, and drawing an edge that +already exists changes nothing. Each returns the reference it was called on, so +``loaded.before(cleaned).before(notified)`` draws both edges from ``loaded``. + +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. + ``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/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index 2903b2b6ac4db..6e5f274b0e1e2 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -127,9 +127,6 @@ declare const RETURN_TYPE: unique symbol; /** * A reference to the result of one task, returned by calling that task. * - * Identity only: the handler and the value are deliberately not exposed. Pass a - * reference as an input of a downstream task to make that task depend on it. - * * `TReturn` is the handler's return type, so a construct that needs a * particular one can ask for it. A reference of a narrower type is usable * wherever a wider one is: a `TaskRef` is a `TaskRef`. @@ -145,8 +142,48 @@ export interface TaskRef { readonly taskId: string; /** @internal Never set; see {@link RETURN_TYPE}. */ readonly [RETURN_TYPE]?: TReturn; + /** + * Run this task before each of `downstream`, carrying no value — the + * TypeScript spelling of Python's `>>`. + * + * ```ts + * loaded.before(cleaned, notified); // loaded >> [cleanup, notify] + * ``` + * + * Variadic, so one call fans out, and it returns its own receiver rather + * 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; + /** + * Run this task after each of `upstream`, carrying no value — Python's `<<`. + * + * ```ts + * cleaned.after(loaded, transformed); // [load, transform] >> cleanup + * ``` + * + * Fan-*in* that carries data is the wiring object instead + * (`summarize({ north: extractNorth(), south: extractSouth() })`), so each + * direction has an answer: named keys when values flow, `after` when only + * order does. + */ + after(...upstream: readonly TaskRef[]): TaskRef; +} + +/** + * An order-only edge of a Dag: upstream task ID, then downstream task ID. + * + * Kept apart from the wiring a factory call records, because an edge that + * carries no value has no argument name to be recorded under. + */ +export interface OrderEdge { + readonly upstream: string; + readonly downstream: string; } +// A task id cannot hold a NUL, so a joined pair cannot collide with one. +const EDGE_KEY_SEPARATOR = "\u0000"; + /** Whether `value` is a TaskRef returned by any copy of this package. */ function isTaskRef(value: unknown): value is TaskRef { return hasBrand(value, "TaskRef"); @@ -277,6 +314,7 @@ export type RecordedInputs = Readonly>; // Dag's private state without public accessors on the class. let taskRecordsOf: (dag: Dag) => ReadonlyMap; let inputsOf: (dag: Dag) => ReadonlyMap; +let orderEdgesOf: (dag: Dag) => readonly OrderEdge[]; let finalizeOf: (dag: Dag) => void; /** Internal: whether `value` is a Dag built by any copy of this package. */ @@ -310,11 +348,15 @@ export class Dag { readonly spec: DagSpec; readonly #tasks = new Map(); readonly #inputs = new Map(); + // Keyed by the two task ids, so declaring an edge twice records it once, and + // insertion-ordered so the serialized Dag reads as written. + readonly #orderEdges = new Map(); #finalized = false; static { taskRecordsOf = (dag) => dag.#tasks; inputsOf = (dag) => dag.#inputs; + orderEdgesOf = (dag) => [...dag.#orderEdges.values()]; finalizeOf = (dag) => dag.#finalize(); } @@ -425,7 +467,7 @@ export class Dag { ); } const spec = this.#taskSpecOf(taskId, options); - const task = createTaskRef(this.dagId, taskId); + const task = this.#createTaskRef(taskId); this.#tasks.set(taskId, { task, // The runtime dispatches every handler through one instantiation, as it @@ -466,6 +508,73 @@ export class Dag { return spec as TaskSpec; } + #createTaskRef(taskId: string): TaskRef { + const task: TaskRef = { + dagId: this.dagId, + taskId, + before: (...downstream) => { + for (const other of downstream) this.#addOrderEdge(task, other, "before"); + return task; + }, + after: (...upstream) => { + for (const other of upstream) this.#addOrderEdge(other, task, "after"); + return task; + }, + }; + brand(task, "TaskRef"); + return Object.freeze(task); + } + + #addOrderEdge(upstream: TaskRef, downstream: TaskRef, 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. + const other = verb === "before" ? downstream : upstream; + this.#validateOwnRef(other, verb); + if (upstream.taskId === downstream.taskId) { + throw new Error( + `${verb}() cannot draw an edge from task "${upstream.taskId}" of Dag "${this.dagId}" to ` + + "itself; an edge orders two different tasks", + ); + } + const key = `${upstream.taskId}${EDGE_KEY_SEPARATOR}${downstream.taskId}`; + // 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 }), + ); + } + } + + #validateOwnRef(ref: TaskRef, verb: string): void { + if (!isTaskRef(ref)) { + throw new Error( + `${verb}() on Dag "${this.dagId}" takes task references returned by calling a task, ` + + "not arbitrary values", + ); + } + if (ref.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`, + ); + } + // 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) { + throw new Error( + `${verb}() was given a reference to "${ref.taskId}" that this Dag did not hand out; ` + + `it comes from another Dag object with the same ID, or ${DUPLICATE_COPY_HINT}`, + ); + } + } + #wire(taskId: string, inputs: unknown): void { if (this.#finalized) { throw new Error( @@ -576,12 +685,6 @@ function readFunctionName(handler: unknown): string | undefined { return typeof name === "string" && name.length > 0 ? name : undefined; } -function createTaskRef(dagId: string, taskId: string): TaskRef { - const task: TaskRef = { dagId, taskId }; - brand(task, "TaskRef"); - return Object.freeze(task); -} - function validateDagSpec(dagId: string, spec: DagSpec): void { const value: unknown = spec; if (!isPlainRecord(value)) { @@ -604,6 +707,11 @@ export function getDagTaskRecords(dag: Dag): ReadonlyMap { return taskRecordsOf(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); +} + /** Internal: what each task of a Dag was called with, keyed by task ID. * A task that has not been called is absent. */ export function getDagTaskInputs(dag: Dag): ReadonlyMap { diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 3ffba6332f999..be48599841390 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -59,8 +59,11 @@ describe("public API", () => { ); const upstream = upstreamTask(); const downstream = downstreamTask({ upstream }); - expect(upstream).toEqual({ dagId: "public_api_dag", taskId: "public_api_task" }); - expect(downstream).toEqual({ dagId: "public_api_dag", taskId: "public_api_downstream" }); + expect(upstream).toMatchObject({ dagId: "public_api_dag", taskId: "public_api_task" }); + expect(downstream).toMatchObject({ + dagId: "public_api_dag", + taskId: "public_api_downstream", + }); expect(dag.taskIds).toEqual(["public_api_task", "public_api_downstream"]); // serve() hands the bundle to the runtime, which needs the supervisor's // socket addresses that Airflow puts on argv. @@ -283,6 +286,11 @@ describe("public API", () => { }); it("keeps the Dag authoring signatures extensible via trailing specs", () => { + // Identity, plus the two order-only edge verbs; the handler and the value + // stay hidden. + expectTypeOf>().toEqualTypeOf< + "dagId" | "taskId" | "before" | "after" + >(); expectTypeOf().toEqualTypeOf(); expectTypeOf().toEqualTypeOf(); // A reference carries its handler's return type, so a construct that needs @@ -290,6 +298,14 @@ describe("public API", () => { // a wider one is, and not the other way round. expectTypeOf>().toMatchTypeOf(); expectTypeOf().not.toMatchTypeOf>(); + // Variadic, and each returns its own receiver rather than its arguments, + // return type included. + expectTypeOf["before"]>().toEqualTypeOf< + (...downstream: readonly TaskRef[]) => TaskRef + >(); + expectTypeOf["after"]>().toEqualTypeOf< + (...upstream: readonly TaskRef[]) => TaskRef + >(); // 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 1bd493b602d7c..612cdc9cdb084 100644 --- a/ts-sdk/tests/sdk/dag.test.ts +++ b/ts-sdk/tests/sdk/dag.test.ts @@ -21,6 +21,7 @@ import { describe, it, expect } from "vitest"; import { Dag, finalizeDag, + getDagOrderEdges, getDagTaskInputs, getDagTaskRecords, type DagSpec, @@ -36,7 +37,7 @@ describe("Dag", () => { expect(typeof myTask).toBe("function"); const ref = myTask(); - expect(ref).toEqual({ dagId: "example_dag", taskId: "my_task" }); + expect(ref).toMatchObject({ dagId: "example_dag", taskId: "my_task" }); expect(Object.isFrozen(ref)).toBe(true); }); @@ -53,9 +54,9 @@ describe("Dag", () => { const transformed = transform({ extracted }); const loaded = load({ transformed }); - expect(extracted).toEqual({ dagId: "chained_dag", taskId: "extract" }); - expect(transformed).toEqual({ dagId: "chained_dag", taskId: "transform" }); - expect(loaded).toEqual({ dagId: "chained_dag", taskId: "load" }); + expect(extracted).toMatchObject({ dagId: "chained_dag", taskId: "extract" }); + expect(transformed).toMatchObject({ dagId: "chained_dag", taskId: "transform" }); + expect(loaded).toMatchObject({ dagId: "chained_dag", taskId: "load" }); const inputs = getDagTaskInputs(dag); expect(inputs.get("extract")).toEqual({}); @@ -86,7 +87,7 @@ describe("Dag", () => { const transform = dag.task("transform", async (_: { upstream: unknown }) => undefined); const lookalike = { dagId: "lookalike_dag", taskId: "ghost" }; - transform({ upstream: lookalike }); + transform({ upstream: lookalike } as unknown as { upstream: TaskRef }); expect(getDagTaskInputs(dag).get("transform")).toEqual({ upstream: lookalike }); }); @@ -511,6 +512,172 @@ describe("Dag", () => { }); }); + describe("order-only edges", () => { + /** A Dag whose tasks are all placed, ready for edges to be drawn on it. */ + function placedDag(dagId: string, ...taskIds: string[]) { + const dag = new Dag(dagId); + const refs = Object.fromEntries( + taskIds.map((taskId) => [taskId, dag.task(taskId, async () => undefined)()]), + ); + return { dag, refs }; + } + + it("draws an edge with before, from the receiver to the argument", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + + refs.load!.before(refs.cleanup!); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "load", downstream: "cleanup" }]); + }); + + it("draws an edge with after, from the argument to the receiver", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + + refs.cleanup!.after(refs.load!); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "load", downstream: "cleanup" }]); + }); + + it("fans out from one before call", () => { + const { dag, refs } = placedDag("d", "load", "cleanup", "notify"); + + refs.load!.before(refs.cleanup!, refs.notify!); + + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "load", downstream: "cleanup" }, + { upstream: "load", downstream: "notify" }, + ]); + }); + + it("fans in from one after call", () => { + const { dag, refs } = placedDag("d", "load", "transform", "cleanup"); + + refs.cleanup!.after(refs.load!, refs.transform!); + + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "load", downstream: "cleanup" }, + { upstream: "transform", downstream: "cleanup" }, + ]); + }); + + it.each([ + ["before", (a: TaskRef, b: TaskRef) => a.before(b)], + ["after", (a: TaskRef, b: TaskRef) => b.after(a)], + ])("returns the receiver from %s, not the arguments", (_verb, draw) => { + const { refs } = placedDag("d", "load", "cleanup"); + + // A fan-out has no single "next" reference, so chaining continues from + // the same task rather than from what was just pointed at. + expect(draw(refs.load!, refs.cleanup!)).toBe(_verb === "before" ? refs.load : refs.cleanup); + }); + + it("records an edge once however many times it is drawn", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + + refs.load!.before(refs.cleanup!); + refs.load!.before(refs.cleanup!); + refs.cleanup!.after(refs.load!); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "load", downstream: "cleanup" }]); + }); + + it("rejects an edge from a task to itself", () => { + 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/, + ); + expect(() => refs.a!.after(refs.a!)).toThrowError( + /after\(\) cannot draw an edge from task "a" of Dag "d" to itself/, + ); + }); + + it("keeps the two directions apart", () => { + const { dag, refs } = placedDag("d", "a", "b"); + + refs.a!.before(refs.b!); + refs.b!.before(refs.a!); + + // Two distinct edges, both recorded: rejecting the cycle they form is a + // Dag-level concern, not an edge-level one. + expect(getDagOrderEdges(dag)).toEqual([ + { upstream: "a", downstream: "b" }, + { upstream: "b", downstream: "a" }, + ]); + }); + + it("leaves a frozen edge that a caller cannot rewrite", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + refs.load!.before(refs.cleanup!); + + expect(Object.isFrozen(getDagOrderEdges(dag)[0])).toBe(true); + }); + + it("carries no value, so it records no input", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + + refs.load!.before(refs.cleanup!); + + expect(getDagTaskInputs(dag).get("cleanup")).toEqual({}); + }); + + it.each([ + ["before", (ref: TaskRef, other: TaskRef) => ref.before(other)], + ["after", (ref: TaskRef, other: TaskRef) => ref.after(other)], + ])("rejects a %s edge to a task of another Dag", (verb, draw) => { + const { refs: here } = placedDag("here", "load"); + const { refs: there } = placedDag("there", "cleanup"); + + expect(() => draw(here.load!, there.cleanup!)).toThrowError( + new RegExp( + `${verb}\\(\\) cannot draw an edge to Dag "there" task "cleanup" from Dag "here"`, + ), + ); + }); + + it("rejects a reference from another Dag object carrying the same Dag ID", () => { + const { refs: first } = placedDag("same_id", "load"); + const { refs: second } = placedDag("same_id", "cleanup"); + + expect(() => first.load!.before(second.cleanup!)).toThrowError( + /before\(\) was given a reference to "cleanup" that this Dag did not hand out/, + ); + }); + + it.each([ + ["a plain object", { dagId: "d", taskId: "cleanup" }], + ["a string", "cleanup"], + ["null", null], + ])("rejects %s where a reference belongs", (_label, value) => { + 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/, + ); + }); + + it("records nothing when an edge is rejected", () => { + const { dag, refs } = placedDag("d", "load"); + + expect(() => refs.load!.before("cleanup" as unknown as TaskRef)).toThrow(); + expect(getDagOrderEdges(dag)).toEqual([]); + }); + + it("rejects an edge drawn after the Dag was read", () => { + const { dag, refs } = placedDag("d", "load", "cleanup"); + finalizeDag(dag); + + expect(() => refs.load!.before(refs.cleanup!)).toThrowError( + /An edge was drawn on Dag "d" after the Dag was read/, + ); + }); + + it("has no edges before any are drawn", () => { + const { dag } = placedDag("d", "load"); + expect(getDagOrderEdges(dag)).toEqual([]); + }); + }); + it("exposes its task IDs in attachment order", () => { const dag = new Dag("ordered_dag"); expect(dag.taskIds).toEqual([]);