Show each native Dag's own source in the Code tab

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:

airflow-core/src/airflow/dag_processing/dagbag.py

@@ -327,17 +327,21 @@ 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))
         else:
             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

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,

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"
 
 

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")
 
 

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)

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")

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")
 

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=[],

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"}
     ]

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)

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")

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:
         """

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",

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

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)
 

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:

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")