Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions airflow-core/src/airflow/dag_processing/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,12 @@ def run(self):
# Related: https://github.com/apache/airflow/pull/57459
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"

Stats.initialize(
is_statsd_datadog_enabled=conf.getboolean("metrics", "statsd_datadog_enabled"),
is_statsd_on=conf.getboolean("metrics", "statsd_on"),
is_otel_on=conf.getboolean("metrics", "otel_on"),
)

self.register_exit_signals()

self.log.info("Processing files using up to %s processes at a time ", self._parallelism)
Expand Down
5 changes: 5 additions & 0 deletions airflow-core/src/airflow/jobs/triggerer_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,11 @@ def _exit_gracefully(self, signum, frame) -> None:
def _execute(self) -> int | None:
self.log.info("Starting the triggerer")
self.register_signals()
Stats.initialize(
is_statsd_datadog_enabled=conf.getboolean("metrics", "statsd_datadog_enabled"),
is_statsd_on=conf.getboolean("metrics", "statsd_on"),
is_otel_on=conf.getboolean("metrics", "otel_on"),
)
try:
# Kick off runner sub-process without DB access
self.trigger_runner = TriggerRunnerSupervisor.start(
Expand Down
14 changes: 14 additions & 0 deletions airflow-core/tests/unit/dag_processing/test_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -1474,3 +1474,17 @@ def test_create_process_passes_bundle_name_to_process_start(
mock_process_start.assert_called_once()
call_kwargs = mock_process_start.call_args.kwargs
assert call_kwargs["bundle_name"] == "testing"

@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):
"""Test that Stats.initialize() is called when DagFileProcessorManager.run() is executed."""
with configure_testing_dag_bundle(tmp_path):
manager = DagFileProcessorManager(max_runs=1)
manager.run()

# Verify Stats.initialize was called with the expected configuration parameters
stats_init_mock.assert_called_once()
call_kwargs = stats_init_mock.call_args.kwargs
assert "is_statsd_datadog_enabled" in call_kwargs
assert "is_statsd_on" in call_kwargs
assert "is_otel_on" in call_kwargs
37 changes: 37 additions & 0 deletions airflow-core/tests/unit/jobs/test_triggerer_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -1201,6 +1201,43 @@ def test_update_triggers_skips_when_ti_has_no_dag_version(session, supervisor_bu
supervisor.stdin.write.assert_not_called()


class TestTriggererJobRunner:
@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):
"""Test that Stats.initialize() is called when TriggererJobRunner._execute() is executed."""
# Setup mock supervisor to immediately stop
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

# We don't need to run the full _execute, just verify Stats.initialize is called
# before TriggerRunnerSupervisor.start
with patch.object(job_runner, "register_signals"):
try:
job_runner._execute()
except Exception:
pass # We expect this to fail since we're mocking

# Verify Stats.initialize was called with the expected configuration parameters
stats_init_mock.assert_called_once()
call_kwargs = stats_init_mock.call_args.kwargs
assert "is_statsd_datadog_enabled" in call_kwargs
assert "is_statsd_on" in call_kwargs
assert "is_otel_on" in call_kwargs


class TestTriggererMessageTypes:
def test_message_types_in_triggerer(self):
"""
Expand Down