From cc042a2b7ed4e869b22ee25060548f4897ea78bd Mon Sep 17 00:00:00 2001 From: Andrey Anshin Date: Fri, 24 Nov 2023 18:35:30 +0400 Subject: [PATCH] Create directories based on `AIRFLOW_CONFIG` path (#35818) --- airflow/configuration.py | 30 +++++++++-- tests/core/test_configuration.py | 88 ++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 5 deletions(-) diff --git a/airflow/configuration.py b/airflow/configuration.py index 1df62c9b9944b..7dca8a44d1dfa 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -1955,9 +1955,29 @@ def create_pre_2_7_defaults() -> ConfigParser: def write_default_airflow_configuration_if_needed() -> AirflowConfigParser: - if not os.path.isfile(AIRFLOW_CONFIG): - log.debug("Creating new Airflow config file in: %s", AIRFLOW_CONFIG) - pathlib.Path(AIRFLOW_HOME).mkdir(parents=True, exist_ok=True) + airflow_config = pathlib.Path(AIRFLOW_CONFIG) + if airflow_config.is_dir(): + msg = ( + "Airflow config expected to be a path to the configuration file, " + f"but got a directory {airflow_config.__fspath__()!r}." + ) + raise IsADirectoryError(msg) + elif not airflow_config.exists(): + log.debug("Creating new Airflow config file in: %s", airflow_config.__fspath__()) + config_directory = airflow_config.parent + if not config_directory.exists(): + # Compatibility with Python 3.8, ``PurePath.is_relative_to`` was added in Python 3.9 + try: + config_directory.relative_to(AIRFLOW_HOME) + except ValueError: + msg = ( + f"Config directory {config_directory.__fspath__()!r} not exists " + f"and it is not relative to AIRFLOW_HOME {AIRFLOW_HOME!r}. " + "Please create this directory first." + ) + raise FileNotFoundError(msg) from None + log.debug("Create directory %r for Airflow config", config_directory.__fspath__()) + config_directory.mkdir(parents=True, exist_ok=True) if conf.get("core", "fernet_key", fallback=None) is None: # We know that FERNET_KEY is not set, so we can generate it, set as global key # and also write it to the config file so that same key will be used next time @@ -1965,7 +1985,7 @@ def write_default_airflow_configuration_if_needed() -> AirflowConfigParser: FERNET_KEY = _generate_fernet_key() conf.remove_option("core", "fernet_key") conf.set("core", "fernet_key", FERNET_KEY) - with open(AIRFLOW_CONFIG, "w") as file: + with open(airflow_config, "w") as file: conf.write( file, include_sources=False, @@ -1974,7 +1994,7 @@ def write_default_airflow_configuration_if_needed() -> AirflowConfigParser: extra_spacing=True, only_defaults=True, ) - make_group_other_inaccessible(AIRFLOW_CONFIG) + make_group_other_inaccessible(airflow_config.__fspath__()) return conf diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py index 2d2a3ea6a0235..e528e50617f92 100644 --- a/tests/core/test_configuration.py +++ b/tests/core/test_configuration.py @@ -39,6 +39,7 @@ get_airflow_home, get_all_expansion_variables, run_command, + write_default_airflow_configuration_if_needed, ) from tests.test_utils.config import conf_vars from tests.test_utils.reset_warning_registry import reset_warning_registry @@ -1649,3 +1650,90 @@ def test_error_when_contributing_to_existing_section(): ): conf.load_providers_configuration() assert conf.get("celery", "celery_app_name") == "test" + + +class TestWriteDefaultAirflowConfigurationIfNeeded: + @pytest.fixture(autouse=True) + def setup_test_cases(self, tmp_path_factory): + self.test_airflow_home = tmp_path_factory.mktemp("airflow_home") + self.test_airflow_config = self.test_airflow_home / "airflow.cfg" + self.test_non_relative_path = tmp_path_factory.mktemp("other") + + with pytest.MonkeyPatch.context() as monkeypatch_ctx: + self.monkeypatch = monkeypatch_ctx + self.patch_airflow_home(self.test_airflow_home) + self.patch_airflow_config(self.test_airflow_config) + yield + + def patch_airflow_home(self, airflow_home): + self.monkeypatch.setattr("airflow.configuration.AIRFLOW_HOME", os.fspath(airflow_home)) + + def patch_airflow_config(self, airflow_config): + self.monkeypatch.setattr("airflow.configuration.AIRFLOW_CONFIG", os.fspath(airflow_config)) + + def test_default(self): + """Test write default config in `${AIRFLOW_HOME}/airflow.cfg`.""" + assert not self.test_airflow_config.exists() + write_default_airflow_configuration_if_needed() + assert self.test_airflow_config.exists() + + @pytest.mark.parametrize( + "relative_to_airflow_home", + [ + pytest.param(True, id="relative-to-airflow-home"), + pytest.param(False, id="non-relative-to-airflow-home"), + ], + ) + def test_config_already_created(self, relative_to_airflow_home): + if relative_to_airflow_home: + test_airflow_config = self.test_airflow_home / "test-existed-config" + else: + test_airflow_config = self.test_non_relative_path / "test-existed-config" + + test_airflow_config.write_text("foo=bar") + write_default_airflow_configuration_if_needed() + assert test_airflow_config.read_text() == "foo=bar" + + def test_config_path_relative(self): + """Test write default config in path relative to ${AIRFLOW_HOME}.""" + test_airflow_config_parent = self.test_airflow_home / "config" + test_airflow_config = test_airflow_config_parent / "test-airflow.config" + self.patch_airflow_config(test_airflow_config) + + assert not test_airflow_config_parent.exists() + assert not test_airflow_config.exists() + write_default_airflow_configuration_if_needed() + assert test_airflow_config.exists() + + def test_config_path_non_relative_directory_exists(self): + """Test write default config in path non-relative to ${AIRFLOW_HOME} and directory exists.""" + test_airflow_config_parent = self.test_non_relative_path + test_airflow_config = test_airflow_config_parent / "test-airflow.cfg" + self.patch_airflow_config(test_airflow_config) + + assert test_airflow_config_parent.exists() + assert not test_airflow_config.exists() + write_default_airflow_configuration_if_needed() + assert test_airflow_config.exists() + + def test_config_path_non_relative_directory_not_exists(self): + """Test raise an error if path to config non-relative to ${AIRFLOW_HOME} and directory not exists.""" + test_airflow_config_parent = self.test_non_relative_path / "config" + test_airflow_config = test_airflow_config_parent / "test-airflow.cfg" + self.patch_airflow_config(test_airflow_config) + + assert not test_airflow_config_parent.exists() + assert not test_airflow_config.exists() + with pytest.raises(FileNotFoundError, match="not exists and it is not relative to"): + write_default_airflow_configuration_if_needed() + assert not test_airflow_config.exists() + assert not test_airflow_config_parent.exists() + + def test_config_paths_is_directory(self): + """Test raise an error if AIRFLOW_CONFIG is a directory.""" + test_airflow_config = self.test_airflow_home / "config-dir" + test_airflow_config.mkdir() + self.patch_airflow_config(test_airflow_config) + + with pytest.raises(IsADirectoryError, match="configuration file, but got a directory"): + write_default_airflow_configuration_if_needed()