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/java/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py index 921f1ce993de3..a695203a8180e 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/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py index abb2aa74f88a0..ac483e1a8d4a2 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/_dag_importer.py @@ -26,6 +26,7 @@ BUNDLE_SUFFIX, has_bundle_layout_prefix, read_bundle_entrypoint_source, + read_bundle_source, ) from airflow.sdk.importers.base import DagSourceCode @@ -55,8 +56,17 @@ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: # 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 it embeds none.""" + 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_entrypoint_source(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/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/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py index 88101baa77ba6..3c262fb8fe451 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_dag_importer.py @@ -132,7 +132,7 @@ def raise_permission_error(self, *args, **kwargs): class TestGetSourceCode: - def test_returns_the_entrypoint_source_for_every_dag(self, importer, tmp_path): + 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' @@ -151,6 +151,38 @@ def test_returns_the_entrypoint_source_for_every_dag(self, importer, tmp_path): 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) 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 497212f0ef717..3fafa9344589c 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")