Remove example_dags support from list_py_file_paths - it's now a bundle!

Author: jedcunningham

airflow/dag_processing/bundles/manager.py

@@ -64,6 +64,23 @@ def parse_config(self) -> None:
                 "Bundle config is not a list. Check config value"
                 " for section `dag_bundles` and key `backends`."
             )
+
+        # example dags!
+        if conf.getboolean("core", "LOAD_EXAMPLES"):
+            from airflow import example_dags
+
+            example_dag_folder = next(iter(example_dags.__path__))
+            backends.append(
+                {
+                    "name": "example_dags",
+                    "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
+                    "kwargs": {
+                        "local_folder": example_dag_folder,
+                        "refresh_interval": conf.getint("scheduler", "dag_dir_list_interval"),
+                    },
+                }
+            )
+
         seen = set()
         for cfg in backends:
             name = cfg["name"]

airflow/models/dagbag.py

@@ -568,11 +568,17 @@ def collect_dags(
 
         # Ensure dag_folder is a str -- it may have been a pathlib.Path
         dag_folder = correct_maybe_zipped(str(dag_folder))
-        for filepath in list_py_file_paths(
-            dag_folder,
-            safe_mode=safe_mode,
-            include_examples=include_examples,
-        ):
+
+        files_to_parse = list_py_file_paths(dag_folder, safe_mode=safe_mode)
+
+        if include_examples:
+            from airflow import example_dags
+
+            example_dag_folder = next(iter(example_dags.__path__))
+
+            files_to_parse.extend(list_py_file_paths(example_dag_folder, safe_mode=safe_mode))
+
+        for filepath in files_to_parse:
             try:
                 file_parse_start_dttm = timezone.utcnow()
                 found_dags = self.process_file(filepath, only_if_updated=only_if_updated, safe_mode=safe_mode)

airflow/utils/file.py

@@ -245,7 +245,6 @@ def find_path_from_directory(
 def list_py_file_paths(
     directory: str | os.PathLike[str] | None,
     safe_mode: bool = conf.getboolean("core", "DAG_DISCOVERY_SAFE_MODE", fallback=True),
-    include_examples: bool | None = None,
 ) -> list[str]:
     """
     Traverse a directory and look for Python files.
@@ -255,23 +254,15 @@ def list_py_file_paths(
         contains Airflow DAG definitions. If not provided, use the
         core.DAG_DISCOVERY_SAFE_MODE configuration setting. If not set, default
         to safe.
-    :param include_examples: include example DAGs
     :return: a list of paths to Python files in the specified directory
     """
-    if include_examples is None:
-        include_examples = conf.getboolean("core", "LOAD_EXAMPLES")
     file_paths: list[str] = []
     if directory is None:
         file_paths = []
     elif os.path.isfile(directory):
         file_paths = [str(directory)]
     elif os.path.isdir(directory):
         file_paths.extend(find_dag_file_paths(directory, safe_mode))
-    if include_examples:
-        from airflow import example_dags
-
-        example_dag_folder = next(iter(example_dags.__path__))
-        file_paths.extend(list_py_file_paths(example_dag_folder, safe_mode, include_examples=False))
     return file_paths
 
 

tests/dag_processing/bundles/test_dag_bundle_manager.py

@@ -68,7 +68,9 @@
 )
 def test_parse_bundle_config(value, expected):
     """Test that bundle_configs are read from configuration."""
-    envs = {"AIRFLOW__DAG_BUNDLES__BACKENDS": value} if value else {}
+    envs = {"AIRFLOW__CORE__LOAD_EXAMPLES": "False"}
+    if value:
+        envs["AIRFLOW__DAG_BUNDLES__BACKENDS"] = value
     cm = nullcontext()
     exp_fail = False
     if isinstance(expected, str):
@@ -133,6 +135,7 @@ def clear_db():
 
 
 @pytest.mark.db_test
+@conf_vars({("core", "LOAD_EXAMPLES"): "False"})
 def test_sync_bundles_to_db(clear_db):
     def _get_bundle_names_and_active():
         with create_session() as session:
@@ -167,3 +170,14 @@ def test_view_url(version):
     with patch.object(BaseDagBundle, "view_url") as view_url_mock:
         bundle_manager.view_url("my-test-bundle", version=version)
     view_url_mock.assert_called_once_with(version=version)
+
+
+def test_example_dags_bundle_added():
+    manager = DagBundlesManager()
+    manager.parse_config()
+    assert "example_dags" in manager._bundle_config
+
+    with conf_vars({("core", "LOAD_EXAMPLES"): "False"}):
+        manager = DagBundlesManager()
+        manager.parse_config()
+        assert "example_dags" not in manager._bundle_config

tests/jobs/test_scheduler_job.py

@@ -107,6 +107,7 @@
 ELASTIC_DAG_FILE = os.path.join(PERF_DAGS_FOLDER, "elastic_dag.py")
 
 TEST_DAG_FOLDER = os.environ["AIRFLOW__CORE__DAGS_FOLDER"]
+EXAMPLE_DAGS_FOLDER = airflow.example_dags.__path__[0]
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 DEFAULT_LOGICAL_DATE = timezone.coerce_datetime(DEFAULT_DATE)
 TRY_NUMBER = 1
@@ -119,12 +120,6 @@ def disable_load_example():
             yield
 
 
-@pytest.fixture
-def load_examples():
-    with conf_vars({("core", "load_examples"): "True"}):
-        yield
-
-
 # Patch the MockExecutor into the dict of known executors in the Loader
 @contextlib.contextmanager
 def _loader_mock(mock_executors):
@@ -3562,21 +3557,7 @@ def test_list_py_file_paths(self):
                 if file_name.endswith((".py", ".zip")):
                     if file_name not in ignored_files:
                         expected_files.add(f"{root}/{file_name}")
-        for file_path in list_py_file_paths(TEST_DAG_FOLDER, include_examples=False):
-            detected_files.add(file_path)
-        assert detected_files == expected_files
-
-        ignored_files = {
-            "helper.py",
-        }
-        example_dag_folder = airflow.example_dags.__path__[0]
-        for root, _, files in os.walk(example_dag_folder):
-            for file_name in files:
-                if file_name.endswith((".py", ".zip")):
-                    if file_name not in ["__init__.py"] and file_name not in ignored_files:
-                        expected_files.add(os.path.join(root, file_name))
-        detected_files.clear()
-        for file_path in list_py_file_paths(TEST_DAG_FOLDER, include_examples=True):
+        for file_path in list_py_file_paths(TEST_DAG_FOLDER):
             detected_files.add(file_path)
         assert detected_files == expected_files
 
@@ -5662,9 +5643,9 @@ def test_find_and_purge_zombies_nothing(self):
             self.job_runner._find_and_purge_zombies()
         executor.callback_sink.send.assert_not_called()
 
-    def test_find_and_purge_zombies(self, load_examples, session, testing_dag_bundle):
-        dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False)
-
+    def test_find_and_purge_zombies(self, session, testing_dag_bundle):
+        dagfile = os.path.join(EXAMPLE_DAGS_FOLDER, "example_branch_operator.py")
+        dagbag = DagBag(dagfile)
         dag = dagbag.get_dag("example_branch_operator")
         DAG.bulk_write_to_db("testing", None, [dag])
         data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE)
