diff --git a/Dockerfile.ci b/Dockerfile.ci index 72484d78588df..bf439cc298a35 100644 --- a/Dockerfile.ci +++ b/Dockerfile.ci @@ -1362,6 +1362,7 @@ function environment_initialization() { export AIRFLOW__CELERY__POOL=${AIRFLOW__CELERY__POOL:-solo} fi export AIRFLOW__CORE__LOAD_EXAMPLES=${LOAD_EXAMPLES} + export AIRFLOW__CORE__LOAD_TESTING_DAGS=${LOAD_TESTING_DAGS} if [[ ${SKIP_ASSETS_COMPILATION:="false"} == "false" ]]; then wait_for_asset_compilation fi diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index b3dc3f133d229..549a04540143e 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -209,6 +209,18 @@ core: type: string example: ~ default: "True" + load_testing_dags: + description: | + Whether to load the testing Dags that ship with Airflow providers. These Dags exist to + exercise Airflow itself (integration tests, smoke tests of a deployment) rather than to + teach Dag authoring, so they are not loaded by default even when ``load_examples`` is + enabled. + + Providers ship them in a ``testing_dags`` folder next to their ``example_dags`` folder. + version_added: "3.4.0" + type: string + example: ~ + default: "False" plugins_folder: description: | Path to the folder containing Airflow plugins diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 66fa318316d89..4a621f7a725d3 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -37,7 +37,7 @@ from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Iterator from sqlalchemy.orm import Session @@ -112,9 +112,9 @@ def _add_example_dag_bundle(bundle_config_list: list[_ExternalBundleConfig]): ) -def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundleConfig]): +def _iter_provider_module_paths() -> Iterator[tuple[str, str]]: """ - Add an ``example_dags`` folder of every installed provider as a bundle. + Yield ``(package_name, module_path)`` for every installed provider. Provider locations are resolved through ``ProvidersManager`` instead of walking ``airflow.providers.__path__`` so that: @@ -124,13 +124,6 @@ def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundl - providers installed outside the ``airflow.providers`` namespace package are discovered via their entry point. """ - # Dedup on the resolved on-disk folder rather than the bundle name: distributions - # under ``airflow.providers.common.*`` use ``pkgutil.extend_path``, so when several - # ``common-*`` packages are installed ``airflow.providers.common.__path__`` has - # multiple entries and the inner loop iterates more than once. Path-based dedup - # only skips when the same folder is seen twice; distinct folders are preserved. - seen: set[str] = set() - for package_name in ProvidersManager().providers: # Heuristic: derive the import path from the canonical # ``apache-airflow-providers-*`` distribution name. Tracked as a follow-up @@ -145,26 +138,40 @@ def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundl module = importlib.import_module(module_name) module_paths = list(getattr(module, "__path__", [])) except Exception: - log.exception("Could not load provider module %s for example DAG discovery", module_name) + log.exception("Could not load provider module %s for DAG discovery", module_name) continue for module_path in module_paths: - example_dag_folder = os.path.join(module_path, "example_dags") - if not os.path.isdir(example_dag_folder): - continue - if example_dag_folder in seen: - continue - seen.add(example_dag_folder) - bundle_name = f"{package_name}-example-dags" - bundle_config_list.append( - _ExternalBundleConfig( - name=bundle_name, - classpath="airflow.dag_processing.bundles.local.LocalDagBundle", - kwargs={ - "path": example_dag_folder, - }, - ) + yield package_name, module_path + + +def _add_provider_dags_to_bundle( + bundle_config_list: list[_ExternalBundleConfig], *, folder_name: str, bundle_name_suffix: str +) -> None: + """Add a ``folder_name`` folder of every installed provider as a bundle.""" + # Dedup on the resolved on-disk folder rather than the bundle name: distributions + # under ``airflow.providers.common.*`` use ``pkgutil.extend_path``, so when several + # ``common-*`` packages are installed ``airflow.providers.common.__path__`` has + # multiple entries and the loop iterates more than once. Path-based dedup + # only skips when the same folder is seen twice; distinct folders are preserved. + seen: set[str] = set() + + for package_name, module_path in _iter_provider_module_paths(): + dag_folder = os.path.join(module_path, folder_name) + if not os.path.isdir(dag_folder): + continue + if dag_folder in seen: + continue + seen.add(dag_folder) + bundle_config_list.append( + _ExternalBundleConfig( + name=f"{package_name}-{bundle_name_suffix}", + classpath="airflow.dag_processing.bundles.local.LocalDagBundle", + kwargs={ + "path": dag_folder, + }, ) + ) def _is_safe_bundle_url(url: str) -> bool: @@ -252,7 +259,13 @@ def parse_config(self) -> None: bundle_config_list = _parse_bundle_config(config_list) if conf.getboolean("core", "LOAD_EXAMPLES"): _add_example_dag_bundle(bundle_config_list) - _add_provider_example_dags_to_bundle(bundle_config_list) + _add_provider_dags_to_bundle( + bundle_config_list, folder_name="example_dags", bundle_name_suffix="example-dags" + ) + if conf.getboolean("core", "LOAD_TESTING_DAGS"): + _add_provider_dags_to_bundle( + bundle_config_list, folder_name="testing_dags", bundle_name_suffix="testing-dags" + ) for bundle_config in bundle_config_list: if bundle_config.team_name and not conf.getboolean("core", "multi_team"): diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index 4b3fc1078abb1..b42d156228a2d 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -467,6 +467,47 @@ def test_example_dags_bundle_added(): assert "example_dags" not in manager._bundle_config +@pytest.fixture +def fake_provider_dag_folders(tmp_path): + """Stand in for an installed provider shipping both ``example_dags`` and ``testing_dags``.""" + module_path = tmp_path / "airflow" / "providers" / "fake" + (module_path / "example_dags").mkdir(parents=True) + (module_path / "testing_dags").mkdir(parents=True) + + with patch( + "airflow.dag_processing.bundles.manager._iter_provider_module_paths", + side_effect=lambda: iter([("apache-airflow-providers-fake", str(module_path))]), + ): + yield module_path + + +def test_provider_testing_dags_bundle_added(fake_provider_dag_folders): + with conf_vars({("core", "LOAD_TESTING_DAGS"): "True"}): + manager = DagBundlesManager() + manager.parse_config() + + bundle_config = manager._bundle_config["apache-airflow-providers-fake-testing-dags"] + assert bundle_config.kwargs["path"] == str(fake_provider_dag_folders / "testing_dags") + + +def test_provider_testing_dags_not_loaded_by_default(fake_provider_dag_folders): + manager = DagBundlesManager() + manager.parse_config() + + assert "apache-airflow-providers-fake-example-dags" in manager._bundle_config + assert "apache-airflow-providers-fake-testing-dags" not in manager._bundle_config + + +def test_provider_testing_dags_load_without_example_dags(fake_provider_dag_folders): + with conf_vars({("core", "LOAD_EXAMPLES"): "False", ("core", "LOAD_TESTING_DAGS"): "True"}): + manager = DagBundlesManager() + manager.parse_config() + + assert "example_dags" not in manager._bundle_config + assert "apache-airflow-providers-fake-example-dags" not in manager._bundle_config + assert "apache-airflow-providers-fake-testing-dags" in manager._bundle_config + + def test_example_dags_name_is_reserved(): reserved_name_config = [{"name": "example_dags", "classpath": "yo face", "kwargs": {}}] with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(reserved_name_config)}): diff --git a/contributing-docs/12_provider_distributions.rst b/contributing-docs/12_provider_distributions.rst index 639d0c014cd8e..769fdb32c0532 100644 --- a/contributing-docs/12_provider_distributions.rst +++ b/contributing-docs/12_provider_distributions.rst @@ -335,6 +335,8 @@ The rules are as follows: * sensors -> sensors are stored here * secrets -> secret backends are stored here * transfers -> transfer operators are stored here + * example_dags -> example Dags teaching how to use the provider, loaded when ``[core] load_examples`` is enabled + * testing_dags -> Dags that exercise Airflow itself rather than teach Dag authoring, loaded when ``[core] load_testing_dags`` is enabled * docs * tests * unit diff --git a/dev/breeze/doc/images/output_shell.svg b/dev/breeze/doc/images/output_shell.svg index 3d183fe886dd7..3bd45b7c2476d 100644 --- a/dev/breeze/doc/images/output_shell.svg +++ b/dev/breeze/doc/images/output_shell.svg @@ -1,4 +1,4 @@ - +