diff --git a/airflow-core/src/airflow/dag_processing/importers/base.py b/airflow-core/src/airflow/dag_processing/importers/base.py index 036dfab98d7a4..183b825adfcd3 100644 --- a/airflow-core/src/airflow/dag_processing/importers/base.py +++ b/airflow-core/src/airflow/dag_processing/importers/base.py @@ -116,7 +116,8 @@ def can_handle(self, file_path: str | Path) -> bool: path = Path(file_path) if isinstance(file_path, str) else file_path return path.suffix.lower() in self.supported_extensions() - def get_relative_path(self, file_path: str | Path, bundle_path: Path | None) -> str: + @staticmethod + def get_relative_path(file_path: str | Path, bundle_path: Path | None) -> str: """Get the relative file path from the bundle root.""" if bundle_path is None: return str(file_path) diff --git a/airflow-core/src/airflow/dag_processing/importers/python_importer.py b/airflow-core/src/airflow/dag_processing/importers/python_importer.py index 2eedc4dff4dec..74ac13c9dd5dd 100644 --- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py +++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py @@ -191,7 +191,7 @@ def import_file( ) # Process imported modules to extract DAGs - self._process_modules(filepath, modules, bundle_name, bundle_path, result) + PythonDagImporter._process_modules(modules, bundle_path, result) return result @@ -331,11 +331,9 @@ def _load_modules_from_zip( del sys.path[0] return mods + @staticmethod def _process_modules( - self, - filepath: str, mods: list[Any], - bundle_name: str | None, bundle_path: Path | None, result: DagImportResult, ) -> None: @@ -353,7 +351,7 @@ def _process_modules( for dag, mod in top_level_dags: dag.fileloc = mod.__file__ - relative_fileloc = self.get_relative_path(dag.fileloc, bundle_path) + relative_fileloc = PythonDagImporter.get_relative_path(dag.fileloc, bundle_path) dag.relative_fileloc = relative_fileloc result.dags.append(dag) diff --git a/airflow-core/tests/unit/dag_processing/importers/test_python_importer.py b/airflow-core/tests/unit/dag_processing/importers/test_python_importer.py new file mode 100644 index 0000000000000..0b5b608588f52 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/importers/test_python_importer.py @@ -0,0 +1,58 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from types import ModuleType + +from airflow.dag_processing.importers import DagImportResult +from airflow.dag_processing.importers.python_importer import PythonDagImporter +from airflow.sdk import DAG + + +def mock_module(name: str, file: str, **kwargs) -> ModuleType: + mod = ModuleType(name) + mod.__file__ = file + for k, v in kwargs.items(): + setattr(mod, k, v) + return mod + + +class TestPythonDagImporter: + def test_process_modules(self): + mods = [ + mock_module("mod1", "bundle/dag1/file1", dag1=DAG(dag_id="hello")), + mock_module("mod2", "bundle/dag2/file2", foo=DAG(dag_id="foo")), + ] + bundle_path = "bundle" + result = DagImportResult(file_path="path") + + PythonDagImporter._process_modules(mods, bundle_path, result) + + assert result.file_path == "path" + assert len(result.dags) == 2 + + # Sort dags for testing + result.dags.sort(key=lambda x: x.dag_id) + + assert result.dags[0].dag_id == "foo" + assert result.dags[0].fileloc == "bundle/dag2/file2" + assert result.dags[0].relative_fileloc == "dag2/file2" + + assert result.dags[1].dag_id == "hello" + assert result.dags[1].fileloc == "bundle/dag1/file1" + assert result.dags[1].relative_fileloc == "dag1/file1"