Skip to content

Commit

Permalink
Splitting syspath preparation into stages (apache#41672)
Browse files Browse the repository at this point in the history
  • Loading branch information
amoghrajesh authored and uranusjr committed Aug 23, 2024
1 parent 5fc3f63 commit 6c24561
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 13 deletions.
17 changes: 10 additions & 7 deletions airflow/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -675,11 +675,8 @@ def configure_action_logging() -> None:
"""Any additional configuration (register callback) for airflow.utils.action_loggers module."""


def prepare_syspath():
"""Ensure certain subfolders of AIRFLOW_HOME are on the classpath."""
if DAGS_FOLDER not in sys.path:
sys.path.append(DAGS_FOLDER)

def prepare_syspath_for_config_and_plugins():
"""Update sys.path for the config and plugins directories."""
# Add ./config/ for loading custom log parsers etc, or
# airflow_local_settings etc.
config_path = os.path.join(AIRFLOW_HOME, "config")
Expand All @@ -690,6 +687,12 @@ def prepare_syspath():
sys.path.append(PLUGINS_FOLDER)


def prepare_syspath_for_dags_folder():
"""Update sys.path to include the DAGs folder."""
if DAGS_FOLDER not in sys.path:
sys.path.append(DAGS_FOLDER)


def get_session_lifetime_config():
"""Get session timeout configs and handle outdated configs gracefully."""
session_lifetime_minutes = conf.get("webserver", "session_lifetime_minutes", fallback=None)
Expand Down Expand Up @@ -771,12 +774,13 @@ def import_local_settings():
def initialize():
"""Initialize Airflow with all the settings from this file."""
configure_vars()
prepare_syspath()
prepare_syspath_for_config_and_plugins()
configure_policy_plugin_manager()
# Load policy plugins _before_ importing airflow_local_settings, as Pluggy uses LIFO and we want anything
# in airflow_local_settings to take precendec
load_policy_plugins(POLICY_PLUGIN_MANAGER)
import_local_settings()
prepare_syspath_for_dags_folder()
global LOGGING_CLASS_PATH
LOGGING_CLASS_PATH = configure_logging()
State.state_color.update(STATE_COLORS)
Expand Down Expand Up @@ -806,7 +810,6 @@ def is_usage_data_collection_enabled() -> bool:
MEGABYTE = KILOBYTE * KILOBYTE
WEB_COLORS = {"LIGHTBLUE": "#4d9de0", "LIGHTORANGE": "#FF9933"}


# Updating serialized DAG can not be faster than a minimum interval to reduce database
# write rate.
MIN_SERIALIZED_DAG_UPDATE_INTERVAL = conf.getint("core", "min_serialized_dag_update_interval", fallback=30)
Expand Down
33 changes: 27 additions & 6 deletions tests/core/test_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,21 +115,42 @@ def teardown_method(self):
for mod in [m for m in sys.modules if m not in self.old_modules]:
del sys.modules[mod]

@mock.patch("airflow.settings.prepare_syspath_for_config_and_plugins")
@mock.patch("airflow.settings.import_local_settings")
@mock.patch("airflow.settings.prepare_syspath")
def test_initialize_order(self, prepare_syspath, import_local_settings):
@mock.patch("airflow.settings.prepare_syspath_for_dags_folder")
def test_initialize_order(
self,
mock_prepare_syspath_for_dags_folder,
mock_import_local_settings,
mock_prepare_syspath_for_config_and_plugins,
):
"""
Tests that import_local_settings is called after prepare_classpath
Tests that import_local_settings is called between prepare_syspath_for_config_and_plugins
and prepare_syspath_for_dags_folder
"""
mock_local_settings = mock.Mock()
mock_local_settings.attach_mock(prepare_syspath, "prepare_syspath")
mock_local_settings.attach_mock(import_local_settings, "import_local_settings")

mock_local_settings.attach_mock(
mock_prepare_syspath_for_config_and_plugins, "prepare_syspath_for_config_and_plugins"
)
mock_local_settings.attach_mock(mock_import_local_settings, "import_local_settings")
mock_local_settings.attach_mock(
mock_prepare_syspath_for_dags_folder, "prepare_syspath_for_dags_folder"
)

import airflow.settings

airflow.settings.initialize()

mock_local_settings.assert_has_calls([call.prepare_syspath(), call.import_local_settings()])
expected_calls = [
call.prepare_syspath_for_config_and_plugins(),
call.import_local_settings(),
call.prepare_syspath_for_dags_folder(),
]

mock_local_settings.assert_has_calls(expected_calls)

assert mock_local_settings.mock_calls == expected_calls

def test_import_with_dunder_all_not_specified(self):
"""
Expand Down

0 comments on commit 6c24561

Please sign in to comment.