Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 43 additions & 5 deletions task-sdk/src/airflow/sdk/coordinators/node/coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

from __future__ import annotations

import base64
import os
import pathlib
from typing import TYPE_CHECKING, Any
Expand All @@ -45,6 +46,38 @@

BUNDLE_FILENAME = "bundle.mjs"
METADATA_FILENAME = "airflow-metadata.yaml"
EMBEDDED_METADATA_MARKER = b"//# airflowMetadata="
EMBEDDED_METADATA_MAX_BYTES = 1024 * 1024


def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any] | None:
Comment thread
jason810496 marked this conversation as resolved.
"""
Read the manifest ``airflow-ts-pack`` embeds in the bundle itself.

The packer prepends the ``airflow-metadata.yaml`` content as a leading
``//# airflowMetadata=<base64>`` line comment, keeping bundle and metadata
a single artifact. Returns ``None`` when the bundle has no such marker.
"""
try:
with bundle_path.open("rb") as bundle_file:
line = bundle_file.readline(EMBEDDED_METADATA_MAX_BYTES + 1)
except OSError as exc:
raise ValueError(f"cannot read {bundle_path.name}: {exc}") from exc

if not line.startswith(EMBEDDED_METADATA_MARKER):
return None
Comment thread
guan404ming marked this conversation as resolved.
if len(line) > EMBEDDED_METADATA_MAX_BYTES:
raise ValueError(
f"embedded airflow metadata exceeds {EMBEDDED_METADATA_MAX_BYTES} bytes; "
f"rebuild {bundle_path.name} with airflow-ts-pack"
)

payload = line[len(EMBEDDED_METADATA_MARKER) :].strip()
try:
decoded = base64.b64decode(payload, validate=True)
except ValueError as exc:
raise ValueError(f"cannot parse embedded airflow metadata: {exc}") from exc
Comment thread
guan404ming marked this conversation as resolved.
return parse_metadata_mapping(decoded, source="embedded airflow metadata")


def _read_bundle_metadata(metadata_path: pathlib.Path) -> dict[str, Any]:
Expand All @@ -63,8 +96,9 @@ def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle:
"""
Locate the ``.mjs`` entry point in *bundles_root*.

Scans each configured directory for ``bundle.mjs`` and reads the sibling
``airflow-metadata.yaml`` for the bundle's supervisor schema version.
Scans each configured directory for ``bundle.mjs`` and reads the bundle's
supervisor schema version from the metadata embedded in the bundle,
falling back to a sibling ``airflow-metadata.yaml`` sidecar.

This is an ordered fallback search, not Dag/task-aware multi-bundle
routing. The first bundle found wins. A future version can use the
Expand All @@ -77,7 +111,9 @@ def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle:
if not candidate.is_file():
continue
try:
metadata = _read_bundle_metadata(root / METADATA_FILENAME)
metadata = _read_embedded_metadata(candidate)
if metadata is None:
metadata = _read_bundle_metadata(root / METADATA_FILENAME)
log.debug("Selected TypeScript bundle", path=candidate, root=root)
return ResolvedBundle(
path=candidate,
Expand Down Expand Up @@ -123,8 +159,10 @@ class NodeCoordinator(SubprocessCoordinator):
``"node"``, which relies on ``$PATH``).
:param bundles_root: Ordered list of directories scanned for a usable
TypeScript bundle. Each bundle directory must contain ``bundle.mjs``
and ``airflow-metadata.yaml``. This is a fallback search path; it does
not yet route different Dag/task pairs to different bundles.
with embedded metadata (as produced by ``airflow-ts-pack``), or
``bundle.mjs`` plus an ``airflow-metadata.yaml`` sidecar. This is a
fallback search path; it does not yet route different Dag/task pairs
to different bundles.
:param task_startup_timeout: Maximum time the coordinator waits for a task
process to start, in seconds. The default is 10 seconds.
"""
Expand Down
86 changes: 68 additions & 18 deletions task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,15 @@

from __future__ import annotations

import base64
import pathlib

import pytest
from uuid6 import uuid7

