Parse native TypeScript Dags with the Node coordinator

airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst

@@ -48,7 +48,8 @@ The SDK is the ``apache-airflow-ts-sdk`` package (ESM-only). It is currently in
 Prerequisites
 -------------
 
-* Node.js 22 or later must be available on the Airflow worker nodes.
+* Node.js 22 or later must be available on the Airflow worker nodes, and on the Dag processor to parse
+  Dags declared in TypeScript.
 * The packed bundle (a single ``bundle.min.mjs`` file, see :ref:`typescript-sdk/build`) must be accessible
   from the worker, under a directory the coordinator scans.
 * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional
@@ -236,16 +237,17 @@ task instance.
   bundles in ``bundles_root``) only need to be present wherever tasks actually execute. With
   ``CeleryExecutor``, setting them on the Celery workers is sufficient. With ``LocalExecutor``, tasks run
   inside the scheduler process, so they must be present where the scheduler can read them. The API server
-  and Dag processor do not need them.
+  does not need them, and the Dag processor needs them only to parse Dags declared in TypeScript,
+  see :ref:`typescript-sdk/native-parsing`.
 
 .. _typescript-sdk/native-dag:
 
 Declaring a Dag in TypeScript
 -----------------------------
 
 A ``Dag`` is declared on this side rather than in Python: its schedule, its tasks, their options and
-the edges between them are all written in TypeScript. The surface is still growing, so a Dag declared
-this way is not served to Airflow yet.
+the edges between them are all written in TypeScript. Airflow parses such a Dag from its packed bundle,
+see :ref:`typescript-sdk/native-parsing`.
 
 ``dag.task(taskId, handler)`` returns a *factory*. A handler takes one object of named arguments, and
 calling the factory names each input, so the call graph is the task graph:
@@ -514,6 +516,42 @@ instance's Dag. The artifact's name does not matter beyond that suffix, so one r
 and a Dag is routed to whichever declares it. If multiple bundles declare the same Dag, the first configured
 root wins, and within a root the first in sorted path order.
 
+.. _typescript-sdk/native-parsing:
+
+Parsing native Dags
+-------------------
+
+A Dag declared in TypeScript ships in its packed bundle, so the bundle goes into a Dag bundle,
+next to any Python Dag files. A coordinator configured without ``bundles_root`` parses the ``*.min.mjs``
+bundles of the Dag bundles it serves: every Dag bundle, or only the one ``dag_bundle_name`` names.
+The same entry runs the tasks:
+
+.. code-block:: ini
+
+    [sdk]
+    coordinators = {
+      "ts": {
+        "classpath": "airflow.sdk.coordinators.node.NodeCoordinator",
+        "kwargs": {"node_executable": "/usr/local/bin/node"}
+      }
+    }
+    queue_to_coordinator = {"typescript": "ts"}
+
+At most one such coordinator can serve a Dag bundle, and a coordinator with ``bundles_root`` parses nothing.
+
+The Dag processor runs each bundle with ``node`` to parse it, so it needs Node.js and this ``[sdk]``
+configuration. It parses only files that end in ``.min.mjs`` and start with the header ``airflow-ts-pack``
+writes. A bundle that fails its integrity check, or whose Dags cannot be serialized or fail validation,
+such as a cycle drawn with ``before`` and ``after``, is reported as an import error.
+
+The Code view shows the bundle's entry module for each of its Dags, as ``airflow-ts-pack`` embeds it
+(see :ref:`typescript-sdk/build`). If the source cannot be read from the bundle, the view shows a short
+notice instead.
+
+A bundle that only registers ``TaskHandler`` objects is parsed too when it sits in a served Dag bundle,
+since its metadata does not say whether it declares Dags. Each parse then launches ``node`` once and
+finds no Dags, so keep such bundles in a ``bundles_root`` directory to avoid that cost.
+
 .. _typescript-sdk/coordinator-config:
 
 :class:`~airflow.sdk.coordinators.node.NodeCoordinator` configuration
@@ -559,12 +597,20 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the
   * Leave both unset (the default) to load the bundle from the **task's own** Dag bundle, pinned to the
     version the run was created with.
 
+  Without ``bundles_root``, the coordinator also parses the native Dags of the Dag bundles it serves,
+  see :ref:`typescript-sdk/native-parsing`.
+
 Limitations
 -----------
 
 * **A Python stub Dag is still required.** The Execution API does not yet carry Dag structure for non-Python
   languages, so task names and dependencies are declared in Python with
   :func:`@task.stub <airflow.sdk.task.stub>`.
+* **Cluster policies do not apply to a native Dag.** ``dag_policy`` and ``task_policy`` are not run on
+  it, so they cannot change or reject it.
+* **Some CLI commands do not take a native Dag.** ``airflow dags test``, ``tasks test``, ``tasks render``
+  and ``tasks list`` refuse it, and ``airflow dags reserialize`` does not store the Dags of
+  ``*.min.mjs`` bundles, which only the Dag processor stores.
 * **Beta status.** The SDK API may change in incompatible ways between releases.
 * **One Node.js subprocess per task instance.** Tasks that need to share in-process state between instances
   should use XCom or an external store instead.

