Introduce BundleDagBag for ephemeral pre-loaded bundle path usage.

- usefor CLI and other one-off DagBag instances 'polluting' sys.path

Author: simi

airflow-core/src/airflow/cli/commands/dag_command.py

@@ -37,7 +37,7 @@
 from airflow.cli.simple_table import AirflowConsole
 from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string
 from airflow.dag_processing.bundles.manager import DagBundlesManager
-from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
+from airflow.dag_processing.dagbag import BundleDagBag, DagBag, sync_bag_to_db
 from airflow.exceptions import AirflowConfigException, AirflowException
 from airflow.jobs.job import Job
 from airflow.models import DagModel, DagRun, TaskInstance
@@ -378,10 +378,12 @@ def dag_list_dags(args, session: Session = NEW_SESSION) -> None:
 
             for bundle in all_bundles:
                 if bundle.name in bundles_to_search:
-                    dagbag = DagBag(bundle.path, bundle_path=bundle.path, bundle_name=bundle.name)
-                    dagbag.collect_dags()
-                    dags_list.extend(list(dagbag.dags.values()))
-                    dagbag_import_errors += len(dagbag.import_errors)
+                    bundle_dagbag = BundleDagBag(
+                        bundle.path, bundle_path=bundle.path, bundle_name=bundle.name
+                    )
+                    bundle_dagbag.collect_dags()
+                    dags_list.extend(list(bundle_dagbag.dags.values()))
+                    dagbag_import_errors += len(bundle_dagbag.import_errors)
         else:
             dagbag = DagBag()
             dagbag.collect_dags()
@@ -474,8 +476,10 @@ def dag_list_import_errors(args, session: Session = NEW_SESSION) -> None:
 
             for bundle in all_bundles:
                 if bundle.name in bundles_to_search:
-                    dagbag = DagBag(bundle.path, bundle_path=bundle.path, bundle_name=bundle.name)
-                    for filename, errors in dagbag.import_errors.items():
+                    bundle_dagbag = BundleDagBag(
+                        bundle.path, bundle_path=bundle.path, bundle_name=bundle.name
+                    )
+                    for filename, errors in bundle_dagbag.import_errors.items():
                         data.append({"bundle_name": bundle.name, "filepath": filename, "error": errors})
         else:
             dagbag = DagBag()
@@ -526,7 +530,9 @@ def dag_report(args) -> None:
         if bundle.name not in bundles_to_reserialize:
             continue
         bundle.initialize()