@@ -5718,68 +5699,70 @@ def test_find_and_purge_zombies(self, load_examples, session, testing_dag_bundle
         assert callback_request.ti.run_id == ti.run_id
         assert callback_request.ti.map_index == ti.map_index
 
-    def test_zombie_message(self, load_examples, testing_dag_bundle):
+    def test_zombie_message(self, testing_dag_bundle, session):
         """
         Check that the zombie message comes out as expected
         """
 
         dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False)
-        with create_session() as session:
-            session.query(Job).delete()
-            dag = dagbag.get_dag("example_branch_operator")
-            DAG.bulk_write_to_db("testing", None, [dag])
-
-            data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE)
-            triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST} if AIRFLOW_V_3_0_PLUS else {}
-            dag_run = dag.create_dagrun(
-                state=DagRunState.RUNNING,
-                logical_date=DEFAULT_DATE,
-                run_type=DagRunType.SCHEDULED,
-                session=session,
-                data_interval=data_interval,
-                **triggered_by_kwargs,
-            )
+        dagfile = os.path.join(EXAMPLE_DAGS_FOLDER, "example_branch_operator.py")
+        dagbag = DagBag(dagfile)
+        dag = dagbag.get_dag("example_branch_operator")
+        DAG.bulk_write_to_db("testing", None, [dag])
 