task-sdk/src/airflow/sdk/coordinators/node/_bundle_reader.py

@@ -40,6 +40,8 @@
 if TYPE_CHECKING:
     from typing import BinaryIO
 
+BUNDLE_SUFFIX = ".min.mjs"
+
 # Format prefixes and whole-line limits must agree with the TypeScript encoder.
 _LAYOUT_COMMENT_PREFIX = b"//# airflowBundle="
 _MAX_LAYOUT_LINE_BYTES = 4096
@@ -143,19 +145,32 @@ def read_bundle_source(bundle_path: pathlib.Path, dag_id: str | None = None) ->
     """
     if dag_id is None:
         return None
-    payloads = _read_verified_payloads(bundle_path)
+    return _read_source(_read_verified_payloads(bundle_path), dag_id)
+
+
+def read_bundle_entrypoint_source(bundle_path: pathlib.Path) -> str | None:
+    """Return the entrypoint source ``airflow-ts-pack`` embedded, or ``None`` when it embeds none."""
+    return _read_source(_read_verified_payloads(bundle_path), None)
+
+
+def _read_source(payloads: _VerifiedPayloads, dag_id: str | None) -> str | None:
     source_path = _resolve_source_path(payloads.metadata, dag_id)
     if source_path is None:
         return None
     try:
         payload = payloads.sources[source_path]
     except KeyError:
-        raise ValueError(
-            f"bundle declares no source region at path {source_path!r} for dag_id {dag_id!r}"
-        ) from None
+        owner = "the entrypoint" if dag_id is None else f"dag_id {dag_id!r}"
+        raise ValueError(f"bundle declares no source region at path {source_path!r} for {owner}") from None
     return _decode_source(payload)
 
 
+def has_bundle_layout_prefix(bundle_path: pathlib.Path) -> bool:
+    """Return whether *bundle_path* starts with the layout comment ``airflow-ts-pack`` writes."""
+    with bundle_path.open("rb") as bundle_file:
+        return bundle_file.read(len(_LAYOUT_COMMENT_PREFIX)) == _LAYOUT_COMMENT_PREFIX
+
+
 @attrs.define(frozen=True)
 class _VerifiedPayloads:
     """The metadata and per-path source payloads of a bundle whose digests have been checked."""
@@ -451,9 +466,9 @@ def _parse_bundle_metadata(payload: bytes) -> BundleMetadata:
     )
 
 
-def _resolve_source_path(payload: bytes, dag_id: str) -> str | None:
+def _resolve_source_path(payload: bytes, dag_id: str | None) -> str | None:
     """
-    Return the embedded source path for *dag_id*, or ``None``.
+    Return the embedded source path for *dag_id*, or for the entrypoint when it is ``None``.
 
     A native TypeScript Dag maps to its own file through ``dag_source_paths``. A Dag that is not
     mapped (one constructed dynamically, so the packer could not attribute a file) falls back to
@@ -466,7 +481,7 @@ def _resolve_source_path(payload: bytes, dag_id: str) -> str | None:
     if not isinstance(metadata, dict):
         raise ValueError("embedded airflow metadata must contain a mapping")
     dag_source_paths = metadata.get("dag_source_paths")
-    if isinstance(dag_source_paths, dict):
+    if dag_id is not None and isinstance(dag_source_paths, dict):
         mapped = dag_source_paths.get(dag_id)
         if isinstance(mapped, str):
             return mapped

task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py

@@ -0,0 +1,71 @@
+#
+# 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 importer that parses native TypeScript Dags from packed bundles."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+import structlog
+
+from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
+from airflow.sdk.coordinators.node._bundle_reader import (
+    BUNDLE_SUFFIX,
+    has_bundle_layout_prefix,
+    read_bundle_entrypoint_source,
+)
+from airflow.sdk.importers.base import DagSourceCode
+
+if TYPE_CHECKING:
+    from structlog.typing import FilteringBoundLogger
+
+    from airflow.sdk.coordinators.node.coordinator import NodeCoordinator
+    from airflow.sdk.importers.base import DagDefinition
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node")
+
+_NO_SOURCE: Final = "// Source code is not available: the bundle embeds no entrypoint source.\n"
+
+
+class NodeDagImporter(CoordinatorDagImporter):
+    """Parse the native Dags of packed ``*.min.mjs`` TypeScript bundles with a :class:`NodeCoordinator`."""
+
+    artifact_suffix = BUNDLE_SUFFIX
+    supported_extensions = [".mjs"]
+    coordinator: NodeCoordinator
+
+    def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool:
+        # The header identifies a packed bundle rather than guessing at content, so safe_mode keeps it.
+        with definition.as_file() as path:
+            try:
+                return has_bundle_layout_prefix(path)
+            except OSError:
+                # Keep the file, so parsing it records why it cannot be read.
+                return True
+
+    def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
+        """Return the embedded entrypoint source of the bundle, or a notice when there is none."""
+        with definition.as_file() as path:
+            try:
+                source = read_bundle_entrypoint_source(path)
+            except (OSError, ValueError) as e:
+                log.warning(
+                    "Cannot read the Dag source embedded in a TypeScript bundle", path=path, reason=str(e)
+                )
+                source = f"// Source code is not available: {e}\n"
+        return DagSourceCode(source_code=source or _NO_SOURCE, language="typescript")

