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 9364305ee2ddd..b07d8ef160f5e 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -49,6 +49,9 @@ Prerequisites ------------- * Node.js 22 or later must be available on the Airflow worker nodes and the Dag processor. + Once a ``NodeCoordinator`` is configured, the Dag processor runs ``node`` on every packed ``*.min.mjs`` bundle + in every Dag bundle, including ones that only register ``TaskHandler`` objects, + and reports an import error for each when ``node`` is missing. * The packed bundle (a single ``bundle.min.mjs`` file, see :ref:`typescript-sdk/build`) must be accessible from the worker and the Dag processor, in the Dag bundle the coordinator scans. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional @@ -241,13 +244,24 @@ located. There is no separate Node.js worker to run: the Airflow worker launches the bundle with ``node`` once per task instance. +A Dag processor with this ``[sdk]`` configuration also parses the ``*.min.mjs`` bundles of every Dag bundle, +and needs Node.js to do so (see :ref:`typescript-sdk/native-parsing`). +The ``ts-task-handlers`` Dag bundle only holds the bundles that the Python stub Dag's tasks run, +so keep the Dag processor from parsing it by listing ``*`` in its ``.airflowignore``: + +.. code-block:: bash + + echo '*' > /opt/airflow/ts-bundles/.airflowignore + .. note:: The ``[sdk]`` config, the packed ``*.min.mjs`` bundles and Node.js must be present wherever tasks execute and on the Dag processor. With ``CeleryExecutor``, tasks execute on the Celery workers; with ``LocalExecutor``, they run inside the scheduler process. The Dag processor checks the stub tasks of each - Python Dag against the task handlers the packed bundles register, so it runs them too. The API server - does not need any of it. Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every + Python Dag against the task handlers the packed bundles register, so it runs them too. + It also runs them to parse the bundles of every Dag bundle, see :ref:`typescript-sdk/native-parsing`. + The API server does not need any of it. + Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker and the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read it is rejected if the name is missing there. @@ -258,8 +272,8 @@ 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: @@ -497,12 +511,11 @@ The code is minified because an integrity digest is only worth taking over an ar read or edit in place. Function names are kept through minification, since a task id defaults to its handler's name. The ``/*! */`` license banners of bundled dependencies are kept. -Because the shipped code is not the code anyone wrote, the packer also embeds each Dag-defining source -file verbatim in its own ``/*# airflowSource: ... #*/`` block comment, verified by its own digest, -so Airflow has something readable to display for each Dag. Each native Dag's file is embedded — the file -its ``new Dag(...)`` constructor ran in — so a bundle that declares its Dags across several files gets one -source region per file, mapped to their ``dag_id`` by the ``dag_source_paths`` field in the manifest. Files -that only supply utilities or types are not embedded. +Because the shipped code is not the code anyone wrote, the packer also embeds source files verbatim, each in +its own ``/*# airflowSource: ... #*/`` block comment verified by its own digest: the entry module, and +each native Dag's file (the file its ``new Dag(...)`` constructor ran in). The ``dag_source_paths`` field in +the manifest maps each ``dag_id`` to its file. Files that only supply utilities or types are not embedded. +The Code view currently shows the embedded entry module for each Dag in the bundle. ``esbuild`` is an optional peer dependency: packing is build-time only, so the runtime install of ``apache-airflow-ts-sdk`` skips it, and it must be installed separately before running ``airflow-ts-pack``. @@ -522,10 +535,71 @@ Deploying ~~~~~~~~~ Copy or mount the bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``. +For a task of a Python stub Dag, :class:`~airflow.sdk.coordinators.node.NodeCoordinator` searches that Dag bundle recursively and launches the first integrity-verified ``*.min.mjs`` bundle whose metadata declares the task instance's Dag. The artifact's name does not matter beyond that suffix, so one Dag bundle can hold several bundles and a Dag is routed to whichever declares it. If multiple bundles declare the same Dag, the first in sorted path order wins. +A task of a native TypeScript Dag runs the bundle its Dag was parsed from, see +:ref:`typescript-sdk/native-parsing`. + +.. _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. +To have Airflow parse it, configure a :class:`~airflow.sdk.coordinators.node.NodeCoordinator`. +The Dag processor runs ``node`` on the bundle to list its Dags, so it needs Node.js, as the workers do: + +.. code-block:: ini + + [sdk] + coordinators = { + "ts": { + "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", + "kwargs": {"node_executable": "/usr/local/bin/node"} + } + } + queue_to_coordinator = {"typescript": "ts"} + +Once a ``NodeCoordinator`` is configured, the Dag processor parses the packed bundles of every Dag bundle, +so it needs this ``[sdk]`` configuration and Node.js. With one ``NodeCoordinator``, it parses them all. With several, +map each Dag bundle that holds native TypeScript Dags to one of them in ``[sdk] dag_bundle_to_coordinator``. +Each packed bundle of a Dag bundle that has no entry, or whose entry names no ``NodeCoordinator``, +fails to parse with an import error: + +.. code-block:: ini + + [sdk] + dag_bundle_to_coordinator = {"dags-folder": "ts"} + +A Dag bundle that holds only the bundles that Python stub Dags' tasks run should list ``*`` in its ``.airflowignore``. +Otherwise, with several Node coordinators, its bundles fail to parse. + +A task of a native Dag runs the bundle file its Dag was parsed from, on the ``NodeCoordinator`` its queue routes to, +which need not be the one that parsed it. For example, a queue can route to a coordinator that uses another Node.js. +A task whose queue routes to another kind of coordinator fails without retries. A task whose bundle is missing, +or that the coordinator cannot run (for example, because the bundle fails its integrity check), +fails and retries while it has retries left. +A task of a Python stub Dag runs the first bundle, in sorted path order, that declares its Dag. + +The Dag processor 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. + +Do not declare a Dag in TypeScript that a Python file in the same bundle also defines. When you move a Dag +such as the Quick start's ``typescript_example`` to ``new Dag(...)``, remove its Python stub, otherwise the +two files overwrite each other's Dag on every parse. + +The Code view currently 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, +since its metadata does not say whether it declares Dags. Each parse then launches ``node`` once and finds no Dags. +To avoid that cost, list such bundles in ``.airflowignore``; the coordinator still finds them to run tasks. .. _typescript-sdk/coordinator-config: @@ -568,12 +642,22 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * If ``task_handler_bundle_name`` is unset, the bundle is read from the **task's own** Dag bundle, pinned to the version the run was created with. + A task of a native TypeScript Dag ignores ``task_handler_bundle_name``: + it runs the bundle of its Dag from the Dag's own bundle, at the version the run was created with. + 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 `. +* **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`` and ``tasks render`` refuse it, + and ``airflow dags reserialize`` does not store the Dags of ``*.min.mjs`` bundles, + which only the Dag processor stores. + ``airflow tasks list`` lists its tasks by running the bundle with ``node``, so it needs Node.js. * **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. diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index ec8889db31845..8471e90c0b07d 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -619,7 +619,7 @@ def update_dag_parsing_results_in_db( :param files_parsed: Set of (bundle_name, relative_fileloc) tuples for all files that were parsed. If None, will be inferred from dags and import_errors. Passing this explicitly ensures that import errors are cleared for files that were parsed but no longer contain DAGs. - :param dag_source_codes: Source code read by the Dag importers, keyed by Dag fileloc. Dags + :param dag_source_codes: Source code read by the Dag importers, keyed by dag_id. Dags without an entry have their source read from ``fileloc``. """ accepted = _reject_other_teams_plugin_classes(bundle_name, dags, import_errors, session=session) @@ -670,7 +670,7 @@ def update_dag_parsing_results_in_db( version_data=version_data, session=session, _prefetched=prefetched_metadata.get(dag.dag_id), - dag_source_code=dag_source_codes.get(dag.fileloc), + dag_source_code=dag_source_codes.get(dag.dag_id), ) ) except OperationalError: diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index f8f2055e8081b..3b68817758dee 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -327,7 +327,11 @@ def _process_definition( self.captured_warnings.pop(fileloc, None) self._import_warnings.pop(fileloc, None) - self.dag_source_codes.pop(fileloc, None) + # Keyed by dag_id, not fileloc: a bundle can embed several Dags' own source under one + # fileloc, so a stale entry is found by which dag_ids this fileloc previously bagged. + for dag_id, dag in list(self.dags.items()): + if dag.fileloc == fileloc: + self.dag_source_codes.pop(dag_id, None) bagged_dags: list[BaggedDAG] if self.parse_lang_sdk_files and is_coordinator_importer(importer): bagged_dags = list(self._import_lang_sdk_definition(definition)) @@ -335,9 +339,9 @@ def _process_definition( bagged_dags = list(self._import_sdk_definition(importer, definition)) self.file_last_changed[fileloc] = freshness_token - if bagged_dags: + for dag in bagged_dags: try: - self.dag_source_codes[fileloc] = importer.get_source_code(definition) + self.dag_source_codes[dag.dag_id] = importer.get_source_code(definition, dag.dag_id) except Exception: self.log.exception("Failed to read source code of %s", fileloc) return bagged_dags diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 50360c5ea6f3d..9a0998dd05109 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -473,7 +473,7 @@ def _apply_team_rules(self, data: dict, dag: SerializedDAG) -> None: def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> dict[str, DagSourceCode]: """ - Read the file's source with its Dag importer, for the fileloc of each Dag. + Read each Dag's own source with its Dag importer, keyed by dag_id. A binary artifact cannot be read as text, so a source that cannot be read is a placeholder. """ @@ -483,11 +483,25 @@ def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> try: if (importer := get_claiming_importer(file, self.bundle_name)) is None: raise RuntimeError(f"No coordinator Dag importer claims {file}") - source = importer.get_source_code(FilesystemDagDefinition(Path(file))) except Exception as e: self.process_log.warning("Cannot read the Dag source", fileloc=file, error=str(e)) - source = DagSourceCode(f"Cannot read the source of {self.dag_file_rel_path}: {e}", "text") - return {dag.data["dag"].get("fileloc", file): source for dag in serialized_dags} + placeholder = DagSourceCode(f"Cannot read the source of {self.dag_file_rel_path}: {e}", "text") + return dict.fromkeys((dag.dag_id for dag in serialized_dags), placeholder) + + sources: dict[str, DagSourceCode] = {} + for dag in serialized_dags: + try: + sources[dag.dag_id] = importer.get_source_code( + FilesystemDagDefinition(Path(file)), dag.dag_id + ) + except Exception as e: + self.process_log.warning( + "Cannot read the Dag source", fileloc=file, dag_id=dag.dag_id, error=str(e) + ) + sources[dag.dag_id] = DagSourceCode( + f"Cannot read the source of {self.dag_file_rel_path}: {e}", "text" + ) + return sources _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[LangSDKDagFileProcessorProcess]]] = { **BaseDagFileProcessorProcess._common_request_handlers, diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index c94e7dde882a3..fb3dd55f43832 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -132,7 +132,7 @@ class DagFileParsingResult(BaseModel): parsed_definitions: list[str] = Field(default_factory=list) """Bundle-relative locations of the Dag definitions imported from ``fileloc``.""" dag_source_codes: dict[str, DagSourceCode] = Field(default_factory=dict) - """Source code of the parsed Dags, keyed by Dag fileloc.""" + """Source code of the parsed Dags, keyed by dag_id.""" type: Literal["DagFileParsingResult"] = "DagFileParsingResult" diff --git a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py index 92c8e596972f6..8e30282b7f264 100644 --- a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py +++ b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py @@ -68,7 +68,7 @@ class FakeCoordinatorDagImporter(CoordinatorDagImporter): artifact_suffix = ".native" supported_extensions = [".native"] - def get_source_code(self, definition) -> DagSourceCode: + def get_source_code(self, definition, dag_id: str | None = None) -> DagSourceCode: return DagSourceCode(definition.read_text(), "fake") diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index d7e5213fad5f1..0f66f067b838d 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -956,7 +956,7 @@ def test_dag_source_codes_are_written_to_dag_code(self, testing_dag_bundle, sess parse_duration=None, warnings=set(), session=session, - dag_source_codes={dag.fileloc: DagSourceCode(source_code="dag_id: yaml_dag\n", language="yaml")}, + dag_source_codes={dag.dag_id: DagSourceCode(source_code="dag_id: yaml_dag\n", language="yaml")}, ) dag_code = DagCode.get_latest_dagcode("yaml_dag", session=session) diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 1cbb0f0471344..4c0f744fa1e2d 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -396,7 +396,7 @@ def import_definition(self, definition, bundle): ], ) - def get_source_code(self, definition): + def get_source_code(self, definition, dag_id: str | None = None): return DagSourceCode(source_code=definition.read_text(), language="text") @@ -775,12 +775,8 @@ def test_dag_source_codes(self, tmp_path): dagbag = DagBag(dag_folder=tmp_path, bundle_path=tmp_path, bundle_name="test_bundle") assert dagbag.dag_source_codes == { - os.fspath(tmp_path / "plain.py"): DagSourceCode( - source_code=_dag_source("plain"), language="python" - ), - os.fspath(tmp_path / "dags.zip" / "member.py"): DagSourceCode( - source_code=_dag_source("member"), language="python" - ), + "plain": DagSourceCode(source_code=_dag_source("plain"), language="python"), + "member": DagSourceCode(source_code=_dag_source("member"), language="python"), } def test_dag_source_codes_drops_stale_source_when_read_fails(self, tmp_path): @@ -808,9 +804,7 @@ def test_sync_bag_to_db_passes_parsed_files_and_source_codes(self, mock_update, kwargs = mock_update.call_args.kwargs assert kwargs["files_parsed"] == {("test_bundle", "dags.zip/member.py"), ("test_bundle", "dags.zip")} assert kwargs["dag_source_codes"] == { - os.fspath(tmp_path / "dags.zip" / "member.py"): DagSourceCode( - source_code=_dag_source("member"), language="python" - ) + "member": DagSourceCode(source_code=_dag_source("member"), language="python") } def test_zip(self, tmp_path, test_zip_path): @@ -1402,7 +1396,7 @@ def test_custom_importer_dags_are_bound_to_their_definition(self, tmp_path): assert dag.fileloc == os.fspath(dag_file) assert dag.relative_fileloc == "nested/text_dag.dagtxt" assert dagbag.dag_source_codes == { - os.fspath(dag_file): DagSourceCode(source_code="text_dag\n", language="text") + "text_dag": DagSourceCode(source_code="text_dag\n", language="text") } assert dagbag.dag_warnings == {DagWarning("text_dag", "test:deprecated_field", "Deprecated field")} assert dagbag.captured_warnings == { @@ -1664,7 +1658,8 @@ def test_native_file_is_bagged_as_serialized_dags(self, mock_parse_dag, tmp_path assert isinstance(dag.task_dict["extract"], SerializedBaseOperator) assert (dag.fileloc, dag.relative_fileloc) == (os.fspath(native), "sub/dags.native") assert dag.bundle_name == "testing" - assert dagbag.dag_source_codes[os.fspath(native)].source_code == "{}" + assert dagbag.dag_source_codes["native_a"].source_code == "{}" + assert dagbag.dag_source_codes["native_b"].source_code == "{}" def test_native_file_is_an_import_error_by_default(self, tmp_path): write_native_file(tmp_path / "dags.native") diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 879e43710edad..3f6284e35c5bc 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -171,6 +171,25 @@ def test_parses_the_dags_the_runtime_returns(self, mock_parse_dag, parse, tmp_pa assert proc._subprocess_schema_version == OLDEST_SCHEMA_VERSION assert "Parsing the bundle" in cap_structlog + @patch.object(FakeCoordinatorDagImporter, "get_source_code", autospec=True) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_each_dag_gets_its_own_source_keyed_by_dag_id(self, mock_parse_dag, mock_get_source_code, parse): + """Two Dags from one file get their own ``get_source_code`` call and their own entry.""" + mock_parse_dag.side_effect = play_runtime( + _reply_with(_serialize_dag("north"), _serialize_dag("south")) + ) + mock_get_source_code.side_effect = lambda self, definition, dag_id=None: DagSourceCode( + f"source for {dag_id}", "fake" + ) + + proc = parse() + + assert proc.parsing_result.dag_source_codes == { + "north": DagSourceCode("source for north", "fake"), + "south": DagSourceCode("source for south", "fake"), + } + assert mock_get_source_code.call_args_list == [call(ANY, ANY, "north"), call(ANY, ANY, "south")] + @patch("airflow.dag_processing.lang_sdk_processor._is_connection_from_pid", autospec=True) @patch.object(FakeCoordinator, "parse_dag", autospec=True) def test_a_connection_is_used_once_it_is_verified(self, mock_parse_dag, mock_owned, tmp_path): @@ -650,7 +669,7 @@ def test_a_dag_source_that_cannot_be_read_is_a_placeholder(): proc._handle_request(DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[dag]), MagicMock(), 1) - source = proc.parsing_result.dag_source_codes[dag.data["dag"]["fileloc"]] + source = proc.parsing_result.dag_source_codes[dag.dag_id] assert source.language == "text" assert source.source_code.startswith("Cannot read the source of dag.native: [Errno 2] No such file") diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a5531ec57b888..2a92835edba39 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2459,7 +2459,7 @@ def test_deactivate_deleted_dags_keeps_dags_of_observed_definitions(self, dag_ma @mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True) def test_persist_parsing_result_passes_parsed_definitions_and_source_codes(self, mock_update): - source_codes = {"/bundle/dags.zip/a.py": DagSourceCode(source_code="src", language="python")} + source_codes = {"a": DagSourceCode(source_code="src", language="python")} parsing_result = DagFileParsingResult( fileloc="/bundle/dags.zip", serialized_dags=[], diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 117b04b07a7c1..9363937693fc7 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -788,9 +788,7 @@ def test_parse_file_reports_definitions_source_codes_and_dag_warnings(mock_dag_w ) assert result.parsed_definitions == ["dags.zip/member.py"] - assert result.dag_source_codes == { - os.fspath(tmp_path / "dags.zip" / "member.py"): DagSourceCode(source_code=source, language="python") - } + assert result.dag_source_codes == {"member": DagSourceCode(source_code=source, language="python")} assert result.warnings == [ {"dag_id": "member", "warning_type": "test:deprecated_field", "message": "Deprecated field"} ] diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index f339546934a69..4a72755d11c74 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -43,6 +43,7 @@ COORDINATOR_DAG_IMPORTERS: Final[tuple[str, ...]] = ( "airflow.sdk.coordinators.java._dag_importer.JavaDagImporter", + "airflow.sdk.coordinators.node._dag_importer.NodeDagImporter", ) """ The classpaths of the :class:`CoordinatorDagImporter` subclasses a Dag bundle's registry may hold. diff --git a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py index dfe81ba9be3be..6a20f3d60182b 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py @@ -73,7 +73,7 @@ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: return True return not wanted or main_class == wanted - def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: """Return the entrypoint source the JAR embeds, or a notice when it embeds none.""" with definition.as_file() as path, zipfile.ZipFile(path) as zf: info = _find_source_entry(zf) diff --git a/task-sdk/src/airflow/sdk/coordinators/node/_bundle_reader.py b/task-sdk/src/airflow/sdk/coordinators/node/_bundle_reader.py index 934bb8f6837b9..8bb39953ac124 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/_bundle_reader.py +++ b/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 diff --git a/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py new file mode 100644 index 0000000000000..ac483e1a8d4a2 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py @@ -0,0 +1,72 @@ +# +# 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, ClassVar, Final + +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, + read_bundle_source, +) +from airflow.sdk.importers.base import DagSourceCode + +if TYPE_CHECKING: + from airflow.sdk.importers.base import DagDefinition + +_NO_SOURCE: Final = "// Source code is not available: the bundle embeds no entrypoint source.\n" + + +class NodeDagImporter(CoordinatorDagImporter): + """ + Claim the native Dags of packed ``*.min.mjs`` TypeScript bundles. + + A :class:`~airflow.sdk.coordinators.node.NodeCoordinator` parses them. + """ + + coordinator_classpath: ClassVar[str] = "airflow.sdk.coordinators.node.NodeCoordinator" + artifact_suffix = BUNDLE_SUFFIX + supported_extensions = [".mjs"] + + 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, dag_id: str | None = None) -> DagSourceCode: + """ + Return the embedded source of *dag_id*'s own file, or a notice when the bundle embeds none. + + Falls back to the entrypoint source when *dag_id* is not given: ``read_bundle_source`` already + falls back to it for a *dag_id* the bundle has no ``dag_source_paths`` entry for. + """ + with definition.as_file() as path: + source = ( + read_bundle_source(path, dag_id) + if dag_id is not None + else read_bundle_entrypoint_source(path) + ) + return DagSourceCode(source_code=source or _NO_SOURCE, language="typescript") diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index 27fc142a3036c..86ed54eb703a3 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/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 @@ -28,7 +28,7 @@ from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle, 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 @@ -40,8 +40,6 @@ 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) @@ -93,7 +91,7 @@ def find(cls, roots: 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:: @@ -115,6 +113,13 @@ class NodeCoordinator(SubprocessCoordinator): ``[dag_processor] dag_bundle_config_list``. If unset, the task's own Dag bundle is used. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. + + A task of a native TypeScript Dag runs the ``*.min.mjs`` bundle its Dag was parsed from. + + The coordinator also parses native TypeScript Dags: every packed ``*.min.mjs`` bundle in a + Dag bundle is run to list its Dags. With one NodeCoordinator configured, it parses the bundles + of every Dag bundle. With several, ``[sdk] dag_bundle_to_coordinator`` picks the one that + parses a Dag bundle. """ node_executable: str = "node" @@ -123,3 +128,15 @@ 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 + + def _build_bundle_command(self, path: pathlib.Path) -> tuple[list[str], str]: + """Return the command that runs the packed bundle at *path*, and its supervisor schema version.""" + return [self.node_executable, os.fspath(path)], read_bundle(path).supervisor_schema_version + + def _build_dag_file_command( + self, *, what: TaskInstance, path: pathlib.Path + ) -> tuple[list[str], str | None]: + return self._build_bundle_command(path) + + def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + return self._build_bundle_command(path) diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index e8eb8d1a884b6..83a707245e1b4 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -264,8 +264,13 @@ def import_definition( """Import DAGs from a DAG definition.""" @abstractmethod - def get_source_code(self, definition: DagDefinition) -> DagSourceCode: - """Retrieve the raw source code and its language identifier for the specified DAG definition.""" + def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: + """ + Retrieve the raw source code and its language identifier for the specified DAG definition. + + :param dag_id: The DAG whose own source is wanted, when *definition* may hold more than one. + An importer that cannot distinguish between the DAGs of one definition ignores it. + """ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: """ diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py b/task-sdk/src/airflow/sdk/importers/python_importer.py index 5a884f707d641..e6fae5b44529a 100644 --- a/task-sdk/src/airflow/sdk/importers/python_importer.py +++ b/task-sdk/src/airflow/sdk/importers/python_importer.py @@ -213,8 +213,8 @@ def import_definition( self._process_modules(modules, result, bundle=bundle) return result - def get_source_code(self, definition: DagDefinition) -> DagSourceCode: - """Retrieve the raw source code for the Python definition.""" + def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: + """Retrieve the raw source code for the Python definition; *dag_id* does not narrow it further.""" if get_file_suffix(definition) == ".pyc": return DagSourceCode( source_code="# Sourceless bytecode (.pyc) — source code not available\n", diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py b/task-sdk/src/airflow/sdk/importers/zip_importer.py index 09b12c14f0419..1b7c53ad104e6 100644 --- a/task-sdk/src/airflow/sdk/importers/zip_importer.py +++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py @@ -268,7 +268,7 @@ def import_definition( return result return importer.import_definition(definition, bundle) - def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: """ Return the source of a single archive member. @@ -281,7 +281,7 @@ def get_source_code(self, definition: DagDefinition) -> DagSourceCode: importer = self._get_internal_importer(definition) if importer is None: raise ValueError(f"No internal importer to read source for {definition!r}") - return importer.get_source_code(definition) + return importer.get_source_code(definition, dag_id) def _register_internal( self, importer: AbstractDagImporter[Any], extensions: list[str] | None = None diff --git a/task-sdk/tests/task_sdk/coordinators/node/_bundle_test_utils.py b/task-sdk/tests/task_sdk/coordinators/node/_bundle_test_utils.py index 8039cbc032cdd..1d222d469c25d 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/_bundle_test_utils.py +++ b/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] diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_bundle_reader.py b/task-sdk/tests/task_sdk/coordinators/node/test_bundle_reader.py index 93440a2118223..29764941fbba9 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_bundle_reader.py +++ b/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, ) @@ -356,14 +357,21 @@ def test_reads_one_source_region_per_dag(self, tmp_path): sources=[ ("sales.ts", b"export const sales = 1;\n"), ("inventory.ts", b"export const inv = 2;\n"), + ("main.ts", b'import "./sales";\n'), ], dag_source_paths={"sales": "sales.ts", "inventory": "inventory.ts"}, - entrypoint_path="sales.ts", + entrypoint_path="main.ts", ) 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) == 'import "./sales";\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 diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 3dc01be33f297..2d2e4ba6b82f9 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -21,20 +21,27 @@ import json import pathlib from unittest import mock +from unittest.mock import MagicMock, patch import pytest from task_sdk.coordinators.node._bundle_test_utils import ( BUNDLE_NAME, mutate_byte, + mutate_section, read_layout, write_bundle, ) from uuid6 import uuid7 -from airflow.sdk.api.datamodels._generated import TaskInstance +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance +from airflow.sdk.coordinators._subprocess import _PopenActivitySubprocess from airflow.sdk.coordinators.node import _bundle_reader as _reader from airflow.sdk.coordinators.node._bundle_reader import _digest_cache from airflow.sdk.coordinators.node.coordinator import NodeCoordinator, _Bundle +from airflow.sdk.execution_time.coordinator import TaskLaunchError +from airflow.sdk.importers import reset_importer_registry + +from tests_common.test_utils.config import conf_vars SCHEMA_VERSION = "2026-06-16" @@ -106,6 +113,111 @@ def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, assert schema_version == SCHEMA_VERSION +@pytest.fixture +def mock_client(make_ti_context): + client = MagicMock() + client.task_instances.start.return_value = make_ti_context() + return client + + +class TestNodeCoordinatorExecuteNativeDag: + """With a NodeCoordinator configured, a task of a native TypeScript Dag runs the bundle of its Dag.""" + + @pytest.fixture(autouse=True) + def _node_coordinator_config(self): + coordinators = {"ts": {"classpath": "airflow.sdk.coordinators.node.NodeCoordinator"}} + reset_importer_registry() + with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}): + yield + reset_importer_registry() + + @pytest.fixture + def mock_start(self, tmp_path): + write_bundle(tmp_path, "sales", name="a.min.mjs", schema_version="2026-06-16") + write_bundle(tmp_path, "sales", name="b.min.mjs", schema_version="2026-10-30") + bundle = MagicMock(path=tmp_path, version="v1") + bundle.name = "dags-folder" + with ( + patch("airflow.sdk.coordinators._subprocess._initialize_pinned_bundle", return_value=bundle), + patch("airflow.sdk.coordinators._subprocess.BundleVersionLock"), + patch.object(_PopenActivitySubprocess, "start") as mock_start, + ): + mock_start.return_value.wait.return_value = 0 + yield mock_start + + def _execute(self, rel_path: str, client): + return NodeCoordinator().execute_task( + what=_make_ti(dag_id="sales"), + dag_rel_path=rel_path, + bundle_info=BundleInfo(name="dags-folder", version="v1"), + client=client, + subprocess_logs_to_stdout=False, + ) + + def test_runs_the_bundle_of_the_dag_and_not_the_first_bundle_by_path( + self, mock_start, mock_client, tmp_path + ): + self._execute("b.min.mjs", mock_client) + + assert mock_start.call_args.kwargs["command"] == ["node", str(tmp_path / "b.min.mjs")] + assert mock_start.call_args.kwargs["subprocess_schema_version"] == "2026-10-30" + + def test_fails_without_starting_node_for_a_bundle_that_fails_its_integrity_check( + self, mock_start, mock_client, tmp_path + ): + mutate_section(tmp_path / "b.min.mjs", "code") + + with pytest.raises(TaskLaunchError, match="code SHA-256 mismatch"): + self._execute("b.min.mjs", mock_client) + + mock_start.assert_not_called() + + +class TestNodeCoordinatorDagFileCommand: + def test_runs_the_bundle_the_dag_was_parsed_from(self, tmp_path): + bundle = write_bundle(tmp_path, "sales", schema_version="2026-10-30") + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") + + command, schema_version = coordinator._build_dag_file_command( + what=_make_ti(dag_id="sales"), path=bundle + ) + + assert command == ["/opt/node/bin/node", str(bundle)] + assert schema_version == "2026-10-30" + + def test_leaves_it_to_the_runtime_to_report_a_dag_the_bundle_does_not_declare(self, tmp_path): + bundle = write_bundle(tmp_path, "inventory") + + command, _ = NodeCoordinator()._build_dag_file_command(what=_make_ti(dag_id="sales"), path=bundle) + + assert command == ["node", str(bundle)] + + def test_tampered_bundle_raises(self, tmp_path): + bundle = write_bundle(tmp_path, "sales") + mutate_section(bundle, "code") + + with pytest.raises(ValueError, match="code SHA-256 mismatch"): + NodeCoordinator()._build_dag_file_command(what=_make_ti(dag_id="sales"), path=bundle) + + +class TestNodeCoordinatorParseDagCommand: + def test_returns_node_and_bundle_schema_version(self, tmp_path): + bundle = write_bundle(tmp_path / "typescript", "native_dag") + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") + + command, schema_version = coordinator._build_parse_dag_command(path=bundle) + + assert command == ["/opt/node/bin/node", str(bundle)] + assert schema_version == SCHEMA_VERSION + + def test_tampered_bundle_raises(self, tmp_path): + bundle = write_bundle(tmp_path, "native_dag") + mutate_section(bundle, "code") + + with pytest.raises(ValueError, match="code SHA-256 mismatch"): + NodeCoordinator()._build_parse_dag_command(path=bundle) + + class TestBundleFind: @pytest.mark.parametrize( "name", diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py new file mode 100644 index 0000000000000..3c262fb8fe451 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py @@ -0,0 +1,275 @@ +# +# 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. +from __future__ import annotations + +import json +import pathlib +from types import SimpleNamespace + +import pytest +from task_sdk.coordinators.node._bundle_test_utils import ( + BUNDLE_NAME, + LAYOUT_PREFIX, + SOURCE_OPEN, + mutate_byte, + mutate_section, + write_bundle, +) + +from airflow.sdk.coordinators._dag_importer import find_claiming_importer +from airflow.sdk.coordinators.java._dag_importer import JavaDagImporter +from airflow.sdk.coordinators.node import NodeCoordinator +from airflow.sdk.coordinators.node._bundle_reader import _digest_cache +from airflow.sdk.coordinators.node._dag_importer import NodeDagImporter +from airflow.sdk.execution_time.coordinator import InvalidCoordinatorError +from airflow.sdk.importers import ( + DagSourceCode, + FilesystemDagDefinition, + get_importer_registry, + reset_importer_registry, +) + +from tests_common.test_utils.config import conf_vars + +NODE_COORDINATOR = "airflow.sdk.coordinators.node.NodeCoordinator" + + +class _NodeCoordinatorSubclass(NodeCoordinator): + pass + + +def _coordinators(*kwargs: dict, mapping: dict[str, str] | None = None) -> dict[tuple[str, str], str]: + """Config with one NodeCoordinator per *kwargs*, keyed "ts-0", "ts-1" and so on.""" + specs = {f"ts-{i}": {"classpath": NODE_COORDINATOR, "kwargs": kw} for i, kw in enumerate(kwargs)} + config = {("sdk", "coordinators"): json.dumps(specs)} + if mapping is not None: + config[("sdk", "dag_bundle_to_coordinator")] = json.dumps(mapping) + return config + + +@pytest.fixture(autouse=True) +def clear_digest_cache(): + _digest_cache.clear() + + +@pytest.fixture(autouse=True) +def _clean_registry(): + reset_importer_registry() + yield + reset_importer_registry() + + +@pytest.fixture +def importer() -> NodeDagImporter: + return NodeDagImporter(bundle_name="dags-folder") + + +def test_claims_packed_bundles_only(importer): + assert importer.supported_extensions == [".mjs"] + assert importer.can_handle("dags/bundle.min.mjs") is True + assert importer.can_handle("dags/helper.mjs") is False + + +def test_lists_only_packed_bundles(importer, tmp_path): + nested = write_bundle(tmp_path / "team", "sales") + tampered = write_bundle(tmp_path, "inventory", name="tampered.min.mjs") + mutate_section(tampered, "code") + write_bundle(tmp_path, "orders", name="plain.mjs") + (tmp_path / "vendor.min.mjs").write_bytes(b"export {};\n") + + definitions = importer.list_dag_definitions(SimpleNamespace(name="dags-folder", path=tmp_path)) + + # A tampered bundle keeps its header, so parsing it reports the integrity failure. + assert sorted(d.path for d in definitions) == sorted([nested, tampered]) + + +class TestMightContainDag: + @pytest.mark.parametrize("safe_mode", [True, False]) + @pytest.mark.parametrize( + ("content", "expected"), + [ + (None, True), + (b"export {};\n", False), + (b"", False), + (LAYOUT_PREFIX[:-1], False), + ], + ids=["bundle", "plain-module", "empty", "truncated-prefix"], + ) + def test_checks_the_layout_header(self, importer, tmp_path, content, expected, safe_mode): + path = write_bundle(tmp_path, "sales") + if content is not None: + path.write_bytes(content) + + assert importer.might_contain_dag(FilesystemDagDefinition(path), safe_mode) is expected + + def test_keeps_an_unreadable_file(self, importer, tmp_path, monkeypatch): + path = write_bundle(tmp_path, "sales") + original_open = pathlib.Path.open + + def raise_permission_error(self, *args, **kwargs): + if self.name == BUNDLE_NAME: + raise PermissionError("denied") + return original_open(self, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "open", raise_permission_error) + + assert importer.might_contain_dag(FilesystemDagDefinition(path), True) is True + + +class TestGetSourceCode: + def test_returns_the_entrypoint_source_without_a_dag_id(self, importer, tmp_path): + sales = 'export const sales = new Dag({ dagId: "sales" });\n' + inventory = 'export const inventory = new Dag({ dagId: "inventory" });\n' + main = '/* the */ import "./sales";\nimport "./inventory";\n' + path = write_bundle( + tmp_path, + "sales", + "inventory", + sources=[ + ("sales.ts", sales.encode()), + ("inventory.ts", inventory.encode()), + ("main.ts", main.encode()), + ], + dag_source_paths={"sales": "sales.ts", "inventory": "inventory.ts"}, + entrypoint_path="main.ts", + ) + + assert importer.get_source_code(FilesystemDagDefinition(path)) == DagSourceCode(main, "typescript") + + def test_returns_each_dags_own_source_when_given_its_dag_id(self, importer, tmp_path): + sales = 'export const sales = new Dag({ dagId: "sales" });\n' + inventory = 'export const inventory = new Dag({ dagId: "inventory" });\n' + main = '/* the */ import "./sales";\nimport "./inventory";\n' + path = write_bundle( + tmp_path, + "sales", + "inventory", + sources=[ + ("sales.ts", sales.encode()), + ("inventory.ts", inventory.encode()), + ("main.ts", main.encode()), + ], + dag_source_paths={"sales": "sales.ts", "inventory": "inventory.ts"}, + entrypoint_path="main.ts", + ) + definition = FilesystemDagDefinition(path) + + assert importer.get_source_code(definition, "sales") == DagSourceCode(sales, "typescript") + assert importer.get_source_code(definition, "inventory") == DagSourceCode(inventory, "typescript") + + def test_falls_back_to_the_entrypoint_for_a_dag_id_with_no_mapped_file(self, importer, tmp_path): + """A Dag constructed dynamically has no ``dag_source_paths`` entry to look up.""" + main = 'export const sales = new Dag({ dagId: "sales" });\n' + path = write_bundle( + tmp_path, "sales", sources=[("main.ts", main.encode())], entrypoint_path="main.ts" + ) + + source_code = importer.get_source_code(FilesystemDagDefinition(path), "sales") + + assert source_code == DagSourceCode(main, "typescript") + + def test_returns_a_notice_without_an_entrypoint_source(self, importer, tmp_path): + path = write_bundle(tmp_path, "sales", entrypoint_path=None) + + source_code = importer.get_source_code(FilesystemDagDefinition(path)) + + assert source_code == DagSourceCode( + "// Source code is not available: the bundle embeds no entrypoint source.\n", "typescript" + ) + + @pytest.mark.parametrize( + ("break_bundle", "reason"), + [ + ( + lambda path: mutate_byte(path, path.read_bytes().index(SOURCE_OPEN) + len(SOURCE_OPEN)), + "source main.ts SHA-256 mismatch", + ), + (lambda path: write_bundle(path.parent, "sales", source=b"x" * (1024 * 1024 + 1)), "exceeds"), + (lambda path: path.unlink(), "cannot read bundle.min.mjs"), + ], + ids=["tampered", "oversize", "missing"], + ) + def test_raises_when_the_source_cannot_be_read(self, importer, tmp_path, break_bundle, reason): + path = write_bundle(tmp_path, "sales") + break_bundle(path) + + with pytest.raises((OSError, ValueError), match=reason): + importer.get_source_code(FilesystemDagDefinition(path)) + + +class TestRegistry: + def test_a_node_coordinator_registers_the_importer_in_every_bundle(self): + bundles = [ + {"name": name, "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}} + for name in ("dags-folder", "ts-bundles") + ] + coordinators = _coordinators({"task_handler_bundle_name": "ts-bundles"}) + with conf_vars({**coordinators, ("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + in_named = get_importer_registry("ts-bundles").get_importer("ts-bundles/bundle.min.mjs") + in_other = get_importer_registry("dags-folder").get_importer("dags/bundle.min.mjs") + + assert isinstance(in_named, NodeDagImporter) + assert isinstance(in_other, NodeDagImporter) + assert in_other.bundle_name == "dags-folder" + + def test_a_subclass_of_the_node_coordinator_registers_the_importer(self): + subclass = {"ts": {"classpath": f"{__name__}._NodeCoordinatorSubclass"}} + with conf_vars({("sdk", "coordinators"): json.dumps(subclass)}): + claiming = find_claiming_importer("dags/bundle.min.mjs", "dags-folder") + + assert isinstance(claiming, NodeDagImporter) + + def test_a_bundle_with_java_and_node_coordinators_sends_each_file_to_its_runtime(self): + coordinators = { + "java": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator"}, + "ts": {"classpath": NODE_COORDINATOR}, + } + with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}): + jar = find_claiming_importer("dags/app.jar", "dags-folder") + bundle = find_claiming_importer("dags/bundle.min.mjs", "dags-folder") + + assert isinstance(jar, JavaDagImporter) + assert isinstance(bundle, NodeDagImporter) + + +class TestGetParsingCoordinator: + def test_is_the_only_node_coordinator(self, importer): + with conf_vars(_coordinators({"node_executable": "/opt/node/bin/node"})): + coordinator = importer.get_parsing_coordinator() + + assert isinstance(coordinator, NodeCoordinator) + assert coordinator.node_executable == "/opt/node/bin/node" + + @pytest.mark.parametrize(("mapped", "executable"), [("ts-0", "/node/18"), ("ts-1", "/node/22")]) + def test_is_the_node_coordinator_mapped_to_its_bundle_among_several(self, importer, mapped, executable): + coordinators = _coordinators( + {"node_executable": "/node/18"}, + {"node_executable": "/node/22"}, + mapping={"dags-folder": mapped}, + ) + with conf_vars(coordinators): + coordinator = importer.get_parsing_coordinator() + + assert coordinator.node_executable == executable + + def test_fails_among_several_without_an_entry_for_its_bundle(self, importer): + with conf_vars(_coordinators({}, {})): + with pytest.raises( + InvalidCoordinatorError, match="Dag bundle 'dags-folder' has 2 NodeCoordinator" + ): + importer.get_parsing_coordinator() diff --git a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py index 53442f0b24b38..c0b323e4f84d5 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py +++ b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py @@ -55,7 +55,7 @@ class _NativeImporter(CoordinatorDagImporter): artifact_suffix = ".min.native" supported_extensions = [".native"] - def get_source_code(self, definition) -> DagSourceCode: + def get_source_code(self, definition, dag_id: str | None = None) -> DagSourceCode: return DagSourceCode("", "native") def might_contain_dag(self, definition, safe_mode: bool) -> bool: diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index d5fd0689064fd..219bcdd8b5169 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1080,7 +1080,7 @@ class _NativeDagImporter(CoordinatorDagImporter): artifact_suffix = ".native" supported_extensions = [".native"] - def get_source_code(self, definition) -> DagSourceCode: + def get_source_code(self, definition, dag_id: str | None = None) -> DagSourceCode: return DagSourceCode("", "native")