from airflow.sdk.api.datamodels._generated import TaskInstance
from airflow.sdk.coordinators.node.coordinator import (
EMBEDDED_METADATA_MAX_BYTES,
NodeCoordinator,
_find_bundle,
)
Expand All @@ -45,27 +47,33 @@ def _make_ti(dag_id: str = "test_dag", queue: str = "ts") -> TaskInstance:
)


def _metadata_yaml(schema_version: str) -> str:
return f"""\
airflow_bundle_metadata_version: "1.0"
sdk:
language: typescript
version: "0.1.0"
supervisor_schema_version: "{schema_version}"
source: src/airflow.ts
dags:
test_dag:
tasks:
- test_task
"""


def write_bundle(root: pathlib.Path, schema_version: str = SCHEMA_VERSION) -> pathlib.Path:
bundle = root / "bundle.mjs"
bundle.write_text("export {};\n", encoding="utf-8")
(root / "airflow-metadata.yaml").write_text(
"\n".join(
[
'airflow_bundle_metadata_version: "1.0"',
"sdk:",
" language: typescript",
' version: "0.1.0"',
f' supervisor_schema_version: "{schema_version}"',
"source: src/airflow.ts",
"dags:",
" test_dag:",
" tasks:",
" - test_task",
"",
]
),
encoding="utf-8",
)
(root / "airflow-metadata.yaml").write_text(_metadata_yaml(schema_version), encoding="utf-8")
return bundle


def write_embedded_bundle(root: pathlib.Path, payload: str | None = None) -> pathlib.Path:
if payload is None:
payload = base64.b64encode(_metadata_yaml(SCHEMA_VERSION).encode("utf-8")).decode("ascii")
bundle = root / "bundle.mjs"
bundle.write_text(f"//# airflowMetadata={payload}\nexport {{}};\n", encoding="utf-8")
return bundle


Expand Down Expand Up @@ -105,6 +113,48 @@ def test_find_bundle_returns_bundle_mjs(self, tmp_path):
assert found.path == bundle
assert found.schema_version == SCHEMA_VERSION

def test_find_bundle_reads_embedded_metadata_without_sidecar(self, tmp_path):
bundle = write_embedded_bundle(tmp_path)

found = _find_bundle([tmp_path])

assert found.path == bundle
assert found.schema_version == SCHEMA_VERSION

def test_find_bundle_prefers_embedded_metadata_over_sidecar(self, tmp_path):
bundle = write_embedded_bundle(tmp_path)
(tmp_path / "airflow-metadata.yaml").write_text("[not-a-mapping]\n", encoding="utf-8")

found = _find_bundle([tmp_path])

assert found.path == bundle
assert found.schema_version == SCHEMA_VERSION

@pytest.mark.parametrize(
("payload", "message"),
[
("not-base64!", "cannot parse embedded airflow metadata"),
(base64.b64encode(b"[not-a-mapping]").decode("ascii"), "must contain a mapping"),
],
)
def test_find_bundle_rejects_invalid_embedded_metadata(self, tmp_path, payload, message):
write_embedded_bundle(tmp_path, payload=payload)

with pytest.raises(FileNotFoundError, match=message):
_find_bundle([tmp_path])

def test_find_bundle_rejects_oversized_embedded_metadata(self, tmp_path):
write_embedded_bundle(tmp_path, payload="A" * EMBEDDED_METADATA_MAX_BYTES)

with pytest.raises(FileNotFoundError, match="embedded airflow metadata exceeds"):
_find_bundle([tmp_path])

def test_find_bundle_rejects_empty_marker(self, tmp_path):
(tmp_path / "bundle.mjs").write_text("//# airflowMetadata=\nexport {};\n", encoding="utf-8")

with pytest.raises(FileNotFoundError, match="must contain a mapping"):
_find_bundle([tmp_path])

