From 78b2668809c0d4ad40242a961fff8b35858d3d5c Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Mon, 15 Jun 2026 16:05:42 -0700 Subject: [PATCH 1/5] Add team_name tag to dag processor metrics for multi-team deployments --- .../src/airflow/dag_processing/manager.py | 69 ++++++- .../src/airflow/dag_processing/processor.py | 16 +- .../tests/unit/dag_processing/test_manager.py | 170 ++++++++++++++++++ .../unit/dag_processing/test_processor.py | 77 ++++++++ 4 files changed, 322 insertions(+), 10 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 92ca10c67fa68..237cb8d87d0c0 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -69,6 +69,7 @@ from airflow.sdk.log import init_log_file, logging_processors from airflow.typing_compat import assert_never from airflow.utils.file import list_py_file_paths, might_contain_dag +from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname from airflow.utils.process_utils import ( @@ -1020,6 +1021,14 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]): utcnow = timezone.utcnow() now = time.monotonic() + if conf.getboolean("core", "multi_team"): + bundle_names = {bundle_name for bundle_name in known_files} + bundle_to_team = { + bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names + } + else: + bundle_to_team = {} + for files in known_files.values(): for file in files: stat = self._file_stats[file] @@ -1040,11 +1049,14 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]): stats.gauge( "dag_processing.last_run.seconds_ago", seconds_ago, - tags={ - "file_path": file.normalized_file_path_for_stats, - "bundle_name": normalize_name_for_stats(file.bundle_name), - "file_name": file_name, - }, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "bundle_name": normalize_name_for_stats(file.bundle_name), + "file_name": file_name, + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) rows.append( @@ -1183,6 +1195,9 @@ def handle_parsing_result( run_duration = time.monotonic() - proc.start_time finish_time = timezone.utcnow() + team_name = ( + DagBundleModel.get_team_name(file.bundle_name) if conf.getboolean("core", "multi_team") else None + ) next_stat = process_parse_results( run_duration=run_duration, finish_time=finish_time, @@ -1191,6 +1206,7 @@ def handle_parsing_result( parsing_result=proc.parsing_result, is_callback_only=is_callback_only, relative_fileloc=str(file.rel_path), + team_name=team_name, ) if proc.parsing_result is not None: @@ -1358,6 +1374,14 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" + if conf.getboolean("core", "multi_team"): + bundle_names = {file.bundle_name for file in self._file_queue} + bundle_to_team = { + bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names + } + else: + bundle_to_team = {} + while self._parallelism > len(self._processors) and self._file_queue: file, _ = self._file_queue.popitem(last=False) # Stop creating duplicate processor i.e. processor with the same filepath @@ -1367,7 +1391,13 @@ def _start_new_processes(self): processor = self._create_process(file) stats.incr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "start"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "start", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) self._processors[file] = processor @@ -1533,6 +1563,15 @@ def _kill_timed_out_processors(self): """Kill any file processors that timeout to defend against process hangs.""" now = time.monotonic() processors_to_remove = [] + + if conf.getboolean("core", "multi_team"): + bundle_names = {file.bundle_name for file in self._processors} + bundle_to_team = { + bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names + } + else: + bundle_to_team = {} + for file, processor in self._processors.items(): duration = now - processor.start_time if duration > self.processor_timeout: @@ -1544,8 +1583,17 @@ def _kill_timed_out_processors(self): self.processor_timeout, ) file_path_tag = file.normalized_file_path_for_stats - stats.decr("dag_processing.processes", tags={"file_path": file_path_tag, "action": "timeout"}) - stats.incr("dag_processing.processor_timeouts", tags={"file_path": file_path_tag}) + team_name = bundle_to_team.get(file.bundle_name) + stats.decr( + "dag_processing.processes", + tags=prune_dict( + {"file_path": file_path_tag, "action": "timeout", "team_name": team_name} + ), + ) + stats.incr( + "dag_processing.processor_timeouts", + tags=prune_dict({"file_path": file_path_tag, "team_name": team_name}), + ) processor.kill(signal.SIGKILL) processors_to_remove.append(file) @@ -1642,6 +1690,7 @@ def process_parse_results( *, is_callback_only: bool = False, relative_fileloc: str | None = None, + team_name: str | None = None, ) -> DagFileStat: """ Create a DagFileStat from parsing results and emit metrics. @@ -1673,7 +1722,9 @@ def process_parse_results( stats.timing( "dag_processing.last_duration", stat.last_duration, - tags={"bundle_name": normalized_bundle, "file_name": file_name}, + tags=prune_dict( + {"bundle_name": normalized_bundle, "file_name": file_name, "team_name": team_name} + ), ) if parsing_result is None: diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index f7b3affd00830..251b7b629747d 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -38,6 +38,7 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.base import BundleVersionLock from airflow.dag_processing.dagbag import BundleDagBag, DagBag +from airflow.models.dag import DagModel from airflow.sdk.exceptions import TaskNotFound from airflow.sdk.execution_time.comms import ( ConnectionResult, @@ -90,6 +91,7 @@ from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.dag_version_inflation_checker import check_dag_file_stability from airflow.utils.file import iter_airflow_imports +from airflow.utils.helpers import prune_dict from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: @@ -386,7 +388,19 @@ def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: Fil callback(context) except Exception: log.exception("Callback failed", dag_id=request.dag_id) - stats.incr("dag.callback_exceptions", tags={"dag_id": request.dag_id}) + stats.incr( + "dag.callback_exceptions", + tags=prune_dict( + { + "dag_id": request.dag_id, + "team_name": ( + DagModel.get_team_name(request.dag_id) + if conf.getboolean("core", "multi_team") + else None + ), + } + ), + ) def _execute_task_callbacks(dagbag: DagBag, request: TaskCallbackRequest, log: FilteringBoundLogger) -> None: diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 0f6c061fc43a4..d9b55f8c9db04 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2922,3 +2922,173 @@ def test_bundle_version_data_stored_after_refresh(self, session): assert manager._bundle_versions["mock_bundle"] == "newhash" assert manager._bundle_version_data["mock_bundle"] == test_data + + +class TestMultiTeamMetrics: + """Tests for team_name tag on dag processing metrics in multi-team mode.""" + + def mock_processor(self, start_time: float | None = None) -> DagFileProcessorProcess: + proc = MagicMock() + proc.wait.return_value = 0 + read_end, write_end = socketpair() + ret = DagFileProcessorProcess( + process_log=MagicMock(), + id=uuid7(), + pid=1234, + process=proc, + stdin=write_end, + logger_filehandle=MagicMock(), + client=MagicMock(), + bundle_name="testing", + dag_file_rel_path="test_dag.py", + ) + if start_time: + ret.start_time = start_time + ret._open_sockets.clear() + return ret + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value="team_alpha") + @mock.patch("airflow.dag_processing.manager.stats.gauge") + def test_log_file_processing_stats_includes_team_name(self, mock_gauge, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_stats[dag_file] = DagFileStat( + last_finish_time=timezone.utcnow() - timedelta(seconds=5), + ) + + manager._log_file_processing_stats({"testing": {dag_file}}) + + mock_gauge.assert_any_call( + "dag_processing.last_run.seconds_ago", + mock.ANY, + tags={ + "file_path": "dag_file.py", + "bundle_name": "testing", + "file_name": "dag_file", + "team_name": "team_alpha", + }, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name") + @mock.patch("airflow.dag_processing.manager.stats.gauge") + def test_log_file_processing_stats_omits_team_name_when_not_multi_team( + self, mock_gauge, mock_get_team_name + ): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_stats[dag_file] = DagFileStat( + last_finish_time=timezone.utcnow() - timedelta(seconds=5), + ) + + manager._log_file_processing_stats({"testing": {dag_file}}) + + mock_get_team_name.assert_not_called() + mock_gauge.assert_any_call( + "dag_processing.last_run.seconds_ago", + mock.ANY, + tags={ + "file_path": "dag_file.py", + "bundle_name": "testing", + "file_name": "dag_file", + }, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value="team_alpha") + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_start_new_processes_includes_team_name(self, mock_incr, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_queue[dag_file] = None + manager._create_process = MagicMock(return_value=self.mock_processor()) + + manager._start_new_processes() + + mock_incr.assert_any_call( + "dag_processing.processes", + tags={"file_path": "dag_file.py", "action": "start", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value="team_alpha") + @mock.patch("airflow.dag_processing.manager.stats.incr") + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_kill_timed_out_processors_includes_team_name(self, mock_decr, mock_incr, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1, processor_timeout=5) + start_time = time.monotonic() - manager.processor_timeout - 1 + processor = self.mock_processor(start_time=start_time) + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + manager._processors = {dag_file: processor} + + with mock.patch.object(type(processor), "kill"): + manager._kill_timed_out_processors() + + mock_decr.assert_called_once_with( + "dag_processing.processes", + tags={"file_path": "dag_file.py", "action": "timeout", "team_name": "team_alpha"}, + ) + mock_incr.assert_any_call( + "dag_processing.processor_timeouts", + tags={"file_path": "dag_file.py", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value="team_alpha") + @mock.patch("airflow.dag_processing.manager.stats.timing") + def test_process_parse_results_includes_team_name(self, mock_timing, mock_get_team_name): + from airflow.dag_processing.manager import process_parse_results + + result = DagFileParsingResult(fileloc="/tmp/dag.py", serialized_dags=[]) + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=result, + relative_fileloc="dag.py", + team_name="team_alpha", + ) + + mock_timing.assert_called_once_with( + "dag_processing.last_duration", + 1.5, + tags={"bundle_name": "testing", "file_name": "dag", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @mock.patch("airflow.dag_processing.manager.stats.timing") + def test_process_parse_results_omits_team_name_when_none(self, mock_timing): + from airflow.dag_processing.manager import process_parse_results + + result = DagFileParsingResult(fileloc="/tmp/dag.py", serialized_dags=[]) + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=result, + relative_fileloc="dag.py", + team_name=None, + ) + + mock_timing.assert_called_once_with( + "dag_processing.last_duration", + 1.5, + tags={"bundle_name": "testing", "file_name": "dag"}, + ) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 2e3d6940cc7eb..599c36d47f695 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2149,3 +2149,80 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): "value": "super-secret-value", "type": "VariableResult", } + + +class TestMultiTeamCallbackMetrics: + """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" + + @conf_vars({("core", "multi_team"): "true"}) + @patch("airflow.dag_processing.processor.DagModel.get_team_name", return_value="team_alpha") + @patch("airflow.dag_processing.processor.stats.incr") + def test_callback_exception_includes_team_name(self, mock_incr, mock_get_team_name, spy_agency): + def failing_callback(context): + raise RuntimeError("boom") + + with DAG(dag_id="test_dag", on_failure_callback=failing_callback) as dag: + BaseOperator(task_id="test_task") + + @spy_agency.spy_for(DagBag.collect_dags, owner=DagBag) + def fake_collect_dags(self, *args, **kwargs): + self.dags[dag.dag_id] = dag + + dagbag = DagBag() + dagbag.collect_dags() + + request = DagCallbackRequest( + filepath="test.py", + dag_id="test_dag", + run_id="test_run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + msg="Test failure", + ) + + log = structlog.get_logger() + _execute_dag_callbacks(dagbag, request, log) + + mock_incr.assert_called_once_with( + "dag.callback_exceptions", + tags={"dag_id": "test_dag", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @patch("airflow.dag_processing.processor.DagModel.get_team_name") + @patch("airflow.dag_processing.processor.stats.incr") + def test_callback_exception_omits_team_name_when_not_multi_team( + self, mock_incr, mock_get_team_name, spy_agency + ): + def failing_callback(context): + raise RuntimeError("boom") + + with DAG(dag_id="test_dag", on_failure_callback=failing_callback) as dag: + BaseOperator(task_id="test_task") + + @spy_agency.spy_for(DagBag.collect_dags, owner=DagBag) + def fake_collect_dags(self, *args, **kwargs): + self.dags[dag.dag_id] = dag + + dagbag = DagBag() + dagbag.collect_dags() + + request = DagCallbackRequest( + filepath="test.py", + dag_id="test_dag", + run_id="test_run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + msg="Test failure", + ) + + log = structlog.get_logger() + _execute_dag_callbacks(dagbag, request, log) + + mock_get_team_name.assert_not_called() + mock_incr.assert_called_once_with( + "dag.callback_exceptions", + tags={"dag_id": "test_dag"}, + ) From dea0d1d3bce9d4154acb68373ea0c196b1c59da1 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Mon, 15 Jun 2026 16:50:28 -0700 Subject: [PATCH 2/5] missed one --- airflow-core/src/airflow/models/dagrun.py | 15 ++++++++- airflow-core/tests/unit/models/test_dagrun.py | 31 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 80e279f149936..ab024142918f7 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1420,6 +1420,7 @@ def execute_dag_callbacks( TaskInstance as TIDataModel, TIRunContext, ) + from airflow.models.dag import DagModel from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance if relevant_ti: @@ -1463,7 +1464,19 @@ def execute_dag_callbacks( callback(context) except Exception: self.log.exception("Callback failed for %s", dag.dag_id) - stats.incr("dag.callback_exceptions", tags={"dag_id": dag.dag_id}) + stats.incr( + "dag.callback_exceptions", + tags=prune_dict( + { + "dag_id": dag.dag_id, + "team_name": ( + DagModel.get_team_name(dag.dag_id) + if airflow_conf.getboolean("core", "multi_team") + else None + ), + } + ), + ) def _get_ready_tis( self, diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index a8d384d948099..9905dd2638652 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -4118,6 +4118,37 @@ def on_failure(context): assert "ti" not in context_received assert context_received["run_id"] == dr.run_id + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + "true", "team_alpha", {"dag_id": "test_dag", "team_name": "team_alpha"}, id="with_team" + ), + pytest.param("false", None, {"dag_id": "test_dag"}, id="without_team"), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats.incr") + def test_callback_exception_team_name_tag( + self, mock_incr, multi_team, team_name, expected_tags, dag_maker, session + ): + def failing_callback(context): + raise RuntimeError("boom") + + with dag_maker("test_dag", session=session, on_failure_callback=failing_callback) as dag: + BashOperator(task_id="test_task", bash_command="echo 1") + + dr = dag_maker.create_dagrun() + dag.on_failure_callback = failing_callback + dag.has_on_failure_callback = True + + with ( + conf_vars({("core", "multi_team"): multi_team}), + mock.patch("airflow.models.dagrun.DagModel.get_team_name", return_value=team_name), + ): + dr.execute_dag_callbacks(dag, success=False) + + mock_incr.assert_any_call("dag.callback_exceptions", tags=expected_tags) + class TestDagRunTracing: """Tests for DagRun OpenTelemetry span behavior.""" From 1d8669fca555aa43c0d4a6c5b74e9e5b0569a5bc Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 16 Jun 2026 16:03:21 -0700 Subject: [PATCH 3/5] fix mock patch target --- airflow-core/tests/unit/models/test_dagrun.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 9905dd2638652..a55b93685d9ee 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -4143,7 +4143,7 @@ def failing_callback(context): with ( conf_vars({("core", "multi_team"): multi_team}), - mock.patch("airflow.models.dagrun.DagModel.get_team_name", return_value=team_name), + mock.patch("airflow.models.dag.DagModel.get_team_name", return_value=team_name), ): dr.execute_dag_callbacks(dag, success=False) From d734f8fb1c03f5d32fdc49c5d60aceb8c396ef49 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 16 Jun 2026 16:20:20 -0700 Subject: [PATCH 4/5] Add team_name to dag_processing.processes stop and terminate --- .../src/airflow/dag_processing/manager.py | 33 ++++++++- .../tests/unit/dag_processing/test_manager.py | 73 +++++++++++++++++++ 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 237cb8d87d0c0..d2f0a0cad9a4a 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1148,6 +1148,15 @@ def remove_orphaned_file_stats(self, present: set[DagFileInfo]): def terminate_orphan_processes(self, present: set[DagFileInfo]): """Stop processors that are working on deleted files.""" present_keys = {file.presence_key for file in present} + + if conf.getboolean("core", "multi_team"): + bundle_names = {file.bundle_name for file in self._processors} + bundle_to_team = { + bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names + } + else: + bundle_to_team = {} + for file in list(self._processors.keys()): if file.presence_key not in present_keys: processor = self._processors.pop(file, None) @@ -1157,7 +1166,13 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]): self.log.warning("Stopping processor for %s", file_name) stats.decr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "stop"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "stop", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) processor.kill(signal.SIGKILL) processor.logger_filehandle.close() @@ -1654,10 +1669,24 @@ def max_runs_reached(self): def terminate(self): """Stop all running processors.""" + if conf.getboolean("core", "multi_team"): + bundle_names = {file.bundle_name for file in self._processors} + bundle_to_team = { + bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names + } + else: + bundle_to_team = {} + for file, processor in self._processors.items(): stats.decr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "terminate"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "terminate", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) # SIGTERM, wait 5s, SIGKILL if still alive processor.kill(signal.SIGTERM, escalation_delay=5.0) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d9b55f8c9db04..e97ac652b59e2 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -3092,3 +3092,76 @@ def test_process_parse_results_omits_team_name_when_none(self, mock_timing): 1.5, tags={"bundle_name": "testing", "file_name": "dag"}, ) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + "true", + "team_alpha", + {"file_path": "dag_file.py", "action": "stop", "team_name": "team_alpha"}, + id="with_team", + ), + pytest.param( + "false", + None, + {"file_path": "dag_file.py", "action": "stop"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_terminate_orphan_processes_includes_team_name( + self, mock_decr, multi_team, team_name, expected_tags + ): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + processor = self.mock_processor() + manager._processors = {dag_file: processor} + + with ( + conf_vars({("core", "multi_team"): multi_team}), + mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value=team_name), + mock.patch.object(type(processor), "kill"), + ): + # Empty "present" set means the file is orphaned, so its processor is stopped. + manager.terminate_orphan_processes(present=set()) + + mock_decr.assert_called_once_with("dag_processing.processes", tags=expected_tags) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + "true", + "team_alpha", + {"file_path": "dag_file.py", "action": "terminate", "team_name": "team_alpha"}, + id="with_team", + ), + pytest.param( + "false", + None, + {"file_path": "dag_file.py", "action": "terminate"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_terminate_includes_team_name(self, mock_decr, multi_team, team_name, expected_tags): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + processor = self.mock_processor() + manager._processors = {dag_file: processor} + + with ( + conf_vars({("core", "multi_team"): multi_team}), + mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value=team_name), + mock.patch.object(type(processor), "kill"), + ): + manager.terminate() + + mock_decr.assert_called_once_with("dag_processing.processes", tags=expected_tags) From 169669e4c56e85529843b27d11f199dfc0d6c9c8 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 16 Jun 2026 20:34:19 -0700 Subject: [PATCH 5/5] Add team_name to dag_processing callback count metrics --- .../src/airflow/dag_processing/manager.py | 9 +++- .../tests/unit/dag_processing/test_manager.py | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index d2f0a0cad9a4a..e43920f2618bc 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -720,7 +720,12 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: ) self._callback_to_execute[file_info].append(request) self._add_files_to_queue([file_info], mode="front") - stats.incr("dag_processing.other_callback_count") + team_name = ( + DagBundleModel.get_team_name(file_info.bundle_name) + if conf.getboolean("core", "multi_team") + else None + ) + stats.incr("dag_processing.other_callback_count", tags=prune_dict({"team_name": team_name})) @provide_session def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION) -> BundleState | None: @@ -1733,7 +1738,7 @@ def process_parse_results( last_duration=run_duration, run_count=run_count, # Don't increment for callback-only processing ) - stats.incr("dag_processing.callback_only_count") + stats.incr("dag_processing.callback_only_count", tags=prune_dict({"team_name": team_name})) else: # Actual DAG parsing or import error stat = DagFileStat( diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index e97ac652b59e2..6d10e5c1a0665 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -3165,3 +3165,52 @@ def test_terminate_includes_team_name(self, mock_decr, multi_team, team_name, ex manager.terminate() mock_decr.assert_called_once_with("dag_processing.processes", tags=expected_tags) + + @pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_process_parse_results_callback_only_count_includes_team_name( + self, mock_incr, team_name, expected_tags + ): + from airflow.dag_processing.manager import process_parse_results + + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=None, + is_callback_only=True, + relative_fileloc="dag.py", + team_name=team_name, + ) + + mock_incr.assert_called_once_with("dag_processing.callback_only_count", tags=expected_tags) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param("true", "team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param("false", None, {}, id="without_team"), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_add_callback_to_queue_includes_team_name(self, mock_incr, multi_team, team_name, expected_tags): + manager = DagFileProcessorManager(max_runs=1) + request = MagicMock(filepath="test_dag.py", bundle_name="testing", bundle_version=None) + bundle = MagicMock(path=TEST_DAGS_FOLDER) + + with ( + conf_vars({("core", "multi_team"): multi_team}), + mock.patch.object(manager, "prepare_callback_bundle", return_value=bundle), + mock.patch.object(manager, "_add_files_to_queue"), + mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name", return_value=team_name), + ): + manager._add_callback_to_queue(request) + + mock_incr.assert_called_once_with("dag_processing.other_callback_count", tags=expected_tags)