jason810496 commented on code in PR #72442:
URL: https://github.com/apache/airflow/pull/72442#discussion_r3930438556
##########
airflow-core/src/airflow/dag_processing/importers/base.py:
##########
@@ -258,9 +256,19 @@ def list_dag_files(
@classmethod
def reset(cls) -> None:
"""Reset the singleton (for testing)."""
- cls._instance = None
+ global _global_registry
+ with _global_registry_lock:
+ _global_registry = None
+
+
+_global_registry: DagImporterRegistry | None = None
+_global_registry_lock = threading.Lock()
def get_importer_registry() -> DagImporterRegistry:
"""Get the global importer registry instance."""
- return DagImporterRegistry()
+ global _global_registry
+ with _global_registry_lock:
+ if _global_registry is None:
+ _global_registry = DagImporterRegistry()
+ return _global_registry
Review Comment:
Would it be better to leverage the `@cache` decorator? Since we're trying to
get rid of all the global instance in the repo. Additionally, the `@cache`
should be thread-safe for the concurrent access.
##########
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:
I feel the whole `_find_files_in_bundle` method should be replaced with
`registry.list_dag_files`.
IIUC, the `registry` should contain all the importer respecting each levels
(bundle, dag-processor, global), so `registry.list_dag_files` should give out
all the possible and valid file path.
##########
airflow-core/src/airflow/dag_processing/bundles/base.py:
##########
@@ -444,6 +446,16 @@ def lock(self):
fcntl.flock(lock_file, LOCK_UN)
self._locked = False
+ @property
+ def importer_registry(self) -> DagImporterRegistry:
+ """Get the DAG importer registry for this bundle."""
+ if self._importer_registry is None:
+ from airflow.dag_processing.bundles.manager import
DagBundlesManager
+
+ manager = DagBundlesManager()
+ self._importer_registry = manager.get_importer_registry(self.name)
+ return self._importer_registry
+
Review Comment:
May I ask why do we need to couple `DagImporterRegistry` with
`BaseDagBundle` interface? Since there isn't any consumer so far in the PR.
##########
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):
+ registry.register(importer, extensions=extensions)
+
+ def get_importer_registry(self, bundle_name: str) -> DagImporterRegistry:
+ """Get the DAG importer registry for a bundle."""
+ if bundle_name not in self._bundle_importers:
+ cfg = self._bundle_config.get(bundle_name)
+ importers_config = cfg.importers if cfg else None
+ self._bundle_importers[bundle_name] =
self.create_importer_registry(bundle_name, importers_config)
+ return self._bundle_importers[bundle_name]
Review Comment:
IIUC, the purpose of coupling `DagImporterRegistry` and `DagBundlesManager`
is re-using the `_bundle_config`.
Would it be better to introduce a common layer for retrieving the configs
like --
https://github.com/apache/airflow/pull/70805/changes#diff-1232041046ad2f8635ae867e2cf30914e6ad4619401897d723f6e57731c05e7bR118-R136.
Or wait until #70805 first, then make the `_load_bundle_config_snapshot` a
public function then we don't need to couple `DagImporterRegistry` and
`DagBundlesManager` together.
--
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]