task-sdk/src/airflow/sdk/coordinators/node/coordinator.py

@@ -15,7 +15,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Node.js runtime coordinator that launches a Node.js subprocess for task execution."""
+"""Node.js runtime coordinator that launches a Node.js subprocess for task execution and Dag parsing."""
 
 from __future__ import annotations
 
@@ -33,7 +33,7 @@
     walk_files,
 )
 from airflow.sdk.coordinators._subprocess import SubprocessCoordinator
-from airflow.sdk.coordinators.node._bundle_reader import read_bundle
+from airflow.sdk.coordinators.node._bundle_reader import BUNDLE_SUFFIX, read_bundle
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -42,11 +42,10 @@
     from typing_extensions import Self
 
     from airflow.sdk.api.datamodels._generated import TaskInstance
+    from airflow.sdk.coordinators.node._dag_importer import NodeDagImporter
 
 log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node")
 
-BUNDLE_SUFFIX = ".min.mjs"
-
 
 def _is_bundle(path: pathlib.Path) -> bool:
     return path.name.endswith(BUNDLE_SUFFIX)
@@ -98,7 +97,7 @@ def find(cls, bundles_root: Sequence[pathlib.Path], dag_id: str) -> Self:
 @attrs.define(kw_only=True)
 class NodeCoordinator(SubprocessCoordinator):
     """
-    Coordinator that launches a Node.js subprocess for task execution.
+    Coordinator that launches a Node.js subprocess for task execution and Dag parsing.
 
     Configuration is taken from the ``[sdk] coordinators`` entry that constructs
     this instance::
@@ -117,7 +116,8 @@ class NodeCoordinator(SubprocessCoordinator):
         ``"node"``, which relies on ``$PATH``).
     :param bundles_root: Directories searched recursively, in order, for the first verified
         ``*.min.mjs`` bundle declaring the task instance's Dag. See :class:`SubprocessCoordinator`
-        for its interaction with ``dag_bundle_name``.
+        for its interaction with ``dag_bundle_name``. Without it, the coordinator also parses the
+        native Dags of the ``*.min.mjs`` bundles in the Dag bundles it serves.
     :param task_startup_timeout: Maximum time the coordinator waits for a task
         process to start, in seconds. The default is 10 seconds.
     """
@@ -134,3 +134,12 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str],
         roots = self._get_scan_roots()
         bundle = _Bundle.find(roots, what.dag_id)
         return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version
+
+    @classmethod
+    def get_dag_importer_class(cls) -> type[NodeDagImporter]:
+        from airflow.sdk.coordinators.node._dag_importer import NodeDagImporter
+
+        return NodeDagImporter
+
+    def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]:
+        return [self.node_executable, os.fspath(path)], read_bundle(path).supervisor_schema_version

task-sdk/tests/task_sdk/coordinators/node/_bundle_test_utils.py

@@ -173,3 +173,8 @@ def mutate_byte(bundle: pathlib.Path, offset: int) -> None:
     contents = bytearray(bundle.read_bytes())
     contents[offset] = ord("A") if contents[offset] != ord("A") else ord("B")
     bundle.write_bytes(contents)
+
+
+def mutate_section(bundle: pathlib.Path, section: str) -> None:
+    """Change the first byte of a layout section, so its digest no longer matches."""
+    mutate_byte(bundle, int(read_layout(bundle)[section]["start"], 16))  # type: ignore[index, call-overload]

task-sdk/tests/task_sdk/coordinators/node/test_bundle_reader.py

@@ -46,6 +46,7 @@
     _digest_cache,
     _hash_region,
     read_bundle,
+    read_bundle_entrypoint_source,
     read_bundle_source,
 )
 
@@ -364,6 +365,12 @@ def test_reads_one_source_region_per_dag(self, tmp_path):
         assert read_bundle(bundle).dag_ids == frozenset({"sales", "inventory"})
         assert read_bundle_source(bundle, "sales") == "export const sales = 1;\n"
         assert read_bundle_source(bundle, "inventory") == "export const inv = 2;\n"
+        assert read_bundle_entrypoint_source(bundle) == "export const sales = 1;\n"
+
+    def test_reads_no_entrypoint_source_when_none_is_declared(self, tmp_path):
+        bundle = write_bundle(tmp_path, "sales", entrypoint_path=None)
+
+        assert read_bundle_entrypoint_source(bundle) is None
 
     def test_returns_none_without_a_dag_id(self, tmp_path):
         # A Dag owned by another language (Python) shows its own source; the caller omits dag_id