From 52759fd4ac3de749859ba7838da9876af42516b4 Mon Sep 17 00:00:00 2001 From: bramhanandlingala Date: Tue, 16 Jun 2026 16:29:08 +0530 Subject: [PATCH] Add bundle_name attribute to DAG and update PythonDagImporter to set it --- .../airflow/dag_processing/importers/python_importer.py | 1 + airflow-core/tests/unit/dag_processing/test_dagbag.py | 7 +++++++ task-sdk/src/airflow/sdk/definitions/dag.py | 1 + 3 files changed, 9 insertions(+) diff --git a/airflow-core/src/airflow/dag_processing/importers/python_importer.py b/airflow-core/src/airflow/dag_processing/importers/python_importer.py index 47f77a86f73ec..41cb6166e6c58 100644 --- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py +++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py @@ -355,6 +355,7 @@ def _process_modules( dag.fileloc = mod.__file__ relative_fileloc = self.get_relative_path(dag.fileloc, bundle_path) dag.relative_fileloc = relative_fileloc + dag.bundle_name = bundle_name result.dags.append(dag) log.debug("Found DAG %s", dag.dag_id) diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 99abd92a59f73..6f69ccafe1c63 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1091,9 +1091,16 @@ def test_dagbag_dag_collection(self): dagbag.collect_dags() assert dagbag.dags + # NEW TEST + for dag in dagbag.dags.values(): + assert dag.bundle_name == "test_collection" + # test that dagbag.dags is not empty if collect_dags is True dagbag = DagBag(dag_folder=TEST_DAGS_FOLDER, bundle_name="test_collection") assert dagbag.dags + + for dag in dagbag.dags.values(): + assert dag.bundle_name == "test_collection" def test_dabgag_captured_warnings(self): dag_file = os.path.join(TEST_DAGS_FOLDER, "test_dag_warnings.py") diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index fd47de467fea4..c8a9b7c027fb1 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -542,6 +542,7 @@ def __rich_repr__(self): fileloc: str = attrs.field(init=False, factory=_default_fileloc) relative_fileloc: str | None = attrs.field(init=False, default=None) + bundle_name: str | None = attrs.field(init=False, default=None) partial: bool = attrs.field(init=False, default=False) edge_info: dict[str, dict[str, EdgeInfoType]] = attrs.field(init=False, factory=dict)