-        dagbag = DagBag(bundle.path, bundle_path=bundle.path, bundle_name=bundle.name, include_examples=False)
+        dagbag = BundleDagBag(
+            bundle.path, bundle_path=bundle.path, bundle_name=bundle.name, include_examples=False
+        )
         all_dagbag_stats.extend(dagbag.dagbag_stats)
 
     AirflowConsole().print_as(
@@ -690,7 +696,7 @@ def dag_reserialize(args, session: Session = NEW_SESSION) -> None:
         if bundle.name not in bundles_to_reserialize:
             continue
         bundle.initialize()
-        dag_bag = DagBag(
+        dag_bag = BundleDagBag(
             bundle.path, bundle_path=bundle.path, bundle_name=bundle.name, include_examples=False
         )
         sync_bag_to_db(dag_bag, bundle.name, bundle_version=bundle.get_current_version(), session=session)

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

@@ -235,13 +235,6 @@ def __init__(
         self.bundle_path = bundle_path
         self.bundle_name = bundle_name
 
-        # Add bundle path to sys.path if provided.
-        # This allows DAG files to import modules from their bundle directory.
-        # No cleanup is performed - this is intentional for ephemeral processes
-        # (dag processor, task runner, CLI) where the process exits after use.
-        if bundle_path and str(bundle_path) not in sys.path:
-            sys.path.append(str(bundle_path))
-
         dag_folder = dag_folder or settings.DAGS_FOLDER
         self.dag_folder = dag_folder
         self.dags: dict[str, DAG] = {}
@@ -702,6 +695,30 @@ def dagbag_report(self):
         return report
 
 
+class BundleDagBag(DagBag):
+    """
+    Bundle-aware DagBag that permanently modifies sys.path.
+
+    This class adds the bundle_path to sys.path permanently to allow DAG files
+    to import modules from their bundle directory. No cleanup is performed.
+
+    WARNING: Only use for one-off usages like CLI commands. Using this in long-running
+    processes will cause sys.path to accumulate entries.
+
+    Same parameters as DagBag, but bundle_path is required.
+    """
+
+    def __init__(self, *args, bundle_path: Path | None = None, **kwargs):
+        if not bundle_path:
+            raise ValueError("bundle_path is required for BundleDagBag")
+
+        if str(bundle_path) not in sys.path:
+            sys.path.append(str(bundle_path))
+
+        kwargs["bundle_path"] = bundle_path
+        super().__init__(*args, **kwargs)
+
+
 @provide_session
 def sync_bag_to_db(
     dagbag: DagBag,

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

@@ -34,7 +34,7 @@
     TaskCallbackRequest,
 )
 from airflow.configuration import conf
-from airflow.dag_processing.dagbag import DagBag
+from airflow.dag_processing.dagbag import BundleDagBag, DagBag
 from airflow.observability.stats import Stats
 from airflow.sdk.exceptions import TaskNotFound
 from airflow.sdk.execution_time.comms import (
@@ -205,7 +205,7 @@ def _parse_file_entrypoint():
 def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileParsingResult | None:
     # TODO: Set known_pool names on DagBag!
 
-    bag = DagBag(
+    bag = BundleDagBag(
         dag_folder=msg.file,
         bundle_path=msg.bundle_path,
         bundle_name=msg.bundle_name,

airflow-core/src/airflow/models/__init__.py

@@ -29,6 +29,7 @@
     "BaseOperator",
     "BaseOperatorLink",
     "BaseXCom",
+    "BundleDagBag",
     "Connection",
     "DagBag",
     "DagWarning",
@@ -98,6 +99,7 @@ def __getattr__(name):
     "BaseOperator": "airflow.sdk",
     "BaseOperatorLink": "airflow.sdk",
     "BaseXCom": "airflow.sdk.bases.xcom",
+    "BundleDagBag": "airflow.dag_processing.dagbag",
     "Callback": "airflow.models.callback",
     "Connection": "airflow.models.connection",
     "DagBag": "airflow.dag_processing.dagbag",
@@ -126,7 +128,7 @@ def __getattr__(name):
 if TYPE_CHECKING:
     # I was unable to get mypy to respect a airflow/models/__init__.pyi, so
     # having to resort back to this hacky method
-    from airflow.dag_processing.dagbag import DagBag
+    from airflow.dag_processing.dagbag import BundleDagBag, DagBag
     from airflow.models.base import ID_LEN, Base
     from airflow.models.callback import Callback
     from airflow.models.connection import Connection

airflow-core/src/airflow/utils/cli.py

@@ -272,14 +272,14 @@ def get_bagged_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | N
     find the correct path (assuming it's a file) and failing that, use the configured
     dags folder.
     """
-    from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
+    from airflow.dag_processing.dagbag import BundleDagBag, sync_bag_to_db
     from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
 
     manager = DagBundlesManager()
     for bundle_name in bundle_names or ():
         bundle = manager.get_bundle(bundle_name)
         with _airflow_parsing_context_manager(dag_id=dag_id):
-            dagbag = DagBag(
+            dagbag = BundleDagBag(
                 dag_folder=dagfile_path or bundle.path,
                 bundle_path=bundle.path,
                 bundle_name=bundle.name,
@@ -292,7 +292,7 @@ def get_bagged_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | N
     for bundle in manager.get_all_dag_bundles():
         bundle.initialize()
         with _airflow_parsing_context_manager(dag_id=dag_id):
-            dagbag = DagBag(
+            dagbag = BundleDagBag(
                 dag_folder=dagfile_path or bundle.path,
                 bundle_path=bundle.path,
                 bundle_name=bundle.name,
@@ -323,7 +323,7 @@ def get_db_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | None
 
 def get_dags(bundle_names: list | None, dag_id: str, use_regex: bool = False, from_db: bool = False):
     """Return DAG(s) matching a given regex or dag_id."""
-    from airflow.dag_processing.dagbag import DagBag
+    from airflow.dag_processing.dagbag import BundleDagBag
 
     bundle_names = bundle_names or []
 
@@ -333,7 +333,7 @@ def get_dags(bundle_names: list | None, dag_id: str, use_regex: bool = False, fr
         return [get_bagged_dag(bundle_names=bundle_names, dag_id=dag_id)]
 
     def _find_dag(bundle):
-        dagbag = DagBag(dag_folder=bundle.path, bundle_path=bundle.path, bundle_name=bundle.name)
+        dagbag = BundleDagBag(dag_folder=bundle.path, bundle_path=bundle.path, bundle_name=bundle.name)
         matched_dags = [dag for dag in dagbag.dags.values() if re.search(dag_id, dag.dag_id)]
         return matched_dags
 

airflow-core/tests/unit/cli/commands/test_dag_command.py

@@ -760,7 +760,7 @@ def test_dag_test_show_dag(self, mock_get_dag, mock_render_dag, stdout_capture):
         mock_render_dag.assert_has_calls([mock.call(mock_get_dag.return_value, tis=[])])
         assert "SOURCE" in output
 
-    @mock.patch("airflow.dag_processing.dagbag.DagBag")
+    @mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
     def test_dag_test_with_bundle_name(self, mock_dagbag, configure_dag_bundles):
         """Test that DAG can be tested using bundle name."""
         mock_dagbag.return_value.get_dag.return_value.test.return_value = DagRun(
@@ -788,7 +788,7 @@ def test_dag_test_with_bundle_name(self, mock_dagbag, configure_dag_bundles):
             include_examples=False,
         )
 
-    @mock.patch("airflow.dag_processing.dagbag.DagBag")
+    @mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
     def test_dag_test_with_dagfile_path(self, mock_dagbag, configure_dag_bundles):
         """Test that DAG can be tested using dagfile path."""
         mock_dagbag.return_value.get_dag.return_value.test.return_value = DagRun(
@@ -810,7 +810,7 @@ def test_dag_test_with_dagfile_path(self, mock_dagbag, configure_dag_bundles):
             include_examples=False,
         )
 
-    @mock.patch("airflow.dag_processing.dagbag.DagBag")
+    @mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
     def test_dag_test_with_both_bundle_and_dagfile_path(self, mock_dagbag, configure_dag_bundles):
         """Test that DAG can be tested using both bundle name and dagfile path."""
         mock_dagbag.return_value.get_dag.return_value.test.return_value = DagRun(

airflow-core/tests/unit/dag_processing/test_dagbag.py

@@ -35,7 +35,12 @@
 from sqlalchemy import select
 
 from airflow import settings
-from airflow.dag_processing.dagbag import DagBag, _capture_with_reraise, _validate_executor_fields
+from airflow.dag_processing.dagbag import (
+    BundleDagBag,
+    DagBag,
+    _capture_with_reraise,
+    _validate_executor_fields,
+)
 from airflow.exceptions import UnknownExecutorException
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.models.dag import DagModel
@@ -1195,10 +1200,10 @@ def test_capture_warnings_with_error_filters(self):
 
 
 class TestBundlePathSysPath:
-    """Tests for bundle_path sys.path handling in DagBag."""
+    """Tests for bundle_path sys.path handling in BundleDagBag."""
 
     def test_bundle_path_added_to_syspath(self, tmp_path):
-        """Test that DagBag adds bundle_path to sys.path when provided."""
+        """Test that BundleDagBag adds bundle_path to sys.path when provided."""
         util_file = tmp_path / "bundle_util.py"
         util_file.write_text('def get_message(): return "Hello from bundle!"')
 
@@ -1220,7 +1225,9 @@ def test_bundle_path_added_to_syspath(self, tmp_path):
 
         assert str(tmp_path) not in sys.path
 
-        dagbag = DagBag(dag_folder=str(dag_file), bundle_path=tmp_path, include_examples=False)
+        dagbag = BundleDagBag(
+            dag_folder=str(dag_file), bundle_path=tmp_path, bundle_name="test-bundle", include_examples=False
+        )
 
         # Check import was successful
         assert len(dagbag.dags) == 1
@@ -1255,16 +1262,18 @@ def test_bundle_path_not_duplicated(self, tmp_path):
         sys.path.append(str(tmp_path))
         count_before = sys.path.count(str(tmp_path))
 
-        DagBag(dag_folder=str(dag_file), bundle_path=tmp_path, include_examples=False)
+        BundleDagBag(
+            dag_folder=str(dag_file), bundle_path=tmp_path, bundle_name="test-bundle", include_examples=False
+        )
 
         # Should not add duplicate
         assert sys.path.count(str(tmp_path)) == count_before
 
         # Cleanup for other tests
         sys.path.remove(str(tmp_path))
 
-    def test_bundle_path_none_no_syspath_modification(self, tmp_path):
-        """Test that no sys.path modification occurs when bundle_path is None."""
+    def test_dagbag_no_bundle_path_no_syspath_modification(self, tmp_path):
+        """Test that no sys.path modification occurs when DagBag is used without bundle_path."""
         dag_file = tmp_path / "simple_dag.py"
         dag_file.write_text(
             textwrap.dedent(
@@ -1280,7 +1289,7 @@ def test_bundle_path_none_no_syspath_modification(self, tmp_path):
             )
         )
         syspath_before = deepcopy(sys.path)
-        dagbag = DagBag(dag_folder=str(dag_file), bundle_path=None, include_examples=False)
+        dagbag = DagBag(dag_folder=str(dag_file), include_examples=False)
         dag = dagbag.get_dag("simple_dag")
 
         assert str(tmp_path) not in dag.description

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -1795,9 +1795,9 @@ def fake_collect_dags(self, *args, **kwargs):
             _execute_email_callbacks(dagbag, request, log)
 
     def test_parse_file_passes_bundle_name_to_dagbag(self):
-        """Test that _parse_file() creates DagBag with correct bundle_name parameter"""
-        # Mock the DagBag constructor to capture its arguments
-        with patch("airflow.dag_processing.processor.DagBag") as mock_dagbag_class:
+        """Test that _parse_file() creates BundleDagBag with correct bundle_name parameter"""
+        # Mock the BundleDagBag constructor to capture its arguments
+        with patch("airflow.dag_processing.processor.BundleDagBag") as mock_dagbag_class:
             # Create a mock instance with proper attributes for Pydantic validation
             mock_dagbag_instance = MagicMock()
             mock_dagbag_instance.dags = {}
@@ -1813,7 +1813,7 @@ def test_parse_file_passes_bundle_name_to_dagbag(self):
 
             _parse_file(request, log=structlog.get_logger())
 
-            # Verify DagBag was called with correct bundle_name
+            # Verify BundleDagBag was called with correct bundle_name
             mock_dagbag_class.assert_called_once()
             call_kwargs = mock_dagbag_class.call_args.kwargs
             assert call_kwargs["bundle_name"] == "test_bundle"

airflow-core/tests/unit/models/test_dag.py

@@ -40,7 +40,7 @@
 from airflow._shared.timezones import timezone
 from airflow._shared.timezones.timezone import datetime as datetime_tz
 from airflow.configuration import conf
-from airflow.dag_processing.dagbag import DagBag
+from airflow.dag_processing.dagbag import BundleDagBag, DagBag
 from airflow.exceptions import AirflowException
 from airflow.models.asset import (
     AssetAliasModel,
@@ -2194,7 +2194,7 @@ def test_relative_fileloc(self, session, testing_dag_bundle):
         rel_path = "test_assets.py"
         bundle_path = TEST_DAGS_FOLDER
         file_path = bundle_path / rel_path
-        bag = DagBag(dag_folder=file_path, bundle_path=bundle_path)
+        bag = BundleDagBag(dag_folder=file_path, bundle_path=bundle_path, bundle_name="testing")
 
         dag = bag.get_dag("dag_with_skip_task")
 

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -1241,7 +1241,7 @@ def test(
             version = DagVersion.get_version(self.dag_id)
             if not version:
                 from airflow.dag_processing.bundles.manager import DagBundlesManager
-                from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
+                from airflow.dag_processing.dagbag import BundleDagBag, sync_bag_to_db
                 from airflow.sdk.definitions._internal.dag_parsing_context import (
                     _airflow_parsing_context_manager,
                 )
@@ -1255,8 +1255,11 @@ def test(
                     if not bundle.is_initialized:
                         bundle.initialize()
                     with _airflow_parsing_context_manager(dag_id=self.dag_id):
-                        dagbag = DagBag(
-                            dag_folder=bundle.path, bundle_path=bundle.path, include_examples=False
+                        dagbag = BundleDagBag(
+                            dag_folder=bundle.path,
+                            bundle_path=bundle.path,
+                            bundle_name=bundle.name,
+                            include_examples=False,
                         )
                         sync_bag_to_db(dagbag, bundle.name, bundle.version)
                     version = DagVersion.get_version(self.dag_id)