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 959a7408362 TS SDK: register mixed-language task handlers with the
TaskHandler class (#73188)
959a7408362 is described below
commit 959a74083628532112101acdc5d7e0e15f219fdc
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Thu Sep 17 13:19:12 2026 +0800
TS SDK: register mixed-language task handlers with the TaskHandler class
(#73188)
A mixed-language Dag is declared in Python: the `@task.stub` per task, the
queue that routes it to the Node coordinator, and the order between them all
live there. The SDK said otherwise. The only way to attach a handler was
`new Dag(dagId)` plus `dag.task(taskId, fn)`, so every mixed-language bundle
constructed a Dag object for a Dag it does not own.
`new TaskHandler(dagId, taskId, handler)` binds a function to the task it
implements and carries nothing else: no schedule, no task order, no dag_id
of
its own. A handler is a value with no call
signature, so wiring one the way a natively declared task is wired is a
compile
error rather than a runtime throw.
`Dag` becomes exclusively the native case, and a dag_id is one or the
other: a
native Dag attaches its own tasks, so registering a handler for one is
rejected
rather than silently becoming a second, disagreeing source for its task
list.
`Registerable` gains its second arm rather than the bundle gaining a second
verb, so one `register` call still lists everything a bundle provides in any
mixture. What that changes underneath is the key: a bundle now holds one
entry
per dag_id and dispatches on the `(dag_id, task_id)` pair, so one bundle can
provide for several Dags and the same task_id under two of them is two
different handlers. The example and the end-to-end test both exercise that,
with `typescript_example` and the new `typescript_taskflow_example` each
declaring a `build_message` served from one bundle.mjs.
---
.../language-sdks/typescript.rst | 35 ++--
.../tests/airflow_e2e_tests/conftest.py | 9 +-
.../ts_sdk_tests/test_ts_sdk_dag.py | 100 ++++++++---
ts-sdk/README.md | 61 ++++---
ts-sdk/api-docs/dag-authoring-api.ts | 2 +-
ts-sdk/docs/index.md | 13 +-
ts-sdk/example/README.md | 10 +-
ts-sdk/example/dags/typescript_taskflow_example.py | 57 ++++++
ts-sdk/example/src/main.ts | 22 ++-
ts-sdk/example/src/taskflow.ts | 75 ++++++++
ts-sdk/scripts/verify-package.mjs | 2 +-
ts-sdk/src/cli/pack.ts | 4 +-
ts-sdk/src/coordinator/runtime.ts | 8 +-
ts-sdk/src/index.ts | 1 +
ts-sdk/src/sdk/bundle.ts | 144 +++++++++++----
ts-sdk/src/sdk/task-handler.ts | 97 +++++++++++
ts-sdk/tests/cli/fixtures/entry.ts | 13 +-
ts-sdk/tests/cli/pack.test.ts | 6 +-
ts-sdk/tests/coordinator/integration.test.ts | 49 +++++-
ts-sdk/tests/public-api.test.ts | 34 +++-
ts-sdk/tests/sdk/bundle.test.ts | 8 +-
ts-sdk/tests/sdk/task-handler.test.ts | 193 +++++++++++++++++++++
22 files changed, 798 insertions(+), 145 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 9bfc50e0f67..7aaf6408ac8 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
@@ -22,7 +22,7 @@ TypeScript SDK
|experimental|
-The TypeScript SDK lets you group task handlers in a ``Dag`` and implement
their logic in TypeScript (or
+The TypeScript SDK lets you register task handlers on a ``Bundle`` and
implement their logic in TypeScript (or
plain JavaScript), running on Node.js. A matching Python stub Dag still
declares the scheduling shape and
dependencies; individual tasks delegate to a Node.js subprocess that is
spawned by
:class:`~airflow.sdk.coordinators.node.NodeCoordinator` for each task instance.
@@ -37,7 +37,7 @@ The SDK is the ``apache-airflow-ts-sdk`` package (ESM-only).
It is currently in
.. seealso::
- For the full TypeScript API reference (``Dag``, ``Bundle``, the task handler
getters,
+ For the full TypeScript API reference (``Bundle``, ``TaskHandler``, ``Dag``,
the task handler getters,
``TaskClient``, supporting types, and exceptions),
see the `TypeScript SDK API reference
<https://airflow.apache.org/docs/ts-sdk/stable/>`__.
@@ -96,13 +96,13 @@ TypeScript implementation
A task is an ordinary (usually ``async``) function taking no arguments:
``getContext()`` and ``getClient()`` reach the runtime from inside the call,
so nothing the SDK supplies is a parameter.
-Create a ``Dag`` with the ``dag_id`` it implements, attach each handler with
``dag.task``,
-register it on a ``Bundle``, then serve it to Airflow with ``bundle.serve()``.
+Create a ``TaskHandler`` per task, binding the function to the ``dag_id`` and
``task_id`` it implements,
+register them on a ``Bundle``, then serve it to Airflow with
``bundle.serve()``.
That top-level ``await`` makes the module a runnable bundle entry point.
.. code-block:: typescript
- import { Bundle, Dag, getClient } from "apache-airflow-ts-sdk";
+ import { Bundle, getClient, TaskHandler } from "apache-airflow-ts-sdk";
export async function buildMessage() {
const client = getClient();
@@ -114,25 +114,22 @@ That top-level ``await`` makes the module a runnable
bundle entry point.
return `${greeting ?? "hello from TypeScript"}; upstream=${upstream ??
"missing"}`;
}
- const dag = new Dag("typescript_example");
- dag.task("build_message", buildMessage);
-
const bundle = new Bundle();
- bundle.register(dag);
+ bundle.register(new TaskHandler("typescript_example", "build_message",
buildMessage));
await bundle.serve();
-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. What
the bundle holds is its
-complete set of Dags; a second ``bundle.serve()`` call is rejected. A Dag left
unregistered is not part of
-the packed bundle, and its tasks are marked removed at runtime.
+The ``dagId`` a handler binds must match the ``dag_id`` of the Python Dag, and
the ``taskId`` a
+``@task.stub`` function in that Dag, including any TaskGroup prefix.
-``register`` is the bundle's one registration verb, and takes any number of
items, so a bundle that
-collects what it provides across several modules can call it repeatedly
instead of passing everything to
-the constructor. Registering holds no sockets and starts nothing, so a unit
test can build a bundle and
-dispatch a handler through ``bundle.getTaskHandler(dagId, taskId)`` without a
coordinator runtime.
+``register`` takes any number of task handlers and ``bundle.serve()`` serves
exactly what is registered,
+so a task left out is not part of the packed bundle and is marked removed at
runtime.
+A second ``bundle.serve()`` call is rejected.
+Registering holds no sockets and starts nothing, so a unit test can build a
bundle and dispatch a handler
+through ``bundle.getTaskHandler(dagId, taskId)`` without a coordinator runtime.
-``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.
+``Dag`` is another interface, for a Dag declared in TypeScript rather than in
Python, and is still a work
+in progress. ``new Dag`` and ``dag.task`` take a trailing options object
(``spec`` on both, plus
+``inputs`` on a task) that is not used yet; do not set them.
.. note::
diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
index 2877a4dc7bb..961f367b907 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
@@ -724,9 +724,10 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir):
ts_bundles_dir.mkdir()
copyfile(TS_SDK_EXAMPLE_PATH / "dist" / "bundle.mjs", ts_bundles_dir /
"bundle.mjs")
- copyfile(
- TS_SDK_EXAMPLE_PATH / "dags" / "typescript_example.py", tmp_dir /
"dags" / "typescript_example.py"
- )
+ # Both of the example bundle's Dags: one bundle.mjs provides for two
dag_ids,
+ # and the tests check that dispatch tells their same-named tasks apart.
+ for dag_file in ("typescript_example.py",
"typescript_taskflow_example.py"):
+ copyfile(TS_SDK_EXAMPLE_PATH / "dags" / dag_file, tmp_dir / "dags" /
dag_file)
coordinator_config = json.dumps(
{
@@ -770,7 +771,7 @@ def spin_up_airflow_environment(tmp_path_factory:
pytest.TempPathFactory):
_E2ETestState.airflow_dags_path = tmp_dir / "dags"
# openlineage sources its dags from the provider system tests (via
_setup_openlineage_integration),
- # so it must not also load the stock e2e dags — the harness triggers every
dag it finds.
+ # so it must not also load the stock e2e dags, since the harness triggers
every dag it finds.
if E2E_TEST_MODE != "openlineage":
console.print(f"[yellow]Copying dags to:[/ {tmp_dir / 'dags'}")
copytree(E2E_DAGS_FOLDER, tmp_dir / "dags", dirs_exist_ok=True)
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
index 3c308100654..5e25ef5801c 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
@@ -21,14 +21,17 @@ Run with::
E2E_TEST_MODE=ts_sdk uv run --project airflow-e2e-tests pytest \\
tests/airflow_e2e_tests/ts_sdk_tests/ -xvs
-The ``typescript_example`` Dag mixes a Python task with ``@task.stub``
-TypeScript tasks whose handlers live in the ``airflow-ts-pack`` bundle built
-by ``conftest._setup_ts_sdk_integration``. Triggered once via the
-module-scoped ``completed_run`` fixture, the run confirms end-to-end that
-``NodeCoordinator`` launches the bundle on the volume-provided Node runtime,
-Variable/Connection reads and Python <-> TypeScript XCom round-trips work
-through the Task Execution API, and coordinator-channel logs reach the
-task-log store.
+Two Dags mix Python tasks with ``@task.stub`` TypeScript tasks whose handlers
live in the single
+``airflow-ts-pack`` bundle built by ``conftest._setup_ts_sdk_integration``.
+Each is triggered once via a module-scoped fixture.
+
+``typescript_example`` confirms that ``NodeCoordinator`` launches the bundle
on the volume-provided
+Node runtime, that Variable/Connection reads and Python <-> TypeScript XCom
round-trips work through
+the Task Execution API, and that coordinator-channel logs reach the task-log
store.
+
+``typescript_taskflow_example`` confirms that one bundle provides for two
``dag_id``s: its
+``build_message`` shares a ``task_id`` with a task in ``typescript_example``,
so a bundle that keyed
+dispatch on the task ID alone would run the wrong handler for one of them.
"""
from __future__ import annotations
@@ -49,26 +52,28 @@ _TS_TASK_TIMEOUT = 600
_LOG_FETCH_TIMEOUT = 120
_DAG_ID = "typescript_example"
+_TASKFLOW_DAG_ID = "typescript_taskflow_example"
@dataclass
class _CompletedRun:
client: AirflowClient
+ dag_id: str
run_id: str
state: str
ti_states: dict[str, str]
def xcom(self, task_id: str, key: str = "return_value"):
- return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id,
run_id=self.run_id, key=key).get(
- "value"
- )
+ return self.client.get_xcom_value(
+ dag_id=self.dag_id, task_id=task_id, run_id=self.run_id, key=key
+ ).get("value")
def logs(self, task_id: str, try_number: int = 1) -> str:
"""Fetch task logs, retrying until present (log upload is async)."""
deadline = time.monotonic() + _LOG_FETCH_TIMEOUT
while True:
resp = self.client.get_task_logs(
- dag_id=_DAG_ID, run_id=self.run_id, task_id=task_id,
try_number=try_number
+ dag_id=self.dag_id, run_id=self.run_id, task_id=task_id,
try_number=try_number
)
text = "\n".join(str(entry) for entry in resp.get("content", []))
if text.strip() or time.monotonic() > deadline:
@@ -76,16 +81,26 @@ class _CompletedRun:
time.sleep(3)
[email protected](scope="module")
-def completed_run() -> _CompletedRun:
- """Trigger ``typescript_example`` once; every test inspects the same
run."""
+def _trigger_and_wait(dag_id: str) -> _CompletedRun:
client = AirflowClient()
- resp = client.trigger_dag(_DAG_ID, json={"logical_date":
datetime.now(timezone.utc).isoformat()})
+ resp = client.trigger_dag(dag_id, json={"logical_date":
datetime.now(timezone.utc).isoformat()})
run_id = resp["dag_run_id"]
- state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id,
timeout=_TS_TASK_TIMEOUT)
- ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id)
+ state = client.wait_for_dag_run(dag_id=dag_id, run_id=run_id,
timeout=_TS_TASK_TIMEOUT)
+ ti_resp = client.get_task_instances(dag_id=dag_id, run_id=run_id)
ti_states = {ti["task_id"]: ti.get("state") for ti in
ti_resp.get("task_instances", [])}
- return _CompletedRun(client=client, run_id=run_id, state=state,
ti_states=ti_states)
+ return _CompletedRun(client=client, dag_id=dag_id, run_id=run_id,
state=state, ti_states=ti_states)
+
+
[email protected](scope="module")
+def completed_run() -> _CompletedRun:
+ """Trigger ``typescript_example`` once; every test inspects the same
run."""
+ return _trigger_and_wait(_DAG_ID)
+
+
[email protected](scope="module")
+def completed_taskflow_run() -> _CompletedRun:
+ """Trigger ``typescript_taskflow_example`` once, from the same bundle."""
+ return _trigger_and_wait(_TASKFLOW_DAG_ID)
def test_dag_run_succeeded(completed_run: _CompletedRun):
@@ -107,8 +122,8 @@ def test_task_states(completed_run: _CompletedRun):
def test_build_message_xcom_round_trip(completed_run: _CompletedRun):
- """``build_message`` combines ``python_start``'s XCom with the Variable,
- pushes it under ``typescript_message``, and returns it."""
+ """``build_message`` combines ``python_start``'s XCom with the Variable,
pushes it under
+ ``typescript_message``, and returns it."""
assert completed_run.xcom("python_start") == "hello from Python"
message = "greetings from e2e; upstream=hello from Python"
@@ -132,3 +147,46 @@ def test_read_connection_xcom(completed_run:
_CompletedRun):
def test_coordinator_logs_reach_task_log_store(completed_run: _CompletedRun):
assert "[ts-sdk.runtime] Coordinator runtime started" in
completed_run.logs("build_message")
+
+
+def test_second_dag_from_the_same_bundle_succeeded(completed_taskflow_run:
_CompletedRun):
+ """One packed bundle provides for both Dags, so the second one also
runs."""
+ assert completed_taskflow_run.state == "success", (
+ f"expected the run to succeed; got {completed_taskflow_run.state!r}. "
+ f"task states: {completed_taskflow_run.ti_states}"
+ )
+ expected = {"make_totals": "success", "summarize": "success",
"build_message": "success"}
+ for task_id, want in expected.items():
+ assert completed_taskflow_run.ti_states.get(task_id) == want, (
+ f"{task_id!r} expected {want!r}. all task states:
{completed_taskflow_run.ti_states}"
+ )
+
+
+def test_summarize_xcom(completed_taskflow_run: _CompletedRun):
+ """``summarize`` reads ``make_totals``'s output and averages it."""
+ assert completed_taskflow_run.xcom("make_totals") == {"orders": 12,
"revenue": 3402.0}
+ value = completed_taskflow_run.xcom("summarize")
+ assert value == {"orders": 12, "averageOrder": 283.5, "currency": "GBP"}, (
+ f"unexpected 'summarize' return_value: {value!r}"
+ )
+
+
+def test_same_task_id_under_two_dags_runs_its_own_handler(
+ completed_run: _CompletedRun, completed_taskflow_run: _CompletedRun
+):
+ """Both Dags have a ``build_message``; each must reach its own handler.
+
+ A bundle that keyed dispatch on the task ID alone would answer both from
+ whichever handler was registered last, and both assertions below would
+ report the same shape.
+ """
+ example_value = completed_run.xcom("build_message")
+ taskflow_value = completed_taskflow_run.xcom("build_message")
+
+ assert set(example_value) == {"message", "upstream"}, (
+ f"unexpected 'typescript_example.build_message' return_value:
{example_value!r}"
+ )
+ assert taskflow_value == {
+ "dagId": _TASKFLOW_DAG_ID,
+ "message": "12 orders averaging 283.5 GBP",
+ }, f"unexpected 'typescript_taskflow_example.build_message' return_value:
{taskflow_value!r}"
diff --git a/ts-sdk/README.md b/ts-sdk/README.md
index 4cc5c9b4e53..c5f60210337 100644
--- a/ts-sdk/README.md
+++ b/ts-sdk/README.md
@@ -35,24 +35,24 @@ npm install [email protected]
## Task Handlers
```ts
-import { Bundle, Dag, getClient, getContext } from "apache-airflow-ts-sdk";
+import { Bundle, getClient, getContext, TaskHandler } from
"apache-airflow-ts-sdk";
export async function sayHello() {
const greeting = await getClient().getVariable("greeting");
return { message: `Hello from ${getContext().taskId}: ${greeting}` };
}
-const dag = new Dag("example_dag");
-dag.task("say_hello", sayHello);
-
const bundle = new Bundle();
-bundle.register(dag);
+bundle.register(new TaskHandler("example_dag", "say_hello", sayHello));
await bundle.serve();
```
A handler is a plain function. `getContext()` and `getClient()` reach the
runtime from inside the call,
so a handler takes no SDK-supplied parameter.
+`new TaskHandler(dagId, taskId, handler)` binds the function to the
Python-owned task it implements.
+The Dag is declared in Python with `@task.stub`, so the TypeScript side only
supplies the task bodies.
+
Non-`undefined` return values are pushed to XCom under the `"return_value"`
key by the active runtime, matching Python `@task` behavior.
@@ -105,7 +105,7 @@ Airflow metadata in the bundle itself.
TypeScript entrypoint:
```ts
-import { Bundle, Dag, getClient } from "apache-airflow-ts-sdk";
+import { Bundle, getClient, TaskHandler } from "apache-airflow-ts-sdk";
export async function extract() {
const client = getClient();
@@ -129,44 +129,41 @@ export async function transform() {
};
}
-const salesPipeline = new Dag("sales_pipeline");
-salesPipeline.task("extract", extract);
-salesPipeline.task("transform", transform);
-
const bundle = new Bundle();
-bundle.register(salesPipeline);
+bundle.register(
+ new TaskHandler("sales_pipeline", "extract", extract),
+ new TaskHandler("sales_pipeline", "transform", transform),
+);
await bundle.serve();
```
The Python stub defines the Dag dependency graph. The TypeScript handler does
-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 `Bundle` holds
-what this bundle process provides, and `bundle.serve()` serves it to Airflow.
-
-`bundle.serve()` is the entrypoint: a Dag left unregistered is not part of the
bundle,
-and its tasks are marked removed at runtime.
+the work and uses `TaskClient` for task-time Airflow data access. The handler
+function is the reusable task implementation; a `TaskHandler` binds it to a
+Python stub task identity, a `Bundle` holds what this bundle process provides,
+and `bundle.serve()` serves it to Airflow.
+
+`bundle.serve()` is the entrypoint: a task left unregistered is not part of
the bundle,
+and one bundle can provide for several `TaskHandler`s.
Registering holds no sockets and starts nothing, so a unit test can build a
bundle
and dispatch through `bundle.getTaskHandler(dagId, taskId)` without any
runtime involved.
-`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:
+Dispatch keys on the `(dagId, taskId)` pair, so the same `taskId` under two
Dags is two different handlers:
```ts
-import { salesDag } from "./sales/dag";
-import { billingDag } from "./billing/dag";
-import { Bundle } from "apache-airflow-ts-sdk";
-
-await new Bundle(salesDag, billingDag).serve();
+import { Bundle, TaskHandler } from "apache-airflow-ts-sdk";
+import { chargeCustomer } from "./billing/tasks";
+import { extract } from "./sales/tasks";
+
+await new Bundle(
+ new TaskHandler("sales_pipeline", "extract", extract),
+ new TaskHandler("billing_pipeline", "extract", chargeCustomer),
+).serve();
```
-`register` is the bundle's one registration verb: a bundle that collects what
it
-provides across several modules can call it repeatedly instead of passing
-everything to the constructor.
+Register `TaskHandler` and `Dag` values with the `register` method, or pass
them to the `Bundle` constructor.
+
+`Dag` is another interface, for a Dag declared natively in TypeScript, and is
still a work in progress.
Airflow launches the bundled entrypoint with `--comm=host:port` and
`--logs=host:port`. `bundle.serve()` connects to those sockets, receives the
diff --git a/ts-sdk/api-docs/dag-authoring-api.ts
b/ts-sdk/api-docs/dag-authoring-api.ts
index b1e92900f56..da90ccb81d4 100644
--- a/ts-sdk/api-docs/dag-authoring-api.ts
+++ b/ts-sdk/api-docs/dag-authoring-api.ts
@@ -19,7 +19,7 @@
/** @module Authoring */
-export { Bundle, Dag, getClient, getContext } from "../src/index.js";
+export { Bundle, Dag, getClient, getContext, TaskHandler } from
"../src/index.js";
export type {
DagSpec,
Registerable,
diff --git a/ts-sdk/docs/index.md b/ts-sdk/docs/index.md
index e53f5b8dc86..689a8299e32 100644
--- a/ts-sdk/docs/index.md
+++ b/ts-sdk/docs/index.md
@@ -35,28 +35,29 @@ Install the beta package from npm:
npm install [email protected]
```
-Define a Dag, register it on a `Bundle`, and serve it.
+Bind a handler to the Python-owned task it implements, register it on a
`Bundle`, and serve it.
A handler is a plain function: `getContext()` returns the `TaskContext` and
`getClient()` the `TaskClient`
for as long as it runs, so neither is a parameter.
Any non-`undefined` return value is pushed to XCom under the `"return_value"`
key by the active runtime,
matching Python `@task` behavior:
```ts
-import { Bundle, Dag, getClient, getContext } from "apache-airflow-ts-sdk";
+import { Bundle, getClient, getContext, TaskHandler } from
"apache-airflow-ts-sdk";
export async function sayHello() {
const greeting = await getClient().getVariable("greeting");
return { message: `Hello from ${getContext().taskId}: ${greeting}` };
}
-const dag = new Dag("example_dag");
-dag.task("say_hello", sayHello);
-
const bundle = new Bundle();
-bundle.register(dag);
+bundle.register(new TaskHandler("example_dag", "say_hello", sayHello));
await bundle.serve();
```
+`register` takes any number of items, so one bundle can provide for several
`TaskHandler`s.
+
+`Dag` is another interface, for a Dag declared in TypeScript rather than in
Python, and is still a work in progress.
+
## Coordinators
Airflow runs TypeScript task bundles through the Python-side `NodeCoordinator`
diff --git a/ts-sdk/example/README.md b/ts-sdk/example/README.md
index d8c690eec90..e9a1e43967b 100644
--- a/ts-sdk/example/README.md
+++ b/ts-sdk/example/README.md
@@ -21,9 +21,10 @@
This example shows the coordinator-mode shape for TypeScript task handlers:
-- `dags/typescript_example.py` declares the Airflow Dag and stub tasks.
-- `src/main.ts` registers TypeScript handlers for the same Dag/task IDs and
- starts the coordinator runtime.
+- `dags/typescript_example.py` and `dags/typescript_taskflow_example.py`
declare two Airflow Dags and their stub tasks.
+- `src/main.ts` and `src/taskflow.ts` register a `TaskHandler` per stub task
and start the coordinator runtime.
+ One bundle provides for both Dags, and both declare a task called
`build_message`.
+ A handler binds the `(dag_id, task_id)` pair, so the two are different tasks
with different bodies.
- `dist/bundle.mjs` is the generated Node.js bundle that Airflow launches.
The build uses the SDK's `airflow-ts-pack` tool, which bundles the entrypoint
@@ -70,7 +71,7 @@ export AIRFLOW__SDK__COORDINATORS='{
export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"typescript": "ts"}'
```
-Copy `dags/typescript_example.py` into your Airflow Dags folder.
+Copy both files in `dags/` into your Airflow Dags folder.
The example also uses one Variable and one Connection:
@@ -87,4 +88,5 @@ Then start Airflow and trigger the Dag:
```bash
airflow dags trigger typescript_example
+airflow dags trigger typescript_taskflow_example
```
diff --git a/ts-sdk/example/dags/typescript_taskflow_example.py
b/ts-sdk/example/dags/typescript_taskflow_example.py
new file mode 100644
index 00000000000..57b0d08ef9b
--- /dev/null
+++ b/ts-sdk/example/dags/typescript_taskflow_example.py
@@ -0,0 +1,57 @@
+# 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.
+
+
+"""
+A second Python-owned Dag served by the same TypeScript bundle.
+
+``typescript_example`` shows the basics; this Dag exists so one bundle
provides for two ``dag_id``s at once.
+Its ``build_message`` stub deliberately shares a ``task_id`` with a task in
``typescript_example``:
+a handler binds the ``(dag_id, task_id)`` pair, so the two are different tasks
with different bodies.
+See ``src/taskflow.ts``.
+"""
+
+from __future__ import annotations
+
+from airflow.sdk import dag, task
+
+
+@task
+def make_totals():
+ return {"orders": 12, "revenue": 3402.0}
+
+
[email protected](queue="typescript")
+def summarize(): ...
+
+
+# Same task_id as `typescript_example.build_message`, on purpose.
[email protected](queue="typescript")
+def build_message(): ...
+
+
+@dag(
+ dag_id="typescript_taskflow_example",
+ schedule=None,
+ catchup=False,
+ tags=["typescript", "example", "taskflow"],
+)
+def typescript_taskflow_example():
+ make_totals() >> summarize() >> build_message()
+
+
+typescript_taskflow_example()
diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts
index 3914148937a..d1644ade4e2 100644
--- a/ts-sdk/example/src/main.ts
+++ b/ts-sdk/example/src/main.ts
@@ -17,9 +17,14 @@
* under the License.
*/
-import { Bundle, Dag, getClient } from "apache-airflow-ts-sdk";
+// The bundle entry point: one bundle, two Python-owned Dags.
+//
+// Both Dags in `dags/` are declared in Python with `@task.stub` tasks routed
to the Node
+// coordinator, so this side only supplies the task bodies.
-const dag = new Dag("typescript_example");
+import { Bundle, getClient, TaskHandler } from "apache-airflow-ts-sdk";
+
+import { buildSummaryMessage, summarize } from "./taskflow.js";
export async function buildMessage() {
const client = getClient();
@@ -50,9 +55,14 @@ export async function readConnection() {
};
}
-dag.task("build_message", buildMessage);
-dag.task("read_connection", readConnection);
-
+// One register call lists everything this bundle provides.
+// `build_message` appears under both Dags: two different handlers, told apart
by the dag_id each
+// is bound to and never by the task_id alone.
const bundle = new Bundle();
-bundle.register(dag);
+bundle.register(
+ new TaskHandler("typescript_example", "build_message", buildMessage),
+ new TaskHandler("typescript_example", "read_connection", readConnection),
+ new TaskHandler("typescript_taskflow_example", "summarize", summarize),
+ new TaskHandler("typescript_taskflow_example", "build_message",
buildSummaryMessage),
+);
await bundle.serve();
diff --git a/ts-sdk/example/src/taskflow.ts b/ts-sdk/example/src/taskflow.ts
new file mode 100644
index 00000000000..54629cab87d
--- /dev/null
+++ b/ts-sdk/example/src/taskflow.ts
@@ -0,0 +1,75 @@
+/*!
+ * 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.
+ */
+
+// Handlers for the `typescript_taskflow_example` Dag.
+//
+// A second Python-owned Dag served by the same bundle, so the pair of ids a
handler binds is what
+// tells its tasks apart from `typescript_example`'s.
+// `build_summary_message` implements a task named `build_message`, exactly as
the other Dag has,
+// and the two share nothing else.
+
+import { getClient, getContext } from "apache-airflow-ts-sdk";
+
+/** What `make_totals` returns on the Python side. */
+export interface Totals {
+ orders: number;
+ revenue: number;
+}
+
+/** What {@link summarize} returns, and what {@link buildSummaryMessage}
reads. */
+export interface Summary {
+ orders: number;
+ averageOrder: number;
+ currency: string;
+}
+
+export async function summarize(): Promise<Summary> {
+ const totals = await getClient().getXCom<Totals>({
+ key: "return_value",
+ taskId: "make_totals",
+ });
+ if (totals === null) {
+ throw new Error(`task ${getContext().taskId} has no totals to summarize`);
+ }
+ const average = totals.orders === 0 ? 0 : totals.revenue / totals.orders;
+
+ return {
+ orders: totals.orders,
+ averageOrder: Number(average.toFixed(2)),
+ currency: "GBP",
+ };
+}
+
+export async function buildSummaryMessage() {
+ const ctx = getContext();
+ const summary = await getClient().getXCom<Summary>({
+ key: "return_value",
+ taskId: "summarize",
+ });
+ if (summary === null) {
+ throw new Error(`task ${ctx.taskId} has no summary to report`);
+ }
+
+ return {
+ // The dag_id is in the return value on purpose: it is how the end-to-end
test tells this task
+ // apart from `typescript_example`'s `build_message`.
+ dagId: ctx.dagId,
+ message: `${summary.orders} orders averaging ${summary.averageOrder}
${summary.currency}`,
+ };
+}
diff --git a/ts-sdk/scripts/verify-package.mjs
b/ts-sdk/scripts/verify-package.mjs
index e8064682da4..463a21e679e 100644
--- a/ts-sdk/scripts/verify-package.mjs
+++ b/ts-sdk/scripts/verify-package.mjs
@@ -34,7 +34,7 @@ export const REQUIRED_ROOT_FILES = ["LICENSE", "NOTICE",
"README.md", "package.j
// The Dag-authoring entrypoints a consumer must be able to reach from the
package root.
// The two getters are here because a handler cannot reach the runtime without
them: a
// published build that dropped them would still import, and fail at the first
task.
-export const REQUIRED_ROOT_EXPORTS = ["Bundle", "Dag", "getClient",
"getContext"];
+export const REQUIRED_ROOT_EXPORTS = ["Bundle", "Dag", "TaskHandler",
"getClient", "getContext"];
function run(command, args, options = {}) {
const result = spawnSync(command, args, {
diff --git a/ts-sdk/src/cli/pack.ts b/ts-sdk/src/cli/pack.ts
index 5acea1b9626..08a3aff008b 100644
--- a/ts-sdk/src/cli/pack.ts
+++ b/ts-sdk/src/cli/pack.ts
@@ -200,7 +200,9 @@ export async function runPack(argv: readonly string[]):
Promise<void> {
const manifest = readBundleManifest(stagingPath);
const dagEntries = Object.entries(manifest.dags);
if (dagEntries.length === 0) {
- throw new Error(`${args.entry} served no Dags; register them with
bundle.register(...)`);
+ throw new Error(
+ `${args.entry} served nothing; register Dags or task handlers with
bundle.register(...)`,
+ );
}
// Warn rather than fail, as airflow-go-pack does: the shared schema
allows a
// Dag with no tasks.
diff --git a/ts-sdk/src/coordinator/runtime.ts
b/ts-sdk/src/coordinator/runtime.ts
index 0e05ec24b58..e34e37c1303 100644
--- a/ts-sdk/src/coordinator/runtime.ts
+++ b/ts-sdk/src/coordinator/runtime.ts
@@ -23,10 +23,10 @@
//
// node my-bundle.mjs --comm=host:port --logs=host:port
//
-// where `my-bundle.mjs` is a user-bundled Node script that imports
-// the SDK, creates `Dag` objects, attaches a handler per task with
-// `dag.task(...)`, registers them on a `Bundle`, then awaits `bundle.serve()`.
-// Each handler runs inside a task scope, which is what `getContext()` and
+// where `my-bundle.mjs` is a user-bundled Node script that imports the SDK,
+// registers what it provides on a `Bundle` (a `TaskHandler` per Python-owned
+// task, a `Dag` per natively declared one), then awaits `bundle.serve()`. Each
+// handler runs inside a task scope, which is what `getContext()` and
// `getClient()` read.
//
// Lifecycle:
diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts
index edf2ceaafc9..935d58f95fd 100644
--- a/ts-sdk/src/index.ts
+++ b/ts-sdk/src/index.ts
@@ -19,6 +19,7 @@
export { Dag } from "./sdk/dag.js";
export { Bundle } from "./sdk/bundle.js";
+export { TaskHandler } from "./sdk/task-handler.js";
export { getClient, getContext } from "./sdk/task.js";
export { ConnectionNotFoundError, VariableNotFoundError } from
"./sdk/client.js";
export { SUPERVISOR_API_VERSION } from "./coordinator/index.js";
diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts
index 615fe971ee9..612adbac8b7 100644
--- a/ts-sdk/src/sdk/bundle.ts
+++ b/ts-sdk/src/sdk/bundle.ts
@@ -21,27 +21,38 @@
import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js";
import { Dag, getDagTaskRecords, isDag, type TaskRef } from "./dag.js";
+import { getTaskHandlerFunction, isTaskHandler, TaskHandler } from
"./task-handler.js";
import type { TaskFunction } from "./task.js";
// Assigned inside Bundle's static block, as Dag does for its tasks.
-let dagsOf: (bundle: Bundle) => ReadonlyMap<string, Dag>;
+let entriesOf: (bundle: Bundle) => ReadonlyMap<string, BundleEntry>;
/**
- * Anything {@link Bundle.register} accepts.
- *
- * A union rather than a base class or an interface: TypeScript's equivalent of
- * the sealed interface the Go SDK uses for the same purpose. Registering gains
- * a kind by gaining an arm here, never a second verb.
+ * What {@link Bundle.register} and the {@link Bundle} constructor accept: a
+ * {@link TaskHandler} for a task that a Python Dag declares, or a {@link Dag}
+ * declared in TypeScript.
*/
-export type Registerable = Dag;
+export type Registerable = Dag | TaskHandler;
+
+// What a bundle holds per dag_id. The two arms are exclusive by construction:
+// a Dag is the native case and owns its own tasks, while task handlers supply
+// bodies for a Dag that Python declares, so one dag_id is never both.
+type BundleEntry =
+ | { readonly kind: "dag"; readonly dag: Dag }
+ | { readonly kind: "handlers"; readonly handlers: Map<string, TaskFunction>
};
+
+function entryTaskIds(entry: BundleEntry): string[] {
+ return entry.kind === "dag" ? [...entry.dag.taskIds] :
[...entry.handlers.keys()];
+}
/** Internal: whether `value` is a Bundle built by any copy of this package. */
export function isBundle(value: unknown): value is Bundle {
return hasBrand(value, "Bundle");
}
-/** Internal: a registered Dag with its task IDs, as {@link listBundleDags}
reports it.
- * A task-less Dag is included, so the bundle manifest keeps it visible. */
+/** Internal: a Dag this bundle provides for, with its task IDs, as
+ * {@link listBundleDags} reports it. A task-less native Dag is included, so
+ * the bundle manifest keeps it visible. */
export interface RegisteredDag {
/** Identifier of the registered Dag. */
readonly dagId: string;
@@ -73,10 +84,12 @@ export interface RegisteredDag {
* a snapshot of its tasks.
*/
export class Bundle {
- readonly #dags = new Map<string, Dag>();
+ // Keyed by dag_id and insertion-ordered, so the bundle manifest lists what
+ // this process provides in the order the entry point registered it.
+ readonly #entries = new Map<string, BundleEntry>();
static {
- dagsOf = (bundle) => bundle.#dags;
+ entriesOf = (bundle) => bundle.#entries;
}
/** Registers `items`, on the same terms as {@link register}. */
@@ -91,28 +104,87 @@ export class Bundle {
* The constructor covers the common case; this is for a bundle that
* collects what it provides across several modules. */
register(...items: Registerable[]): void {
- const incoming = new Set<string>();
+ // Validated against what is already held *and* against this call, in full,
+ // before anything is written: a call that throws registers none of its
+ // items, so a bundle never half-provides what its author listed once.
+ const incomingDags = new Set<string>();
+ const incomingTaskHandlers = new Set<string>();
for (const item of items) {
// Typed as Registerable, so narrowing it would collapse to never; these
// guard callers reaching this from plain JavaScript.
const candidate: unknown = item;
- // Another copy's Dag cannot be registered, since lookups read a private
- // task map keyed to this copy's class, so it is rejected by its cause.
- if (!(candidate instanceof Dag)) {
+ // Another copy's value cannot be registered, since both kinds read
+ // private state keyed to this copy's class, so it is rejected by its
cause.
+ if (candidate instanceof Dag) {
+ this.#checkDag(candidate, incomingDags);
+ } else if (candidate instanceof TaskHandler) {
+ this.#checkTaskHandler(candidate, incomingDags, incomingTaskHandlers);
+ } else if (isDag(candidate)) {
+ throw new Error(`Dag "${(candidate as Dag).dagId}"
${DUPLICATE_COPY_HINT}`);
+ } else if (isTaskHandler(candidate)) {
+ const foreign = candidate as TaskHandler;
throw new Error(
- isDag(candidate)
- ? `Dag "${candidate.dagId}" ${DUPLICATE_COPY_HINT}`
- : "only Dag instances can be registered",
+ `Task handler for Dag "${foreign.dagId}" task "${foreign.taskId}"
${DUPLICATE_COPY_HINT}`,
);
+ } else {
+ throw new Error("only Dag and TaskHandler instances can be
registered");
}
- if (this.#dags.has(item.dagId) || incoming.has(item.dagId)) {
- throw new Error(`Dag "${item.dagId}" is already registered`);
- }
- incoming.add(item.dagId);
}
for (const item of items) {
- this.#dags.set(item.dagId, item);
+ if (item instanceof Dag) {
+ this.#entries.set(item.dagId, { kind: "dag", dag: item });
+ } else {
+ const existing = this.#entries.get(item.dagId);
+ const handlers =
+ existing?.kind === "handlers" ? existing.handlers : new Map<string,
TaskFunction>();
+ handlers.set(item.taskId, getTaskHandlerFunction(item));
+ if (existing === undefined) {
+ this.#entries.set(item.dagId, { kind: "handlers", handlers });
+ }
+ }
+ }
+ }
+
+ #checkDag(dag: Dag, incomingDags: Set<string>): void {
+ if (this.#entries.get(dag.dagId)?.kind === "handlers") {
+ throw new Error(
+ `Dag "${dag.dagId}" already has registered task handlers; a Dag
declared in ` +
+ "TypeScript owns its own tasks, so one Dag ID cannot have both",
+ );
+ }
+ if (this.#entries.has(dag.dagId) || incomingDags.has(dag.dagId)) {
+ throw new Error(`Dag "${dag.dagId}" is already registered`);
+ }
+ incomingDags.add(dag.dagId);
+ }
+
+ #checkTaskHandler(
+ handler: TaskHandler,
+ incomingDags: Set<string>,
+ incomingTaskHandlers: Set<string>,
+ ): void {
+ // A native Dag attaches its tasks with dag.task(...), so a task handler
for
+ // the same Dag ID would be a second, disagreeing source for its task list.
+ if (this.#entries.get(handler.dagId)?.kind === "dag" ||
incomingDags.has(handler.dagId)) {
+ throw new Error(
+ `Dag "${handler.dagId}" is declared in TypeScript; attach its tasks
with ` +
+ "dag.task(...) rather than registering task handlers for them",
+ );
+ }
+ const entry = this.#entries.get(handler.dagId);
+ // Keyed on the pair, not the task ID: one bundle serves several Dags, and
+ // the same task_id under two of them is two different handlers.
+ const pair = `${handler.dagId}\u0000${handler.taskId}`;
+ if (
+ entry?.kind === "handlers"
+ ? entry.handlers.has(handler.taskId)
+ : incomingTaskHandlers.has(pair)
+ ) {
+ throw new Error(
+ `A handler for Dag "${handler.dagId}" task "${handler.taskId}" is
already registered`,
+ );
}
+ incomingTaskHandlers.add(pair);
}
/**
@@ -143,8 +215,11 @@ export class Bundle {
/** Look up a registered handler, the way the runtime dispatches a task.
* Returns `undefined` when no handler exists. */
getTaskHandler(dagId: string, taskId: string): TaskFunction | undefined {
- const dag = this.#dags.get(dagId);
- return dag ? getDagTaskRecords(dag).get(taskId)?.fn : undefined;
+ const entry = this.#entries.get(dagId);
+ if (entry === undefined) return undefined;
+ return entry.kind === "dag"
+ ? getDagTaskRecords(entry.dag).get(taskId)?.fn
+ : entry.handlers.get(taskId);
}
}
@@ -159,18 +234,21 @@ export function validateOwnBundle(value: unknown,
accessor: string): asserts val
);
}
-/** Internal: the task handles across a bundle's Dags. Not re-exported from the
- * package root: enumerating what the runtime dispatches is the runtime's
job. */
+/** Internal: every task handle this bundle can dispatch, across both kinds.
Not
+ * re-exported from the package root: enumerating what the runtime dispatches
+ * is the runtime's job. */
export function listBundleTasks(bundle: Bundle): TaskRef[] {
- return [...dagsOf(bundle).values()].flatMap((dag) =>
- [...getDagTaskRecords(dag).values()].map((record) => record.task),
+ return [...entriesOf(bundle)].flatMap(([dagId, entry]) =>
+ entryTaskIds(entry).map((taskId) => ({ dagId, taskId })),
);
}
-/** Internal: every registered Dag with its task IDs, empty Dags included. */
+/** Internal: every Dag this bundle provides for, with its task IDs. A native
+ * Dag with no tasks is included; a Dag known only through task handlers
always
+ * has at least one, since a handler is what put it here. */
export function listBundleDags(bundle: Bundle): RegisteredDag[] {
- return [...dagsOf(bundle).values()].map((dag) => ({
- dagId: dag.dagId,
- tasks: [...dag.taskIds],
+ return [...entriesOf(bundle)].map(([dagId, entry]) => ({
+ dagId,
+ tasks: entryTaskIds(entry),
}));
}
diff --git a/ts-sdk/src/sdk/task-handler.ts b/ts-sdk/src/sdk/task-handler.ts
new file mode 100644
index 00000000000..a682b586d1b
--- /dev/null
+++ b/ts-sdk/src/sdk/task-handler.ts
@@ -0,0 +1,97 @@
+/*!
+ * 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 mixed-language authoring surface: a TypeScript function bound to the
+// Python-owned task it implements.
+
+import { brand, hasBrand } from "./brand.js";
+import type { TaskFunction } from "./task.js";
+
+// Assigned inside TaskHandler's static block, as Dag does for its tasks.
+let functionOf: (handler: TaskHandler) => TaskFunction;
+
+/** Internal: whether `value` is a TaskHandler built by any copy of this
package. */
+export function isTaskHandler(value: unknown): value is TaskHandler {
+ return hasBrand(value, "TaskHandler");
+}
+
+function requireId(label: string, value: string): string {
+ // Typed as string, so plain JavaScript is what these catch, along with the
+ // empty string, which types cannot rule out and which no Airflow id can be.
+ const candidate: unknown = value;
+ if (typeof candidate !== "string" || candidate.length === 0) {
+ throw new Error(`${label} for a task handler must be a non-empty string`);
+ }
+ return value;
+}
+
+/**
+ * A TypeScript function bound to the Python-owned task it implements.
+ *
+ * A mixed-language Dag declares its structure in Python, with a `@task.stub`
+ * per task routed to the Node coordinator. TypeScript supplies the task bodies
+ * and nothing else: no dag_id of its own, no schedule, no task order.
+ *
+ * ```ts
+ * const bundle = new Bundle();
+ * bundle.register(new TaskHandler("etl", "transform", transform));
+ * await bundle.serve();
+ * ```
+ *
+ * `dagId` must match the Python Dag's `dag_id` and `taskId` a `@task.stub` in
+ * it, including any TaskGroup prefix. Both are written out: nothing is derived
+ * from the handler's function name, which the build step is free to rename.
+ *
+ * Identity and a body, and no more. A handler has no factory to call, so
wiring
+ * one the way a natively declared task is wired is a compile error rather than
+ * a runtime throw. For a native Dag, use {@link Dag} instead.
+ */
+export class TaskHandler<TReturn = unknown> {
+ /** Identifier of the Python Dag this task belongs to. */
+ readonly dagId: string;
+ /** Airflow task ID this handler implements, including any TaskGroup prefix.
*/
+ readonly taskId: string;
+ readonly #handler: TaskFunction<TReturn>;
+
+ static {
+ functionOf = (handler) => handler.#handler as TaskFunction;
+ }
+
+ constructor(dagId: string, taskId: string, handler: TaskFunction<TReturn>) {
+ this.dagId = requireId("dagId", dagId);
+ this.taskId = requireId("taskId", taskId);
+ if (typeof handler !== "function") {
+ throw new Error(`handler for Dag "${dagId}" task "${taskId}" must be a
function`);
+ }
+ brand(this, "TaskHandler");
+ this.#handler = handler;
+ }
+}
+
+/**
+ * Internal: the function a TaskHandler carries, for bundle dispatch.
+ *
+ * Not re-exported from the package root, and the package `"exports"` map
blocks
+ * deep imports, so this is unreachable from outside the SDK. The field is
+ * private for the same reason `TaskRef` does not expose its handler: what a
+ * handler binds is identity, and reaching the body is the runtime's business.
+ */
+export function getTaskHandlerFunction(handler: TaskHandler): TaskFunction {
+ return functionOf(handler);
+}
diff --git a/ts-sdk/tests/cli/fixtures/entry.ts
b/ts-sdk/tests/cli/fixtures/entry.ts
index bd388974a2e..607c62cdf0b 100644
--- a/ts-sdk/tests/cli/fixtures/entry.ts
+++ b/ts-sdk/tests/cli/fixtures/entry.ts
@@ -17,12 +17,15 @@
* under the License.
*/
-import { Bundle, Dag } from "../../../src/index.js";
+import { Bundle, Dag, TaskHandler } from "../../../src/index.js";
-const fixtureDag = new Dag("fixture_dag");
-fixtureDag.task("extract", async () => "extracted");
-fixtureDag.task("transform", async () => "transformed");
+// A mixed bundle: task handlers for a Python-owned Dag, plus a natively
+// declared one, so packing covers both paths into the manifest.
const otherDag = new Dag("other_dag");
otherDag.task("solo", async () => undefined);
-await new Bundle(fixtureDag, otherDag).serve();
+await new Bundle(
+ new TaskHandler("fixture_dag", "extract", async () => "extracted"),
+ new TaskHandler("fixture_dag", "transform", async () => "transformed"),
+ otherDag,
+).serve();
diff --git a/ts-sdk/tests/cli/pack.test.ts b/ts-sdk/tests/cli/pack.test.ts
index cd5dc1ca0b6..34d4b3a3d86 100644
--- a/ts-sdk/tests/cli/pack.test.ts
+++ b/ts-sdk/tests/cli/pack.test.ts
@@ -286,10 +286,12 @@ describe("runPack", () => {
expect(existsSync(path.join(outdir,
"bundle.pack-staging.mjs"))).toBe(false);
});
- it("leaves no bundle behind when the entry serves no Dags", async () => {
+ it("leaves no bundle behind when the entry serves nothing", async () => {
outdir = mkdtempSync(path.join(tmpdir(), "ts-pack-"));
- await expect(runPack([EMPTY_ENTRY, "--outdir",
outdir])).rejects.toThrow("served no Dags");
+ await expect(runPack([EMPTY_ENTRY, "--outdir", outdir])).rejects.toThrow(
+ "served nothing; register Dags or task handlers",
+ );
expect(existsSync(path.join(outdir, "bundle.mjs"))).toBe(false);
expect(existsSync(path.join(outdir,
"bundle.pack-staging.mjs"))).toBe(false);
});
diff --git a/ts-sdk/tests/coordinator/integration.test.ts
b/ts-sdk/tests/coordinator/integration.test.ts
index e6294154187..fabd39a089c 100644
--- a/ts-sdk/tests/coordinator/integration.test.ts
+++ b/ts-sdk/tests/coordinator/integration.test.ts
@@ -37,13 +37,15 @@ import {
} from "../../src/coordinator/runtime.js";
import { Dag } from "../../src/sdk/dag.js";
import { Bundle } from "../../src/sdk/bundle.js";
+import { TaskHandler } from "../../src/sdk/task-handler.js";
import { getClient, getContext } from "../../src/sdk/task.js";
const testDag = new Dag("test_dag");
const otherDag = new Dag("other_dag");
// The bundle the runtime dispatches through. startCoordinator() is driven
// directly rather than through bundle.serve(), so these tests can supply mock
-// socket addresses.
+// socket addresses. Both authoring kinds are registered in one call, since the
+// runtime dispatches through one lookup regardless of which put a task there.
const bundle = new Bundle(testDag, otherDag);
interface MockResult {
@@ -356,6 +358,51 @@ describe("coordinator runtime integration", () => {
});
});
+ it("dispatches to a task handler registered for a Python-owned Dag", async
() => {
+ // The mixed-language path: no Dag object exists for `py_dag` on this side,
+ // only a handler bound to its dag_id and task_id.
+ let observedCtx: unknown = null;
+ bundle.register(
+ new TaskHandler("py_dag", "transform", async () => {
+ observedCtx = getContext();
+ return "transformed";
+ }),
+ );
+
+ const result = await driveSupervisor(makeStartupDetails("transform",
"py_dag"));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+ expect(observedCtx).toMatchObject({ dagId: "py_dag", taskId: "transform"
});
+ const setXComReqs = result.runtimeRequests.filter((r) => r.type ===
"SetXCom");
+ expect(setXComReqs[0]!.body).toMatchObject({
+ key: "return_value",
+ value: "transformed",
+ dag_id: "py_dag",
+ task_id: "transform",
+ });
+ });
+
+ it("keeps two Dags' same-named tasks apart when dispatching", async () => {
+ // Registration keys on the pair, and so must dispatch: the runtime is
+ // handed dag_id and task_id together and must not answer from the wrong
one.
+ bundle.register(
+ new TaskHandler("pair_dag_a", "shared", async () => "from a"),
+ new TaskHandler("pair_dag_b", "shared", async () => "from b"),
+ );
+
+ for (const [dagId, expected] of [
+ ["pair_dag_a", "from a"],
+ ["pair_dag_b", "from b"],
+ ]) {
+ const result = await driveSupervisor(makeStartupDetails("shared",
dagId));
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(result.runtimeRequests.find((r) => r.type ===
"SetXCom")!.body).toMatchObject({
+ value: expected,
+ dag_id: dagId,
+ });
+ }
+ });
+
it("aborts the context signal on SIGTERM and reports a thrown task error",
async () => {
let sawAbort = false;
testDag.task("aborted_then_failed", async () => {
diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts
index ad94cd389a0..dae26146e1a 100644
--- a/ts-sdk/tests/public-api.test.ts
+++ b/ts-sdk/tests/public-api.test.ts
@@ -41,6 +41,7 @@ import {
getClient,
getContext,
SUPERVISOR_API_VERSION,
+ TaskHandler,
VariableNotFoundError,
} from "../src/index.js";
@@ -155,6 +156,37 @@ describe("public API", () => {
expectTypeOf<typeof sdk>().not.toHaveProperty("startCoordinator");
});
+ it("exports TaskHandler as the mixed-language authoring surface", () => {
+ const transform = async () => "transformed";
+ const bundle = new Bundle(new TaskHandler("py_etl", "transform",
transform));
+
+ expect(bundle.getTaskHandler("py_etl", "transform")).toBe(transform);
+ // Identity and a body, and no more: no schedule, no task order, no dag_id
+ // of its own to declare.
+ expectTypeOf<keyof TaskHandler>().toEqualTypeOf<"dagId" | "taskId">();
+ expectTypeOf<ConstructorParameters<typeof TaskHandler>>().toEqualTypeOf<
+ [string, string, TaskFunction<unknown>]
+ >();
+ });
+
+ it("does not let a task handler be wired the way a native task is", () => {
+ // The guarantee an earlier draft's separate MixedLangDag class existed to
+ // provide: a handler has no factory to call, so calling one is a compile
+ // error rather than a runtime throw.
+ const rejectsFactoryMisuse = () => {
+ const handler = new TaskHandler("py_etl", "transform", async () =>
undefined);
+ // @ts-expect-error a task handler is a value, not a callable task
factory.
+ handler();
+ // @ts-expect-error dagId and taskId are positional, not an options
object.
+ new TaskHandler({ dagId: "py_etl", taskId: "transform" }, async () =>
undefined);
+ // @ts-expect-error the task_id is always written out, never derived.
+ new TaskHandler("py_etl", async () => undefined);
+ // @ts-expect-error a handler does not expose the function it carries.
+ void handler.handler;
+ };
+ void rejectsFactoryMisuse;
+ });
+
describe("the task-handler getters", () => {
it("throw outside a handler, naming the accessor", () => {
// The full scope behaviour is covered in tests/sdk/task-scope.test.ts;
@@ -194,7 +226,7 @@ describe("public API", () => {
expectTypeOf<Bundle["serve"]>().toEqualTypeOf<() => Promise<void>>();
expectTypeOf<Bundle["register"]>().toEqualTypeOf<(...items:
Registerable[]) => void>();
expectTypeOf<ConstructorParameters<typeof
Bundle>>().toEqualTypeOf<Registerable[]>();
- expectTypeOf<Registerable>().toEqualTypeOf<Dag>();
+ expectTypeOf<Registerable>().toEqualTypeOf<Dag | TaskHandler>();
for (const name of ["serveDags", "DagRegistry"]) {
expect(name in sdk).toBe(false);
}
diff --git a/ts-sdk/tests/sdk/bundle.test.ts b/ts-sdk/tests/sdk/bundle.test.ts
index e5725cb7810..37f6f68b663 100644
--- a/ts-sdk/tests/sdk/bundle.test.ts
+++ b/ts-sdk/tests/sdk/bundle.test.ts
@@ -49,9 +49,9 @@ describe("Bundle", () => {
);
});
- it("rejects constructor values that are not Dag instances", () => {
+ it("rejects constructor values that are neither a Dag nor a task handler",
() => {
expect(() => new Bundle({ dagId: "example_dag" } as unknown as
Dag)).toThrowError(
- /only Dag instances can be registered/,
+ /only Dag and TaskHandler instances can be registered/,
);
});
@@ -111,10 +111,10 @@ describe("Bundle", () => {
expect(listBundleTasks(bundle)).toEqual([]);
});
- it("rejects values that are not Dag instances", () => {
+ it("rejects values that are neither a Dag nor a task handler", () => {
const bundle = new Bundle();
expect(() => bundle.register({ dagId: "example_dag" } as unknown as
Dag)).toThrowError(
- /only Dag instances can be registered/,
+ /only Dag and TaskHandler instances can be registered/,
);
});
diff --git a/ts-sdk/tests/sdk/task-handler.test.ts
b/ts-sdk/tests/sdk/task-handler.test.ts
new file mode 100644
index 00000000000..78bb1541802
--- /dev/null
+++ b/ts-sdk/tests/sdk/task-handler.test.ts
@@ -0,0 +1,193 @@
+/*!
+ * 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, expect, it } from "vitest";
+
+import { Bundle, listBundleDags, listBundleTasks } from
"../../src/sdk/bundle.js";
+import { Dag } from "../../src/sdk/dag.js";
+import { getTaskHandlerFunction, TaskHandler } from
"../../src/sdk/task-handler.js";
+
+describe("TaskHandler", () => {
+ it("binds a function to the Python-owned task it implements", () => {
+ const transform = async () => "transformed";
+ const handler = new TaskHandler("etl", "transform", transform);
+
+ expect(handler.dagId).toBe("etl");
+ expect(handler.taskId).toBe("transform");
+ expect(getTaskHandlerFunction(handler)).toBe(transform);
+ });
+
+ it("derives nothing from the function's name", () => {
+ // The build step is free to rename or inline a function, so the task_id is
+ // always written out and an anonymous handler is perfectly ordinary.
+ const handler = new TaskHandler("etl", "transform", async () => undefined);
+ expect(handler.taskId).toBe("transform");
+ });
+
+ it.each([
+ ["an empty dagId", "", "transform", /dagId for a task handler must be a
non-empty string/],
+ ["an empty taskId", "etl", "", /taskId for a task handler must be a
non-empty string/],
+ ])("rejects %s", (_label, dagId, taskId, message) => {
+ expect(() => new TaskHandler(dagId, taskId, async () =>
undefined)).toThrowError(message);
+ });
+
+ it("rejects a handler that is not a function", () => {
+ expect(
+ () => new TaskHandler("etl", "transform", "not a function" as unknown as
() => void),
+ ).toThrowError(/handler for Dag "etl" task "transform" must be a
function/);
+ });
+
+ it("does not expose the function it carries", () => {
+ // As with TaskRef: what a handler binds is identity. Reaching the body is
+ // the runtime's business, through an accessor the package root never
ships.
+ const handler = new TaskHandler("etl", "transform", async () => undefined);
+ for (const name of ["handler", "fn", "run", "call"]) {
+ expect(name in handler).toBe(false);
+ }
+ });
+});
+
+describe("a bundle of task handlers", () => {
+ it("registers Dags and task handlers in one call", () => {
+ const nativeDag = new Dag("native_etl");
+ nativeDag.task("extract", async () => undefined);
+ const transform = async () => "transformed";
+
+ const bundle = new Bundle();
+ bundle.register(nativeDag, new TaskHandler("py_etl", "transform",
transform));
+
+ expect(bundle.getTaskHandler("py_etl", "transform")).toBe(transform);
+ expect(bundle.getTaskHandler("native_etl", "extract")).toBeDefined();
+ });
+
+ it("dispatches on the Dag/task pair, not the task ID alone", () => {
+ // The property the flattened map can get wrong and the old per-Dag map
+ // could not: one bundle serves several Dags, and the same task_id under
+ // two of them is two different handlers.
+ const first = async () => "from etl";
+ const second = async () => "from reporting";
+
+ const bundle = new Bundle(
+ new TaskHandler("etl", "build_message", first),
+ new TaskHandler("reporting", "build_message", second),
+ );
+
+ expect(bundle.getTaskHandler("etl", "build_message")).toBe(first);
+ expect(bundle.getTaskHandler("reporting", "build_message")).toBe(second);
+ expect(bundle.getTaskHandler("unknown", "build_message")).toBeUndefined();
+ expect(bundle.getTaskHandler("etl", "unknown")).toBeUndefined();
+ });
+
+ it("lists every Dag it provides for, in registration order", () => {
+ const nativeDag = new Dag("native_etl");
+ nativeDag.task("extract", async () => undefined);
+
+ const bundle = new Bundle(
+ new TaskHandler("py_etl", "transform", async () => undefined),
+ nativeDag,
+ new TaskHandler("py_etl", "report", async () => undefined),
+ );
+
+ // A Dag registered through handlers keeps its place from the first handler
+ // that named it, so a later one does not reorder the manifest.
+ expect(listBundleDags(bundle)).toEqual([
+ { dagId: "py_etl", tasks: ["transform", "report"] },
+ { dagId: "native_etl", tasks: ["extract"] },
+ ]);
+ expect(listBundleTasks(bundle)).toEqual([
+ { dagId: "py_etl", taskId: "transform" },
+ { dagId: "py_etl", taskId: "report" },
+ { dagId: "native_etl", taskId: "extract" },
+ ]);
+ });
+
+ it("accumulates handlers for one Dag across several calls", () => {
+ const bundle = new Bundle();
+ bundle.register(new TaskHandler("py_etl", "transform", async () =>
undefined));
+ bundle.register(new TaskHandler("py_etl", "report", async () =>
undefined));
+
+ expect(listBundleDags(bundle)).toEqual([{ dagId: "py_etl", tasks:
["transform", "report"] }]);
+ });
+
+ it("rejects a second handler for the same Dag and task", () => {
+ const bundle = new Bundle(new TaskHandler("etl", "transform", async () =>
undefined));
+ expect(() =>
+ bundle.register(new TaskHandler("etl", "transform", async () =>
undefined)),
+ ).toThrowError(/A handler for Dag "etl" task "transform" is already
registered/);
+ });
+
+ it("rejects a duplicate Dag and task within a single call", () => {
+ expect(
+ () =>
+ new Bundle(
+ new TaskHandler("etl", "transform", async () => undefined),
+ new TaskHandler("etl", "transform", async () => undefined),
+ ),
+ ).toThrowError(/A handler for Dag "etl" task "transform" is already
registered/);
+ });
+
+ it("registers none of its items when a call throws", () => {
+ const bundle = new Bundle();
+ expect(() =>
+ bundle.register(
+ new TaskHandler("etl", "transform", async () => undefined),
+ new TaskHandler("etl", "transform", async () => undefined),
+ ),
+ ).toThrowError(/already registered/);
+ expect(listBundleDags(bundle)).toEqual([]);
+ expect(bundle.getTaskHandler("etl", "transform")).toBeUndefined();
+ });
+
+ it("rejects a task handler for a Dag declared in TypeScript", () => {
+ // A native Dag attaches its tasks with dag.task(...), so a handler for the
+ // same Dag ID would be a second, disagreeing source for its task list.
+ const bundle = new Bundle(new Dag("native_etl"));
+ expect(() =>
+ bundle.register(new TaskHandler("native_etl", "transform", async () =>
undefined)),
+ ).toThrowError(/is declared in TypeScript; attach its tasks with
dag\.task/);
+ });
+
+ it("rejects a task handler for a Dag declared in TypeScript in the same
call", () => {
+ expect(
+ () =>
+ new Bundle(
+ new Dag("native_etl"),
+ new TaskHandler("native_etl", "transform", async () => undefined),
+ ),
+ ).toThrowError(/is declared in TypeScript; attach its tasks with
dag\.task/);
+ });
+
+ it("rejects a Dag whose ID already has task handlers", () => {
+ const bundle = new Bundle(new TaskHandler("py_etl", "transform", async ()
=> undefined));
+ expect(() => bundle.register(new Dag("py_etl"))).toThrowError(
+ /already has registered task handlers/,
+ );
+ });
+
+ it("names the duplicate-copy cause for a handler carrying the brand but not
this class", () => {
+ // Stands in for a handler from a second resolved copy: same brand, other
+ // class. Its private function field is unreadable here, so the point is
+ // only that it says why.
+ const foreign = { dagId: "etl", taskId: "transform" };
+ Object.defineProperty(foreign, Symbol.for("airflow.ts-sdk.TaskHandler"), {
value: true });
+ expect(() => new Bundle(foreign as unknown as TaskHandler)).toThrowError(
+ /Task handler for Dag "etl" task "transform" comes from a different
copy/,
+ );
+ });
+});