-            scheduler_job = Job(executor=MockExecutor())
-            self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
-            self.job_runner.processor_agent = mock.MagicMock()
+        session.query(Job).delete()
 
-            # We will provision 2 tasks so we can check we only find zombies from this scheduler
-            tasks_to_setup = ["branching", "run_this_first"]
+        data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE)
+        triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST} if AIRFLOW_V_3_0_PLUS else {}
+        dag_run = dag.create_dagrun(
+            state=DagRunState.RUNNING,
+            logical_date=DEFAULT_DATE,
+            run_type=DagRunType.SCHEDULED,
+            session=session,
+            data_interval=data_interval,
+            **triggered_by_kwargs,
+        )
 
-            for task_id in tasks_to_setup:
-                task = dag.get_task(task_id=task_id)
-                ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING)
-                ti.queued_by_job_id = 999
+        scheduler_job = Job(executor=MockExecutor())
+        self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
+        self.job_runner.processor_agent = mock.MagicMock()
 
-                session.add(ti)
-                session.flush()
+        # We will provision 2 tasks so we can check we only find zombies from this scheduler
+        tasks_to_setup = ["branching", "run_this_first"]
 
-            assert task.task_id == "run_this_first"  # Make sure we have the task/ti we expect
+        for task_id in tasks_to_setup:
+            task = dag.get_task(task_id=task_id)
+            ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING)
+            ti.queued_by_job_id = 999
 
-            ti.queued_by_job_id = scheduler_job.id
+            session.add(ti)
             session.flush()
 
-            zombie_message = self.job_runner._generate_zombie_message_details(ti)
-            assert zombie_message == {
-                "DAG Id": "example_branch_operator",
-                "Task Id": "run_this_first",
-                "Run Id": "scheduled__2016-01-01T00:00:00+00:00",
-            }
-
-            ti.hostname = "10.10.10.10"
-            ti.map_index = 2
-            ti.external_executor_id = "abcdefg"
-
-            zombie_message = self.job_runner._generate_zombie_message_details(ti)
-            assert zombie_message == {
-                "DAG Id": "example_branch_operator",
-                "Task Id": "run_this_first",
-                "Run Id": "scheduled__2016-01-01T00:00:00+00:00",
-                "Hostname": "10.10.10.10",
-                "Map Index": 2,
-                "External Executor Id": "abcdefg",
-            }
+        assert task.task_id == "run_this_first"  # Make sure we have the task/ti we expect
+
+        ti.queued_by_job_id = scheduler_job.id
+        session.flush()
+
+        zombie_message = self.job_runner._generate_zombie_message_details(ti)
+        assert zombie_message == {
+            "DAG Id": "example_branch_operator",
+            "Task Id": "run_this_first",
+            "Run Id": "scheduled__2016-01-01T00:00:00+00:00",
+        }
+
+        ti.hostname = "10.10.10.10"
+        ti.map_index = 2
+        ti.external_executor_id = "abcdefg"
+
+        zombie_message = self.job_runner._generate_zombie_message_details(ti)
+        assert zombie_message == {
+            "DAG Id": "example_branch_operator",
+            "Task Id": "run_this_first",
+            "Run Id": "scheduled__2016-01-01T00:00:00+00:00",
+            "Hostname": "10.10.10.10",
+            "Map Index": 2,
+            "External Executor Id": "abcdefg",
+        }
 
     def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_processor(
         self, testing_dag_bundle