This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 14ab7252212 TS SDK: replace registerTask with Dag and registerDags
(#71144)
14ab7252212 is described below
commit 14ab7252212f05ef28273fc014b8a9017ddd5389
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Wed Aug 19 15:17:47 2026 +0800
TS SDK: replace registerTask with Dag and registerDags (#71144)
* TS SDK: replace registerTask with Dag and registerDags
registerTask bound one handler at a time to a Dag/task pair declared in a
Python stub file, and left no object that could later carry a natively
declared TypeScript Dag. Reworking only the authoring interface now means
native Dag support and TaskFlow-style data passing can arrive without a
second breaking change for users: a Dag instance keeps its spec and every
task's handler, spec and declared inputs, which is what a future
serialize() needs to emit the serialized Dag JSON.
Producing that JSON stays out of scope, so Dag parsing still answers with
no serialized Dags. The coordinator wire protocol and the bundle manifest
shape are unchanged.
* TS SDK: make registerDags the only Dag bundle entrypoint
A Dag author has no reason to know Airflow's coordinator exists, and the
three-step shape left two steps that failed quietly when one was missed: a
Dag
that was built but never registered was dropped from the packed bundle, and
its
task instances showed up as *removed* at runtime with only a warning in the
task
log. Folding the runtime handoff into registerDags removes one of those
steps
outright; the packer now names the Dags that fell through the other.
Reserved options were the second quiet failure. `inputs` read as though it
declared a dependency, so a user could conclude the value was wired and call
getXCom without a taskId, which reads the running task's own XCom and
returns
null without failing the task. Saying plainly that these fields are inert,
and
accepting only `{}` until they are real, turns a wrong value into a compile
error. Generated specs will be all-optional types that `{}` still
satisfies, so
filling them in later cannot break a call site.
Empty-Dag handling followed airflow-go-pack rather than diverging from it:
one
placeholder Dag should not block a bundle build in TypeScript and not in Go.
* TS SDK: serve Dags from an explicit DagRegistry
An entrypoint named for registration gave no hint that it also connects
sockets and blocks until Airflow's supervisor has its terminal frame, so
the name now says what it does. Naming it after serving also matches the
Java and Go SDKs, where the socket-owning half is kept separate from the
collection of Dags it serves.
Making that collection an object the author builds removes the
process-wide state the previous shape needed: constructing a Dag no
longer registers it as a side effect, and the runtime dispatches through
the registry it is handed rather than a module-level singleton. A
registry owns no sockets and starts nothing, so a bundle's handlers can
be dispatched from a unit test with no runtime in scope. Leaving it in
the authoring layer also keeps the dependency between the two pointing
one way; a serve() method on the registry would have made it circular.
Warning about a Dag that was built but never served was the only thing
the process-wide tracking bought, so it goes too. As in the Go and Java
SDKs, a Dag left out of the bundle is simply absent from it, and its
tasks are marked removed at runtime.
* TS SDK: let a failed serveDags be retried
A serve that never got off the ground is not a serve, but the one-shot
latch was set before the runtime ran and never released, so a bundle that
failed on its socket arguments reported "already called" on the next
attempt and hid the real cause.
The latch also has to outlive a duplicated dependency: what must not
happen twice is one process connecting two pairs of sockets, and a
workspace that resolves two copies of the package gives each its own
module state, so both would believe they were first. Keying it globally
rather than per module makes them agree.
Objects from a second copy still cannot be used, because both classes read
private state keyed to the class that declared it. They can be recognised
though, so a genuine Dag or registry from another copy is now reported as
that rather than as the wrong type.
* TS SDK: reject bundle metadata the manifest schema would not accept
Nothing downstream re-validates the metadata line: the reader in
_bundle_metadata.py only checks that the document is a mapping, so a
schema version that is merely truthy rather than a non-empty string was
rendered into the bundle verbatim and travelled all the way to Airflow.
A bundle that printed a bare null crashed the packer with a raw TypeError
from reading a field off it, which reads as a bug in airflow-ts-pack
rather than as a report about the bundle it was given.
* TS SDK: keep registry enumeration out of the public API
Listing a registry's tasks and Dags is how the runtime dispatches and how
the packer builds a manifest; a Dag author declares tasks and serves them,
and never needs to ask a registry what is in it. Exporting DagRegistry
made both methods public API that has to be kept working forever, for no
one's benefit.
Serving another copy's registry cannot work either, for the reason
registering another copy's Dag cannot: the lookups read private state
keyed to the class object of the copy that defined it. The check says that
now instead of reporting a genuine registry as the wrong type.
* TS SDK: rewrite the reserved-option docs as documentation
These read as replies to review threads rather than as something useful in
an editor hover: the rationale came first, the plain statement of what the
option does came last or not at all, and each declaration ordered it
differently. "Inert" was doing work that "not used yet" does plainly.
Every one now says what it is, what it does not do yet, then what it is
for, in that order.
* TS SDK: give the copy-detection brand its own module
The brand is one mechanism with two users and had no owner between them:
the shared error text sat in dag.ts, where the coordinator had to import
it to describe a registry, and registry.ts carried a "for the reason given
on Dag's" pointer that would rot away from what it points at.
The comments introduced alongside it are cut back to what the code does
not already say.
* TS SDK: check Dag/task key format only during packing
validateKey ran inside the Dag constructor and dag.task(), part of the
bundle module's top level that also re-runs on every task-execution-runtime
startup, so a bundle's keys were re-checked on every task attempt instead
of once. Move the check into buildBundleManifest, which only runs for
airflow-ts-pack's --airflow-metadata query, so it happens once at pack time.
* Preserve valid Dag IDs in TypeScript bundle metadata
* Reject invalid TypeScript Dag option objects
* TS SDK: Report invalid Dag and task IDs clearly
* TS SDK: Report the latest bundle error
---
.../language-sdks/typescript.rst | 27 ++-
ts-sdk/README.md | 57 +++--
ts-sdk/docs/index.md | 14 +-
ts-sdk/example/src/main.ts | 10 +-
ts-sdk/src/cli/pack.ts | 65 +++++-
ts-sdk/src/coordinator/index.ts | 5 +-
ts-sdk/src/coordinator/manifest.ts | 46 +++-
ts-sdk/src/coordinator/runtime.ts | 92 +++++++-
ts-sdk/src/index.ts | 8 +-
ts-sdk/src/sdk/brand.ts | 43 ++++
ts-sdk/src/sdk/dag.ts | 245 +++++++++++++++++++++
ts-sdk/src/sdk/registry.ts | 147 +++++++------
ts-sdk/tests/cli/fixtures/empty-entry.ts | 4 +-
ts-sdk/tests/cli/fixtures/entry.ts | 12 +-
ts-sdk/tests/cli/fixtures/noisy-entry.ts | 7 +-
ts-sdk/tests/cli/pack.test.ts | 159 ++++++++++++-
ts-sdk/tests/coordinator/integration.test.ts | 42 ++--
ts-sdk/tests/coordinator/protocol.test.ts | 5 +-
ts-sdk/tests/coordinator/public-api.test.ts | 20 +-
ts-sdk/tests/coordinator/runtime-manifest.test.ts | 91 +++++++-
ts-sdk/tests/public-api.test.ts | 201 ++++++++++++++---
ts-sdk/tests/sdk/dag.test.ts | 237 ++++++++++++++++++++
ts-sdk/tests/sdk/registry.test.ts | 197 +++++++++--------
23 files changed, 1426 insertions(+), 308 deletions(-)
diff --git
a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
index 3571d7bf519..b93be7b38ce 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<string>({
@@ -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 73cfdcab641..45ea2bc31a9 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 6fe81326553..edda33fe48d 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 1c480b8056b..d08530615e5 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<string>({
@@ -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 0d863fbe898..3268ec9aab6 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<BundleManifest>;
+ 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<void> {
});
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 815e0c849b9..f06214a4c2f 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 32e503dce7a..6baca7f4d7b 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<string, { tasks: string[] }>;
}
-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 a3a057b4c01..5f8c1afaf25 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<symbol, boolean | undefined> {
+ return globalThis as unknown as Record<symbol, boolean | undefined>;
+}
+
+/**
+ * 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<void> {
+ // 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<void> {
+/** 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<void> {
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<RuntimeSucceedTask | RuntimeRetryTask | RuntimeTaskState> {
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 0a63cb263f4..a7c8f4960b0 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 00000000000..b9fab223862
--- /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 00000000000..b00492ffd87
--- /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<string, unknown> {
+ 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<string, never>;
+
+/**
+ * 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<string, never>;
+
+/**
+ * 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<Record<string, TaskRef>>;
+
+/**
+ * Named options for `dag.task()`.
+ *
+ * Keyword-only so neither field has to be positioned around the other, and so
+ * future fields can be added without a new parameter. Unknown keys are
+ * rejected, so a typo fails at import time rather than being ignored.
+ */
+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<string, TaskRecord>;
+
+/** 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<string, TaskRecord>();
+
+ 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<TReturn = unknown>(
+ taskId: string,
+ handler: TaskHandler<TReturn>,
+ 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<string, TaskRecord> {
+ return taskRecordsOf(dag);
+}
diff --git a/ts-sdk/src/sdk/registry.ts b/ts-sdk/src/sdk/registry.ts
index 122be56e02d..ee4286de926 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<string, Dag>;
-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<string, Map<string, TaskHandler>>();
+/**
+ * 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<string, Dag>();
- /**
- * 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<TReturn = unknown>(registration: TaskRegistration, handler:
TaskHandler<TReturn>): 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<string, TaskHandler>();
- 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<string>();
+ 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<TReturn = unknown>(
- registration: TaskRegistration,
- handler: TaskHandler<TReturn>,
-): 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 3e2b582b224..c48c6c24cf8 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 8d53c43aee5..8549fb56a32 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 3b9675619d3..dcd7b0d8c0b 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 4a5c2d572bf..6de652e7f28 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 a580159073e..040de00fa96 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<string, unknown>();
let observedGreeting: string | null = "<unset>";
- 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 = "<unset>";
- 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 330f94b62b0..8089a595a5d 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 218fa42508e..16e91d4b31f 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<typeof startCoordinator>().toEqualTypeOf<
- (opts?: StartCoordinatorOptions) => Promise<void>
- >();
- expectTypeOf<StartCoordinatorOptions>().toEqualTypeOf<{
- commAddr?: string;
- logsAddr?: string;
- argv?: readonly string[];
- }>();
+ it("exposes serveDags, not the coordinator itself, from the coordinator
subpath", () => {
+ expectTypeOf<typeof serveDags>().toEqualTypeOf<(registry: DagRegistry) =>
Promise<void>>();
+ expect("startCoordinator" in coordinator).toBe(false);
+ expectTypeOf<typeof coordinator>().not.toHaveProperty("startCoordinator");
});
});
diff --git a/ts-sdk/tests/coordinator/runtime-manifest.test.ts
b/ts-sdk/tests/coordinator/runtime-manifest.test.ts
index f6aea6f2fd5..9c8beec94e1 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 4e7a49ddb4d..a9b2f888019 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, unknown>)[
+ 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<keyof DagRegistry>().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<typeof sdk>().not.toHaveProperty("registerTask");
+ expectTypeOf<typeof sdk>().not.toHaveProperty("listRegisteredTasks");
+ expectTypeOf<typeof sdk>().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<typeof sdk>().not.toHaveProperty("defaultRegistry");
+ expectTypeOf<typeof sdk>().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<typeof startCoordinator>().toEqualTypeOf<
- (opts?: StartCoordinatorOptions) => Promise<void>
- >();
- expectTypeOf<StartCoordinatorOptions>().toEqualTypeOf<{
- commAddr?: string;
- logsAddr?: string;
- argv?: readonly string[];
- }>();
+ it("reaches the runtime only through serveDags, which takes a registry", ()
=> {
+ expectTypeOf<typeof serveDags>().toEqualTypeOf<(registry: DagRegistry) =>
Promise<void>>();
+ expectTypeOf<ConstructorParameters<typeof
DagRegistry>>().toEqualTypeOf<Dag[]>();
expectTypeOf(SUPERVISOR_API_VERSION).toMatchTypeOf<string>();
});
- it("uses idiomatic TypeScript names for public client types", () => {
- expectTypeOf<TaskRegistration>().toEqualTypeOf<{
+ it("keeps the Dag authoring signatures extensible via trailing specs", () =>
{
+ expectTypeOf<TaskRef>().toEqualTypeOf<{
readonly dagId: string;
readonly taskId: string;
}>();
+ expectTypeOf<TaskInputs>().toEqualTypeOf<Readonly<Record<string,
TaskRef>>>();
+ expectTypeOf<TaskOptions>().toEqualTypeOf<{
+ readonly inputs?: TaskInputs;
+ readonly spec?: TaskSpec;
+ }>();
+ expectTypeOf<ConstructorParameters<typeof Dag>>().toEqualTypeOf<[string,
DagSpec?]>();
+ expectTypeOf<Dag["task"]>().toEqualTypeOf<
+ <TReturn = unknown>(
+ taskId: string,
+ handler: TaskHandler<TReturn>,
+ options?: TaskOptions,
+ ) => TaskRef
+ >();
+ expectTypeOf<Dag["taskIds"]>().toEqualTypeOf<readonly string[]>();
+ // 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<DagSpec>().toEqualTypeOf<Record<string, never>>();
+ expectTypeOf<TaskSpec>().toEqualTypeOf<Record<string, never>>();
+ });
+
+ it("uses idiomatic TypeScript names for public client types", () => {
expectTypeOf<TaskContext>().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<ConnectionResult>().toEqualTypeOf<{ connId: string; connType:
string }>();
// @ts-expect-error public TaskContext does not expose the raw
task-instance id.
expectTypeOf<TaskContext>().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<TaskRef>().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 00000000000..e76c24ab4c1
--- /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<string, TaskRef> = { 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<string,
never>)).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<string, never>,
+ }),
+ ).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<string, never>),
+ ).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<string, never>),
+ ).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<string, TaskRef>,
+ }),
+ ).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<unknown>)).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 4597f1d0845..289ffd5935b 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<unknown>,
- ),
- ).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",
+ });
});
});
