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(); });