Skip to content

Commit

Permalink
Create directories based on AIRFLOW_CONFIG path (#35818)
Browse files Browse the repository at this point in the history
  • Loading branch information
Taragolis committed Nov 24, 2023
1 parent 0b23d56 commit cc042a2
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 5 deletions.
30 changes: 25 additions & 5 deletions airflow/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -1955,17 +1955,37 @@ 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
global FERNET_KEY
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,
Expand All @@ -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


Expand Down
88 changes: 88 additions & 0 deletions tests/core/test_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

0 comments on commit cc042a2

Please sign in to comment.