def test_find_bundle_checks_multiple_roots(self, tmp_path):
first = tmp_path / "first"
second = tmp_path / "second"
Expand Down
34 changes: 30 additions & 4 deletions ts-sdk/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,10 @@ coordinators = {
queue_to_coordinator = {"typescript": "ts"}
```

Each configured bundle directory must contain `bundle.mjs` and
`airflow-metadata.yaml`.
Each configured bundle directory must contain a `bundle.mjs` built with
`airflow-ts-pack` (see [Packing bundles](#packing-bundles)), which embeds the
Airflow metadata in the bundle itself. A `bundle.mjs` without embedded
metadata is also accepted alongside an `airflow-metadata.yaml` sidecar.

TypeScript entrypoint:

Expand Down Expand Up @@ -147,8 +149,32 @@ Airflow launches the bundled entrypoint with `--comm=host:port` and
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 builds a
`bundle.mjs` with `esbuild` and uses a Python stub Dag.
See [`example/`](example/) for a coordinator-runtime example that packs a
bundle with `airflow-ts-pack` and uses a Python stub Dag.

## Packing bundles

`airflow-ts-pack` produces everything `NodeCoordinator` needs in one command.
Packing is build-time only, so `esbuild` is an optional peer dependency the
runtime install skips:

```bash
npm install --save-dev esbuild
airflow-ts-pack src/main.ts --outdir dist
```

It bundles the entrypoint into `dist/bundle.mjs` with esbuild, runs the
bundle with `--airflow-metadata` so the bundle reports its own registered
Dag/task pairs and supervisor schema version, and embeds that manifest in the
bundle as a leading `//# airflowMetadata=<base64>` comment. The result is a
single deployable file whose metadata cannot drift from its code; no
hand-written sidecar is needed.

Options:

- `--outdir <dir>` — output directory (default `dist`)
- `--source <name>` — display name of the primary source file shown in the
Airflow UI (default: entry basename)

## TaskClient

Expand Down
19 changes: 4 additions & 15 deletions ts-sdk/example/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,9 @@ This example shows the coordinator-mode shape for TypeScript task handlers:
starts the coordinator runtime.
- `dist/bundle.mjs` is the generated Node.js bundle that Airflow launches.

The TypeScript SDK does not include a packer yet, so this example builds the
bundle with `esbuild` and writes the Airflow metadata file manually.
The build uses the SDK's `airflow-ts-pack` tool, which bundles the entrypoint
with esbuild and embeds the Airflow metadata generated from the bundle's
registered tasks, producing a single deployable file.

## Build

Expand All @@ -39,31 +40,19 @@ pnpm install
pnpm run build
```

Build the example bundle:
Build the example bundle and its metadata:

```bash
cd ts-sdk/example
pnpm install
pnpm run build
```

Create the metadata file next to the generated bundle:

```bash
node --input-type=module > dist/airflow-metadata.yaml <<'EOF'
import { SUPERVISOR_API_VERSION } from "@apache-airflow/ts-sdk";

console.log(`sdk:
supervisor_schema_version: "${SUPERVISOR_API_VERSION}"`);
EOF
```

The coordinator expects this layout:

```text
ts-sdk/example/dist/
bundle.mjs
airflow-metadata.yaml
```

## Airflow Configuration
Expand Down
2 changes: 1 addition & 1 deletion ts-sdk/example/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
"type": "module",
"license": "Apache-2.0",
"scripts": {
"build": "esbuild src/main.ts --bundle --platform=node --format=esm --target=node22 --outfile=dist/bundle.mjs",
"build": "airflow-ts-pack src/main.ts --outdir dist",
Comment thread
jason810496 marked this conversation as resolved.
"typecheck": "tsc --noEmit"
},
"dependencies": {
Expand Down
12 changes: 12 additions & 0 deletions ts-sdk/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
"type": "module",
"main": "./dist/index.js",
"types": "./dist/index.d.ts",
"bin": {
"airflow-ts-pack": "./dist/cli/main.js"
},
"homepage": "https://airflow.apache.org",
"repository": {
"type": "git",
Expand Down Expand Up @@ -59,8 +62,17 @@
"dependencies": {
"@msgpack/msgpack": "^3.1.2"
},
"peerDependencies": {
"esbuild": "^0.28.1"
Comment thread
guan404ming marked this conversation as resolved.
},
"peerDependenciesMeta": {
"esbuild": {
"optional": true
}
},
"devDependencies": {
"@eslint/js": "^10.0.1",
"esbuild": "^0.28.1",
"@types/node": "^22.19.17",
"eslint": "^10.4.0",
"json-schema-to-typescript": "^15.0.4",
Expand Down
Loading