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
29 changes: 6 additions & 23 deletions task-sdk/src/airflow/sdk/coordinators/node/coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,18 +45,17 @@
log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node")

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:
def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any]:
"""
Read the manifest ``airflow-ts-pack`` embeds in the bundle itself.

The packer prepends the ``airflow-metadata.yaml`` content as a leading
The packer prepends the metadata as a leading
``//# airflowMetadata=<base64>`` line comment, keeping bundle and metadata
a single artifact. Returns ``None`` when the bundle has no such marker.
a single artifact. Raises ``ValueError`` when the bundle has no such marker.
"""
try:
with bundle_path.open("rb") as bundle_file:
Expand All @@ -65,7 +64,7 @@ def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any] | None:
raise ValueError(f"cannot read {bundle_path.name}: {exc}") from exc

if not line.startswith(EMBEDDED_METADATA_MARKER):
return None
raise ValueError(f"{bundle_path.name} has no embedded airflow metadata; rebuild with airflow-ts-pack")
if len(line) > EMBEDDED_METADATA_MAX_BYTES:
raise ValueError(
f"embedded airflow metadata exceeds {EMBEDDED_METADATA_MAX_BYTES} bytes; "
Expand All @@ -80,25 +79,12 @@ def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any] | None:
return parse_metadata_mapping(decoded, source="embedded airflow metadata")


def _read_bundle_metadata(metadata_path: pathlib.Path) -> dict[str, Any]:
if not metadata_path.is_file():
raise ValueError(f"missing {METADATA_FILENAME}")

try:
content = metadata_path.read_text(encoding="utf-8")
except OSError as exc:
raise ValueError(f"cannot read {METADATA_FILENAME}: {exc}") from exc

return parse_metadata_mapping(content, source=METADATA_FILENAME)


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 bundle's
supervisor schema version from the metadata embedded in the bundle,
falling back to a sibling ``airflow-metadata.yaml`` sidecar.
supervisor schema version from the metadata embedded in the bundle.

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 @@ -112,8 +98,6 @@ def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle:
continue
try:
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 @@ -159,8 +143,7 @@ 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``
with embedded metadata (as produced by ``airflow-ts-pack``), or
``bundle.mjs`` plus an ``airflow-metadata.yaml`` sidecar. This is a
with embedded metadata (as produced by ``airflow-ts-pack``). 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
Expand Down
62 changes: 15 additions & 47 deletions task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,16 +62,11 @@ def _metadata_yaml(schema_version: str) -> str:
"""


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(_metadata_yaml(schema_version), encoding="utf-8")
return bundle


def write_embedded_bundle(root: pathlib.Path, payload: str | None = None) -> pathlib.Path:
def write_bundle(
root: pathlib.Path, schema_version: str = SCHEMA_VERSION, payload: str | None = None
) -> pathlib.Path:
if payload is None:
payload = base64.b64encode(_metadata_yaml(SCHEMA_VERSION).encode("utf-8")).decode("ascii")
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 @@ -113,38 +108,26 @@ 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"),
(base64.b64encode(b"sdk: [not-a-mapping]").decode("ascii"), "missing sdk metadata mapping"),
(
base64.b64encode(b"sdk:\n language: typescript").decode("ascii"),
"missing or invalid sdk.supervisor_schema_version",
),
],
)
def test_find_bundle_rejects_invalid_embedded_metadata(self, tmp_path, payload, message):
write_embedded_bundle(tmp_path, payload=payload)
write_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)
write_bundle(tmp_path, payload="A" * EMBEDDED_METADATA_MAX_BYTES)

with pytest.raises(FileNotFoundError, match="embedded airflow metadata exceeds"):
_find_bundle([tmp_path])
Expand Down Expand Up @@ -176,36 +159,21 @@ def test_find_bundle_ignores_other_mjs_names(self, tmp_path):
def test_find_bundle_rejects_bundle_without_metadata(self, tmp_path):
(tmp_path / "bundle.mjs").write_text("export {};\n", encoding="utf-8")

with pytest.raises(FileNotFoundError, match="missing airflow-metadata.yaml"):
_find_bundle([tmp_path])

@pytest.mark.parametrize(
("metadata", "message"),
[
("[not-a-mapping]\n", "airflow-metadata.yaml must contain a mapping"),
("sdk: [not-a-mapping]\n", "missing sdk metadata mapping"),
("sdk:\n language: typescript\n", "missing or invalid sdk.supervisor_schema_version"),
],
)
def test_find_bundle_rejects_invalid_metadata_shape(self, tmp_path, metadata, message):
(tmp_path / "bundle.mjs").write_text("export {};\n", encoding="utf-8")
(tmp_path / "airflow-metadata.yaml").write_text(metadata, encoding="utf-8")

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

def test_find_bundle_reports_unreadable_metadata(self, tmp_path, monkeypatch):
def test_find_bundle_reports_unreadable_bundle(self, tmp_path, monkeypatch):
write_bundle(tmp_path)

def raise_os_error(self, *args, **kwargs):
if self.name == "airflow-metadata.yaml":
if self.name == "bundle.mjs":
raise PermissionError("denied")
return original_open(self, *args, **kwargs)

original_open = pathlib.Path.open
monkeypatch.setattr(pathlib.Path, "open", raise_os_error)

with pytest.raises(FileNotFoundError, match="cannot read airflow-metadata.yaml"):
with pytest.raises(FileNotFoundError, match="cannot read bundle.mjs"):
_find_bundle([tmp_path])

def test_find_bundle_rejects_invalid_schema_version(self, tmp_path):
Expand Down
3 changes: 1 addition & 2 deletions ts-sdk/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,7 @@ queue_to_coordinator = {"typescript": "ts"}

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.
Airflow metadata in the bundle itself.

TypeScript entrypoint:

Expand Down