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 3571d7bf519c4..b93be7b38cefc 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -86,13 +86,14 @@ value routes the task to the Node.js coordinator. TypeScript implementation ~~~~~~~~~~~~~~~~~~~~~~~~~ -A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Register it with the -``dag_id`` and ``task_id`` it implements, then start the coordinator runtime; the registrations and the -top-level ``await startCoordinator()`` make the module a runnable bundle entry point. +A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Create a ``Dag`` with +the ``dag_id`` it implements, attach each handler with ``dag.task``, collect the Dags in a ``DagRegistry``, +then serve them to Airflow with ``serveDags``; that top-level ``await`` makes the module a runnable bundle +entry point. .. code-block:: typescript - import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; + import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; export async function buildMessage({ ctx, client }: TaskHandlerArgs) { const upstream = await client.getXCom({ @@ -103,12 +104,22 @@ top-level ``await startCoordinator()`` make the module a runnable bundle entry p return `${greeting ?? "hello from TypeScript"}; upstream=${upstream ?? "missing"}`; } - registerTask({ dagId: "typescript_example", taskId: "build_message" }, buildMessage); + const dag = new Dag("typescript_example"); + dag.task("build_message", buildMessage); - await startCoordinator(); + await serveDags(new DagRegistry(dag)); -The ``dagId`` passed to ``registerTask`` must match the ``dag_id`` of the Python Dag, and each ``taskId`` -must match a ``@task.stub`` function in that Dag. +The ``dagId`` passed to ``new Dag(...)`` must match the ``dag_id`` of the Python Dag, and each ``taskId`` +passed to ``dag.task`` must match a ``@task.stub`` function in that Dag. The registry passed to +``serveDags`` is the bundle's complete set of Dags; a second ``serveDags`` call is rejected. A Dag left out +of the registry is not part of the packed bundle, and its tasks are marked removed at runtime. + +``DagRegistry`` holds no sockets and starts nothing, so a unit test can build one and dispatch a handler +through ``registry.getTaskHandler(dagId, taskId)`` without a coordinator runtime. A bundle that collects +its Dags across several modules can add them incrementally with ``registry.register(...)``. + +``new Dag`` and ``dag.task`` take a trailing options object — ``spec`` on both, plus ``inputs`` on a task. +These are not used yet; do not set them. Any other key is rejected. .. note:: diff --git a/ts-sdk/README.md b/ts-sdk/README.md index 73cfdcab64153..45ea2bc31a98a 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -33,14 +33,17 @@ runtime used to execute registered TypeScript handlers from Airflow. ## Task Handlers ```ts -import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; +import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; export async function sayHello({ ctx, client }: TaskHandlerArgs) { const greeting = await client.getVariable("greeting"); return { message: `Hello from ${ctx.taskId}: ${greeting}` }; } -registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello); +const dag = new Dag("example_dag"); +dag.task("say_hello", sayHello); + +await serveDags(new DagRegistry(dag)); ``` Non-`undefined` return values are pushed to XCom under the `"return_value"` @@ -95,7 +98,7 @@ Airflow metadata in the bundle itself. TypeScript entrypoint: ```ts -import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; +import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; export async function extract({ client }: TaskHandlerArgs) { const connection = await client.getConnection("sales_db"); @@ -118,33 +121,49 @@ export async function transform({ client }: TaskHandlerArgs) { }; } -registerTask({ dagId: "sales_pipeline", taskId: "extract" }, extract); -registerTask({ dagId: "sales_pipeline", taskId: "transform" }, transform); +const salesPipeline = new Dag("sales_pipeline"); +salesPipeline.task("extract", extract); +salesPipeline.task("transform", transform); -await startCoordinator(); +await serveDags(new DagRegistry(salesPipeline)); ``` The Python stub defines the Dag dependency graph. The TypeScript handler does -the work and uses `TaskClient` for task-time Airflow data access. Register each -handler with the Python Dag's `dag_id` and the stub task's `task_id`. The -handler function is the reusable task implementation; `registerTask` binds that -handler to a Python stub Dag/task identity for coordinator mode. +the work and uses `TaskClient` for task-time Airflow data access. Create a +`Dag` with the Python Dag's `dag_id` and attach each handler with the stub +task's `task_id`. The handler function is the reusable task implementation; +`dag.task` binds that handler to a Python stub task identity, a `DagRegistry` +collects the Dags this bundle can execute, and `serveDags` serves them to +Airflow. + +`serveDags` is the entrypoint, and the registry it is given is the whole bundle: +a Dag left out of the registry is not part of the bundle, and its tasks are +marked removed at runtime. The registry itself holds no sockets and starts +nothing, so a unit test can build one and dispatch through +`registry.getTaskHandler(dagId, taskId)` without any runtime involved. -For larger projects, keep one Airflow entrypoint that imports every module that -registers tasks, then starts the coordinator: +`new Dag` and `dag.task` take a trailing options object — `spec` on both, plus +`inputs` on a task. These are not used yet; do not set them. + +For larger projects, declare each Dag in its own module and keep one Airflow +entrypoint that serves them all: ```ts -import "./sales/tasks"; -import "./billing/tasks"; -import { startCoordinator } from "@apache-airflow/ts-sdk"; +import { salesDag } from "./sales/dag"; +import { billingDag } from "./billing/dag"; +import { DagRegistry, serveDags } from "@apache-airflow/ts-sdk"; -await startCoordinator(); +await serveDags(new DagRegistry(salesDag, billingDag)); ``` +A bundle that collects its Dags across several modules can add them +incrementally with `registry.register(...)` instead of passing them all to the +constructor. + Airflow launches the bundled entrypoint with `--comm=host:port` and -`--logs=host:port`. `startCoordinator()` connects to those sockets, receives -the task startup message, finds the registered handler for the Dag/task pair, -and reports the terminal task state back to Airflow. +`--logs=host:port`. `serveDags()` connects to those sockets, receives the task +startup message, finds the registered handler for the Dag/task pair, and +reports the terminal task state back to Airflow. See [`example/`](example/) for a coordinator-runtime example that packs a bundle with `airflow-ts-pack` and uses a Python stub Dag. diff --git a/ts-sdk/docs/index.md b/ts-sdk/docs/index.md index 6fe813265533b..edda33fe48dc1 100644 --- a/ts-sdk/docs/index.md +++ b/ts-sdk/docs/index.md @@ -34,19 +34,23 @@ The SDK is currently distributed as source in the `ts-sdk/` directory of the Apache Airflow repository. Build it there and add it as a local dependency of your task bundle; it is not yet published to a public npm registry. -Register a task handler. Handlers receive a `TaskContext` and a `TaskClient`; -any non-`undefined` return value is pushed to XCom under the `"return_value"` -key by the active runtime, matching Python `@task` behavior: +Define a Dag and register its task handlers. Handlers receive a `TaskContext` +and a `TaskClient`; any non-`undefined` return value is pushed to XCom under +the `"return_value"` key by the active runtime, matching Python `@task` +behavior: ```ts -import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; +import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; export async function sayHello({ ctx, client }: TaskHandlerArgs) { const greeting = await client.getVariable("greeting"); return { message: `Hello from ${ctx.taskId}: ${greeting}` }; } -registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello); +const dag = new Dag("example_dag"); +dag.task("say_hello", sayHello); + +await serveDags(new DagRegistry(dag)); ``` ## Coordinators diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts index 1c480b8056b85..d08530615e57f 100644 --- a/ts-sdk/example/src/main.ts +++ b/ts-sdk/example/src/main.ts @@ -17,9 +17,9 @@ * under the License. */ -import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; +import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; -const DAG_ID = "typescript_example"; +const dag = new Dag("typescript_example"); export async function buildMessage({ client }: TaskHandlerArgs) { const upstream = await client.getXCom({ @@ -49,7 +49,7 @@ export async function readConnection({ client }: TaskHandlerArgs) { }; } -registerTask({ dagId: DAG_ID, taskId: "build_message" }, buildMessage); -registerTask({ dagId: DAG_ID, taskId: "read_connection" }, readConnection); +dag.task("build_message", buildMessage); +dag.task("read_connection", readConnection); -await startCoordinator(); +await serveDags(new DagRegistry(dag)); diff --git a/ts-sdk/src/cli/pack.ts b/ts-sdk/src/cli/pack.ts index 0d863fbe89885..3268ec9aab611 100644 --- a/ts-sdk/src/cli/pack.ts +++ b/ts-sdk/src/cli/pack.ts @@ -129,9 +129,16 @@ function readBundleManifest(bundlePath: string): BundleManifest { maxBuffer: MANIFEST_MAX_BUFFER_BYTES, }); } catch (error) { - throw new Error(`Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`, { - cause: error, - }); + const stderr = (error as { stderr?: string }).stderr ?? ""; + const reported = stderr + .split("\n") + .reverse() + .find((line) => /^\w*Error: /.test(line.trim())) + ?.trim(); + throw new Error( + reported ?? `Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`, + { cause: error }, + ); } // Import-time logging from user code lands on stdout too; pick the sentinel line. @@ -143,20 +150,52 @@ function readBundleManifest(bundlePath: string): BundleManifest { throw new Error(`Bundle produced no ${AIRFLOW_METADATA_FLAG} output`); } - let manifest: BundleManifest; + let parsed: unknown; try { - manifest = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length)) as BundleManifest; + parsed = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length)); } catch (error) { throw new Error(`Bundle produced invalid ${AIRFLOW_METADATA_FLAG} output: ${String(error)}`, { cause: error, }); } - if (!manifest.supervisor_schema_version || !manifest.dags || typeof manifest.dags !== "object") { + if (!isBundleManifest(parsed)) { throw new Error(`Bundle produced incomplete ${AIRFLOW_METADATA_FLAG} output`); } + const manifest = parsed; + // The line is whatever the bundle printed and nothing downstream re-validates + // it, so check each Dag entry down to the task-id element. + for (const [dagId, dag] of Object.entries(manifest.dags)) { + if (dag == null || !isTaskIdList(dag.tasks)) { + throw new Error( + `Bundle produced ${AIRFLOW_METADATA_FLAG} output with a malformed entry for Dag "${dagId}"`, + ); + } + } return manifest; } +// The document is checked before anything is read off it: JSON.parse also yields +// null and primitives, and `null.supervisor_schema_version` would surface as a +// raw TypeError rather than a report about the bundle. +function isBundleManifest(value: unknown): value is BundleManifest { + if (typeof value !== "object" || value === null || Array.isArray(value)) return false; + const { supervisor_schema_version: version, dags } = value as Partial; + return ( + // Rendered into the manifest verbatim, where the schema requires a non-empty + // string, so a truthy number or boolean would travel to Airflow as-is. + typeof version === "string" && + version.length > 0 && + typeof dags === "object" && + dags !== null && + // An array would pass the typeof check and yield Dags named "0", "1", ... + !Array.isArray(dags) + ); +} + +function isTaskIdList(value: unknown): value is string[] { + return Array.isArray(value) && value.every((item) => typeof item === "string" && item.length > 0); +} + // esbuild keeps an entry hashbang as line 1, where the metadata comment must go; // NodeCoordinator always runs the bundle through `node`, so drop it. function stripShebang(bundle: string): string { @@ -193,10 +232,16 @@ export async function runPack(argv: readonly string[]): Promise { }); const manifest = readBundleManifest(stagingPath); - if (Object.keys(manifest.dags).length === 0) { - throw new Error( - `${args.entry} registered no tasks; call registerTask(...) before startCoordinator()`, - ); + const dagEntries = Object.entries(manifest.dags); + if (dagEntries.length === 0) { + throw new Error(`${args.entry} served no Dags; pass them to serveDags(new DagRegistry(...))`); + } + // Warn rather than fail, as airflow-go-pack does: the shared schema allows a + // Dag with no tasks. + for (const [dagId, dag] of dagEntries) { + if (dag.tasks.length === 0) { + process.stderr.write(`warning: dag ${JSON.stringify(dagId)} has no tasks\n`); + } } const metadataYaml = renderMetadataYaml({ diff --git a/ts-sdk/src/coordinator/index.ts b/ts-sdk/src/coordinator/index.ts index 815e0c849b972..f06214a4c2f4e 100644 --- a/ts-sdk/src/coordinator/index.ts +++ b/ts-sdk/src/coordinator/index.ts @@ -21,8 +21,11 @@ // // TaskClient and related types are exported from the package root. This // barrel only exports coordinator-specific entry points. +// +// `startCoordinator` is deliberately not exported: Dag authors reach the +// runtime through `serveDags()`, and never name the coordinator itself. -export { startCoordinator, type StartCoordinatorOptions } from "./runtime.js"; +export { serveDags } from "./runtime.js"; /** Cadwyn schema version this SDK was generated against. Not sent on * the wire — exposed so callers can read it for bundle metadata, * health checks, or to confirm which schema their build is pinned to. */ diff --git a/ts-sdk/src/coordinator/manifest.ts b/ts-sdk/src/coordinator/manifest.ts index 32e503dce7a5f..6baca7f4d7b3a 100644 --- a/ts-sdk/src/coordinator/manifest.ts +++ b/ts-sdk/src/coordinator/manifest.ts @@ -18,27 +18,57 @@ */ import { SUPERVISOR_API_VERSION } from "./protocol.js"; -import { listRegisteredTasks, type TaskRegistration } from "../sdk/registry.js"; +import { listRegistryDags, type DagRegistry } from "../sdk/registry.js"; export const AIRFLOW_METADATA_FLAG = "--airflow-metadata"; /** Marks the manifest line on stdout, which import-time logging may also reach. */ export const AIRFLOW_METADATA_SENTINEL = "__AIRFLOW_METADATA__ "; +// Mirrors the Python task-SDK KEY_REGEX and validate_key in airflow.sdk.definitions._internal.node. +// Checked here rather than in Dag()/task() registration: that runs on every +// bundle module load, including at task-execution-runtime startup, whereas +// this only runs once, when the bundle is packed. +const KEY_REGEX = /^[\p{L}\p{N}_.-]+$/u; +const MAX_KEY_LENGTH = 250; + +function validateKey(label: string, value: string): void { + if (typeof value !== "string" || !KEY_REGEX.test(value)) { + throw new Error( + `${label} must be made of alphanumeric characters, dashes, dots, and underscores`, + ); + } + if (value.length > MAX_KEY_LENGTH) { + throw new Error(`${label} must be less than ${MAX_KEY_LENGTH} characters, not ${value.length}`); + } +} + /** Bundle manifest fields only the built bundle itself knows: the schema * version it was compiled against and the Dag/task pairs it registered. - * `airflow-ts-pack` runs `node bundle.mjs --airflow-metadata` to read this. */ + * Registered Dags without tasks appear with an empty `tasks` list so + * `airflow-ts-pack` (which runs `node bundle.mjs --airflow-metadata` to + * read this) can warn about them instead of silently dropping them. */ export interface BundleManifest { supervisor_schema_version: string; dags: Record; } -export function buildBundleManifest( - registrations: readonly TaskRegistration[] = listRegisteredTasks(), -): BundleManifest { +export function buildBundleManifest(registry: DagRegistry): BundleManifest { const dags: BundleManifest["dags"] = {}; - for (const { dagId, taskId } of registrations) { - (dags[dagId] ??= { tasks: [] }).tasks.push(taskId); + for (const { dagId, tasks } of listRegistryDags(registry)) { + validateKey(`Dag "${dagId}"`, dagId); + for (const taskId of tasks) { + validateKey(`Task "${taskId}" of Dag "${dagId}"`, taskId); + } + Object.defineProperty(dags, dagId, { + configurable: true, + enumerable: true, + value: { tasks: [...tasks] }, + writable: true, + }); } - return { supervisor_schema_version: SUPERVISOR_API_VERSION, dags }; + return { + supervisor_schema_version: SUPERVISOR_API_VERSION, + dags, + }; } diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index a3a057b4c0176..5f8c1afaf25e1 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -24,8 +24,9 @@ // node my-bundle.mjs --comm=host:port --logs=host:port // // where `my-bundle.mjs` is a user-bundled Node script that imports -// the SDK, calls `registerTask(...)` for each handler, then calls -// `startCoordinator()`. +// the SDK, creates `Dag` objects, attaches a handler per task with +// `dag.task(...)`, collects them in a `DagRegistry`, then awaits +// `serveDags(registry)`. // // Lifecycle: // 1. Parse --comm / --logs from argv @@ -51,13 +52,68 @@ import { type RuntimeTaskState, type StartupDetails, } from "./protocol.js"; -import { getRegisteredTask, listRegisteredTasks } from "../sdk/registry.js"; +import { DagRegistry, isDagRegistry, listRegistryTasks } from "../sdk/registry.js"; +import { DUPLICATE_COPY_HINT } from "../sdk/brand.js"; import type { TaskContext, TaskHandlerArgs } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; export const ABORT_GRACE_PERIOD_MS = 30_000; export const COORDINATOR_RESPONSE_TIMEOUT_MS = 30_000; +// What must not happen twice is one process connecting two pairs of sockets, so +// this is keyed globally rather than held in a module variable: two resolved +// copies of the package would each get their own, and both would be first. +const SERVED = Symbol.for("airflow.ts-sdk.served"); + +function serveLatch(): Record { + return globalThis as unknown as Record; +} + +/** + * Serve a bundle's Dags to Airflow. The entry point of a TypeScript Dag bundle. + * + * Build the Dags, attach a handler per task with `dag.task(...)`, collect them + * in a {@link DagRegistry}, then await this at module top level: + * + * ```ts + * const dag = new Dag("my_dag"); + * dag.task("extract", extractFn); + * await serveDags(new DagRegistry(dag)); + * ``` + * + * The registry is the bundle's complete set of Dags: this process serves one + * supervisor request, so a second call — which would connect a second pair of + * sockets — is rejected. A call that fails is not a serve, and may be retried. + * Resolves when Airflow's supervisor has been sent the terminal frame for the + * work this process was started for; the same call also answers the build-time + * `--airflow-metadata` query `airflow-ts-pack` makes. + */ +export async function serveDags(registry: DagRegistry): Promise { + // Checked here rather than at first use: a bad argument otherwise surfaces + // only after the sockets are up, as a missing method deep in the runtime. + if (!(registry instanceof DagRegistry)) { + throw new Error( + isDagRegistry(registry) + ? `The registry passed to serveDags(...) ${DUPLICATE_COPY_HINT}` + : "serveDags(...) takes a DagRegistry; build one with new DagRegistry(...dags)", + ); + } + const latch = serveLatch(); + if (latch[SERVED]) { + throw new Error("serveDags(...) was already called; serve every Dag from a single registry"); + } + // Set before the first await, so two concurrent calls cannot both pass. + latch[SERVED] = true; + try { + await startCoordinator(registry); + } catch (err) { + // startCoordinator closes both sockets on its way out, so a failed serve + // holds nothing open and may be retried. + latch[SERVED] = undefined; + throw err; + } +} + /** Options for `startCoordinator()`. */ export interface StartCoordinatorOptions { /** Comm socket address (host:port). Must be supplied together with `logsAddr`; otherwise parsed from argv. */ @@ -108,12 +164,21 @@ export function parseArgs(argv: readonly string[]): ParsedArgs { return { commAddr, logsAddr }; } -/** Start the coordinator runtime. Resolves when the subprocess has - * delivered its terminal frame and closed both sockets. */ -export async function startCoordinator(opts: StartCoordinatorOptions = {}): Promise { +/** Start the coordinator runtime, dispatching to `registry`. Resolves when the + * subprocess has delivered its terminal frame and closed both sockets. + * + * Internal: `serveDags()` is the entry point Dag authors call, so this is + * deliberately absent from the package `"exports"` map. Tests drive it + * directly to supply explicit socket addresses. */ +export async function startCoordinator( + registry: DagRegistry, + opts: StartCoordinatorOptions = {}, +): Promise { const argv = opts.argv ?? process.argv; if (argv.includes(AIRFLOW_METADATA_FLAG)) { - process.stdout.write(`${AIRFLOW_METADATA_SENTINEL}${JSON.stringify(buildBundleManifest())}\n`); + process.stdout.write( + `${AIRFLOW_METADATA_SENTINEL}${JSON.stringify(buildBundleManifest(registry))}\n`, + ); return; } const parsed = @@ -135,7 +200,7 @@ export async function startCoordinator(opts: StartCoordinatorOptions = {}): Prom const runtimeLogs = logs.child("runtime"); runtimeLogs.debug("Connecting log socket", { logs_addr: parsed.logsAddr }); await logs.connect(parsed.logsAddr); - const tasks = listRegisteredTasks(); + const tasks = listRegistryTasks(registry); runtimeLogs.info("Coordinator runtime started", { registered_tasks: tasks, count: tasks.length, @@ -157,7 +222,7 @@ export async function startCoordinator(opts: StartCoordinatorOptions = {}): Prom file: body.file, bundle_path: body.bundle_path, }); - const response = handleParse(body, runtimeLogs); + const response = handleParse(body, registry, runtimeLogs); await sendSupervisorResponse(firstFrame.id, response, comm, runtimeLogs); } else if (body.type === "StartupDetails") { runtimeLogs.info("Received task startup details", { @@ -170,6 +235,7 @@ export async function startCoordinator(opts: StartCoordinatorOptions = {}): Prom }); const response = await handleTask( body, + registry, comm, runtimeLogs, logs.child("client"), @@ -250,12 +316,13 @@ export function createRuntimeAbort( function handleParse( request: { file: string; bundle_path: string }, + registry: DagRegistry, logs: LogChannel, ): RuntimeDagFileParsingResult { // TypeScript-native Dag parsing is not yet supported. // Respond with an empty result so the Python-stub-Dag workflow works. logs.info("Parse-mode response (TS Dag parsing not yet supported)", { - registered_tasks: listRegisteredTasks(), + registered_tasks: listRegistryTasks(registry), }); const response: RuntimeDagFileParsingResult = { type: "DagFileParsingResult", @@ -267,19 +334,20 @@ function handleParse( async function handleTask( details: StartupDetails, + registry: DagRegistry, comm: CommChannel, logs: LogChannel, clientLogs: LogChannel, signal: AbortSignal, ): Promise { const ti = details.ti; - const handler = getRegisteredTask(ti.dag_id, ti.task_id); + const handler = registry.getTaskHandler(ti.dag_id, ti.task_id); if (!handler) { logs.warning("No handler registered for task", { dag_id: ti.dag_id, task_id: ti.task_id, - available: listRegisteredTasks(), + available: listRegistryTasks(registry), }); // 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 0a63cb263f400..a7c8f4960b0ea 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -17,11 +17,11 @@ * under the License. */ -export { registerTask, listRegisteredTasks } from "./sdk/registry.js"; +export { Dag } from "./sdk/dag.js"; +export { DagRegistry } from "./sdk/registry.js"; export { ConnectionNotFoundError, VariableNotFoundError } from "./sdk/client.js"; -export { startCoordinator, SUPERVISOR_API_VERSION } from "./coordinator/index.js"; +export { serveDags, SUPERVISOR_API_VERSION } from "./coordinator/index.js"; +export type { DagSpec, 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 { StartCoordinatorOptions } from "./coordinator/index.js"; -export type { TaskRegistration } from "./sdk/registry.js"; export type { TaskContext, TaskHandler, TaskHandlerArgs } from "./sdk/task.js"; diff --git a/ts-sdk/src/sdk/brand.ts b/ts-sdk/src/sdk/brand.ts new file mode 100644 index 0000000000000..b9fab223862f4 --- /dev/null +++ b/ts-sdk/src/sdk/brand.ts @@ -0,0 +1,43 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Internal: tells an SDK object apart from a stray one, across copies. +// +// Brands live in the cross-realm symbol registry, which two resolved copies of +// this package share even though each gets its own class object. That does not +// make another copy's objects usable — `Dag` and `DagRegistry` read private +// state keyed to the class that declared it — so callers still guard with +// `instanceof` and use these only to tell the two failures apart. + +const PREFIX = "airflow.ts-sdk."; + +/** Mark `target` as built by this package. Not a declared field, so it stays + * out of the public type. */ +export function brand(target: object, name: string): void { + Object.defineProperty(target, Symbol.for(PREFIX + name), { value: true }); +} + +/** Whether `value` carries `name`'s brand, from any copy of this package. */ +export function hasBrand(value: unknown, name: string): boolean { + return typeof value === "object" && value !== null && Symbol.for(PREFIX + name) in value; +} + +/** Tail shared by the errors reporting a second resolved copy. */ +export const DUPLICATE_COPY_HINT = + "comes from a different copy of @apache-airflow/ts-sdk; deduplicate the dependency so one copy is resolved"; diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts new file mode 100644 index 0000000000000..b00492ffd877b --- /dev/null +++ b/ts-sdk/src/sdk/dag.ts @@ -0,0 +1,245 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// The Dag authoring surface: `new Dag(dagId)` plus `dag.task(taskId, handler)`. + +import { brand, hasBrand } from "./brand.js"; +import type { TaskHandler } from "./task.js"; + +function isPlainRecord(value: unknown): value is Record { + if (typeof value !== "object" || value === null || Array.isArray(value)) return false; + const prototype = Object.getPrototypeOf(value); + return prototype === Object.prototype || prototype === null; +} + +function validateEmptySpec(name: string, value: unknown): void { + if (!isPlainRecord(value) || Reflect.ownKeys(value).length > 0) { + throw new Error(`${name} must be an empty object`); + } +} + +/** + * Dag-level options. + * + * No fields yet, so only `{}` is accepted: a field that would be silently + * dropped — `new Dag("d", { schedule: "@daily" })` — is a compile error. + * + * Native Dag declaration will add optional fields here, generated from the + * serialized-Dag JSON schema as `src/generated/supervisor.ts` is. + */ +export type DagSpec = Record; + +/** + * Task-level options. + * + * No fields yet, so only `{}` is accepted, as with {@link DagSpec}. + * + * Future task fields (retries, ...) will land here. + */ +export type TaskSpec = Record; + +/** + * A reference to a task registered on a {@link Dag}, returned by `dag.task(...)`. + * + * Identity only: the handler is deliberately not exposed. + * + * References are what `inputs` accepts, and what native Dag declaration will + * use to wire dependencies. + */ +export interface TaskRef { + /** Identifier of the Dag this task belongs to. */ + readonly dagId: string; + /** Airflow task ID, including any TaskGroup prefix. */ + readonly taskId: string; +} + +/** + * Task references keyed by input name. + * + * 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. + */ +export type TaskInputs = Readonly>; + +/** + * 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. + */ +export interface TaskOptions { + /** + * References to the upstream tasks this task consumes. + * + * Not used yet: a handler receives `{ctx, client}` only, and the Python stub + * Dag defines task order. Read an upstream return value explicitly instead — + * `client.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. + */ + readonly inputs?: TaskInputs; + /** Task-level options. Stored, but not used yet — see {@link TaskSpec}. */ + readonly spec?: TaskSpec; +} + +/** Per-task record a Dag retains: the handle, the handler, its spec, and the + * upstream handles feeding it. */ +export interface TaskRecord { + readonly task: TaskRef; + readonly handler: TaskHandler; + 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. +let taskRecordsOf: (dag: Dag) => ReadonlyMap; + +/** Internal: whether `value` is a Dag built by any copy of this package. */ +export function isDag(value: unknown): value is Dag { + return hasBrand(value, "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. + * + * Constructing a Dag has no effect beyond the instance itself. Collect the ones + * a bundle should serve in a `DagRegistry` and pass it to `serveDags(...)`. + */ +export class Dag { + /** Identifier of this Dag. Must match the Python Dag's `dag_id`. */ + readonly dagId: string; + /** Dag-level options this instance was constructed with, copied and frozen. */ + readonly spec: DagSpec; + readonly #tasks = new Map(); + + static { + taskRecordsOf = (dag) => dag.#tasks; + } + + 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. + // Shallow — a nested value in a future generated spec stays mutable. + this.spec = Object.freeze({ ...spec }); + } + + /** Task IDs attached to this Dag, in attachment order. */ + get taskIds(): readonly string[] { + return [...this.#tasks.keys()]; + } + + /** + * Register a TypeScript handler for a task of this Dag. + * + * `taskId` must match the Dag-side operator's `task_id` exactly, including + * any TaskGroup prefix. Returns this task's handle. + */ + task( + taskId: string, + handler: TaskHandler, + options: TaskOptions = {}, + ): TaskRef { + 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}"`); + } + this.#validateOptions(taskId, options); + const { inputs = {}, spec = {} } = options; + validateEmptySpec(`spec for Dag "${this.dagId}" task "${taskId}"`, spec); + this.#validateInputs(taskId, inputs); + const task: TaskRef = Object.freeze({ dagId: this.dagId, taskId }); + this.#tasks.set(taskId, { + task, + handler: handler as TaskHandler, + spec: Object.freeze({ ...spec }), + inputs: Object.freeze({ ...inputs }), + }); + return task; + } + + // 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") { + 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`); + } + for (const [name, upstream] of Object.entries(inputs)) { + if ( + upstream == null || + typeof upstream.dagId !== "string" || + typeof upstream.taskId !== "string" + ) { + throw new Error( + `Input "${name}" of task "${taskId}" must be a task handle returned by dag.task(...)`, + ); + } + if (upstream.dagId !== this.dagId) { + throw new Error( + `Input "${name}" of task "${taskId}" comes from Dag "${upstream.dagId}", not "${this.dagId}"`, + ); + } + // 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)) { + throw new Error( + `Input "${name}" of task "${taskId}" refers to unregistered task "${upstream.taskId}"`, + ); + } + } + } +} + +/** + * Internal: the task records of a Dag, for registry lookups. + * + * Not re-exported from the package root, and the package `"exports"` map + * blocks deep imports, so this is unreachable from outside the SDK. + */ +export function getDagTaskRecords(dag: Dag): ReadonlyMap { + return taskRecordsOf(dag); +} diff --git a/ts-sdk/src/sdk/registry.ts b/ts-sdk/src/sdk/registry.ts index 122be56e02dfa..ee4286de926fd 100644 --- a/ts-sdk/src/sdk/registry.ts +++ b/ts-sdk/src/sdk/registry.ts @@ -17,85 +17,108 @@ * under the License. */ +import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; +import { Dag, getDagTaskRecords, isDag, type TaskRef } from "./dag.js"; import type { TaskHandler } from "./task.js"; -// Mirrors the Python task-SDK KEY_REGEX and validate_key in airflow.sdk.definitions._internal.node. -const KEY_REGEX = /^[\p{L}\p{N}_.-]+$/u; -const MAX_KEY_LENGTH = 250; +// Assigned inside DagRegistry's static block, as Dag does for its tasks. +let dagsOf: (registry: DagRegistry) => ReadonlyMap; -function validateKey(name: string, value: string): void { - if (typeof value !== "string" || !KEY_REGEX.test(value)) { - throw new Error( - `${name} must be made of alphanumeric characters, dashes, dots, and underscores`, - ); - } - if (value.length > MAX_KEY_LENGTH) { - throw new Error(`${name} must be less than ${MAX_KEY_LENGTH} characters, not ${value.length}`); - } +/** Internal: whether `value` is a DagRegistry built by any copy of this package. */ +export function isDagRegistry(value: unknown): value is DagRegistry { + return hasBrand(value, "DagRegistry"); } -/** Identifies the Airflow task handled by a TypeScript function. */ -export interface TaskRegistration { - /** Identifier of the Dag containing this task. */ +/** Internal: a registered Dag with its task IDs, as {@link listRegistryDags} reports it. + * A task-less Dag is included, so the bundle manifest keeps it visible. */ +export interface RegisteredDag { + /** Identifier of the registered Dag. */ readonly dagId: string; - /** Airflow task ID, including any TaskGroup prefix. */ - readonly taskId: string; + /** Airflow task IDs, including any TaskGroup prefix. */ + readonly tasks: string[]; } -/** Registry of TypeScript task handlers keyed by Dag ID and task ID. */ -export class TaskRegistry { - readonly #tasks = new Map>(); +/** + * The Dags a bundle process can execute, keyed by Dag ID. + * + * This is what a bundle entry point builds and hands to `serveDags(registry)`: + * + * ```ts + * const dag = new Dag("my_dag"); + * dag.task("extract", extractFn); + * await serveDags(new DagRegistry(dag)); + * ``` + * + * It holds no sockets and starts nothing, so a test can build one and invoke a + * handler through {@link getTaskHandler} without any runtime in scope. + * + * Lookups delegate live to each Dag's task map, so tasks added to a Dag + * after registration are visible — the registry records Dag identity, not + * a snapshot of its tasks. + */ +export class DagRegistry { + readonly #dags = new Map(); - /** - * Register a TypeScript handler for an Airflow task. - * - * `dagId` must match the Python Dag's `dag_id`. `taskId` must match the - * Dag-side operator's `task_id` exactly, including any TaskGroup prefix. - */ - register(registration: TaskRegistration, handler: TaskHandler): void { - const { dagId, taskId } = registration; - validateKey("dagId", dagId); - validateKey("taskId", taskId); - if (typeof handler !== "function") { - throw new Error(`handler for Dag "${dagId}" task "${taskId}" must be a function`); - } - const dagTasks = this.#tasks.get(dagId) ?? new Map(); - if (dagTasks.has(taskId)) { - throw new Error(`Task "${taskId}" is already registered for Dag "${dagId}"`); - } - dagTasks.set(taskId, handler as TaskHandler); - this.#tasks.set(dagId, dagTasks); + static { + dagsOf = (registry) => registry.#dags; } - /** Look up a registered handler. Returns `undefined` when no handler exists. */ - get(dagId: string, taskId: string): TaskHandler | undefined { - return this.#tasks.get(dagId)?.get(taskId); + /** Registers `dags`, on the same terms as {@link register}. */ + constructor(...dags: Dag[]) { + brand(this, "DagRegistry"); + this.register(...dags); } - /** List all registered tasks. */ - list(): TaskRegistration[] { - return [...this.#tasks.entries()].flatMap(([dagId, tasks]) => - [...tasks.keys()].map((taskId) => ({ dagId, taskId })), - ); + /** Register Dags. Registering an already-registered `dagId` throws, + * and a call that throws registers none of its Dags. + * + * The constructor covers the common case; this is for a bundle that + * collects its Dags across several modules. */ + register(...dags: Dag[]): void { + const incoming = new Set(); + for (const dag of dags) { + // Typed as Dag, so narrowing it would collapse to never; these guard + // callers reaching this from plain JavaScript. + const candidate: unknown = dag; + // Another copy's Dag cannot be registered — lookups read a private task + // map keyed to this copy's class — so it is rejected, but by its cause. + if (!(candidate instanceof Dag)) { + throw new Error( + isDag(candidate) + ? `Dag "${candidate.dagId}" ${DUPLICATE_COPY_HINT}` + : "only Dag instances can be registered", + ); + } + if (this.#dags.has(dag.dagId) || incoming.has(dag.dagId)) { + throw new Error(`Dag "${dag.dagId}" is already registered`); + } + incoming.add(dag.dagId); + } + for (const dag of dags) { + this.#dags.set(dag.dagId, dag); + } } -} -const defaultRegistry = new TaskRegistry(); - -/** Register a TypeScript handler in the default task registry. */ -export function registerTask( - registration: TaskRegistration, - handler: TaskHandler, -): void { - defaultRegistry.register(registration, handler); + /** Look up a registered handler, the way the runtime dispatches a task. + * Returns `undefined` when no handler exists. */ + getTaskHandler(dagId: string, taskId: string): TaskHandler | undefined { + const dag = this.#dags.get(dagId); + return dag ? getDagTaskRecords(dag).get(taskId)?.handler : undefined; + } } -/** Look up a registered handler. Returns `undefined` when no handler exists. */ -export function getRegisteredTask(dagId: string, taskId: string): TaskHandler | undefined { - return defaultRegistry.get(dagId, taskId); +/** Internal: the task handles across a registry's Dags. Not re-exported from the + * package root — enumerating what the runtime dispatches is the runtime's job. */ +export function listRegistryTasks(registry: DagRegistry): TaskRef[] { + return [...dagsOf(registry).values()].flatMap((dag) => + [...getDagTaskRecords(dag).values()].map((record) => record.task), + ); } -/** List all registered tasks. */ -export function listRegisteredTasks(): TaskRegistration[] { - return defaultRegistry.list(); +/** Internal: every registered Dag with its task IDs, empty Dags included. */ +export function listRegistryDags(registry: DagRegistry): RegisteredDag[] { + return [...dagsOf(registry).values()].map((dag) => ({ + dagId: dag.dagId, + tasks: [...dag.taskIds], + })); } diff --git a/ts-sdk/tests/cli/fixtures/empty-entry.ts b/ts-sdk/tests/cli/fixtures/empty-entry.ts index 3e2b582b224cd..c48c6c24cf8fc 100644 --- a/ts-sdk/tests/cli/fixtures/empty-entry.ts +++ b/ts-sdk/tests/cli/fixtures/empty-entry.ts @@ -17,6 +17,6 @@ * under the License. */ -import { startCoordinator } from "../../../src/index.js"; +import { DagRegistry, serveDags } from "../../../src/index.js"; -await startCoordinator(); +await serveDags(new DagRegistry()); diff --git a/ts-sdk/tests/cli/fixtures/entry.ts b/ts-sdk/tests/cli/fixtures/entry.ts index 8d53c43aee500..8549fb56a321e 100644 --- a/ts-sdk/tests/cli/fixtures/entry.ts +++ b/ts-sdk/tests/cli/fixtures/entry.ts @@ -17,10 +17,12 @@ * under the License. */ -import { registerTask, startCoordinator } from "../../../src/index.js"; +import { Dag, DagRegistry, serveDags } from "../../../src/index.js"; -registerTask({ dagId: "fixture_dag", taskId: "extract" }, async () => "extracted"); -registerTask({ dagId: "fixture_dag", taskId: "transform" }, async () => "transformed"); -registerTask({ dagId: "other_dag", taskId: "solo" }, async () => undefined); +const fixtureDag = new Dag("fixture_dag"); +fixtureDag.task("extract", async () => "extracted"); +fixtureDag.task("transform", async () => "transformed"); +const otherDag = new Dag("other_dag"); +otherDag.task("solo", async () => undefined); -await startCoordinator(); +await serveDags(new DagRegistry(fixtureDag, otherDag)); diff --git a/ts-sdk/tests/cli/fixtures/noisy-entry.ts b/ts-sdk/tests/cli/fixtures/noisy-entry.ts index 3b9675619d3c6..dcd7b0d8c0b4d 100644 --- a/ts-sdk/tests/cli/fixtures/noisy-entry.ts +++ b/ts-sdk/tests/cli/fixtures/noisy-entry.ts @@ -18,10 +18,11 @@ * under the License. */ -import { registerTask, startCoordinator } from "../../../src/index.js"; +import { Dag, DagRegistry, serveDags } from "../../../src/index.js"; console.log("noise from an import-time dependency"); -registerTask({ dagId: "noisy_dag", taskId: "only" }, async () => undefined); +const noisyDag = new Dag("noisy_dag"); +noisyDag.task("only", async () => undefined); -await startCoordinator(); +await serveDags(new DagRegistry(noisyDag)); diff --git a/ts-sdk/tests/cli/pack.test.ts b/ts-sdk/tests/cli/pack.test.ts index 4a5c2d572bfb3..6de652e7f2830 100644 --- a/ts-sdk/tests/cli/pack.test.ts +++ b/ts-sdk/tests/cli/pack.test.ts @@ -22,7 +22,7 @@ import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "no import { tmpdir } from "node:os"; import path from "node:path"; import { fileURLToPath } from "node:url"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { EMBEDDED_METADATA_PREFIX, @@ -95,10 +95,26 @@ describe("renderMetadataYaml", () => { }); }); +function readEmbeddedMetadata(bundlePath: string): string { + const firstLine = readFileSync(bundlePath, "utf-8").split("\n")[0]!; + return Buffer.from(firstLine.slice(EMBEDDED_METADATA_PREFIX.length), "base64").toString("utf-8"); +} + +/** Collect what runPack writes to stderr; returns a reader for the text so far. */ +function captureStderr(): () => string { + const chunks: string[] = []; + vi.spyOn(process.stderr, "write").mockImplementation((chunk: string | Uint8Array) => { + chunks.push(typeof chunk === "string" ? chunk : Buffer.from(chunk).toString("utf-8")); + return true; + }); + return () => chunks.join(""); +} + describe("runPack", () => { let outdir: string; afterEach(() => { + vi.restoreAllMocks(); if (outdir) rmSync(outdir, { recursive: true, force: true }); }); @@ -167,9 +183,10 @@ describe("runPack", () => { writeFileSync( entry, [ - `import { registerTask, startCoordinator } from ${JSON.stringify(SDK_INDEX)};`, - 'for (let i = 0; i < 4000; i += 1) registerTask({ dagId: "big_dag", taskId: String(i).padStart(240, "t") }, async () => undefined);', - "await startCoordinator();", + `import { Dag, DagRegistry, serveDags } from ${JSON.stringify(SDK_INDEX)};`, + 'const bigDag = new Dag("big_dag");', + 'for (let i = 0; i < 4000; i += 1) bigDag.task(String(i).padStart(240, "t"), async () => undefined);', + "await serveDags(new DagRegistry(bigDag));", ].join("\n"), ); @@ -180,11 +197,141 @@ describe("runPack", () => { expect(existsSync(path.join(outdir, "bundle.pack-staging.mjs"))).toBe(false); }); - it("leaves no bundle behind when the entry registers no tasks", async () => { + it("leaves no bundle behind when the entry serves no Dags", async () => { + outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + + await expect(runPack([EMPTY_ENTRY, "--outdir", outdir])).rejects.toThrow("served no Dags"); + expect(existsSync(path.join(outdir, "bundle.mjs"))).toBe(false); + expect(existsSync(path.join(outdir, "bundle.pack-staging.mjs"))).toBe(false); + }); + + it.each([ + { + label: "Dag", + dagId: "bad id!", + taskId: "valid_task", + expected: + 'Error: Dag "bad id!" must be made of alphanumeric characters, dashes, dots, and underscores', + }, + { + label: "task", + dagId: "valid_dag", + taskId: "bad id!", + expected: + 'Error: Task "bad id!" of Dag "valid_dag" must be made of alphanumeric characters, dashes, dots, and underscores', + }, + ])( + "reports an invalid $label ID without a staging-bundle stack", + async ({ dagId, taskId, expected }) => { + outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + const entry = path.join(outdir, "invalid-id-entry.ts"); + writeFileSync( + entry, + [ + `import { Dag, DagRegistry, serveDags } from ${JSON.stringify(SDK_INDEX)};`, + `const invalidDag = new Dag(${JSON.stringify(dagId)});`, + `invalidDag.task(${JSON.stringify(taskId)}, async () => undefined);`, + "await serveDags(new DagRegistry(invalidDag));", + ].join("\n"), + ); + + await expect(runPack([entry, "--outdir", outdir])).rejects.toHaveProperty( + "message", + expected, + ); + expect(existsSync(path.join(outdir, "bundle.mjs"))).toBe(false); + expect(existsSync(path.join(outdir, "bundle.pack-staging.mjs"))).toBe(false); + }, + ); + + it("reports the last error from a failed bundle", async () => { outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + const entry = path.join(outdir, "multiple-errors-entry.ts"); + writeFileSync( + entry, + ['console.error("Error: earlier failure");', 'throw new Error("final failure");'].join("\n"), + ); - await expect(runPack([EMPTY_ENTRY, "--outdir", outdir])).rejects.toThrow("registered no tasks"); + await expect(runPack([entry, "--outdir", outdir])).rejects.toHaveProperty( + "message", + "Error: final failure", + ); expect(existsSync(path.join(outdir, "bundle.mjs"))).toBe(false); expect(existsSync(path.join(outdir, "bundle.pack-staging.mjs"))).toBe(false); }); + + // A bundle can print the sentinel itself, so nothing on that line is trusted. + it.each([ + ['{ supervisor_schema_version: "1", dags: { broken_dag: {} } }', "malformed entry"], + [ + '{ supervisor_schema_version: "1", dags: { broken_dag: { tasks: ["ok", 7] } } }', + "malformed entry", + ], + [ + '{ supervisor_schema_version: "1", dags: { broken_dag: { tasks: [""] } } }', + "malformed entry", + ], + ['{ supervisor_schema_version: "1", dags: [{ tasks: ["a"] }] }', "incomplete"], + // Was read off before the document itself was checked, so it surfaced as a + // raw TypeError. + ["null", "incomplete"], + // Truthy, but not the non-empty string the schema requires. + ['{ supervisor_schema_version: true, dags: { d: { tasks: ["a"] } } }', "incomplete"], + ['{ supervisor_schema_version: 20260616, dags: { d: { tasks: ["a"] } } }', "incomplete"], + ])("rejects the metadata line %s", async (manifest, message) => { + outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + const entry = path.join(outdir, "malformed-entry.ts"); + writeFileSync( + entry, + `console.log(${JSON.stringify(AIRFLOW_METADATA_SENTINEL)} + JSON.stringify(${manifest}));`, + ); + + await expect(runPack([entry, "--outdir", outdir])).rejects.toThrow(message); + expect(existsSync(path.join(outdir, "bundle.mjs"))).toBe(false); + expect(existsSync(path.join(outdir, "bundle.pack-staging.mjs"))).toBe(false); + }); + + it("warns but still packs a registered Dag with no tasks, as airflow-go-pack does", async () => { + outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + const entry = path.join(outdir, "mixed-entry.ts"); + writeFileSync( + entry, + [ + `import { Dag, DagRegistry, serveDags } from ${JSON.stringify(SDK_INDEX)};`, + 'const salesDag = new Dag("sales_dag");', + 'salesDag.task("extract", async () => undefined);', + 'await serveDags(new DagRegistry(salesDag, new Dag("empty_dag")));', + ].join("\n"), + ); + const stderr = captureStderr(); + + await runPack([entry, "--outdir", outdir]); + + expect(stderr()).toContain('warning: dag "empty_dag" has no tasks\n'); + expect(readEmbeddedMetadata(path.join(outdir, "bundle.mjs"))).toContain( + ' "empty_dag":\n tasks: []', + ); + }); + + it("packs only the Dags the served registry holds", async () => { + outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-")); + const entry = path.join(outdir, "forgotten-entry.ts"); + writeFileSync( + entry, + [ + `import { Dag, DagRegistry, serveDags } from ${JSON.stringify(SDK_INDEX)};`, + 'const salesDag = new Dag("sales_dag");', + 'salesDag.task("extract", async () => undefined);', + 'const billingDag = new Dag("billing_dag");', + 'billingDag.task("charge", async () => undefined);', + "await serveDags(new DagRegistry(salesDag));", + ].join("\n"), + ); + + await runPack([entry, "--outdir", outdir]); + + const metadata = readEmbeddedMetadata(path.join(outdir, "bundle.mjs")); + expect(metadata).toContain(' "sales_dag":'); + expect(metadata).not.toContain("billing_dag"); + }); }); diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index a580159073ec2..040de00fa9611 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -35,7 +35,15 @@ import { COORDINATOR_RESPONSE_TIMEOUT_MS, startCoordinator, } from "../../src/coordinator/runtime.js"; -import { registerTask } from "../../src/sdk/registry.js"; +import { Dag } from "../../src/sdk/dag.js"; +import { DagRegistry } from "../../src/sdk/registry.js"; + +const testDag = new Dag("test_dag"); +const otherDag = new Dag("other_dag"); +// The registry the runtime dispatches through. startCoordinator() is driven +// directly rather than through serveDags(), so these tests can supply mock +// socket addresses. +const registry = new DagRegistry(testDag, otherDag); interface MockResult { firstResponse: { id: number; body: unknown; isResponse: boolean } | null; @@ -141,7 +149,7 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr const commAccept = acceptOne(comm.server); const logsAccept = acceptOne(logs.server); - const runtimeDone = startCoordinator({ + const runtimeDone = startCoordinator(registry, { commAddr: `127.0.0.1:${comm.port}`, logsAddr: `127.0.0.1:${logs.port}`, argv: [], @@ -212,7 +220,7 @@ describe("coordinator runtime integration", () => { const logsSockPromise = acceptOne(logs.server); const commSockPromise = acceptOne(comm.server); - const runtimeDone = startCoordinator({ + const runtimeDone = startCoordinator(registry, { commAddr: `127.0.0.1:${comm.port}`, logsAddr: `127.0.0.1:${logs.port}`, argv: [], @@ -231,7 +239,7 @@ describe("coordinator runtime integration", () => { it("dispatches StartupDetails to a registered handler and emits SucceedTask", async () => { let observedCtx: unknown = null; - registerTask({ dagId: "test_dag", taskId: "say_hello" }, async ({ ctx }) => { + testDag.task("say_hello", async ({ ctx }) => { observedCtx = ctx; return "ok"; }); @@ -287,8 +295,8 @@ describe("coordinator runtime integration", () => { const commAccept = acceptOne(comm.server); const logsAccept = acceptOne(logs.server); - registerTask({ dagId: "test_dag", taskId: "terminal_timeout" }, async () => undefined); - const runtimeDone = startCoordinator({ + testDag.task("terminal_timeout", async () => undefined); + const runtimeDone = startCoordinator(registry, { commAddr: `127.0.0.1:${comm.port}`, logsAddr: `127.0.0.1:${logs.port}`, argv: [], @@ -317,7 +325,7 @@ describe("coordinator runtime integration", () => { }); it("returns TaskState=failed when the handler throws", async () => { - registerTask({ dagId: "test_dag", taskId: "boom" }, async () => { + testDag.task("boom", async () => { throw new Error("boom"); }); @@ -330,7 +338,7 @@ describe("coordinator runtime integration", () => { }); it("returns RetryTask when the handler throws and Airflow says the failure is retryable", async () => { - registerTask({ dagId: "test_dag", taskId: "boom_retry" }, async () => { + testDag.task("boom_retry", async () => { throw new Error("boom"); }); @@ -349,7 +357,7 @@ describe("coordinator runtime integration", () => { it("aborts ctx.signal on SIGTERM and reports a thrown task error", async () => { let sawAbort = false; - registerTask({ dagId: "test_dag", taskId: "aborted_then_failed" }, async ({ ctx }) => { + testDag.task("aborted_then_failed", async ({ ctx }) => { process.emit("SIGTERM"); sawAbort = ctx.signal.aborted; throw new Error("interrupted"); @@ -367,7 +375,7 @@ describe("coordinator runtime integration", () => { it("returns RetryTask with the thrown error when a task fails after SIGTERM", async () => { let sawAbort = false; - registerTask({ dagId: "test_dag", taskId: "aborted_then_failed_retry" }, async ({ ctx }) => { + testDag.task("aborted_then_failed_retry", async ({ ctx }) => { process.emit("SIGTERM"); sawAbort = ctx.signal.aborted; throw new Error("interrupted"); @@ -390,7 +398,7 @@ describe("coordinator runtime integration", () => { it("does not discard a completed task result after SIGTERM", async () => { let sawAbort = false; - registerTask({ dagId: "test_dag", taskId: "completed_after_sigterm" }, async ({ ctx }) => { + testDag.task("completed_after_sigterm", async ({ ctx }) => { process.emit("SIGTERM"); sawAbort = ctx.signal.aborted; return "completed"; @@ -426,7 +434,7 @@ describe("coordinator runtime integration", () => { const xcomStore = new Map(); let observedGreeting: string | null = ""; - registerTask({ dagId: "test_dag", taskId: "say_hello_client" }, async ({ ctx, client }) => { + testDag.task("say_hello_client", async ({ ctx, client }) => { // The coordinator-mode handler MUST receive a client. if (!client) throw new Error("client missing in coordinator mode"); @@ -495,7 +503,7 @@ describe("coordinator runtime integration", () => { it("returns null from getVariable when the supervisor signals NOT_FOUND", async () => { let observed: string | null = ""; - registerTask({ dagId: "test_dag", taskId: "missing_variable" }, async ({ client }) => { + testDag.task("missing_variable", async ({ client }) => { observed = await client.getVariable("missing_key"); }); @@ -520,10 +528,10 @@ describe("coordinator runtime integration", () => { it("looks up handlers by exact Dag and task id", async () => { let calledFirstDag = false; let calledSecondDag = false; - registerTask({ dagId: "test_dag", taskId: "shared_task" }, async () => { + testDag.task("shared_task", async () => { calledFirstDag = true; }); - registerTask({ dagId: "other_dag", taskId: "shared_task" }, async () => { + otherDag.task("shared_task", async () => { calledSecondDag = true; }); @@ -548,7 +556,7 @@ describe("coordinator runtime integration", () => { }); it("auto-pushes return_value XCom when handler returns a value", async () => { - registerTask({ dagId: "test_dag", taskId: "pusher" }, async () => "my-result"); + testDag.task("pusher", async () => "my-result"); const responder: Responder = (msgType, _body) => { if (msgType === "SetXCom") return { body: null }; @@ -568,7 +576,7 @@ describe("coordinator runtime integration", () => { }); it("does NOT push return_value XCom when handler returns undefined", async () => { - registerTask({ dagId: "test_dag", taskId: "void_task" }, async () => { + testDag.task("void_task", async () => { // no return value }); diff --git a/ts-sdk/tests/coordinator/protocol.test.ts b/ts-sdk/tests/coordinator/protocol.test.ts index 330f94b62b05d..8089a595a5dfb 100644 --- a/ts-sdk/tests/coordinator/protocol.test.ts +++ b/ts-sdk/tests/coordinator/protocol.test.ts @@ -20,6 +20,7 @@ import { describe, expect, it } from "vitest"; import { asMsgFromSupervisor } from "../../src/coordinator/protocol.js"; import { parseArgs, startCoordinator } from "../../src/coordinator/runtime.js"; +import { DagRegistry } from "../../src/sdk/registry.js"; describe("protocol decode", () => { it("accepts StartupDetails", () => { @@ -85,14 +86,14 @@ describe("runtime arg parser", () => { it("requires commAddr and logsAddr overrides to be supplied together", async () => { await expect( - startCoordinator({ + startCoordinator(new DagRegistry(), { commAddr: "127.0.0.1:5001", argv: ["node", "bundle.mjs", "--logs=127.0.0.1:5002"], }), ).rejects.toThrow(/Missing --comm/); await expect( - startCoordinator({ + startCoordinator(new DagRegistry(), { logsAddr: "127.0.0.1:5002", argv: ["node", "bundle.mjs", "--comm=127.0.0.1:5001"], }), diff --git a/ts-sdk/tests/coordinator/public-api.test.ts b/ts-sdk/tests/coordinator/public-api.test.ts index 218fa42508ea4..16e91d4b31f17 100644 --- a/ts-sdk/tests/coordinator/public-api.test.ts +++ b/ts-sdk/tests/coordinator/public-api.test.ts @@ -17,19 +17,15 @@ * under the License. */ -import { describe, expectTypeOf, it } from "vitest"; -import type { StartCoordinatorOptions } from "../../src/coordinator/index.js"; -import { startCoordinator } from "../../src/coordinator/index.js"; +import { describe, expect, expectTypeOf, it } from "vitest"; +import type { DagRegistry } from "../../src/sdk/registry.js"; +import * as coordinator from "../../src/coordinator/index.js"; +import { serveDags } from "../../src/coordinator/index.js"; describe("coordinator public API", () => { - it("exports the coordinator runtime entrypoint from the coordinator subpath", () => { - expectTypeOf().toEqualTypeOf< - (opts?: StartCoordinatorOptions) => Promise - >(); - expectTypeOf().toEqualTypeOf<{ - commAddr?: string; - logsAddr?: string; - argv?: readonly string[]; - }>(); + it("exposes serveDags, not the coordinator itself, from the coordinator subpath", () => { + expectTypeOf().toEqualTypeOf<(registry: DagRegistry) => Promise>(); + expect("startCoordinator" in coordinator).toBe(false); + expectTypeOf().not.toHaveProperty("startCoordinator"); }); }); diff --git a/ts-sdk/tests/coordinator/runtime-manifest.test.ts b/ts-sdk/tests/coordinator/runtime-manifest.test.ts index f6aea6f2fd501..9c8beec94e108 100644 --- a/ts-sdk/tests/coordinator/runtime-manifest.test.ts +++ b/ts-sdk/tests/coordinator/runtime-manifest.test.ts @@ -22,16 +22,21 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { AIRFLOW_METADATA_SENTINEL, buildBundleManifest } from "../../src/coordinator/manifest.js"; import { startCoordinator } from "../../src/coordinator/runtime.js"; import { SUPERVISOR_API_VERSION } from "../../src/coordinator/protocol.js"; +import { Dag } from "../../src/sdk/dag.js"; +import { DagRegistry } from "../../src/sdk/registry.js"; + +function buildDag(dagId: string, ...taskIds: string[]): Dag { + const dag = new Dag(dagId); + for (const taskId of taskIds) { + dag.task(taskId, async () => undefined); + } + return dag; +} describe("buildBundleManifest", () => { - it("groups registrations by Dag ID under the SDK's schema version", () => { - expect( - buildBundleManifest([ - { dagId: "dag_a", taskId: "t1" }, - { dagId: "dag_b", taskId: "t2" }, - { dagId: "dag_a", taskId: "t3" }, - ]), - ).toEqual({ + it("maps a registry's Dags to their tasks under the SDK's schema version", () => { + const registry = new DagRegistry(buildDag("dag_a", "t1", "t3"), buildDag("dag_b", "t2")); + expect(buildBundleManifest(registry)).toEqual({ supervisor_schema_version: SUPERVISOR_API_VERSION, dags: { dag_a: { tasks: ["t1", "t3"] }, @@ -39,6 +44,70 @@ describe("buildBundleManifest", () => { }, }); }); + + it("keeps a registered Dag without tasks visible in the manifest", () => { + expect(buildBundleManifest(new DagRegistry(buildDag("empty_dag"))).dags).toEqual({ + empty_dag: { tasks: [] }, + }); + }); + + it("keeps a Dag named __proto__ visible in serialized metadata", () => { + const manifest = buildBundleManifest(new DagRegistry(buildDag("__proto__", "task"))); + const serializedDags = JSON.parse(JSON.stringify(manifest)).dags; + + expect(Object.keys(serializedDags)).toEqual(["__proto__"]); + expect(serializedDags["__proto__"]).toEqual({ tasks: ["task"] }); + }); + + it("reports only the Dags the registry was given", () => { + const registry = new DagRegistry(buildDag("dag_a", "t1")); + buildDag("dag_b", "t2"); + expect(Object.keys(buildBundleManifest(registry).dags)).toEqual(["dag_a"]); + }); + + it("rejects an empty dagId", () => { + const registry = new DagRegistry(buildDag("")); + expect(() => buildBundleManifest(registry)).toThrowError(/must be made of alphanumeric/); + }); + + it("rejects an empty taskId", () => { + const registry = new DagRegistry(buildDag("example_dag", "")); + expect(() => buildBundleManifest(registry)).toThrowError(/must be made of alphanumeric/); + }); + + it.each([" ", "\t", "my dag", "a/b", "task@1"])( + "rejects a dagId with characters no Python dag_id allows: %j", + (dagId) => { + const registry = new DagRegistry(buildDag(dagId)); + expect(() => buildBundleManifest(registry)).toThrowError(/must be made of alphanumeric/); + }, + ); + + it.each([" ", "\t", "my task", "a/b", "task@1"])( + "rejects a taskId with characters no Python task_id allows: %j", + (taskId) => { + const registry = new DagRegistry(buildDag("example_dag", taskId)); + expect(() => buildBundleManifest(registry)).toThrowError(/must be made of alphanumeric/); + }, + ); + + it("rejects a dagId longer than 250 characters", () => { + const registry = new DagRegistry(buildDag("d".repeat(251))); + expect(() => buildBundleManifest(registry)).toThrowError( + /must be less than 250 characters, not 251/, + ); + }); + + it("rejects a taskId longer than 250 characters", () => { + const registry = new DagRegistry(buildDag("example_dag", "t".repeat(251))); + expect(() => buildBundleManifest(registry)).toThrowError( + /must be less than 250 characters, not 251/, + ); + }); + + it("does not validate key format when Dags are only registered, not packed", () => { + expect(() => new DagRegistry(buildDag("bad dag id", "bad task id"))).not.toThrow(); + }); }); describe("startCoordinator --airflow-metadata", () => { @@ -49,13 +118,15 @@ describe("startCoordinator --airflow-metadata", () => { it("dumps the manifest to stdout and returns without connecting", async () => { const write = vi.spyOn(process.stdout, "write").mockReturnValue(true); - await startCoordinator({ argv: ["node", "bundle.mjs", "--airflow-metadata"] }); + await startCoordinator(new DagRegistry(buildDag("metadata_dag", "only")), { + argv: ["node", "bundle.mjs", "--airflow-metadata"], + }); expect(write).toHaveBeenCalledTimes(1); const written = String(write.mock.calls[0]![0]); expect(written.startsWith(AIRFLOW_METADATA_SENTINEL)).toBe(true); const payload = JSON.parse(written.slice(AIRFLOW_METADATA_SENTINEL.length)); expect(payload.supervisor_schema_version).toBe(SUPERVISOR_API_VERSION); - expect(payload.dags).toEqual({}); + expect(payload.dags).toEqual({ metadata_dag: { tasks: ["only"] } }); }); }); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 4e7a49ddb4df2..a9b2f888019d3 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -17,30 +17,140 @@ * under the License. */ -import { describe, expect, expectTypeOf, it } from "vitest"; +import { afterEach, describe, expect, expectTypeOf, it, vi } from "vitest"; +import { AIRFLOW_METADATA_FLAG } from "../src/coordinator/manifest.js"; import type { ConnectionResult, + DagSpec, GetXComOpts, SetXComOpts, - StartCoordinatorOptions, TaskClient, TaskContext, - TaskRegistration, + TaskHandler, + TaskInputs, + TaskOptions, + TaskRef, + TaskSpec, } from "../src/index.js"; +import * as sdk from "../src/index.js"; import { ConnectionNotFoundError, - listRegisteredTasks, - registerTask, - startCoordinator, + Dag, + DagRegistry, + serveDags, SUPERVISOR_API_VERSION, VariableNotFoundError, } from "../src/index.js"; describe("public API", () => { - it("exports task registration helpers", () => { - const registration = { dagId: "public_api_dag", taskId: "public_api_task" }; - registerTask(registration, async () => undefined); - expect(listRegisteredTasks()).toContainEqual(registration); + 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 }, + }); + 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"]); + // serveDags hands the registry to the runtime, which needs the supervisor's + // socket addresses that Airflow puts on argv. + await expect(serveDags(new DagRegistry(dag))).rejects.toThrow("Missing --comm"); + }); + + // The registry guard runs before the already-served latch, so this holds + // regardless of whether another test in this file already served a registry. + it.each([ + ["a bare Dag", new Dag("not_a_registry_dag")], + ["a plain object", { register: () => {} }], + ["null", null], + ])("rejects a %s in place of a registry", async (_label, value) => { + await expect(serveDags(value as unknown as DagRegistry)).rejects.toThrow( + /serveDags\(\.\.\.\) takes a DagRegistry/, + ); + }); + + it("names the duplicate-copy cause for a registry built by another copy", async () => { + // Stands in for a registry from a second resolved copy: same brand, other + // class. It still cannot be served, so the point is only that it says why. + const foreign = {}; + Object.defineProperty(foreign, Symbol.for("airflow.ts-sdk.DagRegistry"), { value: true }); + await expect(serveDags(foreign as unknown as DagRegistry)).rejects.toThrow( + /different copy of @apache-airflow\/ts-sdk/, + ); + }); + + describe("the one-shot serve latch", () => { + // Global by design, so it outlives the test that trips it. + afterEach(() => { + delete (globalThis as unknown as Record)[ + Symbol.for("airflow.ts-sdk.served") + ]; + }); + + it("rejects a second call once a serve has completed", async () => { + const argv = process.argv; + // The one path that completes without sockets. + process.argv = [...argv, AIRFLOW_METADATA_FLAG]; + vi.spyOn(process.stdout, "write").mockReturnValue(true); + try { + await serveDags(new DagRegistry(new Dag("served_dag"))); + await expect(serveDags(new DagRegistry(new Dag("second_call_dag")))).rejects.toThrow( + /serveDags\(\.\.\.\) was already called/, + ); + } finally { + process.argv = argv; + vi.restoreAllMocks(); + } + }); + + it("releases the latch when a serve fails, so the call can be retried", async () => { + await expect(serveDags(new DagRegistry(new Dag("first_try")))).rejects.toThrow( + "Missing --comm", + ); + // The retry reports why it actually failed, not "already called". + await expect(serveDags(new DagRegistry(new Dag("second_try")))).rejects.toThrow( + "Missing --comm", + ); + }); + }); + + it("exports DagRegistry as the Dag collection a bundle serves", () => { + const dag = new Dag("registry_api_dag"); + const handler = async () => "hello"; + dag.task("extract", handler); + // Building a registry starts nothing, so a test can dispatch through it + // exactly as the runtime does, with no sockets in scope. + const registry = new DagRegistry(dag); + expect(registry.getTaskHandler("registry_api_dag", "extract")).toBe(handler); + registry.register(new Dag("late_dag")); + expect(registry.getTaskHandler("late_dag", "extract")).toBeUndefined(); + }); + + it("keeps registry enumeration out of the public surface", () => { + const registry = new DagRegistry(); + for (const name of ["listTasks", "listDags"]) { + expect(name in registry).toBe(false); + } + expectTypeOf().toEqualTypeOf<"register" | "getTaskHandler">(); + }); + + it("does not export the removed registerTask surface or the coordinator itself", () => { + for (const name of [ + "registerTask", + "listRegisteredTasks", + "registerDags", + "defaultRegistry", + "startCoordinator", + ]) { + expect(name in sdk).toBe(false); + } + expectTypeOf().not.toHaveProperty("registerTask"); + expectTypeOf().not.toHaveProperty("listRegisteredTasks"); + expectTypeOf().not.toHaveProperty("registerDags"); + // The runtime reads the registry it is handed, so there is no process-wide + // registry for a Dag constructor to write into. + expectTypeOf().not.toHaveProperty("defaultRegistry"); + expectTypeOf().not.toHaveProperty("startCoordinator"); }); it("exports public error classes", () => { @@ -55,23 +165,39 @@ describe("public API", () => { expect(connErr.connId).toBe("missing_conn"); }); - it("exports the coordinator runtime entrypoint", () => { - expectTypeOf().toEqualTypeOf< - (opts?: StartCoordinatorOptions) => Promise - >(); - expectTypeOf().toEqualTypeOf<{ - commAddr?: string; - logsAddr?: string; - argv?: readonly string[]; - }>(); + it("reaches the runtime only through serveDags, which takes a registry", () => { + expectTypeOf().toEqualTypeOf<(registry: DagRegistry) => Promise>(); + expectTypeOf>().toEqualTypeOf(); expectTypeOf(SUPERVISOR_API_VERSION).toMatchTypeOf(); }); - it("uses idiomatic TypeScript names for public client types", () => { - expectTypeOf().toEqualTypeOf<{ + it("keeps the Dag authoring signatures extensible via trailing specs", () => { + expectTypeOf().toEqualTypeOf<{ readonly dagId: string; readonly taskId: string; }>(); + expectTypeOf().toEqualTypeOf>>(); + expectTypeOf().toEqualTypeOf<{ + readonly inputs?: TaskInputs; + readonly spec?: TaskSpec; + }>(); + expectTypeOf>().toEqualTypeOf<[string, DagSpec?]>(); + expectTypeOf().toEqualTypeOf< + ( + taskId: string, + handler: TaskHandler, + options?: TaskOptions, + ) => TaskRef + >(); + expectTypeOf().toEqualTypeOf(); + // Reserved with no fields yet, so only `{}` is expressible. Generated specs + // will be all-optional (weak) types, and `{}` stays assignable to those, so + // filling these in later cannot break a call site. + expectTypeOf().toEqualTypeOf>(); + expectTypeOf().toEqualTypeOf>(); + }); + + it("uses idiomatic TypeScript names for public client types", () => { expectTypeOf().toEqualTypeOf<{ readonly dagId: string; readonly taskId: string; @@ -120,7 +246,6 @@ describe("public API", () => { it("rejects wire-format names and non-JSON XCom values", () => { function acceptsGetXComOpts(_opts: GetXComOpts): void {} function acceptsSetXComOpts(_opts: SetXComOpts): void {} - function acceptsTaskRegistration(_registration: TaskRegistration): void {} acceptsGetXComOpts({ key: "result", @@ -149,10 +274,34 @@ describe("public API", () => { expectTypeOf().toEqualTypeOf<{ connId: string; connType: string }>(); // @ts-expect-error public TaskContext does not expose the raw task-instance id. expectTypeOf().toHaveProperty("taskInstanceId"); - // @ts-expect-error task registration requires explicit dagId/taskId fields. - acceptsTaskRegistration("public_api_task"); - // @ts-expect-error task registration uses dagId, not dag_id. - acceptsTaskRegistration({ dag_id: "example", taskId: "extract" }); + // Never invoked: these constructor/method misuses also throw at runtime. + const rejectsPositionalMisuse = () => { + // @ts-expect-error dagId is positional, not an options object. + new Dag({ dagId: "example" }); + // @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. + 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 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 serveDags takes the registry, not a bare Dag. + serveDags(dag); + // @ts-expect-error a registry is built from Dags, not from task handles. + new DagRegistry(upstream); + }; + void rejectsPositionalMisuse; + // @ts-expect-error the TaskRef handle is opaque and does not expose the handler. + expectTypeOf().toHaveProperty("handler"); // @ts-expect-error XCom values must be JSON-compatible. acceptsSetXComOpts({ key: "result", value: new Date() }); }); diff --git a/ts-sdk/tests/sdk/dag.test.ts b/ts-sdk/tests/sdk/dag.test.ts new file mode 100644 index 0000000000000..e76c24ab4c1cd --- /dev/null +++ b/ts-sdk/tests/sdk/dag.test.ts @@ -0,0 +1,237 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, it, expect } from "vitest"; +import { Dag, getDagTaskRecords, type TaskRef } from "../../src/sdk/dag.js"; +import { DagRegistry } from "../../src/sdk/registry.js"; + +describe("Dag", () => { + it("returns a frozen TaskRef handle 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); + }); + + it("chains upstream handles 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: {} }); + + 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 }); + }); + + 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 } }); + + expect(getDagTaskRecords(dag).get("transform")?.inputs).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 inputs: Record = { extracted }; + dag.task("transform", async () => undefined, { inputs }); + + inputs.sneaky = extracted; + const recorded = getDagTaskRecords(dag).get("transform")!.inputs; + expect(recorded).toEqual({ extracted }); + expect(Object.isFrozen(recorded)).toBe(true); + }); + + 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( + /Input "extracted" of task "transform" comes from Dag "first_dag", not "second_dag"/, + ); + }); + + 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) => { + 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\(\.\.\.\)/, + ); + expect(getDagTaskRecords(dag).has("transform")).toBe(false); + }); + + it("rejects an input referring to a task that is not registered yet", () => { + 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"/); + }); + + it("retains its spec and each task's handler and spec, copied and frozen", () => { + const dagSpec = {}; + const taskSpec = {}; + const handler = async () => "hello"; + const dag = new Dag("example_dag", dagSpec); + dag.task("my_task", handler, { spec: taskSpec }); + + expect(dag.dagId).toBe("example_dag"); + expect(dag.spec).toEqual(dagSpec); + expect(Object.isFrozen(dag.spec)).toBe(true); + const record = getDagTaskRecords(dag).get("my_task"); + expect(record?.handler).toBe(handler); + expect(record?.spec).toEqual(taskSpec); + expect(Object.isFrozen(record!.spec)).toBe(true); + }); + + it.each([ + ["a populated object", { schedule: "@daily" }], + ["null", null], + ["an array", []], + ["a non-plain object", new Date()], + ])("rejects a Dag spec that is not an empty object: %s", (_label, spec) => { + expect(() => new Dag("example_dag", spec as unknown as Record)).toThrowError( + /spec for Dag "example_dag" must be an empty object/, + ); + }); + + it.each([ + ["a populated object", { retries: 2 }], + ["null", null], + ["an array", []], + ["a non-plain object", new Date()], + ])("rejects a task spec that is not an empty object: %s", (_label, spec) => { + const dag = new Dag("example_dag"); + expect(() => + dag.task("transform", async () => undefined, { + spec: spec as unknown as Record, + }), + ).toThrowError(/spec for Dag "example_dag" task "transform" must be an empty object/); + expect(dag.taskIds).toEqual([]); + }); + + it("exposes its task IDs in attachment order", () => { + const dag = new Dag("ordered_dag"); + expect(dag.taskIds).toEqual([]); + dag.task("extract", async () => undefined); + dag.task("transform", async () => undefined); + expect(dag.taskIds).toEqual(["extract", "transform"]); + }); + + it.each([ + ["a misspelled inputs key", { input: {} }], + ["a misspelled spec key", { specs: {} }], + ["an upstream handle passed positionally", { upstream: { dagId: "d", taskId: "t" } }], + ])("rejects %s in the task options", (_label, options) => { + const dag = new Dag("example_dag"); + expect(() => + dag.task("transform", async () => undefined, options as unknown as Record), + ).toThrowError(/Unknown option ".+" for Dag "example_dag" task "transform"/); + expect(dag.taskIds).toEqual([]); + }); + + it.each([ + ["null", null], + ["an array", []], + ["a string", "inputs"], + ["a non-plain object", new Date()], + ])("rejects task options that are not an options object: %s", (_label, options) => { + const dag = new Dag("example_dag"); + expect(() => + dag.task("transform", async () => undefined, options as unknown as Record), + ).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); + expect(() => dag.task("dup", async () => undefined)).toThrowError(/already registered/); + }); + + it("allows the same taskId in different Dags", () => { + const first = async () => "first"; + const second = async () => "second"; + const firstDag = new Dag("first_dag"); + const secondDag = new Dag("second_dag"); + firstDag.task("extract", first); + secondDag.task("extract", second); + + const registry = new DagRegistry(); + registry.register(firstDag, secondDag); + expect(registry.getTaskHandler("first_dag", "extract")).toBe(first); + expect(registry.getTaskHandler("second_dag", "extract")).toBe(second); + }); + + 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); + const registry = new DagRegistry(); + registry.register(dag); + expect(registry.getTaskHandler("café_dag", "任務")).toBe(handler); + }); + + it("rejects non-function handlers", () => { + const dag = new Dag("example_dag"); + expect(() => dag.task("x", "not a function" as unknown as () => Promise)).toThrowError( + /must be a function/, + ); + }); + + it("treats a dotted TaskGroup taskId as a single taskId (group.task)", () => { + const dag = new Dag("example_dag"); + dag.task("transforms.normalize", async () => "ok"); + const registry = new DagRegistry(); + registry.register(dag); + expect(registry.getTaskHandler("example_dag", "transforms.normalize")).toBeDefined(); + // Should NOT accidentally match the prefix alone + expect(registry.getTaskHandler("example_dag", "transforms")).toBeUndefined(); + expect(registry.getTaskHandler("example_dag", "normalize")).toBeUndefined(); + }); +}); diff --git a/ts-sdk/tests/sdk/registry.test.ts b/ts-sdk/tests/sdk/registry.test.ts index 4597f1d08458e..289ffd5935b26 100644 --- a/ts-sdk/tests/sdk/registry.test.ts +++ b/ts-sdk/tests/sdk/registry.test.ts @@ -18,124 +18,139 @@ */ import { describe, it, expect } from "vitest"; -import { TaskRegistry } from "../../src/sdk/registry.js"; +import { Dag } from "../../src/sdk/dag.js"; +import { DagRegistry, listRegistryDags, listRegistryTasks } from "../../src/sdk/registry.js"; -describe("registry", () => { - it("registers and retrieves a handler", async () => { - const registry = new TaskRegistry(); +describe("DagRegistry", () => { + it("registers a Dag and retrieves its handlers", () => { const handler = async () => "hello"; - registry.register({ dagId: "example_dag", taskId: "my_task" }, handler); - const got = registry.get("example_dag", "my_task"); - expect(got).toBe(handler); + const dag = new Dag("example_dag"); + dag.task("my_task", handler); + const registry = new DagRegistry(); + registry.register(dag); + expect(registry.getTaskHandler("example_dag", "my_task")).toBe(handler); }); - it("returns undefined for unknown taskIds", () => { - const registry = new TaskRegistry(); - expect(registry.get("example_dag", "nope")).toBeUndefined(); - expect(registry.get("unknown_dag", "my_task")).toBeUndefined(); + it("registers the Dags passed to its constructor", () => { + const handler = async () => "hello"; + const dagA = new Dag("dag_a"); + dagA.task("a", handler); + const registry = new DagRegistry(dagA, new Dag("dag_b")); + expect(registry.getTaskHandler("dag_a", "a")).toBe(handler); + expect(listRegistryDags(registry)).toEqual([ + { dagId: "dag_a", tasks: ["a"] }, + { dagId: "dag_b", tasks: [] }, + ]); + }); + + it("rejects duplicate dagIds passed to the constructor", () => { + expect(() => new DagRegistry(new Dag("example_dag"), new Dag("example_dag"))).toThrowError( + /already registered/, + ); + }); + + it("rejects constructor values that are not Dag instances", () => { + expect(() => new DagRegistry({ dagId: "example_dag" } as unknown as Dag)).toThrowError( + /only Dag instances can be registered/, + ); + }); + + it("returns undefined for unknown taskIds and dagIds", () => { + const registry = new DagRegistry(); + const dag = new Dag("example_dag"); + dag.task("my_task", async () => undefined); + registry.register(dag); + expect(registry.getTaskHandler("example_dag", "nope")).toBeUndefined(); + expect(registry.getTaskHandler("unknown_dag", "my_task")).toBeUndefined(); }); - it("returns an empty list when no tasks are registered", () => { - const registry = new TaskRegistry(); - expect(registry.list()).toEqual([]); + it("returns an empty list when no Dags are registered", () => { + const registry = new DagRegistry(); + expect(listRegistryTasks(registry)).toEqual([]); }); - it("lists registered tasks", () => { - const registry = new TaskRegistry(); - registry.register({ dagId: "dag_a", taskId: "a" }, async () => undefined); - registry.register({ dagId: "dag_b", taskId: "b" }, async () => undefined); - const registered = registry.list(); + it("lists tasks across registered Dags", () => { + const dagA = new Dag("dag_a"); + dagA.task("a", async () => undefined); + const dagB = new Dag("dag_b"); + dagB.task("b", async () => undefined); + const registry = new DagRegistry(); + registry.register(dagA, dagB); + const registered = listRegistryTasks(registry); expect(registered).toHaveLength(2); expect(registered).toContainEqual({ dagId: "dag_a", taskId: "a" }); expect(registered).toContainEqual({ dagId: "dag_b", taskId: "b" }); }); - it("rejects duplicate registration within a Dag", () => { - const registry = new TaskRegistry(); - registry.register({ dagId: "example_dag", taskId: "dup" }, async () => undefined); - expect(() => - registry.register({ dagId: "example_dag", taskId: "dup" }, async () => undefined), - ).toThrowError(/already registered/); + it("rejects registering the same dagId in separate calls", () => { + const registry = new DagRegistry(); + registry.register(new Dag("example_dag")); + expect(() => registry.register(new Dag("example_dag"))).toThrowError(/already registered/); }); - it("allows the same taskId in different Dags", () => { - const registry = new TaskRegistry(); - const first = async () => "first"; - const second = async () => "second"; - registry.register({ dagId: "first_dag", taskId: "extract" }, first); - registry.register({ dagId: "second_dag", taskId: "extract" }, second); - - expect(registry.get("first_dag", "extract")).toBe(first); - expect(registry.get("second_dag", "extract")).toBe(second); + it("rejects duplicate dagIds within a single call", () => { + const registry = new DagRegistry(); + expect(() => registry.register(new Dag("example_dag"), new Dag("example_dag"))).toThrowError( + /already registered/, + ); }); - it("rejects an empty dagId", () => { - const registry = new TaskRegistry(); - expect(() => - registry.register({ dagId: "", taskId: "my_task" }, async () => undefined), - ).toThrowError(/dagId must be made of alphanumeric/); + it("rejects registering the same Dag instance twice", () => { + const registry = new DagRegistry(); + const dag = new Dag("example_dag"); + registry.register(dag); + expect(() => registry.register(dag)).toThrowError(/already registered/); }); - it("rejects an empty taskId", () => { - const registry = new TaskRegistry(); - expect(() => - registry.register({ dagId: "example_dag", taskId: "" }, async () => undefined), - ).toThrowError(/taskId must be made of alphanumeric/); + it("registers none of the Dags when a call throws", () => { + const registry = new DagRegistry(); + const dag = new Dag("dag_a"); + dag.task("a", async () => undefined); + expect(() => registry.register(dag, new Dag("dag_a"))).toThrowError(/already registered/); + expect(registry.getTaskHandler("dag_a", "a")).toBeUndefined(); + expect(listRegistryTasks(registry)).toEqual([]); }); - it.each([" ", "\t", "my dag", "a/b", "task@1"])( - "rejects a dagId with characters no Python dag_id allows: %j", - (dagId) => { - const registry = new TaskRegistry(); - expect(() => - registry.register({ dagId, taskId: "my_task" }, async () => undefined), - ).toThrowError(/dagId must be made of alphanumeric/); - }, - ); - - it.each([" ", "\t", "my task", "a/b", "task@1"])( - "rejects a taskId with characters no Python task_id allows: %j", - (taskId) => { - const registry = new TaskRegistry(); - expect(() => - registry.register({ dagId: "example_dag", taskId }, async () => undefined), - ).toThrowError(/taskId must be made of alphanumeric/); - }, - ); - - it.each([ - ["dagId", { dagId: "d".repeat(251), taskId: "my_task" }], - ["taskId", { dagId: "example_dag", taskId: "t".repeat(251) }], - ])("rejects a %s longer than 250 characters", (name, registration) => { - const registry = new TaskRegistry(); - expect(() => registry.register(registration, async () => undefined)).toThrowError( - new RegExp(`${name} must be less than 250 characters, not 251`), + it("rejects values that are not Dag instances", () => { + const registry = new DagRegistry(); + expect(() => registry.register({ dagId: "example_dag" } as unknown as Dag)).toThrowError( + /only Dag instances can be registered/, ); }); - it("accepts a Unicode dagId that Python's word-character rule allows", () => { - const registry = new TaskRegistry(); - const handler = async () => undefined; - registry.register({ dagId: "café_dag", taskId: "任務" }, handler); - expect(registry.get("café_dag", "任務")).toBe(handler); + it("names the duplicate-copy cause when a Dag carries the brand but not this class", () => { + // Stands in for a Dag from a second resolved copy: same brand, other class. + const foreign = { dagId: "foreign_dag" }; + Object.defineProperty(foreign, Symbol.for("airflow.ts-sdk.Dag"), { value: true }); + expect(() => new DagRegistry(foreign as unknown as Dag)).toThrowError( + /different copy of @apache-airflow\/ts-sdk/, + ); }); - it("rejects non-function handlers", () => { - const registry = new TaskRegistry(); - expect(() => - registry.register( - { dagId: "example_dag", taskId: "x" }, - "not a function" as unknown as () => Promise, - ), - ).toThrowError(/must be a function/); + 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); + const registry = new DagRegistry(); + registry.register(dagA, new Dag("empty_dag")); + expect(listRegistryDags(registry)).toEqual([ + { dagId: "dag_a", tasks: ["a1", "a2"] }, + { dagId: "empty_dag", tasks: [] }, + ]); }); - it("treats a dotted TaskGroup taskId as a single taskId (group.task)", () => { - const registry = new TaskRegistry(); - registry.register({ dagId: "example_dag", taskId: "transforms.normalize" }, async () => "ok"); - expect(registry.get("example_dag", "transforms.normalize")).toBeDefined(); - // Should NOT accidentally match the prefix alone - expect(registry.get("example_dag", "transforms")).toBeUndefined(); - expect(registry.get("example_dag", "normalize")).toBeUndefined(); + it("sees tasks added to a Dag after registration", () => { + const registry = new DagRegistry(); + const dag = new Dag("example_dag"); + registry.register(dag); + expect(listRegistryTasks(registry)).toEqual([]); + + const handler = async () => "late"; + dag.task("late_task", handler); + expect(registry.getTaskHandler("example_dag", "late_task")).toBe(handler); + expect(listRegistryTasks(registry)).toContainEqual({ + dagId: "example_dag", + taskId: "late_task", + }); }); });