From 3daa850a002982419f49c926c2c433f4f8fdd067 Mon Sep 17 00:00:00 2001 From: Aaryan Mahajan Date: Thu, 2 Jul 2026 23:10:50 +0530 Subject: [PATCH 1/3] Fix custom stats metrics not being sent from plugins/listeners Airflow keeps two independent Stats singletons: one internal to airflow-core, and the task-sdk copy that airflow.sdk.observability.stats (and the deprecated airflow.stats shim) expose to plugins and listener hooks. Long-running components only initialized the internal copy at startup, so a plugin's DAG-run listener hook running in the scheduler or API server silently got NoStatsLogger regardless of StatsD config, even though native Airflow metrics worked fine. closes: #69172 --- airflow-core/src/airflow/api_fastapi/app.py | 1 + .../src/airflow/dag_processing/manager.py | 1 + .../src/airflow/executors/base_executor.py | 1 + .../src/airflow/jobs/scheduler_job_runner.py | 1 + .../src/airflow/jobs/triggerer_job_runner.py | 1 + .../observability/metrics/stats_utils.py | 18 ++++++++++++ .../tests/unit/api_fastapi/test_app.py | 11 ++++++- .../tests/unit/dag_processing/test_manager.py | 8 ++++- .../unit/executors/test_base_executor.py | 8 +++++ .../tests/unit/jobs/test_scheduler_job.py | 12 ++++++++ .../tests/unit/jobs/test_triggerer_job.py | 8 ++++- .../unit/observability/metrics/test_stats.py | 29 +++++++++++++++++++ 12 files changed, 96 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/app.py b/airflow-core/src/airflow/api_fastapi/app.py index e4d1e40efbb3d..7add5299b78d0 100644 --- a/airflow-core/src/airflow/api_fastapi/app.py +++ b/airflow-core/src/airflow/api_fastapi/app.py @@ -86,6 +86,7 @@ def _initialize_api_server_stats() -> None: factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) + stats_utils.initialize_sdk_stats_backend() except Exception: log.warning( "Failed to initialize API server Stats in the API server; metrics emitted through the " diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index f523ca6c3e1e2..294d4ccea1978 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -396,6 +396,7 @@ def prepare_process_context(self) -> None: factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) + stats_utils.initialize_sdk_stats_backend() def prepare_bundles(self) -> None: """Sync bundle configuration to the DB and load bundles for parsing.""" diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index ebd850ad944b1..b163b6b95b4ce 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -215,6 +215,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None) factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) + stats_utils.initialize_sdk_stats_backend() super().__init__() # Ensure we set this now, so that each subprocess gets the same value from airflow.api_fastapi.auth.tokens import get_signing_args diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index cb4c8f645520c..ee02b3542008b 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1674,6 +1674,7 @@ def _execute(self) -> int | None: factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) + stats_utils.initialize_sdk_stats_backend() self._run_scheduler_loop() diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 759735242ab1d..12c9058cf8eb8 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -266,6 +266,7 @@ def _execute(self) -> int | None: factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) + stats_utils.initialize_sdk_stats_backend() self.trigger_runner = None try: # Kick off runner sub-process without DB access diff --git a/airflow-core/src/airflow/observability/metrics/stats_utils.py b/airflow-core/src/airflow/observability/metrics/stats_utils.py index 9e0f1263bff66..c76b1f6ae4b78 100644 --- a/airflow-core/src/airflow/observability/metrics/stats_utils.py +++ b/airflow-core/src/airflow/observability/metrics/stats_utils.py @@ -37,3 +37,21 @@ def get_stats_factory() -> Callable: return otel_logger.get_otel_logger return NoStatsLogger + + +def initialize_sdk_stats_backend() -> None: + """ + Initialize the task-sdk ``Stats`` singleton with this process's configured backend. + + Plugins and listener hooks commonly get ``Stats`` via ``airflow.sdk.observability.stats`` + (or the deprecated ``airflow.stats`` shim, which re-exports the same module). That module + is a separate singleton from the one this process initializes for its own internal use, so + it must be initialized here too, or plugin code silently falls back to ``NoStatsLogger`` in + any long-running component other than the task runner. + """ + from airflow.sdk.observability import stats as sdk_stats # noqa: SDK001 + + sdk_stats.initialize( + factory=get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) diff --git a/airflow-core/tests/unit/api_fastapi/test_app.py b/airflow-core/tests/unit/api_fastapi/test_app.py index 9fd9a3edad1ae..e6ac8cb911fc5 100644 --- a/airflow-core/tests/unit/api_fastapi/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/test_app.py @@ -24,6 +24,7 @@ import airflow.api_fastapi.app as app_module import airflow.plugins_manager as plugins_manager +import airflow.sdk.observability.stats pytestmark = pytest.mark.db_test @@ -180,6 +181,7 @@ def test_initializes_api_server_stats_with_factory(self): sentinel_factory = object() with ( mock.patch("airflow._shared.observability.metrics.stats") as mock_stats, + mock.patch.object(airflow.sdk.observability.stats, "initialize") as mock_sdk_initialize, mock.patch( "airflow.observability.metrics.stats_utils.get_stats_factory", return_value=sentinel_factory, @@ -187,12 +189,19 @@ def test_initializes_api_server_stats_with_factory(self): ): app_module._initialize_api_server_stats() - mock_get_factory.assert_called_once_with() + assert mock_get_factory.call_count == 2 mock_stats.initialize.assert_called_once() _, kwargs = mock_stats.initialize.call_args assert kwargs["factory"] is sentinel_factory assert isinstance(kwargs["export_legacy_names"], bool) + # Plugins/listeners get Stats via the task-sdk-backed singleton, a separate one + # from the API server's own — it must be initialized too, see initialize_sdk_stats_backend. + mock_sdk_initialize.assert_called_once() + _, sdk_kwargs = mock_sdk_initialize.call_args + assert sdk_kwargs["factory"] is sentinel_factory + assert isinstance(sdk_kwargs["export_legacy_names"], bool) + def test_stats_failure_does_not_block_startup(self): """A metrics misconfiguration must not prevent the API server from starting.""" with ( diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index c61e29c591a2a..faf807e8738bf 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2592,8 +2592,11 @@ def test_create_process_passes_bundle_name_to_process_start( call_kwargs = mock_process_start.call_args.kwargs assert call_kwargs["bundle_name"] == "testing" + @mock.patch("airflow.dag_processing.manager.stats_utils.initialize_sdk_stats_backend") @mock.patch("airflow.dag_processing.manager.stats.initialize") - def test_stats_initialize_called_on_run(self, stats_init_mock, tmp_path, configure_testing_dag_bundle): + def test_stats_initialize_called_on_run( + self, stats_init_mock, sdk_stats_init_mock, tmp_path, configure_testing_dag_bundle + ): """Test that stats.initialize() is called when DagFileProcessorManager.run() is executed.""" with configure_testing_dag_bundle(tmp_path): manager = DagFileProcessorManager(max_runs=1) @@ -2604,6 +2607,9 @@ def test_stats_initialize_called_on_run(self, stats_init_mock, tmp_path, configu call_kwargs = stats_init_mock.call_args.kwargs assert "factory" in call_kwargs + # The task-sdk Stats singleton (what plugins/listeners use) must be initialized too. + sdk_stats_init_mock.assert_called_once() + def test_run_invokes_before_and_after_hooks(self, tmp_path, configure_testing_dag_bundle): """`run()` should call `before_run` then `after_run`, even if the loop raises.""" with configure_testing_dag_bundle(tmp_path): diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 8bc7cfd73db23..96a355a6f74bd 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -72,6 +72,14 @@ def test_invalid_slotspool(): BaseExecutor(0) +@mock.patch("airflow.executors.base_executor.stats_utils.initialize_sdk_stats_backend") +def test_init_initializes_sdk_stats_backend(sdk_stats_init_mock): + """The task-sdk Stats singleton (what plugins/listeners use) must be initialized too.""" + BaseExecutor() + + sdk_stats_init_mock.assert_called_once() + + def test_get_task_log(): executor = BaseExecutor() ti = TaskInstance(task=SerializedBaseOperator(task_id="dummy"), dag_version_id=mock.MagicMock(spec=UUID)) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 782bcf81c4aa3..e12d73106f42b 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9472,6 +9472,18 @@ def test_process_expired_deadlines(self, mock_handle_miss, session, dag_maker): # Assert that all deadlines which are both expired and unhandled get processed. assert mock_handle_miss.call_count == 2 + @mock.patch("airflow.jobs.scheduler_job_runner.stats_utils.initialize_sdk_stats_backend") + def test_execute_initializes_sdk_stats_backend(self, sdk_stats_init_mock, session): + """The task-sdk Stats singleton (what plugins/listeners use) must be initialized too.""" + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1, executors=[MockExecutor()]) + + self.job_runner._execute() + + # BaseExecutor.__init__ also calls this (same shared stats_utils module), so + # constructing MockExecutor() above contributes a call in addition to _execute()'s own. + assert sdk_stats_init_mock.called + @mock.patch("airflow.models.Deadline.handle_miss") def test_process_expired_deadlines_no_deadlines_found(self, mock_handle_miss, session): """Test handling when there are no deadlines to process.""" diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index fa4d68f5fff6a..a1dbb4e4fe484 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2482,9 +2482,12 @@ def test_update_triggers_skips_when_ti_has_no_dag_version(session, supervisor_bu class TestTriggererJobRunner: + @patch("airflow.jobs.triggerer_job_runner.stats_utils.initialize_sdk_stats_backend") @patch("airflow.jobs.triggerer_job_runner.stats.initialize") @patch.object(TriggerRunnerSupervisor, "start") - def test_stats_initialize_called_on_execute(self, mock_supervisor_start, stats_init_mock, session): + def test_stats_initialize_called_on_execute( + self, mock_supervisor_start, stats_init_mock, sdk_stats_init_mock, session + ): """Test that stats.initialize() is called when TriggererJobRunner._execute() is executed.""" # Setup mock supervisor to immediately stop mock_supervisor = MagicMock() @@ -2514,6 +2517,9 @@ def test_stats_initialize_called_on_execute(self, mock_supervisor_start, stats_i call_kwargs = stats_init_mock.call_args.kwargs assert "factory" in call_kwargs + # The task-sdk Stats singleton (what plugins/listeners use) must be initialized too. + sdk_stats_init_mock.assert_called_once() + @patch.object(TriggerRunnerSupervisor, "start") def test_execute_sets_server_process_context(self, mock_supervisor_start, session, monkeypatch): """_execute marks triggerer as server context for secrets backend detection.""" diff --git a/airflow-core/tests/unit/observability/metrics/test_stats.py b/airflow-core/tests/unit/observability/metrics/test_stats.py index 426e87a11991d..de71a9bca1244 100644 --- a/airflow-core/tests/unit/observability/metrics/test_stats.py +++ b/airflow-core/tests/unit/observability/metrics/test_stats.py @@ -18,11 +18,14 @@ import importlib import re +from unittest import mock import pytest import airflow import airflow.observability.stats +import airflow.sdk.observability.stats +from airflow.configuration import conf from airflow.observability.metrics import stats_utils from tests_common.test_utils.config import conf_vars @@ -100,3 +103,29 @@ def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self assert isinstance(backend.dogstatsd, DogStatsd) assert not hasattr(backend, "statsd") importlib.reload(airflow.observability.stats) + + +class TestInitializeSdkStatsBackend: + """ + Plugins/listeners typically get ``Stats`` via ``airflow.sdk.observability.stats`` (or the + deprecated ``airflow.stats`` shim, which re-exports it), a singleton separate from the one + this module initializes for internal use. ``initialize_sdk_stats_backend`` keeps that + singleton configured too, so plugin code works outside the task runner as well. + """ + + def test_configures_sdk_stats_singleton_with_this_process_backend(self): + with conf_vars({("metrics", "statsd_on"): "True"}): + importlib.reload(airflow.sdk.observability.stats) + stats_utils.initialize_sdk_stats_backend() + backend = airflow.sdk._shared.observability.metrics.stats._get_backend() + assert hasattr(backend, "statsd") + importlib.reload(airflow.sdk.observability.stats) + + def test_uses_same_factory_and_legacy_names_flag_as_this_process(self): + with mock.patch("airflow.sdk.observability.stats.initialize") as mock_sdk_initialize: + stats_utils.initialize_sdk_stats_backend() + + mock_sdk_initialize.assert_called_once_with( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) From a0fcf7be041c87be8ea8f5015d6f501c0382f09f Mon Sep 17 00:00:00 2001 From: Aaryan Mahajan Date: Thu, 2 Jul 2026 23:20:05 +0530 Subject: [PATCH 2/3] Add newsfragment for plugin StatsD metrics fix --- airflow-core/newsfragments/69270.bugfix.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 airflow-core/newsfragments/69270.bugfix.rst diff --git a/airflow-core/newsfragments/69270.bugfix.rst b/airflow-core/newsfragments/69270.bugfix.rst new file mode 100644 index 0000000000000..14f49388a2bc8 --- /dev/null +++ b/airflow-core/newsfragments/69270.bugfix.rst @@ -0,0 +1 @@ +Fixed custom StatsD/Datadog/OpenTelemetry metrics not being emitted from plugins when using listener hooks (e.g. ``on_dag_run_success``) that run in the scheduler or API server process. From 396cb29e2fcae8ddc637dceaa63a3fd208f5a3bc Mon Sep 17 00:00:00 2001 From: Aaryan Mahajan Date: Fri, 24 Jul 2026 16:10:49 +0400 Subject: [PATCH 3/3] Fix custom stats metrics not being sent from plugins/listeners Plugins and listeners typically reach Stats through the task-sdk import path, a separate singleton from the one each component initializes for its own internal use, so plugin code kept using an unconfigured NoStatsLogger even after the component's own Stats was set up. Propagate the configuration to any sibling copy of the stats module already loaded under a different distribution path instead of having core import from airflow.sdk directly to reach it. Metrics initialization failures no longer prevent the API server, scheduler, Dag processor, triggerer, or executor from starting. --- airflow-core/newsfragments/69270.bugfix.rst | 1 - airflow-core/src/airflow/api_fastapi/app.py | 1 - .../src/airflow/dag_processing/manager.py | 15 ++++--- .../src/airflow/executors/base_executor.py | 15 ++++--- .../src/airflow/jobs/scheduler_job_runner.py | 15 ++++--- .../src/airflow/jobs/triggerer_job_runner.py | 15 ++++--- .../observability/metrics/stats_utils.py | 18 -------- .../tests/unit/api_fastapi/test_app.py | 11 +---- .../tests/unit/dag_processing/test_manager.py | 15 ++++--- .../unit/executors/test_base_executor.py | 10 ++--- .../tests/unit/jobs/test_scheduler_job.py | 11 ++--- .../tests/unit/jobs/test_triggerer_job.py | 33 +++++++++++--- .../unit/observability/metrics/test_stats.py | 26 +++++------ .../observability/metrics/stats.py | 32 +++++++++++++ .../tests/observability/metrics/test_stats.py | 45 +++++++++++++++++++ 15 files changed, 174 insertions(+), 89 deletions(-) delete mode 100644 airflow-core/newsfragments/69270.bugfix.rst diff --git a/airflow-core/newsfragments/69270.bugfix.rst b/airflow-core/newsfragments/69270.bugfix.rst deleted file mode 100644 index 14f49388a2bc8..0000000000000 --- a/airflow-core/newsfragments/69270.bugfix.rst +++ /dev/null @@ -1 +0,0 @@ -Fixed custom StatsD/Datadog/OpenTelemetry metrics not being emitted from plugins when using listener hooks (e.g. ``on_dag_run_success``) that run in the scheduler or API server process. diff --git a/airflow-core/src/airflow/api_fastapi/app.py b/airflow-core/src/airflow/api_fastapi/app.py index 7add5299b78d0..e4d1e40efbb3d 100644 --- a/airflow-core/src/airflow/api_fastapi/app.py +++ b/airflow-core/src/airflow/api_fastapi/app.py @@ -86,7 +86,6 @@ def _initialize_api_server_stats() -> None: factory=stats_utils.get_stats_factory(), export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), ) - stats_utils.initialize_sdk_stats_backend() except Exception: log.warning( "Failed to initialize API server Stats in the API server; metrics emitted through the " diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 294d4ccea1978..e53624ccd96db 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -392,11 +392,16 @@ def prepare_process_context(self) -> None: # selector implementations. Also see _StubSelector documentation. self.selector = selectors.DefaultSelector() - stats.initialize( - factory=stats_utils.get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) - stats_utils.initialize_sdk_stats_backend() + try: + stats.initialize( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) + except Exception: + self.log.warning( + "Failed to initialize Stats in the Dag processor; metrics will not be recorded.", + exc_info=True, + ) def prepare_bundles(self) -> None: """Sync bundle configuration to the DB and load bundles for parsing.""" diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index b163b6b95b4ce..bb269be26d2cf 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -211,11 +211,16 @@ def jwt_generator(self) -> JWTGenerator: return generator def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None): - stats.initialize( - factory=stats_utils.get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) - stats_utils.initialize_sdk_stats_backend() + try: + stats.initialize( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) + except Exception: + self.log.warning( + "Failed to initialize Stats in the executor; metrics will not be recorded.", + exc_info=True, + ) super().__init__() # Ensure we set this now, so that each subprocess gets the same value from airflow.api_fastapi.auth.tokens import get_signing_args diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index ee02b3542008b..00560c916bf84 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1670,11 +1670,16 @@ def _execute(self) -> int | None: # local import due to type_checking. - stats.initialize( - factory=stats_utils.get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) - stats_utils.initialize_sdk_stats_backend() + try: + stats.initialize( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) + except Exception: + self.log.warning( + "Failed to initialize Stats in the scheduler; metrics will not be recorded.", + exc_info=True, + ) self._run_scheduler_loop() diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 12c9058cf8eb8..69cf4cfcbdc0b 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -262,11 +262,16 @@ def _execute(self) -> int | None: os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server" self.log.info("Starting the triggerer") self.register_signals() - stats.initialize( - factory=stats_utils.get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) - stats_utils.initialize_sdk_stats_backend() + try: + stats.initialize( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) + except Exception: + self.log.warning( + "Failed to initialize Stats in the triggerer; metrics will not be recorded.", + exc_info=True, + ) self.trigger_runner = None try: # Kick off runner sub-process without DB access diff --git a/airflow-core/src/airflow/observability/metrics/stats_utils.py b/airflow-core/src/airflow/observability/metrics/stats_utils.py index c76b1f6ae4b78..9e0f1263bff66 100644 --- a/airflow-core/src/airflow/observability/metrics/stats_utils.py +++ b/airflow-core/src/airflow/observability/metrics/stats_utils.py @@ -37,21 +37,3 @@ def get_stats_factory() -> Callable: return otel_logger.get_otel_logger return NoStatsLogger - - -def initialize_sdk_stats_backend() -> None: - """ - Initialize the task-sdk ``Stats`` singleton with this process's configured backend. - - Plugins and listener hooks commonly get ``Stats`` via ``airflow.sdk.observability.stats`` - (or the deprecated ``airflow.stats`` shim, which re-exports the same module). That module - is a separate singleton from the one this process initializes for its own internal use, so - it must be initialized here too, or plugin code silently falls back to ``NoStatsLogger`` in - any long-running component other than the task runner. - """ - from airflow.sdk.observability import stats as sdk_stats # noqa: SDK001 - - sdk_stats.initialize( - factory=get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) diff --git a/airflow-core/tests/unit/api_fastapi/test_app.py b/airflow-core/tests/unit/api_fastapi/test_app.py index e6ac8cb911fc5..27ce822eac665 100644 --- a/airflow-core/tests/unit/api_fastapi/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/test_app.py @@ -24,7 +24,6 @@ import airflow.api_fastapi.app as app_module import airflow.plugins_manager as plugins_manager -import airflow.sdk.observability.stats pytestmark = pytest.mark.db_test @@ -181,7 +180,6 @@ def test_initializes_api_server_stats_with_factory(self): sentinel_factory = object() with ( mock.patch("airflow._shared.observability.metrics.stats") as mock_stats, - mock.patch.object(airflow.sdk.observability.stats, "initialize") as mock_sdk_initialize, mock.patch( "airflow.observability.metrics.stats_utils.get_stats_factory", return_value=sentinel_factory, @@ -189,19 +187,12 @@ def test_initializes_api_server_stats_with_factory(self): ): app_module._initialize_api_server_stats() - assert mock_get_factory.call_count == 2 + assert mock_get_factory.call_count == 1 mock_stats.initialize.assert_called_once() _, kwargs = mock_stats.initialize.call_args assert kwargs["factory"] is sentinel_factory assert isinstance(kwargs["export_legacy_names"], bool) - # Plugins/listeners get Stats via the task-sdk-backed singleton, a separate one - # from the API server's own — it must be initialized too, see initialize_sdk_stats_backend. - mock_sdk_initialize.assert_called_once() - _, sdk_kwargs = mock_sdk_initialize.call_args - assert sdk_kwargs["factory"] is sentinel_factory - assert isinstance(sdk_kwargs["export_legacy_names"], bool) - def test_stats_failure_does_not_block_startup(self): """A metrics misconfiguration must not prevent the API server from starting.""" with ( diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index faf807e8738bf..81e6f0f94db0d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2592,11 +2592,8 @@ def test_create_process_passes_bundle_name_to_process_start( call_kwargs = mock_process_start.call_args.kwargs assert call_kwargs["bundle_name"] == "testing" - @mock.patch("airflow.dag_processing.manager.stats_utils.initialize_sdk_stats_backend") @mock.patch("airflow.dag_processing.manager.stats.initialize") - def test_stats_initialize_called_on_run( - self, stats_init_mock, sdk_stats_init_mock, tmp_path, configure_testing_dag_bundle - ): + def test_stats_initialize_called_on_run(self, stats_init_mock, tmp_path, configure_testing_dag_bundle): """Test that stats.initialize() is called when DagFileProcessorManager.run() is executed.""" with configure_testing_dag_bundle(tmp_path): manager = DagFileProcessorManager(max_runs=1) @@ -2607,8 +2604,14 @@ def test_stats_initialize_called_on_run( call_kwargs = stats_init_mock.call_args.kwargs assert "factory" in call_kwargs - # The task-sdk Stats singleton (what plugins/listeners use) must be initialized too. - sdk_stats_init_mock.assert_called_once() + @mock.patch("airflow.dag_processing.manager.stats") + def test_stats_init_failure_does_not_block_run(self, mock_stats, tmp_path, configure_testing_dag_bundle): + """A metrics misconfiguration must not prevent the Dag processor from running.""" + mock_stats.initialize.side_effect = RuntimeError("boom") + + with configure_testing_dag_bundle(tmp_path): + manager = DagFileProcessorManager(max_runs=1) + manager.run() def test_run_invokes_before_and_after_hooks(self, tmp_path, configure_testing_dag_bundle): """`run()` should call `before_run` then `after_run`, even if the loop raises.""" diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 96a355a6f74bd..27dd4e74a1c3b 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -72,12 +72,12 @@ def test_invalid_slotspool(): BaseExecutor(0) -@mock.patch("airflow.executors.base_executor.stats_utils.initialize_sdk_stats_backend") -def test_init_initializes_sdk_stats_backend(sdk_stats_init_mock): - """The task-sdk Stats singleton (what plugins/listeners use) must be initialized too.""" - BaseExecutor() +@mock.patch("airflow.executors.base_executor.stats") +def test_init_stats_failure_does_not_block_construction(mock_stats): + """A metrics misconfiguration must not prevent the executor from being constructed.""" + mock_stats.initialize.side_effect = RuntimeError("boom") - sdk_stats_init_mock.assert_called_once() + BaseExecutor() def test_get_task_log(): diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index e12d73106f42b..76a948c02daec 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9472,18 +9472,15 @@ def test_process_expired_deadlines(self, mock_handle_miss, session, dag_maker): # Assert that all deadlines which are both expired and unhandled get processed. assert mock_handle_miss.call_count == 2 - @mock.patch("airflow.jobs.scheduler_job_runner.stats_utils.initialize_sdk_stats_backend") - def test_execute_initializes_sdk_stats_backend(self, sdk_stats_init_mock, session): - """The task-sdk Stats singleton (what plugins/listeners use) must be initialized too.""" + @mock.patch("airflow.jobs.scheduler_job_runner.stats") + def test_execute_stats_init_failure_does_not_block_run(self, mock_stats, session): + """A metrics misconfiguration must not prevent the scheduler from running.""" + mock_stats.initialize.side_effect = RuntimeError("boom") scheduler_job = Job() self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1, executors=[MockExecutor()]) self.job_runner._execute() - # BaseExecutor.__init__ also calls this (same shared stats_utils module), so - # constructing MockExecutor() above contributes a call in addition to _execute()'s own. - assert sdk_stats_init_mock.called - @mock.patch("airflow.models.Deadline.handle_miss") def test_process_expired_deadlines_no_deadlines_found(self, mock_handle_miss, session): """Test handling when there are no deadlines to process.""" diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index a1dbb4e4fe484..398c3477e920a 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2482,12 +2482,9 @@ def test_update_triggers_skips_when_ti_has_no_dag_version(session, supervisor_bu class TestTriggererJobRunner: - @patch("airflow.jobs.triggerer_job_runner.stats_utils.initialize_sdk_stats_backend") @patch("airflow.jobs.triggerer_job_runner.stats.initialize") @patch.object(TriggerRunnerSupervisor, "start") - def test_stats_initialize_called_on_execute( - self, mock_supervisor_start, stats_init_mock, sdk_stats_init_mock, session - ): + def test_stats_initialize_called_on_execute(self, mock_supervisor_start, stats_init_mock, session): """Test that stats.initialize() is called when TriggererJobRunner._execute() is executed.""" # Setup mock supervisor to immediately stop mock_supervisor = MagicMock() @@ -2517,8 +2514,32 @@ def test_stats_initialize_called_on_execute( call_kwargs = stats_init_mock.call_args.kwargs assert "factory" in call_kwargs - # The task-sdk Stats singleton (what plugins/listeners use) must be initialized too. - sdk_stats_init_mock.assert_called_once() + @patch("airflow.jobs.triggerer_job_runner.stats") + @patch.object(TriggerRunnerSupervisor, "start") + def test_stats_init_failure_does_not_block_execute(self, mock_supervisor_start, mock_stats, session): + """A metrics misconfiguration must not prevent the triggerer from running.""" + mock_stats.initialize.side_effect = RuntimeError("boom") + + mock_supervisor = MagicMock() + mock_supervisor.stop = False + mock_supervisor._exit_code = None + mock_supervisor.is_alive.return_value = True + mock_supervisor.run.side_effect = lambda: setattr(mock_supervisor, "stop", True) + mock_supervisor_start.return_value = mock_supervisor + + job = Job() + session.add(job) + session.flush() + + job_runner = TriggererJobRunner(job) + job_runner.trigger_runner = mock_supervisor + mock_supervisor.stop = True # Stop immediately + + with patch.object(job_runner, "register_signals"): + with contextlib.suppress(Exception): + job_runner._execute() + + mock_stats.initialize.assert_called_once() @patch.object(TriggerRunnerSupervisor, "start") def test_execute_sets_server_process_context(self, mock_supervisor_start, session, monkeypatch): diff --git a/airflow-core/tests/unit/observability/metrics/test_stats.py b/airflow-core/tests/unit/observability/metrics/test_stats.py index de71a9bca1244..d1488354825a3 100644 --- a/airflow-core/tests/unit/observability/metrics/test_stats.py +++ b/airflow-core/tests/unit/observability/metrics/test_stats.py @@ -18,7 +18,6 @@ import importlib import re -from unittest import mock import pytest @@ -105,27 +104,24 @@ def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self importlib.reload(airflow.observability.stats) -class TestInitializeSdkStatsBackend: +class TestInitializePropagatesToSdkSingleton: """ Plugins/listeners typically get ``Stats`` via ``airflow.sdk.observability.stats`` (or the - deprecated ``airflow.stats`` shim, which re-exports it), a singleton separate from the one - this module initializes for internal use. ``initialize_sdk_stats_backend`` keeps that - singleton configured too, so plugin code works outside the task runner as well. + deprecated ``airflow.stats`` shim, which re-exports it) — a separate singleton from the one + this process initializes for its own internal use, because the two are loaded from the same + symlinked source under different module names. The shared ``initialize()`` implementation + propagates the configuration to that sibling singleton, so plugin code gets a working + backend outside the task runner too. """ def test_configures_sdk_stats_singleton_with_this_process_backend(self): with conf_vars({("metrics", "statsd_on"): "True"}): importlib.reload(airflow.sdk.observability.stats) - stats_utils.initialize_sdk_stats_backend() + airflow.observability.stats.initialize( + factory=stats_utils.get_stats_factory(), + export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), + ) backend = airflow.sdk._shared.observability.metrics.stats._get_backend() assert hasattr(backend, "statsd") + importlib.reload(airflow.observability.stats) importlib.reload(airflow.sdk.observability.stats) - - def test_uses_same_factory_and_legacy_names_flag_as_this_process(self): - with mock.patch("airflow.sdk.observability.stats.initialize") as mock_sdk_initialize: - stats_utils.initialize_sdk_stats_backend() - - mock_sdk_initialize.assert_called_once_with( - factory=stats_utils.get_stats_factory(), - export_legacy_names=conf.getboolean("metrics", "legacy_names_on"), - ) diff --git a/shared/observability/src/airflow_shared/observability/metrics/stats.py b/shared/observability/src/airflow_shared/observability/metrics/stats.py index 7b51e580cd036..8acad3f97ebb9 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/stats.py +++ b/shared/observability/src/airflow_shared/observability/metrics/stats.py @@ -20,6 +20,7 @@ import os import re import socket +import sys from collections.abc import Callable, Iterable from typing import TYPE_CHECKING, Any @@ -105,6 +106,37 @@ def initialize( _factory = factory _backend = None _export_legacy_names = export_legacy_names + _propagate_to_sibling_modules(factory=factory, export_legacy_names=export_legacy_names) + + +def _propagate_to_sibling_modules( + *, + factory: Callable[[], StatsLogger | NoStatsLogger], + export_legacy_names: bool, +) -> None: + """ + Apply the same configuration to other loaded copies of this module. + + This source file is symlinked into multiple distributions (e.g. ``airflow-core`` and + ``task-sdk``), each importing it under a different module name (``airflow._shared...`` vs + ``airflow.sdk._shared...``). Python treats each as a distinct module object with its own + module-level globals, so a process that has both loaded (e.g. the scheduler, which also runs + plugin/listener code that reaches ``Stats`` through the task-sdk path) would otherwise end up + with one singleton initialized and the other silently defaulting to ``NoStatsLogger``. + ``__file__`` differs per symlinked path, so siblings are identified by resolved real path + instead. Attributes are set via ``setattr`` (mypy has no visibility into another module's + globals) rather than by calling the sibling's own ``initialize()``, to avoid re-entering + this function. + """ + this_file = os.path.realpath(__file__) + for name, module in list(sys.modules.items()): + if name == __name__ or module is None: + continue + if os.path.realpath(getattr(module, "__file__", "") or "") != this_file: + continue + setattr(module, "_factory", factory) + setattr(module, "_backend", None) + setattr(module, "_export_legacy_names", export_legacy_names) def _get_backend() -> StatsLogger | NoStatsLogger: diff --git a/shared/observability/tests/observability/metrics/test_stats.py b/shared/observability/tests/observability/metrics/test_stats.py index 9bdb4c473121c..c6a163d87513c 100644 --- a/shared/observability/tests/observability/metrics/test_stats.py +++ b/shared/observability/tests/observability/metrics/test_stats.py @@ -19,7 +19,9 @@ import importlib import logging +import sys import time +import types from collections.abc import Callable from unittest import mock from unittest.mock import Mock @@ -826,3 +828,46 @@ def test_build_dag_metric_tags(tag_names: list[str], expected: dict[str, str]) - def test_build_dag_metric_tags_accepts_generator() -> None: assert build_dag_metric_tags(name for name in ["env:prod"]) == {"env": "prod"} + + +class TestInitializePropagatesToSiblingModules: + """ + This source file is symlinked into both ``airflow-core`` and ``task-sdk``, each importing it + under a different module name. ``initialize()`` must configure any such sibling copy already + loaded in the process too, or the two singletons drift out of sync. + """ + + def setup_method(self): + self.sibling = types.ModuleType("fake_distribution._shared.observability.metrics.stats") + self.sibling.__file__ = airflow_shared.observability.metrics.stats.__file__ + self.sibling._factory = None + self.sibling._backend = "stale" + self.sibling._export_legacy_names = True + sys.modules[self.sibling.__name__] = self.sibling + + def teardown_method(self): + del sys.modules[self.sibling.__name__] + importlib.reload(airflow_shared.observability.metrics.stats) + + def test_initialize_configures_sibling_module(self): + factory = get_statsd_logger_factory(stats_class=statsd.StatsClient) + + airflow_shared.observability.metrics.stats.initialize(factory=factory, export_legacy_names=False) + + assert self.sibling._factory is factory + assert self.sibling._backend is None + assert self.sibling._export_legacy_names is False + + def test_initialize_does_not_touch_unrelated_modules(self): + unrelated = types.ModuleType("unrelated_module") + unrelated.__file__ = "/some/other/path/stats.py" + unrelated._factory = "untouched" + sys.modules["unrelated_module"] = unrelated + try: + airflow_shared.observability.metrics.stats.initialize( + factory=get_statsd_logger_factory(stats_class=statsd.StatsClient), + export_legacy_names=True, + ) + assert unrelated._factory == "untouched" + finally: + del sys.modules["unrelated_module"]