SameerMesiah97 commented on code in PR #72442:
URL: https://github.com/apache/airflow/pull/72442#discussion_r3926634151
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -1791,3 +1792,76 @@ def test_skipped_row_is_recoverable_after_operator_fix(
assert restored.is_stale is False
assert restored.bundle_name == "configured-bundle"
assert restored.relative_fileloc == "legacy.py"
+
+
+class TestDagBundlesManagerImporters:
+ """Tests for DAG importer registry integration in DagBundlesManager."""
+
+ def test_create_importer_registry_precedence_and_overrides(self, caplog):
Review Comment:
This test currently uses PythonDagImporter at every level, so it does not
demonstrate that the bundle-level importer actually wins. What I would suggest
is to create 2 custom classes at module level:
```
class GlobalDagImporter(PythonDagImporter):
pass
class BundleDagImporter(PythonDagImporter):
pass
```
then you can use them in your test like this:
```
def test_create_importer_registry_precedence_and_overrides(self, caplog):
"""Bundle importers override global importers, which override
defaults."""
global_config = [
{
"classpath": f"{__name__}.GlobalDagImporter",
"extensions": [".py", ".custom"],
}
]
bundle_config = [
{
"classpath": f"{__name__}.BundleDagImporter",
"extensions": [".custom"],
}
]
with (
caplog.at_level(logging.WARNING),
conf_vars(
{
("dag_processor", "dag_importer_configs"): json.dumps(
global_config
)
}
),
):
manager = DagBundlesManager()
registry = manager.create_importer_registry(
"bundle_override",
bundle_config,
)
# Global configuration overrides the default .py importer.
assert isinstance(registry.get_importer("dag.py"), GlobalDagImporter)
# Bundle configuration overrides the global .custom importer.
assert isinstance(
registry.get_importer("dag.custom"),
BundleDagImporter,
)
assert any(
"Extension '.py' already registered" in record.message
for record in caplog.records
)
assert any(
"Extension '.custom' already registered" in record.message
for record in caplog.records
)
```
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -1791,3 +1792,76 @@ def test_skipped_row_is_recoverable_after_operator_fix(
assert restored.is_stale is False
assert restored.bundle_name == "configured-bundle"
assert restored.relative_fileloc == "legacy.py"
+
+
+class TestDagBundlesManagerImporters:
+ """Tests for DAG importer registry integration in DagBundlesManager."""
+
+ def test_create_importer_registry_precedence_and_overrides(self, caplog):
+ """Bundle importer overrides global importer, which overrides default
importers."""
+ from airflow.dag_processing.importers.python_importer import
PythonDagImporter
+
+ global_config = [
+ {
+ "classpath":
"airflow.dag_processing.importers.python_importer.PythonDagImporter",
+ "extensions": [".py", ".custom"],
+ }
+ ]
+ bundle_config = [
+ {
+ "classpath":
"airflow.dag_processing.importers.python_importer.PythonDagImporter",
+ "extensions": [".custom"],
+ }
+ ]
+ with (
+ caplog.at_level(logging.WARNING),
+ conf_vars({("dag_processor", "dag_importer_configs"):
json.dumps(global_config)}),
+ ):
+ manager = DagBundlesManager()
+ registry = manager.create_importer_registry("bundle_override",
bundle_config)
+
+ assert registry.can_handle("dag.py")
+ assert registry.can_handle("dag.custom")
+ importer = registry.get_importer("test.custom")
+ assert isinstance(importer, PythonDagImporter)
+ assert any("Extension '.py' already registered" in r.message for r in
caplog.records)
+ assert any("Extension '.custom' already registered" in r.message for r
in caplog.records)
+
+ @pytest.mark.parametrize(
+ ("global_cfg", "bundle_cfg", "match"),
+ [
+ (json.dumps({"invalid": "object"}), None, "key
`dag_importer_configs` must be a list"),
+ (None, [{"extensions": [".py"]}], "Missing required 'classpath'"),
+ (None, [{"classpath": "invalid.path"}], "Failed to load DAG
importer"),
+ ],
+ )
+ def test_create_importer_registry_invalid_configs(self, global_cfg,
bundle_cfg, match):
+ """Invalid global or bundle configurations raise
AirflowConfigException."""
+ manager = DagBundlesManager()
+ with (
+ conf_vars({("dag_processor", "dag_importer_configs"): global_cfg}
if global_cfg else {}),
+ pytest.raises(AirflowConfigException, match=match),
+ ):
+ manager.create_importer_registry("test", bundle_cfg)
+
+ def test_create_importer_registry_nested_internal_importers(self):
+ """Composite importers normalize nested internal_importers."""
+ manager = DagBundlesManager()
+ importers_config = [
+ {
+ "classpath": "airflow.sdk.importers.zip_importer.ZipImporter",
+ "extensions": [".zip"],
+ "kwargs": {
+ "internal_importers": [
+ {
+ "classpath":
"airflow.sdk.importers.python_importer.PythonDagImporter",
+ "extensions": [".py"],
+ }
+ ]
+ },
+ }
+ ]
+ registry = manager.create_importer_registry("test_bundle",
importers_config)
+ importer = registry.get_importer("archive.zip")
+ assert importer is not None
+ assert ".py" in importer._internal_importers
Review Comment:
I would add a test to cover caching behavior:
```
def test_get_importer_registry_is_cached(self):
manager = DagBundlesManager()
registry = manager.get_importer_registry("test_bundle")
assert manager.get_importer_registry("test_bundle") is registry
```
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -958,12 +958,10 @@ def _refresh_dag_bundles(self, known_files: dict[str,
set[DagFileInfo]]):
def _find_files_in_bundle(self, bundle: BaseDagBundle) -> list[Path]:
"""Get relative paths for dag files from bundle dir."""
- # Build up a list of Python files that could contain DAGs
self.log.info("Searching for files in %s at %s", bundle.name,
bundle.path)
- rel_paths = [
- Path(x).relative_to(bundle.path)
- for x in list_py_file_paths(bundle.path,
safe_mode=self.dag_discovery_safe_mode)
- ]
+ importer_registry = bundle.importer_registry
+ dag_files = importer_registry.list_dag_files(bundle.path,
safe_mode=self.dag_discovery_safe_mode)
Review Comment:
Correct me if I am wrong (I am not as familiar with the DAG processor), but
it looks like explicitly configured extensions are only used for registry
lookup, not file discovery. For example, configring `PythonDegImporter` for
"custom" makes `get_importer("dag.custom")` succeed, but `list_dag_files()`
delegates to `PythonDagImporter.list_dag_files()`, which only discovers Python
files and ZIP archives. Would custom files therefore never reach parsing?
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -661,12 +672,53 @@ def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]:
"""
for name, cfg in self._bundle_config.items():
try:
- yield cfg.bundle_class(name=name, version=None, **cfg.kwargs)
+ bundle = cfg.bundle_class(name=name, version=None,
**cfg.kwargs)
+ bundle._importer_registry = self.get_importer_registry(name)
+ yield bundle
except Exception as e:
self.log.exception("Error creating bundle '%s': %s", name, e)
# Skip this bundle and continue with others
continue
+ def create_importer_registry(
+ self, bundle_name: str, importers_config: list[dict[str, Any]] | None
+ ) -> DagImporterRegistry:
+ """Create and configure a DagImporterRegistry for a bundle with 3-tier
precedence."""
+ from airflow.dag_processing.importers import DagImporterRegistry
+
+ registry = DagImporterRegistry()
+
+ # Global configuration
+ global_importers = conf.getjson("dag_processor",
"dag_importer_configs", fallback=None)
+ if global_importers:
+ if not isinstance(global_importers, list):
+ raise AirflowConfigException(
+ "Section `dag_processor` key `dag_importer_configs` must
be a list "
+ f"but got {global_importers.__class__.__name__}"
+ )
+ self._load_importers_into_registry(registry, global_importers,
context="global configuration")
+
+ # Bundle explicit mapping
+ if importers_config:
+ self._load_importers_into_registry(registry, importers_config,
context=f"bundle '{bundle_name}'")
+
+ return registry
+
+ def _load_importers_into_registry(
+ self, registry: DagImporterRegistry, configs: list[dict[str, Any]],
context: str
+ ) -> None:
+ """Dynamically load and register custom DAG importers."""
+ for importer, extensions in load_dag_importers(configs,
context=context):
Review Comment:
Nit: I would add some validation here. Please see the below:
```
def _load_importers_into_registry(
self,
registry: DagImporterRegistry,
configs: list[dict[str, Any]],
context: str,
) -> None:
"""Dynamically load and register custom DAG importers."""
from airflow.dag_processing.importers import AbstractDagImporter
for importer, extensions in load_dag_importers(configs, context=context):
if not isinstance(importer, AbstractDagImporter):
raise AirflowConfigException(
f"Configured DAG importer {type(importer).__module__}."
f"{type(importer).__qualname__} for {context} must inherit "
"from AbstractDagImporter."
)
registry.register(importer, extensions=extensions)
```
I would ensure there is test coverage for this validation as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
