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 b954ff9cd6c6b..aad721a1a4f4d 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -121,16 +121,12 @@ That top-level ``await`` makes the module a runnable bundle entry point. The ``dagId`` a handler binds must match the ``dag_id`` of the Python Dag, and the ``taskId`` a ``@task.stub`` function in that Dag, including any TaskGroup prefix. -``register`` takes any number of task handlers and ``bundle.serve()`` serves exactly what is registered, -so a task left out is not part of the packed bundle and is marked removed at runtime. +``register`` takes any number of task handlers and Dags, and ``bundle.serve()`` serves exactly what is +registered, so a task left out is not part of the packed bundle and is marked removed at runtime. A second ``bundle.serve()`` call is rejected. Registering holds no sockets and starts nothing, so a unit test can build a bundle and dispatch a handler through ``bundle.getTaskHandler(dagId, taskId)`` without a coordinator runtime. -``Dag`` is another interface, for a Dag declared in TypeScript rather than in Python, and is still a work -in progress. ``new Dag`` and ``dag.task`` take a trailing options object (``spec`` on both, plus -``inputs`` on a task) that is not used yet; do not set them. - TaskFlow arguments ~~~~~~~~~~~~~~~~~~ @@ -242,6 +238,47 @@ task instance. inside the scheduler process, so they must be present where the scheduler can read them. The API server and Dag processor do not need them. +.. _typescript-sdk/native-dag: + +Declaring a Dag in TypeScript +----------------------------- + +A ``Dag`` is declared on this side rather than in Python: its tasks and the edges between them are +written in TypeScript. The surface is still growing, so a Dag declared this way is not served to +Airflow yet. + +``dag.task(taskId, handler)`` returns a *factory*. Calling it places the task in the Dag and supplies the +handler's arguments, so the call graph is the task graph: + +.. code-block:: typescript + + import { Dag } from "apache-airflow-ts-sdk"; + + const dag = new Dag("ts_etl"); + + const extract = dag.task("extract", async (): Promise => 42); + const transform = dag.task("transform", async (rows: number, region: string) => rows * 2); + const load = dag.task("load", async (total: number) => {}); + + load(transform(extract(), "us")); + +Arguments are passed in the order the handler declares them. A handler that declares a single object of +named arguments can also be called with that object, which names each input instead of ordering it: + +.. code-block:: typescript + + const store = dag.task("store", async ({ total }: { total: number }) => {}); + + store({ total: extract() }); + +Each argument takes either an upstream reference or a literal JSON value. A reference has to be the +argument itself: one buried inside an array or an object is a literal, and draws no edge. + +Every task has to be called exactly once. An uncalled task fails when the Dag is read, so none can be +left out of the graph by accident. + +``new Dag`` and ``dag.task`` also take a trailing ``spec`` object that is not used yet; do not set it. + Writing tasks ------------- diff --git a/ts-sdk/adr/0002-native-dag-interface.md b/ts-sdk/adr/0002-native-dag-interface.md index 2c2a1a3df1be7..eca1f5e2ea294 100644 --- a/ts-sdk/adr/0002-native-dag-interface.md +++ b/ts-sdk/adr/0002-native-dag-interface.md @@ -28,11 +28,15 @@ Proposed. Revised after the review on #72047. 1. **`dag.task(handler)` returns a factory, and the task id is optional.** With no id the task takes the handler's function name (`dag.task(extract)` → task `"extract"`); `dag.task(taskId, handler)` sets it explicitly, which an anonymous handler must do. Calling the factory both places the task in - the Dag and supplies its arguments by name — the shape Python TaskFlow uses for + the Dag and supplies its arguments, in the order the handler declares them + (`load(transform(extract(), "us"))`). A handler that declares a single object of named arguments can + also be called with that object — the shape Python TaskFlow uses for `load(transformed=transform(...))`. 2. **The call graph is the task graph.** `tsc` checks every wired key against the handler's own parameter type, and a `TaskRef` exists only once its producing call has returned, so a cycle - through arguments is unrepresentable rather than rejected by a validator. + through arguments is unrepresentable rather than rejected by a validator. A reference passed by + position is checked against the argument's own type, which is what tells the two call shapes apart + when a handler declares a single argument. 3. **Every task is called exactly once.** An uncalled task fails when the Dag is read, so none can be silently left out of the graph. 4. **`before` and `after` draw order-only edges** — the TypeScript pair for `>>` and `<<`, both @@ -145,9 +149,9 @@ convention. by design. Native declaration is what fills them, generated from the serialized-Dag JSON schema the way `src/generated/supervisor.ts` is. This ADR does not choose those fields; it fixes where an author writes them. -- `TaskOptions` collapses into `TaskSpec`. The shipped third argument to `dag.task` is - `{ inputs, spec }`; with wiring moved to the factory call, `inputs` is no longer an option and the - third argument is the spec itself. +- `TaskOptions` carries the spec and the handler's positional argument names, which the packer fills in + from the parameter list so the Dag names each argument as its handler does. With wiring moved to the + factory call, `inputs` is no longer an option. - `TaskHandlerArgs` is removed from the public API, `DagRegistry` becomes `Bundle`, and `serveDags(registry)` becomes `bundle.serve()`, which breaks 0.1.0-beta1 authors; see [ADR-0001](0001-mixed-lang-dag-interface.md) for the shipped call sites @@ -155,11 +159,10 @@ convention. ## Alternatives -- **Positional wiring** (`load(transform(extract()))`), which becomes expressible once data no longer - shares an object with `ctx`/`client`, since a handler can then take its arguments positionally and - `Parameters` is a real tuple. Rejected: it removes the key names from every call - site, and those names are what keeps flat, one-statement-per-task wiring readable at twenty tasks. - A handler may still take several positional arguments; only the *wiring* stays named. +- **Named-only wiring**, rejected in the review on #73435: naming every input reads well at twenty + tasks but forces an object around a single argument, and positional calls are what TypeScript + authors write. Both are offered, and the handler's own parameter list decides which one a task can + use. - **Injected `ctx`/`client` arguments**, mimicking the Python signature. Rejected, per the above and because feeling native to TypeScript matters more than matching Python's parameter list. @@ -175,9 +178,12 @@ convention. the same object, which forced every typed handler to declare `TArgs & TaskHandlerArgs` and left the top-level argument namespace open to collisions with an author's own parameter names. Getters close both. -- **The spec argument already has its slot.** `dag.task(taskId, handler, options)` reads - `{ inputs = {}, spec = {} }` and runs `validateEmptySpec` on the spec today - (`ts-sdk/src/sdk/dag.ts`), so task fields land on a path that exists rather than a new one. +- **The spec argument already has its slot.** `dag.task(taskId, handler, options)` reads `{ spec = {} }` + and runs `validateEmptySpec` on it (`ts-sdk/src/sdk/dag.ts`), so task fields land on a path that + exists rather than a new one. +- **A positional argument binds by order, and its name is a label.** The serialized Dag names each + argument, so the packer reads the names from the handler's parameter list; `arg0`, `arg1` and so on + stand in for a name it cannot see, without changing which value reaches which argument. - **A `TaskRef` is inert** — a handle for wiring, not a promise. Nothing in a Dag file executes a task body. - **A defaulted task id is resolved at pack time, not read at runtime.** esbuild renames function diff --git a/ts-sdk/src/coordinator/manifest.ts b/ts-sdk/src/coordinator/manifest.ts index 97661257efdce..523385236a370 100644 --- a/ts-sdk/src/coordinator/manifest.ts +++ b/ts-sdk/src/coordinator/manifest.ts @@ -18,7 +18,7 @@ */ import { SUPERVISOR_API_VERSION } from "./protocol.js"; -import { listBundleDags, type Bundle } from "../sdk/bundle.js"; +import { bundleDagTaskIds, finalizeBundleDags, type Bundle } from "../sdk/bundle.js"; export const AIRFLOW_METADATA_FLAG = "--airflow-metadata"; @@ -36,7 +36,10 @@ export interface BundleManifest { export function buildBundleManifest(bundle: Bundle): BundleManifest { const taskHandlers: BundleManifest["task_handlers"] = {}; - for (const { dagId, tasks } of listBundleDags(bundle)) { + // The manifest is the bundle reporting what it provides, so this is where its + // Dags are finalized: a Dag missing an edge is reported here rather than packed. + finalizeBundleDags(bundle); + for (const [dagId, tasks] of bundleDagTaskIds(bundle)) { if (typeof dagId !== "string") { throw new Error("Dag ID must be a string"); } diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index deb0cb485a005..bf9764f3e6618 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -55,7 +55,7 @@ import { type StartupDetails, } from "./protocol.js"; import { getArgNames } from "../sdk/arg-names.js"; -import { listBundleTasks, type Bundle } from "../sdk/bundle.js"; +import { bundleDagTaskIds, type Bundle } from "../sdk/bundle.js"; import { runInTaskScope, type TaskContext } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; @@ -148,10 +148,10 @@ export async function startCoordinator( const runtimeLogs = logs.child("runtime"); runtimeLogs.debug("Connecting log socket", { logs_addr: parsed.logsAddr }); await logs.connect(parsed.logsAddr); - const tasks = listBundleTasks(bundle); + const byDag = bundleDagTaskIds(bundle); runtimeLogs.info("Coordinator runtime started", { - registered_tasks: tasks, - count: tasks.length, + registered_tasks: Object.fromEntries(byDag), + count: [...byDag.values()].reduce((total, tasks) => total + tasks.length, 0), // Cadwyn schema version this SDK was generated against. Logged // for operator visibility; not sent on the wire. supervisor_api_version: SUPERVISOR_API_VERSION, @@ -270,7 +270,7 @@ function handleParse( // 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: listBundleTasks(bundle), + registered_tasks: Object.fromEntries(bundleDagTaskIds(bundle)), }); const response: RuntimeDagFileParsingResult = { type: "DagFileParsingResult", @@ -295,7 +295,7 @@ async function handleTask( logs.warning("No handler registered for task", { dag_id: ti.dag_id, task_id: ti.task_id, - available: listBundleTasks(bundle), + available: Object.fromEntries(bundleDagTaskIds(bundle)), }); // A missing handler means this bundle cannot run the task, so retrying the // same bundle/configuration mismatch would not help. diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index e6d162155e31e..d63bdf24186c8 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -26,7 +26,16 @@ export { ConnectionNotFoundError, VariableNotFoundError } from "./sdk/client.js" export { SUPERVISOR_API_VERSION } from "./coordinator/index.js"; export type { ArgNameMap } from "./sdk/arg-names.js"; export type { Registerable } from "./sdk/bundle.js"; -export type { DagSpec, TaskInputs, TaskOptions, TaskRef, TaskSpec } from "./sdk/dag.js"; +export type { + DagSpec, + PositionalInputs, + TaskFactory, + TaskInput, + TaskInputs, + TaskOptions, + TaskRef, + TaskSpec, +} from "./sdk/dag.js"; export type { TaskClient } from "./sdk/client.js"; export type { ConnectionResult, GetXComOpts, JsonValue, SetXComOpts } from "./sdk/client-types.js"; export type { TaskContext, TaskFunction } from "./sdk/task.js"; diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index ad99d0c8f344a..2aad2aa03bcc3 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -20,12 +20,13 @@ // The bundle: what a TypeScript bundle process provides, and how it serves it. import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; -import { Dag, getDagTaskRecords, isDag, type TaskRef } from "./dag.js"; +import { Dag, finalizeDag, getDagTaskRecords, isDag } from "./dag.js"; import { getTaskHandlerFunction, isTaskHandler, TaskHandler } from "./task-handler.js"; import type { TaskFunction } from "./task.js"; // Assigned inside Bundle's static block, as Dag does for its tasks. -let entriesOf: (bundle: Bundle) => ReadonlyMap; +let dagsOf: (bundle: Bundle) => ReadonlyMap; +let taskHandlersOf: (bundle: Bundle) => ReadonlyMap>; /** * What {@link Bundle.register} and the {@link Bundle} constructor accept: a @@ -37,32 +38,11 @@ let entriesOf: (bundle: Bundle) => ReadonlyMap; // is assignable to, `TaskHandler` included. export type Registerable = Dag | TaskHandler; -// What a bundle holds per dag_id. The two arms are exclusive by construction: -// a Dag is the native case and owns its own tasks, while task handlers supply -// bodies for a Dag that Python declares, so one dag_id is never both. -type BundleEntry = - | { readonly kind: "dag"; readonly dag: Dag } - | { readonly kind: "handlers"; readonly handlers: Map }; - -function entryTaskIds(entry: BundleEntry): string[] { - return entry.kind === "dag" ? [...entry.dag.taskIds] : [...entry.handlers.keys()]; -} - /** Internal: whether `value` is a Bundle built by any copy of this package. */ export function isBundle(value: unknown): value is Bundle { return hasBrand(value, "Bundle"); } -/** Internal: a Dag this bundle provides for, with its task IDs, as - * {@link listBundleDags} reports it. A task-less native Dag is included, so - * the bundle manifest keeps it visible. */ -export interface RegisteredDag { - /** Identifier of the registered Dag. */ - readonly dagId: string; - /** Airflow task IDs, including any TaskGroup prefix. */ - readonly tasks: string[]; -} - /** * What a bundle process provides to Airflow, and the thing that serves it. * @@ -87,12 +67,15 @@ export interface RegisteredDag { * a snapshot of its tasks. */ export class Bundle { - // Keyed by dag_id and insertion-ordered, so the bundle manifest lists what - // this process provides in the order the entry point registered it. - readonly #entries = new Map(); + // Dags declared in TypeScript, keyed by dag_id in registration order. + #dags = new Map(); + // Handlers for the tasks a Python Dag declares, keyed by dag_id and then by + // task_id, both in registration order. + #taskHandlers = new Map>(); static { - entriesOf = (bundle) => bundle.#entries; + dagsOf = (bundle) => bundle.#dags; + taskHandlersOf = (bundle) => bundle.#taskHandlers; } /** Registers `items`, on the same terms as {@link register}. */ @@ -107,11 +90,13 @@ export class Bundle { * The constructor covers the common case; this is for a bundle that * collects what it provides across several modules. */ register(...items: Registerable[]): void { - // Validated against what is already held *and* against this call, in full, - // before anything is written: a call that throws registers none of its - // items, so a bundle never half-provides what its author listed once. - const incomingDags = new Set(); - const incomingTaskHandlers = new Set(); + // Staged on copies and committed once every item is accepted, so a call + // that throws registers none of its items and a bundle never + // half-provides what its author listed in one call. + const dags = new Map(this.#dags); + const taskHandlers = new Map( + [...this.#taskHandlers].map(([dagId, handlers]) => [dagId, new Map(handlers)] as const), + ); for (const item of items) { // Typed as Registerable, so narrowing it would collapse to never; these // guard callers reaching this from plain JavaScript. @@ -119,9 +104,9 @@ export class Bundle { // Another copy's value cannot be registered, since both kinds read // private state keyed to this copy's class, so it is rejected by its cause. if (candidate instanceof Dag) { - this.#checkDag(candidate, incomingDags); + stageDag(dags, taskHandlers, candidate); } else if (candidate instanceof TaskHandler) { - this.#checkTaskHandler(candidate, incomingDags, incomingTaskHandlers); + stageTaskHandler(dags, taskHandlers, candidate); } else if (isDag(candidate)) { throw new Error(`Dag "${(candidate as Dag).dagId}" ${DUPLICATE_COPY_HINT}`); } else if (isTaskHandler(candidate)) { @@ -133,61 +118,8 @@ export class Bundle { throw new Error("only Dag and TaskHandler instances can be registered"); } } - for (const item of items) { - if (item instanceof Dag) { - this.#entries.set(item.dagId, { kind: "dag", dag: item }); - } else { - const existing = this.#entries.get(item.dagId); - const handlers = - existing?.kind === "handlers" ? existing.handlers : new Map(); - handlers.set(item.taskId, getTaskHandlerFunction(item)); - if (existing === undefined) { - this.#entries.set(item.dagId, { kind: "handlers", handlers }); - } - } - } - } - - #checkDag(dag: Dag, incomingDags: Set): void { - if (this.#entries.get(dag.dagId)?.kind === "handlers") { - throw new Error( - `Dag "${dag.dagId}" already has registered task handlers; a Dag declared in ` + - "TypeScript owns its own tasks, so one Dag ID cannot have both", - ); - } - if (this.#entries.has(dag.dagId) || incomingDags.has(dag.dagId)) { - throw new Error(`Dag "${dag.dagId}" is already registered`); - } - incomingDags.add(dag.dagId); - } - - #checkTaskHandler( - handler: TaskHandler, - incomingDags: Set, - incomingTaskHandlers: Set, - ): void { - // A native Dag attaches its tasks with dag.task(...), so a task handler for - // the same Dag ID would be a second, disagreeing source for its task list. - if (this.#entries.get(handler.dagId)?.kind === "dag" || incomingDags.has(handler.dagId)) { - throw new Error( - `Dag "${handler.dagId}" is declared in TypeScript; attach its tasks with ` + - "dag.task(...) rather than registering task handlers for them", - ); - } - const entry = this.#entries.get(handler.dagId); - // Keyed on the pair, not the task ID: one bundle serves several Dags, and - // the same task_id under two of them is two different handlers. - const pair = `${handler.dagId}\u0000${handler.taskId}`; - if ( - entry?.kind === "handlers" - ? entry.handlers.has(handler.taskId) - : incomingTaskHandlers.has(pair) - ) { - throw new Error( - `A handler for Dag "${handler.dagId}" task "${handler.taskId}" is already registered`, - ); - } - incomingTaskHandlers.add(pair); + this.#dags = dags; + this.#taskHandlers = taskHandlers; } /** @@ -218,12 +150,54 @@ export class Bundle { /** Look up a registered handler, the way the runtime dispatches a task. * Returns `undefined` when no handler exists. */ getTaskHandler(dagId: string, taskId: string): TaskFunction | undefined { - const entry = this.#entries.get(dagId); - if (entry === undefined) return undefined; - return entry.kind === "dag" - ? getDagTaskRecords(entry.dag).get(taskId)?.fn - : entry.handlers.get(taskId); + const dag = this.#dags.get(dagId); + return dag === undefined + ? this.#taskHandlers.get(dagId)?.get(taskId) + : getDagTaskRecords(dag).get(taskId)?.fn; + } +} + +// A Dag declared in TypeScript owns its own tasks, so a dag_id it holds and a +// dag_id task handlers hold are two disagreeing sources for one task list. +// Each kind is therefore rejected for a dag_id the other one already holds, +// whichever order they were registered in. + +function stageDag( + dags: Map, + taskHandlers: ReadonlyMap>, + dag: Dag, +): void { + if (taskHandlers.has(dag.dagId)) { + throw new Error( + `Dag "${dag.dagId}" already has registered task handlers; a Dag declared in ` + + "TypeScript owns its own tasks, so one Dag ID cannot have both", + ); + } + if (dags.has(dag.dagId)) { + throw new Error(`Dag "${dag.dagId}" is already registered`); + } + dags.set(dag.dagId, dag); +} + +function stageTaskHandler( + dags: ReadonlyMap, + taskHandlers: Map>, + handler: TaskHandler, +): void { + if (dags.has(handler.dagId)) { + throw new Error( + `Dag "${handler.dagId}" is declared in TypeScript; attach its tasks with ` + + "dag.task(...) rather than registering task handlers for them", + ); } + const forDag = taskHandlers.get(handler.dagId) ?? new Map(); + if (forDag.has(handler.taskId)) { + throw new Error( + `A handler for Dag "${handler.dagId}" task "${handler.taskId}" is already registered`, + ); + } + forDag.set(handler.taskId, getTaskHandlerFunction(handler)); + taskHandlers.set(handler.dagId, forDag); } /** Internal: reject a `this` that is not a Bundle built by this copy, naming @@ -237,21 +211,24 @@ export function validateOwnBundle(value: unknown, accessor: string): asserts val ); } -/** Internal: every task handle this bundle can dispatch, across both kinds. Not - * re-exported from the package root: enumerating what the runtime dispatches - * is the runtime's job. */ -export function listBundleTasks(bundle: Bundle): TaskRef[] { - return [...entriesOf(bundle)].flatMap(([dagId, entry]) => - entryTaskIds(entry).map((taskId) => ({ dagId, taskId })), - ); +/** Internal: finalize every Dag this bundle declared in TypeScript, so no task + * can be added or wired afterwards. */ +export function finalizeBundleDags(bundle: Bundle): void { + for (const dag of dagsOf(bundle).values()) { + finalizeDag(dag); + } } -/** Internal: every Dag this bundle provides for, with its task IDs. A native - * Dag with no tasks is included; a Dag known only through task handlers always - * has at least one, since a handler is what put it here. */ -export function listBundleDags(bundle: Bundle): RegisteredDag[] { - return [...entriesOf(bundle)].map(([dagId, entry]) => ({ - dagId, - tasks: entryTaskIds(entry), - })); +/** Internal: the task IDs this bundle can dispatch, per Dag: the Dags declared + * in TypeScript first, then the Dags its task handlers name, each in + * registration order. A Dag declared in TypeScript with no tasks is included. */ +export function bundleDagTaskIds(bundle: Bundle): Map { + const byDag = new Map(); + for (const [dagId, dag] of dagsOf(bundle)) { + byDag.set(dagId, [...dag.taskIds]); + } + for (const [dagId, handlers] of taskHandlersOf(bundle)) { + byDag.set(dagId, [...handlers.keys()]); + } + return byDag; } diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index 1da149473bb1e..94e2a45ec6364 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -17,9 +17,13 @@ * under the License. */ -// The Dag authoring surface: `new Dag(dagId)` plus `dag.task(taskId, handler)`. +// The Dag authoring surface: `new Dag(dagId)`, `dag.task(taskId, handler)`, and +// the factory it returns. Calling a factory declares the task's place in the +// Dag and supplies its arguments, the way calling a TaskFlow function does in +// Python. -import { brand, hasBrand } from "./brand.js"; +import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; +import type { JsonValue } from "./client-types.js"; import type { TaskFunction } from "./task.js"; function isPlainRecord(value: unknown): value is Record { @@ -54,66 +58,195 @@ export type DagSpec = Record; */ export type TaskSpec = Record; +// Carries a reference's return type without carrying a value. Not exported, so +// the property cannot be read or written from outside; it exists only so +// `TaskRef` and `TaskRef` are different types to the compiler. +declare const RETURN_TYPE: unique symbol; + /** - * A reference to a task registered on a {@link Dag}, returned by `dag.task(...)`. + * 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. * - * Identity only: the handler is deliberately not exposed. + * `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`. * - * References are what `inputs` accepts, and what native Dag declaration will - * use to wire dependencies. + * A reference carries a hidden brand, so a hand-written `{dagId, taskId}` + * 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 { +export interface TaskRef { /** Identifier of the Dag this task belongs to. */ readonly dagId: string; /** Airflow task ID, including any TaskGroup prefix. */ readonly taskId: string; + /** @internal Never set; see {@link RETURN_TYPE}. */ + readonly [RETURN_TYPE]?: TReturn; +} + +/** Whether `value` is a TaskRef returned by any copy of this package. */ +function isTaskRef(value: unknown): value is TaskRef { + return hasBrand(value, "TaskRef"); +} + +/** + * The first reference reachable inside `value`, or `undefined` when there is + * none. + * + * A reference is an edge, and an edge has to be an input in its own right: one + * buried in an array or an object is data the serializer would write into the + * Dag verbatim, leaving the task running without the upstream it was given. + * TypeScript rejects that for a well-typed argument, so this catches the `any`, + * the cast and the plain-JavaScript caller. + */ +function findNestedTaskRef(value: unknown, seen: WeakSet): TaskRef | undefined { + if (typeof value !== "object" || value === null) return undefined; + if (isTaskRef(value)) return value; + // A literal that refers back to itself is not JSON either, but it reaches + // here before anything else has rejected it. + if (seen.has(value)) return undefined; + seen.add(value); + for (const nested of Array.isArray(value) ? value : Object.values(value)) { + const found = findNestedTaskRef(nested, seen); + if (found) return found; + } + return undefined; } /** - * Task references keyed by input name. + * The part of `T` a literal can express, matched structurally. + * + * `Extract` would do for a type literal but collapses an + * `interface` to `never`, since an interface gets no implicit index signature + * and so never matches `JsonValue`'s object arm. An object survives only if + * every property does, which leaves types JSON cannot carry — a `Date`, a + * method — as `never`, so such an argument can only be given a reference. + */ +type JsonCompatible = T extends JsonValue + ? T + : // A function is an object with no keys, so it would otherwise map to `{}` + // and take every method of a class along with it. + T extends (...args: never[]) => unknown + ? never + : T extends readonly (infer TElement)[] + ? readonly JsonCompatible[] + : T extends object + ? T extends { [K in keyof T]: JsonCompatible } + ? { [K in keyof T]: JsonCompatible } + : never + : never; + +/** + * One input of a task: the upstream task that produces the value, or the value. + * + * A literal is restricted to the JSON-compatible part of the argument's type, + * because it has to survive the trip through the serialized Dag. An argument + * that cannot be expressed as JSON at all, such as a `Date`, can only be given + * a reference. + */ +export type TaskInput = TaskRef | JsonCompatible; + +/** The inputs of a task that declares several arguments, in declaration order. */ +export type PositionalInputs = { + [K in keyof TParams]: TaskInput; +}; + +/** The inputs of a task that declares one object of named arguments, by name. */ +export type TaskInputs = { + [K in keyof TArgs]: TaskRef | JsonCompatible; +}; + +// Offered only where it means something: `TaskInputs` would map over +// `number`'s own methods and accept `{ toFixed: ... }`. +type NamedInputs = [TOnly] extends [object] ? TaskInputs : never; + +/** + * What `dag.task(...)` returns: call it to declare where the task sits in the Dag. + * + * Pass one input per argument the handler declares, in order. An input is + * either another task's reference, which makes this task wait for that task and + * receive its result, or a literal value: + * + * ```ts + * const extract = dag.task("extract", async (): Promise => 42); + * const transform = dag.task("transform", async (rows: number, region: string) => rows); + * const load = dag.task("load", async (total: number) => {}); + * + * load(transform(extract(), "us")); + * ``` + * + * A handler that declares a single object of named arguments can also be called + * with that object, which names each input instead of ordering it: * - * Stored and validated, but they do not create dependencies or pass values yet - * (see {@link TaskOptions.inputs}). Each must identify an earlier task in the - * same Dag. Literal values are not supported. + * ```ts + * const store = dag.task("store", async ({ total }: { total: number }) => {}); + * + * store({ total: extract() }); + * ``` + * + * The compiler checks that every argument is supplied and that each literal + * matches its argument's type. A reference passed by position is checked against + * the argument's type as well, which is what tells the two call shapes apart + * when a handler declares a single argument. */ -export type TaskInputs = Readonly>; +export type TaskFactory = [TParams] extends [ + readonly [], +] + ? () => TaskRef + : TParams extends readonly [infer TOnly] + ? (input: TaskInput | NamedInputs) => TaskRef + : (...inputs: PositionalInputs) => TaskRef; /** * Named options for `dag.task()`. * - * Keyword-only so neither field has to be positioned around the other, and so - * future fields can be added without a new parameter. Unknown keys are - * rejected, so a typo fails at import time rather than being ignored. + * Keyword-only so future fields can be added without a new parameter. Unknown + * keys are rejected, so a typo fails at import time rather than being ignored. */ export interface TaskOptions { + /** Task-level options. Stored, but not used yet — see {@link TaskSpec}. */ + readonly spec?: TaskSpec; /** - * References to the upstream tasks this task consumes. + * Names for the handler's positional arguments, in declaration order. * - * Not used yet: a handler takes no arguments, and the Python stub Dag - * defines task order. Read an upstream return value explicitly instead, with - * `getClient().getXCom({ key: "return_value", taskId: "extract" })`, where - * omitting `taskId` reads the *running* task's own XCom, not the upstream. - * - * In the future these will declare dependencies in native TypeScript Dags. + * `airflow-ts-pack` fills this in from the handler's parameter list, so the + * Dag names each argument as its handler does. Positional inputs bind by + * order, so a name left out only costs the label: `arg0`, `arg1` and so on + * stand in for it. */ - readonly inputs?: TaskInputs; - /** Task-level options. Stored, but not used yet (see {@link TaskSpec}). */ - readonly spec?: TaskSpec; + readonly argNames?: readonly string[]; } -/** Per-task record a Dag retains: the handle, the function, its spec, and the - * upstream handles feeding it. */ +/** 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; - /** Upstream handles keyed by input name; empty when the task has no inputs. */ - readonly inputs: TaskInputs; } -// Assigned inside Dag's static block: gives package-internal code read access -// to the #tasks private field without a public accessor on the Dag class. +/** + * Internal: what one call to a task factory recorded, by argument name. + * + * A {@link TaskRef} is an edge from the upstream task; anything else is a + * constant argument. + * + * Recorded for the serializer that will turn a native Dag into serialized Dag + * JSON; nothing reads them at execution time. What a running task is called + * with comes from the supervisor's `arg_bindings`, which are derived from the + * serialized Dag and resolved per task instance — see decision G of + * `airflow-core/adr/lang-sdk/0007-taskflow-across-language-boundary.md`. Those + * bindings therefore win over anything recorded here, and the two never have to + * be reconciled: a native Dag's wiring is what produced its bindings. + */ +export type RecordedInputs = Readonly>; + +// Assigned inside Dag's static block: gives package-internal code access to +// Dag's private state without public accessors on the class. let taskRecordsOf: (dag: Dag) => ReadonlyMap; +let inputsOf: (dag: Dag) => ReadonlyMap; +let finalizeOf: (dag: Dag) => void; /** Internal: whether `value` is a Dag built by any copy of this package. */ export function isDag(value: unknown): value is Dag { @@ -123,33 +256,44 @@ export function isDag(value: unknown): value is Dag { /** * A Dag declared in TypeScript. * - * Today the Dag structure itself is still declared by a Python stub file; a - * `Dag` instance binds TypeScript handlers to that stub's Dag/task IDs. The - * instance retains its `spec` and every task's `(taskId, handler, spec)` so a - * future `serialize()` can produce the serialized Dag JSON for native - * TypeScript Dag declaration. + * Declare the tasks with `dag.task(taskId, handler)`, then call what each + * returns to lay the Dag out: + * + * ```ts + * const extract = dag.task("extract", async (): Promise => 42); + * const load = dag.task("load", async ({ rows }: { rows: number }) => {}); + * load({ rows: extract() }); + * ``` + * + * Every task has to be called exactly once, so none can be left out of the Dag + * by accident. A TypeScript handler for a task that a *Python* Dag declares is + * a `TaskHandler` instead, not a task on a `Dag`. * * Constructing a Dag has no effect beyond the instance itself. Collect the ones * a bundle should serve on a `Bundle` and await `bundle.serve()`. */ export class Dag { - /** Identifier of this Dag. Must match the Python Dag's `dag_id`. */ + /** Identifier of this Dag. */ readonly dagId: string; /** Dag-level options this instance was constructed with, copied and frozen. */ readonly spec: DagSpec; readonly #tasks = new Map(); + readonly #inputs = new Map(); + #finalized = false; static { taskRecordsOf = (dag) => dag.#tasks; + inputsOf = (dag) => dag.#inputs; + finalizeOf = (dag) => dag.#finalize(); } constructor(dagId: string, spec: DagSpec = {}) { validateEmptySpec(`spec for Dag "${dagId}"`, spec); brand(this, "Dag"); this.dagId = dagId; - // Copied and frozen, as task specs and inputs are: nothing reads a spec - // until the bundle manifest is built, long after the user's module has run, - // so a later mutation of their object would silently change what is packed. + // Copied and frozen, as task specs are: nothing reads a spec until the + // bundle manifest is built, long after the user's module has run, so a + // later mutation of their object would silently change what is packed. // Shallow, so a nested value in a future generated spec stays mutable. this.spec = Object.freeze({ ...spec }); } @@ -160,82 +304,215 @@ export class Dag { } /** - * Register a TypeScript handler for a task of this Dag. + * Declare a task of this Dag, and return the factory that places it. * - * `taskId` must match the Dag-side operator's `task_id` exactly, including - * any TaskGroup prefix. Returns this task's handle. + * Every argument the handler declares becomes an input of the returned + * {@link TaskFactory}; `getContext()` and `getClient()` reach the runtime + * from inside the call, so neither is an argument. */ - task( + task( taskId: string, - handler: TaskFunction, + handler: (...args: TParams) => TReturn | Promise, options: TaskOptions = {}, - ): TaskRef { + ): TaskFactory { if (typeof handler !== "function") { throw new Error(`handler for Dag "${this.dagId}" task "${taskId}" must be a function`); } 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) { + throw new Error( + `Task "${taskId}" cannot be added to Dag "${this.dagId}" after the Dag was read; ` + + "declare every task while the module is loading", + ); + } this.#validateOptions(taskId, options); - const { inputs = {}, spec = {} } = options; + const { spec = {} } = options; validateEmptySpec(`spec for Dag "${this.dagId}" task "${taskId}"`, spec); - this.#validateInputs(taskId, inputs); - const task: TaskRef = Object.freeze({ dagId: this.dagId, taskId }); + const argNames = this.#validateArgNames(taskId, options.argNames); + const task = createTaskRef(this.dagId, taskId); this.#tasks.set(taskId, { task, - fn: handler as TaskFunction, + // The runtime dispatches every handler through one instantiation, as it + // does a registered TaskHandler; a positional one is wrapped at wiring. + fn: handler as unknown as TaskFunction, spec: Object.freeze({ ...spec }), - inputs: Object.freeze({ ...inputs }), }); - return task; + return ((...inputs: unknown[]) => { + this.#wire(taskId, inputs, argNames); + return task; + }) as TaskFactory; } - // TypeScript is bypassable from plain JavaScript or an `as TaskOptions` - // cast, so an unknown key is rejected rather than silently ignored. + // TypeScript is bypassable — from plain JavaScript, or an `as TaskOptions` + // cast — so an unknown key is rejected rather than silently ignored. #validateOptions(taskId: string, options: TaskOptions): void { const value: unknown = options; if (!isPlainRecord(value)) { throw new Error(`options for Dag "${this.dagId}" task "${taskId}" must be an object`); } for (const key of Object.keys(value)) { - if (key !== "inputs" && key !== "spec") { + if (key !== "spec" && key !== "argNames") { throw new Error(`Unknown option "${key}" for Dag "${this.dagId}" task "${taskId}"`); } } } - #validateInputs(taskId: string, inputs: TaskInputs): void { - if (!isPlainRecord(inputs)) { - throw new Error(`inputs for Dag "${this.dagId}" task "${taskId}" must be an object`); + // TypeScript is bypassable, and these names become the keys the arguments are + // recorded under, so an integer-like one would reorder what it labels. + #validateArgNames(taskId: string, names: unknown): readonly string[] | undefined { + if (names === undefined) return undefined; + const describe = `argNames for Dag "${this.dagId}" task "${taskId}"`; + if (!Array.isArray(names)) throw new Error(`${describe} must be an array of names`); + const seen = new Set(); + for (const name of names as unknown[]) { + if (typeof name !== "string" || name.length === 0 || /^\d+$/.test(name)) { + throw new Error( + `${describe} holds ${JSON.stringify(name)}; each name must be a non-empty ` + + "string that is not a number", + ); + } + if (seen.has(name)) throw new Error(`${describe} names "${name}" twice`); + seen.add(name); + } + return Object.freeze([...(names as string[])]); + } + + #wire(taskId: string, inputs: readonly unknown[], argNames: readonly string[] | undefined): void { + if (this.#finalized) { + throw new Error( + `Task "${taskId}" of Dag "${this.dagId}" was called after the Dag was read; ` + + "call every task while the module is loading", + ); + } + if (this.#inputs.has(taskId)) { + throw new Error( + `Task "${taskId}" of Dag "${this.dagId}" was already called; a task holds one place ` + + "in a Dag, so call it once and reuse the reference", + ); } - for (const [name, upstream] of Object.entries(inputs)) { - if ( - upstream == null || - typeof upstream.dagId !== "string" || - typeof upstream.taskId !== "string" - ) { + const positional = !isNamedCall(inputs); + const recorded = this.#checkInputs( + taskId, + positional ? positionalInputs(inputs, argNames) : (inputs[0] as Record), + ); + this.#inputs.set(taskId, recorded); + if (positional && inputs.length > 0) { + const record = this.#tasks.get(taskId)!; + this.#tasks.set(taskId, { ...record, fn: spreadArgs(record.fn) }); + } + } + + #checkInputs(taskId: string, inputs: Record): RecordedInputs { + // Every own key, not just the enumerable string ones: a symbol key would be + // copied by the spread below and so has to be checked, not stepped over. + for (const key of Reflect.ownKeys(inputs)) { + if (typeof key === "symbol") { + throw new Error( + `Input "${String(key)}" of task "${taskId}" is keyed by a symbol; ` + + "an argument name is a string", + ); + } + const value = inputs[key]; + // Anything unbranded is a literal argument, including a look-alike + // `{dagId, taskId}` object: only a real reference makes an edge. + if (!isTaskRef(value)) { + const nested = findNestedTaskRef(value, new WeakSet()); + if (nested) { + throw new Error( + `Input "${key}" of task "${taskId}" holds a reference to "${nested.taskId}" inside a ` + + "literal value, which draws no edge; pass the reference as the input itself, or " + + "give each upstream its own argument", + ); + } + continue; + } + if (value.dagId !== this.dagId) { + throw new Error( + `Input "${key}" of task "${taskId}" comes from Dag "${value.dagId}", not "${this.dagId}"`, + ); + } + const upstream = this.#tasks.get(value.taskId); + if (!upstream) { throw new Error( - `Input "${name}" of task "${taskId}" must be a task handle returned by dag.task(...)`, + `Input "${key}" of task "${taskId}" refers to unregistered task "${value.taskId}"`, ); } - if (upstream.dagId !== this.dagId) { + // Identity, not just the ID pair: two Dag objects can carry the same + // dagId, and a second resolved copy of this package brands its own + // references, so matching IDs do not make a reference this Dag handed out. + if (upstream.task !== value) { throw new Error( - `Input "${name}" of task "${taskId}" comes from Dag "${upstream.dagId}", not "${this.dagId}"`, + `Input "${key}" of task "${taskId}" was not returned by this Dag's "${value.taskId}"; ` + + `it comes from another Dag object with the same ID, or ${DUPLICATE_COPY_HINT}`, ); } - // An input can only name a task registered earlier on this Dag, which - // makes self-references and cycles unrepresentable. - if (!this.#tasks.has(upstream.taskId)) { + } + return Object.freeze({ ...inputs }) as RecordedInputs; + } + + #finalize(): void { + if (this.#finalized) return; + for (const taskId of this.#tasks.keys()) { + if (!this.#inputs.has(taskId)) { throw new Error( - `Input "${name}" of task "${taskId}" refers to unregistered task "${upstream.taskId}"`, + `Task "${taskId}" of Dag "${this.dagId}" is never called, so it has no place in the ` + + "Dag; call the factory dag.task(...) returned", ); } } + // Last, so a Dag that failed the check reports that same failure again + // rather than reporting itself as already read. + this.#finalized = true; } } /** - * Internal: the task records of a Dag, for registry lookups. + * Whether a call named its inputs rather than ordering them. + * + * One plain object is the named form. A reference is a plain object too, so it + * is ruled out first: `load(extract())` is one positional input, not a map of + * argument names. + */ +function isNamedCall(inputs: readonly unknown[]): boolean { + return inputs.length === 1 && !isTaskRef(inputs[0]) && isPlainRecord(inputs[0]); +} + +function positionalInputs( + inputs: readonly unknown[], + argNames: readonly string[] | undefined, +): Record { + const byName: Record = {}; + inputs.forEach((value, index) => { + byName[argNames?.[index] ?? `arg${index}`] = value; + }); + return byName; +} + +/** + * Dispatch a positional handler through the one call shape the runtime uses. + * + * A task is called with its bound arguments as a single object, in the order + * they were recorded, so spreading its values back restores the argument list + * the handler declared. + */ +function spreadArgs(fn: TaskFunction): TaskFunction { + const handler = fn as unknown as (...args: unknown[]) => unknown; + return ((args: Record) => + handler(...Object.values(args ?? {}))) as unknown as TaskFunction; +} + +function createTaskRef(dagId: string, taskId: string): TaskRef { + const task: TaskRef = { dagId, taskId }; + brand(task, "TaskRef"); + return Object.freeze(task); +} + +/** + * Internal: the task records of a Dag, for bundle lookups. * * Not re-exported from the package root, and the package `"exports"` map * blocks deep imports, so this is unreachable from outside the SDK. @@ -243,3 +520,19 @@ export class Dag { export function getDagTaskRecords(dag: Dag): ReadonlyMap { return taskRecordsOf(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 { + return inputsOf(dag); +} + +/** + * Internal: check that `dag` is fully laid out, then finalize it against further wiring. + * + * Idempotent, and never part of the public surface: a Dag is finished when its + * module has finished loading, so there is nothing for an author to call. + */ +export function finalizeDag(dag: Dag): void { + finalizeOf(dag); +} diff --git a/ts-sdk/tests/cli/fixtures/entry.ts b/ts-sdk/tests/cli/fixtures/entry.ts index 607c62cdf0b66..77fe93fcd9aaf 100644 --- a/ts-sdk/tests/cli/fixtures/entry.ts +++ b/ts-sdk/tests/cli/fixtures/entry.ts @@ -22,7 +22,7 @@ import { Bundle, Dag, TaskHandler } from "../../../src/index.js"; // A mixed bundle: task handlers for a Python-owned Dag, plus a natively // declared one, so packing covers both paths into the manifest. const otherDag = new Dag("other_dag"); -otherDag.task("solo", async () => undefined); +otherDag.task("solo", async () => undefined)(); await new Bundle( new TaskHandler("fixture_dag", "extract", async () => "extracted"), diff --git a/ts-sdk/tests/cli/fixtures/noisy-entry.ts b/ts-sdk/tests/cli/fixtures/noisy-entry.ts index 370022287c7d3..ed7ab0683c3ce 100644 --- a/ts-sdk/tests/cli/fixtures/noisy-entry.ts +++ b/ts-sdk/tests/cli/fixtures/noisy-entry.ts @@ -23,6 +23,6 @@ import { Bundle, Dag } from "../../../src/index.js"; console.log("noise from an import-time dependency"); const noisyDag = new Dag("noisy_dag"); -noisyDag.task("only", async () => undefined); +noisyDag.task("only", async () => undefined)(); await new Bundle(noisyDag).serve(); diff --git a/ts-sdk/tests/cli/pack.test.ts b/ts-sdk/tests/cli/pack.test.ts index fad4d50a8a80b..c9d8e6a9ea63d 100644 --- a/ts-sdk/tests/cli/pack.test.ts +++ b/ts-sdk/tests/cli/pack.test.ts @@ -427,7 +427,7 @@ describe("runPack", () => { [ `import { Bundle, Dag } from ${JSON.stringify(SDK_INDEX)};`, 'const bigDag = new Dag("big_dag");', - 'for (let i = 0; i < 5000; i += 1) bigDag.task(String(i).padStart(240, "t"), async () => undefined);', + 'for (let i = 0; i < 5000; i += 1) bigDag.task(String(i).padStart(240, "t"), async () => undefined)();', "await new Bundle(bigDag).serve();", ].join("\n"), ); @@ -472,7 +472,7 @@ describe("runPack", () => { [ `import { Bundle, Dag } from ${JSON.stringify(SDK_INDEX)};`, `const suspiciousDag = new Dag(${JSON.stringify(dagId)});`, - `suspiciousDag.task(${JSON.stringify(taskId)}, async () => undefined);`, + `suspiciousDag.task(${JSON.stringify(taskId)}, async () => undefined)();`, "await new Bundle(suspiciousDag).serve();", ].join("\n"), ); @@ -542,7 +542,7 @@ describe("runPack", () => { [ `import { Bundle, Dag } from ${JSON.stringify(SDK_INDEX)};`, 'const salesDag = new Dag("sales_dag");', - 'salesDag.task("extract", async () => undefined);', + 'salesDag.task("extract", async () => undefined)();', 'await new Bundle(salesDag, new Dag("empty_dag")).serve();', ].join("\n"), ); @@ -565,9 +565,9 @@ describe("runPack", () => { [ `import { Bundle, Dag } from ${JSON.stringify(SDK_INDEX)};`, 'const salesDag = new Dag("sales_dag");', - 'salesDag.task("extract", async () => undefined);', + 'salesDag.task("extract", async () => undefined)();', 'const billingDag = new Dag("billing_dag");', - 'billingDag.task("charge", async () => undefined);', + 'billingDag.task("charge", async () => undefined)();', "await new Bundle(salesDag).serve();", ].join("\n"), ); diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index c5a6ffc1eaf7e..d0b6d53f83fff 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -28,7 +28,7 @@ // No Python, no Airflow install, but exercises the same wire format // the real coordinator speaks. -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { createServer, Socket as NetSocket, type Server, type Socket } from "node:net"; import { encode, decode } from "@msgpack/msgpack"; import { @@ -49,13 +49,22 @@ interface BoundTransformArgs { dryRun: boolean; } -const testDag = new Dag("test_dag"); -const otherDag = new Dag("other_dag"); // The bundle the runtime dispatches through. startCoordinator() is driven // directly rather than through bundle.serve(), so these tests can supply mock // socket addresses. Both authoring kinds are registered in one call, since the // runtime dispatches through one lookup regardless of which put a task there. -const bundle = new Bundle(testDag, otherDag); +// +// Rebuilt per test: dispatching reads the bundle, which finalizes every native +// Dag in it against further tasks, so one Dag cannot collect a task per test. +let testDag: Dag; +let otherDag: Dag; +let bundle: Bundle; + +beforeEach(() => { + testDag = new Dag("test_dag"); + otherDag = new Dag("other_dag"); + bundle = new Bundle(testDag, otherDag); +}); interface MockResult { firstResponse: { id: number; body: unknown; isResponse: boolean } | null; @@ -254,7 +263,7 @@ describe("coordinator runtime integration", () => { testDag.task("say_hello", async () => { observedCtx = getContext(); return "ok"; - }); + })(); const result = await driveSupervisor(makeStartupDetails("say_hello")); @@ -307,7 +316,7 @@ describe("coordinator runtime integration", () => { const commAccept = acceptOne(comm.server); const logsAccept = acceptOne(logs.server); - testDag.task("terminal_timeout", async () => undefined); + testDag.task("terminal_timeout", async () => undefined)(); const runtimeDone = startCoordinator(bundle, { commAddr: `127.0.0.1:${comm.port}`, logsAddr: `127.0.0.1:${logs.port}`, @@ -339,7 +348,7 @@ describe("coordinator runtime integration", () => { it("returns TaskState=failed when the handler throws", async () => { testDag.task("boom", async () => { throw new Error("boom"); - }); + })(); const result = await driveSupervisor(makeStartupDetails("boom")); @@ -352,7 +361,7 @@ describe("coordinator runtime integration", () => { it("returns RetryTask when the handler throws and Airflow says the failure is retryable", async () => { testDag.task("boom_retry", async () => { throw new Error("boom"); - }); + })(); const result = await driveSupervisor( makeStartupDetails("boom_retry", "test_dag", "r1", { @@ -608,7 +617,7 @@ describe("coordinator runtime integration", () => { process.emit("SIGTERM"); sawAbort = getContext().signal.aborted; throw new Error("interrupted"); - }); + })(); const result = await driveSupervisor(makeStartupDetails("aborted_then_failed")); @@ -626,7 +635,7 @@ describe("coordinator runtime integration", () => { process.emit("SIGTERM"); sawAbort = getContext().signal.aborted; throw new Error("interrupted"); - }); + })(); const result = await driveSupervisor( makeStartupDetails("aborted_then_failed_retry", "test_dag", "r1", { @@ -649,7 +658,7 @@ describe("coordinator runtime integration", () => { process.emit("SIGTERM"); sawAbort = getContext().signal.aborted; return "completed"; - }); + })(); const result = await driveSupervisor(makeStartupDetails("completed_after_sigterm")); @@ -703,7 +712,7 @@ describe("coordinator runtime integration", () => { if (back !== `node says: ${observedGreeting}`) { throw new Error(`xcom round-trip mismatch: ${back}`); } - }); + })(); const responder: Responder = (msgType, body) => { if (msgType === "GetVariable") { @@ -753,7 +762,7 @@ describe("coordinator runtime integration", () => { let observed: string | null = ""; testDag.task("missing_variable", async () => { observed = await getClient().getVariable("missing_key"); - }); + })(); const responder: Responder = (msgType) => { if (msgType === "GetVariable") { @@ -778,10 +787,10 @@ describe("coordinator runtime integration", () => { let calledSecondDag = false; testDag.task("shared_task", async () => { calledFirstDag = true; - }); + })(); otherDag.task("shared_task", async () => { calledSecondDag = true; - }); + })(); await driveSupervisor(makeStartupDetails("shared_task")); @@ -804,7 +813,7 @@ describe("coordinator runtime integration", () => { }); it("auto-pushes return_value XCom when handler returns a value", async () => { - testDag.task("pusher", async () => "my-result"); + testDag.task("pusher", async () => "my-result")(); const responder: Responder = (msgType, _body) => { if (msgType === "SetXCom") return { body: null }; @@ -826,7 +835,7 @@ describe("coordinator runtime integration", () => { it("does NOT push return_value XCom when handler returns undefined", async () => { testDag.task("void_task", async () => { // no return value - }); + })(); const result = await driveSupervisor(makeStartupDetails("void_task")); diff --git a/ts-sdk/tests/coordinator/runtime-manifest.test.ts b/ts-sdk/tests/coordinator/runtime-manifest.test.ts index b7e030f43e7c5..afca723af2106 100644 --- a/ts-sdk/tests/coordinator/runtime-manifest.test.ts +++ b/ts-sdk/tests/coordinator/runtime-manifest.test.ts @@ -28,7 +28,7 @@ import { Bundle } from "../../src/sdk/bundle.js"; function buildDag(dagId: string, ...taskIds: string[]): Dag { const dag = new Dag(dagId); for (const taskId of taskIds) { - dag.task(taskId, async () => undefined); + dag.task(taskId, async () => undefined)(); } return dag; } @@ -87,7 +87,7 @@ describe("buildBundleManifest", () => { it("rejects a non-string dagId before object-key coercion hides it", () => { const dag = new Dag(123 as unknown as string); - dag.task("t1", async () => undefined); + dag.task("t1", async () => undefined)(); expect(() => buildBundleManifest(new Bundle(dag))).toThrowError(/Dag ID must be a string/); }); }); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 2a86f09d66913..b62075579f1d4 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -27,8 +27,11 @@ import type { SetXComOpts, TaskClient, Registerable, + PositionalInputs, TaskContext, + TaskFactory, TaskFunction, + TaskInput, TaskInputs, TaskOptions, TaskRef, @@ -50,10 +53,13 @@ import { describe("public API", () => { it("exports the Dag authoring surface", async () => { const dag = new Dag("public_api_dag"); - const upstream = dag.task("public_api_task", async () => undefined); - const downstream = dag.task("public_api_downstream", async () => undefined, { - inputs: { upstream }, - }); + const upstreamTask = dag.task("public_api_task", async () => undefined); + const downstreamTask = dag.task( + "public_api_downstream", + async (_: { upstream: undefined }) => undefined, + ); + 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(dag.taskIds).toEqual(["public_api_task", "public_api_downstream"]); @@ -278,22 +284,41 @@ describe("public API", () => { }); it("keeps the Dag authoring signatures extensible via trailing specs", () => { - expectTypeOf().toEqualTypeOf<{ - readonly dagId: string; - readonly taskId: string; - }>(); - expectTypeOf().toEqualTypeOf>>(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + // A reference carries its handler's return type, so a construct that needs + // a particular one can ask for it: a narrower reference is usable wherever + // a wider one is, and not the other way round. + expectTypeOf>().toMatchTypeOf(); + expectTypeOf().not.toMatchTypeOf>(); + // Wiring moved to the factory call, so `inputs` is no longer an option and + // the only remaining one is the spec. expectTypeOf().toEqualTypeOf<{ - readonly inputs?: TaskInputs; readonly spec?: TaskSpec; + readonly argNames?: readonly string[]; }>(); + // Each named argument takes any upstream reference or a literal of its own type. + expectTypeOf>().toEqualTypeOf<{ rows: TaskRef | number }>(); + // A positional argument takes a literal or a reference of the argument's own + // type, which is what tells a one-argument positional call from a named one. + expectTypeOf>().toEqualTypeOf | number>(); + expectTypeOf>().toEqualTypeOf< + [TaskRef | number, TaskRef | string] + >(); + // A handler with no arguments is called with none; one with several is + // called with a value per argument, in order. + expectTypeOf>().toEqualTypeOf<() => TaskRef>(); + expectTypeOf>().toEqualTypeOf<() => TaskRef>(); + expectTypeOf>().toEqualTypeOf< + (...inputs: [TaskRef | number, TaskRef | string]) => TaskRef + >(); expectTypeOf>().toEqualTypeOf<[string, DagSpec?]>(); expectTypeOf().toEqualTypeOf< - ( + ( taskId: string, - handler: TaskFunction, + handler: (...args: TParams) => TReturn | Promise, options?: TaskOptions, - ) => TaskRef + ) => TaskFactory >(); expectTypeOf().toEqualTypeOf(); // Reserved with no fields yet, so only `{}` is expressible. Generated specs @@ -391,21 +416,43 @@ describe("public API", () => { // @ts-expect-error a task handler is required. new Dag("example").task("extract"); const dag = new Dag("example"); - const upstream = dag.task("extract", async () => undefined); - // @ts-expect-error inputs must be task handles, not arbitrary values. + const extract = dag.task("extract", async () => undefined); + // @ts-expect-error wiring belongs to the factory call, not the options. dag.task("transform", async () => undefined, { inputs: { count: 1 } }); - // @ts-expect-error inputs and spec are keyword-only, not positional. - dag.task("transform2", async () => undefined, { upstream }); + // @ts-expect-error the spec is keyword-only, not positional. + dag.task("transform2", async () => undefined, { extract }); // @ts-expect-error a Dag spec is an options object, not a primitive. new Dag("spec_dag", 42); // @ts-expect-error DagSpec has no fields yet, so a schedule cannot be declared here. new Dag("spec_dag", { schedule: "@daily" }); // @ts-expect-error TaskSpec has no fields yet, so retries cannot be declared here. dag.task("transform3", async () => undefined, { spec: { retries: 2 } }); - // @ts-expect-error the TaskRef handle is data, not callable. - upstream(); - // @ts-expect-error a bundle is built from Dags, not from task handles. - new Bundle(upstream); + // @ts-expect-error a handler with no arguments is called with none. + extract({ rows: 1 }); + const transform = dag.task("transform4", async (_: { rows: number }) => undefined); + // @ts-expect-error every argument the handler declares has to be supplied. + transform({}); + // @ts-expect-error a literal has to match its argument's type. + transform({ rows: "many" }); + const totals = dag.task("totals", async (): Promise<{ rows: number }> => ({ rows: 1 })); + // A single object of named arguments is given by name or by position. + transform({ rows: 1 }); + transform(totals()); + // @ts-expect-error a positional reference has to return the argument's type. + transform(extract()); + const pair = dag.task("pair", async (rows: number, region: string) => `${region}${rows}`); + pair(1, "us"); + // @ts-expect-error a positional argument cannot be skipped. + pair(1); + // @ts-expect-error each positional literal has to match its own argument. + pair("many", "us"); + const stamp = dag.task("stamp", async (_: { at: Date }) => undefined); + // @ts-expect-error a Date cannot survive the serialized Dag, so only a reference will do. + stamp({ at: new Date() }); + // The reference an upstream call returns is always accepted. + stamp({ at: extract() }); + // @ts-expect-error a bundle is built from Dags, not from task factories. + new Bundle(extract); // @ts-expect-error serve() takes nothing; the bundle already holds it all. new Bundle(dag).serve(dag); }; diff --git a/ts-sdk/tests/sdk/bundle.test.ts b/ts-sdk/tests/sdk/bundle.test.ts index 37f6f68b663d5..24998153c278b 100644 --- a/ts-sdk/tests/sdk/bundle.test.ts +++ b/ts-sdk/tests/sdk/bundle.test.ts @@ -19,13 +19,13 @@ import { describe, it, expect } from "vitest"; import { Dag } from "../../src/sdk/dag.js"; -import { Bundle, listBundleDags, listBundleTasks } from "../../src/sdk/bundle.js"; +import { Bundle, bundleDagTaskIds, finalizeBundleDags } from "../../src/sdk/bundle.js"; describe("Bundle", () => { it("registers a Dag and retrieves its handlers", () => { const handler = async () => "hello"; const dag = new Dag("example_dag"); - dag.task("my_task", handler); + dag.task("my_task", handler)(); const bundle = new Bundle(); bundle.register(dag); expect(bundle.getTaskHandler("example_dag", "my_task")).toBe(handler); @@ -34,13 +34,15 @@ describe("Bundle", () => { it("registers the Dags passed to its constructor", () => { const handler = async () => "hello"; const dagA = new Dag("dag_a"); - dagA.task("a", handler); + dagA.task("a", handler)(); const bundle = new Bundle(dagA, new Dag("dag_b")); expect(bundle.getTaskHandler("dag_a", "a")).toBe(handler); - expect(listBundleDags(bundle)).toEqual([ - { dagId: "dag_a", tasks: ["a"] }, - { dagId: "dag_b", tasks: [] }, - ]); + expect(bundleDagTaskIds(bundle)).toEqual( + new Map([ + ["dag_a", ["a"]], + ["dag_b", []], + ]), + ); }); it("rejects duplicate dagIds passed to the constructor", () => { @@ -58,28 +60,30 @@ describe("Bundle", () => { it("returns undefined for unknown taskIds and dagIds", () => { const bundle = new Bundle(); const dag = new Dag("example_dag"); - dag.task("my_task", async () => undefined); + dag.task("my_task", async () => undefined)(); bundle.register(dag); expect(bundle.getTaskHandler("example_dag", "nope")).toBeUndefined(); expect(bundle.getTaskHandler("unknown_dag", "my_task")).toBeUndefined(); }); - it("returns an empty list when no Dags are registered", () => { + it("returns nothing when no Dags are registered", () => { const bundle = new Bundle(); - expect(listBundleTasks(bundle)).toEqual([]); + expect(bundleDagTaskIds(bundle)).toEqual(new Map()); }); it("lists tasks across registered Dags", () => { const dagA = new Dag("dag_a"); - dagA.task("a", async () => undefined); + dagA.task("a", async () => undefined)(); const dagB = new Dag("dag_b"); - dagB.task("b", async () => undefined); + dagB.task("b", async () => undefined)(); const bundle = new Bundle(); bundle.register(dagA, dagB); - const registered = listBundleTasks(bundle); - expect(registered).toHaveLength(2); - expect(registered).toContainEqual({ dagId: "dag_a", taskId: "a" }); - expect(registered).toContainEqual({ dagId: "dag_b", taskId: "b" }); + expect(bundleDagTaskIds(bundle)).toEqual( + new Map([ + ["dag_a", ["a"]], + ["dag_b", ["b"]], + ]), + ); }); it("rejects registering the same dagId in separate calls", () => { @@ -105,10 +109,10 @@ describe("Bundle", () => { it("registers none of the Dags when a call throws", () => { const bundle = new Bundle(); const dag = new Dag("dag_a"); - dag.task("a", async () => undefined); + dag.task("a", async () => undefined)(); expect(() => bundle.register(dag, new Dag("dag_a"))).toThrowError(/already registered/); expect(bundle.getTaskHandler("dag_a", "a")).toBeUndefined(); - expect(listBundleTasks(bundle)).toEqual([]); + expect(bundleDagTaskIds(bundle)).toEqual(new Map()); }); it("rejects values that are neither a Dag nor a task handler", () => { @@ -129,14 +133,16 @@ describe("Bundle", () => { it("lists every registered Dag with its tasks, empty Dags included", () => { const dagA = new Dag("dag_a"); - dagA.task("a1", async () => undefined); - dagA.task("a2", async () => undefined); + dagA.task("a1", async () => undefined)(); + dagA.task("a2", async () => undefined)(); const bundle = new Bundle(); bundle.register(dagA, new Dag("empty_dag")); - expect(listBundleDags(bundle)).toEqual([ - { dagId: "dag_a", tasks: ["a1", "a2"] }, - { dagId: "empty_dag", tasks: [] }, - ]); + expect(bundleDagTaskIds(bundle)).toEqual( + new Map([ + ["dag_a", ["a1", "a2"]], + ["empty_dag", []], + ]), + ); }); it("accepts a call that registers nothing", () => { @@ -144,7 +150,7 @@ describe("Bundle", () => { // guard at the call site. const bundle = new Bundle(); expect(() => bundle.register()).not.toThrow(); - expect(listBundleDags(bundle)).toEqual([]); + expect(bundleDagTaskIds(bundle)).toEqual(new Map()); }); it("carries the brand its own serve guard reads", () => { @@ -153,18 +159,30 @@ describe("Bundle", () => { expect(Symbol.for("airflow.ts-sdk.Bundle") in new Bundle()).toBe(true); }); - it("sees tasks added to a Dag after registration", () => { + it("sees tasks added to a Dag between registration and the first read", () => { + // Registration records Dag identity rather than a snapshot of its tasks, so + // a Dag assembled across several modules is still complete when it is read. const bundle = new Bundle(); const dag = new Dag("example_dag"); bundle.register(dag); - expect(listBundleTasks(bundle)).toEqual([]); const handler = async () => "late"; - dag.task("late_task", handler); + dag.task("late_task", handler)(); expect(bundle.getTaskHandler("example_dag", "late_task")).toBe(handler); - expect(listBundleTasks(bundle)).toContainEqual({ - dagId: "example_dag", - taskId: "late_task", - }); + expect(bundleDagTaskIds(bundle).get("example_dag")).toContain("late_task"); + }); + + it("rejects a task added to a Dag the bundle has already reported", () => { + const dag = new Dag("example_dag"); + dag.task("extract", async () => undefined)(); + const bundle = new Bundle(dag); + // Reporting what the bundle provides is what finalizes a native Dag; + // enumerating what it can dispatch does not. + expect(bundleDagTaskIds(bundle)).toEqual(new Map([["example_dag", ["extract"]]])); + finalizeBundleDags(bundle); + + expect(() => dag.task("late_task", async () => "late")).toThrowError( + /Task "late_task" cannot be added to Dag "example_dag" after the Dag was read/, + ); }); }); diff --git a/ts-sdk/tests/sdk/dag.test.ts b/ts-sdk/tests/sdk/dag.test.ts index 4dfe1652652ed..5c772eb57c8ea 100644 --- a/ts-sdk/tests/sdk/dag.test.ts +++ b/ts-sdk/tests/sdk/dag.test.ts @@ -18,53 +18,104 @@ */ import { describe, it, expect } from "vitest"; -import { Dag, getDagTaskRecords, type TaskRef } from "../../src/sdk/dag.js"; -import { Bundle } from "../../src/sdk/bundle.js"; +import { + Dag, + finalizeDag, + getDagTaskInputs, + getDagTaskRecords, + type TaskRef, +} from "../../src/sdk/dag.js"; +import { Bundle, finalizeBundleDags } from "../../src/sdk/bundle.js"; describe("Dag", () => { - it("returns a frozen TaskRef handle with the Dag and task identity", () => { + it("returns a factory whose call yields a frozen TaskRef with the Dag and task identity", () => { const dag = new Dag("example_dag"); - const task = dag.task("my_task", async () => "hello"); - expect(task).toEqual({ dagId: "example_dag", taskId: "my_task" }); - expect(Object.isFrozen(task)).toBe(true); + const myTask = dag.task("my_task", async () => "hello"); + expect(typeof myTask).toBe("function"); + + const ref = myTask(); + expect(ref).toEqual({ dagId: "example_dag", taskId: "my_task" }); + expect(Object.isFrozen(ref)).toBe(true); }); - it("chains upstream handles into downstream task inputs", () => { + it("chains upstream references into downstream task inputs", () => { const dag = new Dag("chained_dag"); - const extracted = dag.task("extract", async () => ({ rows: 1 })); - const transformed = dag.task("transform", async () => undefined, { inputs: { extracted } }); - const loaded = dag.task("load", async () => undefined, { inputs: { transformed }, spec: {} }); + const extract = dag.task("extract", async () => ({ rows: 1 })); + const transform = dag.task( + "transform", + async (_: { extracted: { rows: number } }) => undefined, + ); + const load = dag.task("load", async (_: { transformed: undefined }) => undefined, { spec: {} }); + + const extracted = extract(); + 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" }); - const records = getDagTaskRecords(dag); - expect(records.get("extract")?.inputs).toEqual({}); - expect(records.get("transform")?.inputs).toEqual({ extracted }); - expect(records.get("load")?.inputs).toEqual({ transformed }); + const inputs = getDagTaskInputs(dag); + expect(inputs.get("extract")).toEqual({}); + expect(inputs.get("transform")).toEqual({ extracted }); + expect(inputs.get("load")).toEqual({ transformed }); }); - it("accepts several named inputs for one task", () => { - const dag = new Dag("fan_in_dag"); - const extracted = dag.task("extract", async () => undefined); - const otherTaskResult = dag.task("other_task", async () => undefined); - dag.task("transform", async () => undefined, { inputs: { extracted, otherTaskResult } }); + it("records a literal argument as a value rather than an edge", () => { + const dag = new Dag("literal_dag"); + const extract = dag.task("extract", async () => 1); + const transform = dag.task( + "transform", + async (_: { extracted: number; regionCode: string; limits: number[] }) => undefined, + ); + + const extracted = extract(); + transform({ extracted, regionCode: "us", limits: [1, 2] }); - expect(getDagTaskRecords(dag).get("transform")?.inputs).toEqual({ + expect(getDagTaskInputs(dag).get("transform")).toEqual({ extracted, - otherTaskResult, + regionCode: "us", + limits: [1, 2], }); }); + it("treats an unbranded look-alike reference as a literal, not an edge", () => { + const dag = new Dag("lookalike_dag"); + const transform = dag.task("transform", async (_: { upstream: unknown }) => undefined); + const lookalike = { dagId: "lookalike_dag", taskId: "ghost" }; + + transform({ upstream: lookalike }); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ upstream: lookalike }); + }); + + it("accepts several named inputs for one task", () => { + const dag = new Dag("fan_in_dag"); + const extract = dag.task("extract", async () => undefined); + const other = dag.task("other_task", async () => undefined); + const transform = dag.task( + "transform", + async (_: { extracted: undefined; otherTaskResult: undefined }) => undefined, + ); + + const extracted = extract(); + const otherTaskResult = other(); + transform({ extracted, otherTaskResult }); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ extracted, otherTaskResult }); + }); + it("records frozen inputs that later mutation of the caller's object cannot change", () => { const dag = new Dag("example_dag"); - const extracted = dag.task("extract", async () => undefined); + const extract = dag.task("extract", async () => undefined); + const transform = dag.task("transform", async (_: { extracted: undefined }) => undefined); + + const extracted = extract(); const inputs: Record = { extracted }; - dag.task("transform", async () => undefined, { inputs }); + transform(inputs as { extracted: TaskRef }); inputs.sneaky = extracted; - const recorded = getDagTaskRecords(dag).get("transform")!.inputs; + const recorded = getDagTaskInputs(dag).get("transform")!; expect(recorded).toEqual({ extracted }); expect(Object.isFrozen(recorded)).toBe(true); }); @@ -72,36 +123,237 @@ describe("Dag", () => { it("rejects an input taken from another Dag", () => { const first = new Dag("first_dag"); const second = new Dag("second_dag"); - const extracted = first.task("extract", async () => undefined); - expect(() => - second.task("transform", async () => undefined, { inputs: { extracted } }), - ).toThrowError( + const extracted = first.task("extract", async () => undefined)(); + const transform = second.task("transform", async (_: { extracted: TaskRef }) => undefined); + + expect(() => transform({ extracted })).toThrowError( /Input "extracted" of task "transform" comes from Dag "first_dag", not "second_dag"/, ); }); + it("rejects a reference from another Dag object carrying the same Dag ID", () => { + const first = new Dag("same_id"); + const second = new Dag("same_id"); + first.task("extract", async () => undefined); + const extracted = second.task("extract", async () => undefined)(); + const transform = first.task("transform", async (_: { extracted: TaskRef }) => undefined); + + expect(() => transform({ extracted })).toThrowError( + /Input "extracted" of task "transform" was not returned by this Dag's "extract"/, + ); + }); + + it("rejects a reference to a task this Dag never registered", () => { + const first = new Dag("same_id"); + const second = new Dag("same_id"); + const ghost = second.task("ghost", async () => undefined)(); + const transform = first.task("transform", async (_: { ghost: TaskRef }) => undefined); + + expect(() => transform({ ghost })).toThrowError( + /Input "ghost" of task "transform" refers to unregistered task "ghost"/, + ); + }); + + it("rejects a reference buried inside a literal input", () => { + // It would draw no edge, so the task would run without the upstream it was + // given. TypeScript rejects it for a well-typed argument, which leaves the + // `any`, the cast and the plain-JavaScript caller to this check. + const dag = new Dag("nested_dag"); + const extract = dag.task("extract", async () => 1); + const other = dag.task("other", async () => 2); + const fan = dag.task("fan", async (_: { sources: unknown }) => undefined); + const sources = [extract(), other()]; + + expect(() => fan({ sources } as unknown as { sources: never })).toThrowError( + /Input "sources" of task "fan" holds a reference to "extract" inside a literal value/, + ); + }); + + it("finds a reference nested several levels down, and survives a self-reference", () => { + const dag = new Dag("deep_dag"); + const extract = dag.task("extract", async () => 1); + const fan = dag.task("fan", async (_: { config: unknown }) => undefined); + const config: Record = { outer: { inner: [{ from: extract() }] } }; + config["self"] = config; + + expect(() => fan({ config } as unknown as { config: never })).toThrowError( + /holds a reference to "extract" inside a literal value/, + ); + }); + + it("rejects an input keyed by a symbol", () => { + const dag = new Dag("symbol_dag"); + const transform = dag.task("transform", async (_: { real: string }) => undefined); + const inputs = { real: "ok", [Symbol("sneaky")]: "value" }; + + expect(() => transform(inputs)).toThrowError( + /Input "Symbol\(sneaky\)" of task "transform" is keyed by a symbol; an argument name is a string/, + ); + }); + + it("records a positional call in argument order", () => { + const dag = new Dag("example_dag"); + const extract = dag.task("extract", async (): Promise => 1); + const transform = dag.task("transform", async (rows: number, region: string) => `${region}`); + + const extracted = extract(); + transform(extracted, "us"); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ arg0: extracted, arg1: "us" }); + expect(Object.keys(getDagTaskInputs(dag).get("transform")!)).toEqual(["arg0", "arg1"]); + }); + + it("names positional arguments from argNames", () => { + const dag = new Dag("example_dag"); + const extract = dag.task("extract", async (): Promise => 1); + const transform = dag.task("transform", async (rows: number, region: string) => `${region}`, { + argNames: ["rows", "region"], + }); + + const extracted = extract(); + transform(extracted, "us"); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ rows: extracted, region: "us" }); + }); + + it("labels the arguments argNames does not reach", () => { + const dag = new Dag("example_dag"); + const transform = dag.task("transform", async (rows: number, region: string) => `${region}`, { + argNames: ["rows"], + }); + + transform(1, "us"); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ rows: 1, arg1: "us" }); + }); + it.each([ - ["a plain string", "extract"], - ["an object without a dagId", { taskId: "extract" }], - ["null", null], - ])("rejects an input that is not a task handle: %s", (_label, value) => { + ["not an array", { argNames: 1 }, /argNames for Dag "d" task "t" must be an array of names/], + ["not a string", { argNames: [1] }, /holds 1; each name must be a non-empty string/], + ["empty", { argNames: [""] }, /holds ""; each name must be a non-empty string/], + ["a number", { argNames: ["0"] }, /holds "0"; each name must be a non-empty string/], + ["a duplicate", { argNames: ["a", "a"] }, /argNames for Dag "d" task "t" names "a" twice/], + ])("rejects argNames that are %s", (_label, options, expected) => { + const dag = new Dag("d"); + + expect(() => dag.task("t", async (a: number) => a, options as never)).toThrowError(expected); + }); + + it("reads a single argument that is not a map of names as one positional input", () => { const dag = new Dag("example_dag"); - const extracted = value as unknown as TaskRef; - expect(() => - dag.task("transform", async () => undefined, { inputs: { extracted } }), - ).toThrowError( - /Input "extracted" of task "transform" must be a task handle returned by dag\.task\(\.\.\.\)/, + const when = new Date(); + const transform = dag.task("transform", async (at: Date) => at); + + transform(when as never); + + expect(getDagTaskInputs(dag).get("transform")).toEqual({ arg0: when }); + }); + + it("reads a single reference as one positional input rather than a map of names", () => { + const dag = new Dag("example_dag"); + const extract = dag.task("extract", async (): Promise<{ rows: number }> => ({ rows: 1 })); + const load = dag.task("load", async (totals: { rows: number }) => totals.rows); + + const extracted = extract(); + load(extracted); + + expect(getDagTaskInputs(dag).get("load")).toEqual({ arg0: extracted }); + }); + + it("spreads a positional task's bound arguments back into its argument list", async () => { + const dag = new Dag("example_dag"); + const seen: unknown[] = []; + const transform = dag.task( + "transform", + async (rows: number, region: string) => { + seen.push(rows, region); + }, + { argNames: ["rows", "region"] }, ); - expect(getDagTaskRecords(dag).has("transform")).toBe(false); + + transform(1, "us"); + // The order the runtime hands the bound arguments over in, which is the + // order they were recorded. + await new Bundle(dag).getTaskHandler("example_dag", "transform")!({ + rows: 1, + region: "us", + } as never); + + expect(seen).toEqual([1, "us"]); }); - it("rejects an input referring to a task that is not registered yet", () => { + it("calls a task declaring one object of named arguments with that object", async () => { const dag = new Dag("example_dag"); - expect(() => - dag.task("transform", async () => undefined, { - inputs: { ghost: { dagId: "example_dag", taskId: "ghost" } }, - }), - ).toThrowError(/Input "ghost" of task "transform" refers to unregistered task "ghost"/); + const seen: unknown[] = []; + const store = dag.task("store", async ({ rows }: { rows: number }) => { + seen.push(rows); + }); + + store({ rows: 1 }); + await new Bundle(dag).getTaskHandler("example_dag", "store")!({ rows: 7 } as never); + + expect(seen).toEqual([7]); + }); + + it("rejects calling the same task twice", () => { + const dag = new Dag("example_dag"); + const extract = dag.task("extract", async () => undefined); + extract(); + + expect(() => extract()).toThrowError( + /Task "extract" of Dag "example_dag" was already called; a task holds one place in a Dag/, + ); + }); + + it("fails when the Dag is read with a task that was never called", () => { + const dag = new Dag("unplaced_dag"); + dag.task("extract", async () => undefined)(); + dag.task("orphan", async () => undefined); + + expect(() => finalizeDag(dag)).toThrowError( + /Task "orphan" of Dag "unplaced_dag" is never called, so it has no place in the Dag/, + ); + }); + + it("reports the same failure on a second read rather than reporting itself as read", () => { + const dag = new Dag("unplaced_dag"); + dag.task("orphan", async () => undefined); + + expect(() => finalizeDag(dag)).toThrowError(/is never called/); + expect(() => finalizeDag(dag)).toThrowError(/is never called/); + }); + + it("surfaces an uncalled task when a bundle reports what it provides", () => { + const dag = new Dag("served_dag"); + dag.task("orphan", async () => undefined); + const bundle = new Bundle(dag); + + expect(() => finalizeBundleDags(bundle)).toThrowError( + /Task "orphan" of Dag "served_dag" is never called/, + ); + }); + + it("rejects a task added after the Dag was read", () => { + const dag = new Dag("closed_dag"); + dag.task("extract", async () => undefined)(); + finalizeDag(dag); + + expect(() => dag.task("late", async () => undefined)).toThrowError( + /Task "late" cannot be added to Dag "closed_dag" after the Dag was read/, + ); + }); + + it("rejects wiring through a factory after the Dag was read", () => { + // A factory outlives the module that built it, so a stray later call has to + // be rejected rather than silently rewiring a Dag Airflow already read. + const dag = new Dag("closed_dag"); + const extract = dag.task("extract", async () => undefined); + extract(); + finalizeDag(dag); + + expect(() => extract()).toThrowError( + /Task "extract" of Dag "closed_dag" was called after the Dag was read/, + ); }); it("retains its spec and each task's handler and spec, copied and frozen", () => { @@ -109,7 +361,7 @@ describe("Dag", () => { const taskSpec = {}; const handler = async () => "hello"; const dag = new Dag("example_dag", dagSpec); - dag.task("my_task", handler, { spec: taskSpec }); + dag.task("my_task", handler, { spec: taskSpec })(); expect(dag.dagId).toBe("example_dag"); expect(dag.spec).toEqual(dagSpec); @@ -155,9 +407,9 @@ describe("Dag", () => { }); it.each([ - ["a misspelled inputs key", { input: {} }], + ["inputs, which the factory call now carries", { inputs: {} }], ["a misspelled spec key", { specs: {} }], - ["an upstream handle passed positionally", { upstream: { dagId: "d", taskId: "t" } }], + ["an upstream reference passed as an option", { upstream: { dagId: "d", taskId: "t" } }], ])("rejects %s in the task options", (_label, options) => { const dag = new Dag("example_dag"); expect(() => @@ -169,7 +421,7 @@ describe("Dag", () => { it.each([ ["null", null], ["an array", []], - ["a string", "inputs"], + ["a string", "spec"], ["a non-plain object", new Date()], ])("rejects task options that are not an options object: %s", (_label, options) => { const dag = new Dag("example_dag"); @@ -178,16 +430,6 @@ describe("Dag", () => { ).toThrowError(/options for Dag "example_dag" task "transform" must be an object/); }); - it("rejects task inputs that are not a plain object", () => { - const dag = new Dag("example_dag"); - expect(() => - dag.task("transform", async () => undefined, { - inputs: new Date() as unknown as Record, - }), - ).toThrowError(/inputs for Dag "example_dag" task "transform" must be an object/); - expect(dag.taskIds).toEqual([]); - }); - it("rejects duplicate taskIds within a Dag", () => { const dag = new Dag("example_dag"); dag.task("dup", async () => undefined); @@ -199,8 +441,8 @@ describe("Dag", () => { const second = async () => "second"; const firstDag = new Dag("first_dag"); const secondDag = new Dag("second_dag"); - firstDag.task("extract", first); - secondDag.task("extract", second); + firstDag.task("extract", first)(); + secondDag.task("extract", second)(); const bundle = new Bundle(); bundle.register(firstDag, secondDag); @@ -211,7 +453,7 @@ describe("Dag", () => { it("accepts a Unicode dagId that Python's word-character rule allows", () => { const handler = async () => undefined; const dag = new Dag("café_dag"); - dag.task("任務", handler); + dag.task("任務", handler)(); const bundle = new Bundle(); bundle.register(dag); expect(bundle.getTaskHandler("café_dag", "任務")).toBe(handler); @@ -226,7 +468,7 @@ describe("Dag", () => { it("treats a dotted TaskGroup taskId as a single taskId (group.task)", () => { const dag = new Dag("example_dag"); - dag.task("transforms.normalize", async () => "ok"); + dag.task("transforms.normalize", async () => "ok")(); const bundle = new Bundle(); bundle.register(dag); expect(bundle.getTaskHandler("example_dag", "transforms.normalize")).toBeDefined(); diff --git a/ts-sdk/tests/sdk/task-handler.test.ts b/ts-sdk/tests/sdk/task-handler.test.ts index 78bb15418028a..bf406321f0ed9 100644 --- a/ts-sdk/tests/sdk/task-handler.test.ts +++ b/ts-sdk/tests/sdk/task-handler.test.ts @@ -19,7 +19,7 @@ import { describe, expect, it } from "vitest"; -import { Bundle, listBundleDags, listBundleTasks } from "../../src/sdk/bundle.js"; +import { Bundle, bundleDagTaskIds } from "../../src/sdk/bundle.js"; import { Dag } from "../../src/sdk/dag.js"; import { getTaskHandlerFunction, TaskHandler } from "../../src/sdk/task-handler.js"; @@ -66,7 +66,7 @@ describe("TaskHandler", () => { describe("a bundle of task handlers", () => { it("registers Dags and task handlers in one call", () => { const nativeDag = new Dag("native_etl"); - nativeDag.task("extract", async () => undefined); + nativeDag.task("extract", async () => undefined)(); const transform = async () => "transformed"; const bundle = new Bundle(); @@ -94,9 +94,9 @@ describe("a bundle of task handlers", () => { expect(bundle.getTaskHandler("etl", "unknown")).toBeUndefined(); }); - it("lists every Dag it provides for, in registration order", () => { + it("lists every Dag it provides for, the ones declared in TypeScript first", () => { const nativeDag = new Dag("native_etl"); - nativeDag.task("extract", async () => undefined); + nativeDag.task("extract", async () => undefined)(); const bundle = new Bundle( new TaskHandler("py_etl", "transform", async () => undefined), @@ -104,16 +104,12 @@ describe("a bundle of task handlers", () => { new TaskHandler("py_etl", "report", async () => undefined), ); - // A Dag registered through handlers keeps its place from the first handler - // that named it, so a later one does not reorder the manifest. - expect(listBundleDags(bundle)).toEqual([ - { dagId: "py_etl", tasks: ["transform", "report"] }, - { dagId: "native_etl", tasks: ["extract"] }, - ]); - expect(listBundleTasks(bundle)).toEqual([ - { dagId: "py_etl", taskId: "transform" }, - { dagId: "py_etl", taskId: "report" }, - { dagId: "native_etl", taskId: "extract" }, + // Compared as entries rather than as a Map, which would ignore the order. + // A Dag its handlers name keeps its place from the first handler that named + // it, so a later one does not reorder the manifest. + expect([...bundleDagTaskIds(bundle)]).toEqual([ + ["native_etl", ["extract"]], + ["py_etl", ["transform", "report"]], ]); }); @@ -122,7 +118,7 @@ describe("a bundle of task handlers", () => { bundle.register(new TaskHandler("py_etl", "transform", async () => undefined)); bundle.register(new TaskHandler("py_etl", "report", async () => undefined)); - expect(listBundleDags(bundle)).toEqual([{ dagId: "py_etl", tasks: ["transform", "report"] }]); + expect(bundleDagTaskIds(bundle)).toEqual(new Map([["py_etl", ["transform", "report"]]])); }); it("rejects a second handler for the same Dag and task", () => { @@ -150,7 +146,7 @@ describe("a bundle of task handlers", () => { new TaskHandler("etl", "transform", async () => undefined), ), ).toThrowError(/already registered/); - expect(listBundleDags(bundle)).toEqual([]); + expect(bundleDagTaskIds(bundle)).toEqual(new Map()); expect(bundle.getTaskHandler("etl", "transform")).toBeUndefined(); }); @@ -180,6 +176,42 @@ describe("a bundle of task handlers", () => { ); }); + it("rejects a Dag whose ID the same call already gave a handler", () => { + // The mirror of the Dag-first case, which the validation pass has to catch + // too: writing the Dag entry would otherwise replace the handler map and + // drop the handler, leaving its task with no body at run time. + const dag = new Dag("py_etl"); + dag.task("extract", async () => undefined)(); + expect( + () => new Bundle(new TaskHandler("py_etl", "legacy", async () => undefined), dag), + ).toThrowError(/already has registered task handlers/); + }); + + it("registers neither kind when a mixed call throws", () => { + const bundle = new Bundle(); + const dag = new Dag("py_etl"); + dag.task("extract", async () => undefined)(); + expect(() => + bundle.register(new TaskHandler("py_etl", "legacy", async () => undefined), dag), + ).toThrowError(/already has registered task handlers/); + + expect(bundleDagTaskIds(bundle)).toEqual(new Map()); + expect(bundle.getTaskHandler("py_etl", "legacy")).toBeUndefined(); + }); + + it("keeps the handlers it already held when a later call throws", () => { + const bundle = new Bundle(new TaskHandler("py_etl", "transform", async () => undefined)); + expect(() => + bundle.register( + new TaskHandler("py_etl", "report", async () => undefined), + new TaskHandler("py_etl", "report", async () => undefined), + ), + ).toThrowError(/already registered/); + + expect(bundleDagTaskIds(bundle)).toEqual(new Map([["py_etl", ["transform"]]])); + expect(bundle.getTaskHandler("py_etl", "report")).toBeUndefined(); + }); + it("names the duplicate-copy cause for a handler carrying the brand but not this class", () => { // Stands in for a handler from a second resolved copy: same brand, other // class. Its private function field is unreadable here, so the point is