diff --git a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py index 9ac25874c3c92..5c582781a45f0 100644 --- a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py +++ b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py @@ -35,6 +35,7 @@ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics from airflow.providers.edge3.models.types import is_callback_execute from airflow.utils.db import DBLocks, create_global_lock +from airflow.utils.helpers import prune_dict from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import TaskInstanceState @@ -340,7 +341,7 @@ def _purge_jobs(self, session: Session) -> bool: @provide_session def sync(self, *, session: Session = NEW_SESSION) -> None: """Sync will get called periodically by the heartbeat method.""" - with Stats.timer("edge_executor.sync.duration"): + with Stats.timer("edge_executor.sync.duration", tags=prune_dict({"team_name": self.team_name})): orphaned = self._update_orphaned_jobs(session) purged = self._purge_jobs(session) liveness = self._check_worker_liveness(session) diff --git a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py index 3554689b794d2..8975eb7ad9443 100644 --- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py +++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py @@ -120,6 +120,30 @@ def test_sync_orphaned_tasks(self, mock_stats_incr): assert jobs[0].task_id == "started_running_orphaned" assert jobs[0].state == TaskInstanceState.REMOVED + @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team_name is only available in Airflow 3.2+") + @pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + ("team_a", {"team_name": "team_a"}), + (None, {}), + ], + ) + @patch(f"{Stats.__module__}.Stats.timer") + def test_sync_duration_metric_tags_team_name(self, mock_stats_timer, team_name, expected_tags): + executor = EdgeExecutor(team_name=team_name) + + with ( + mock.patch.object(executor, "_update_orphaned_jobs", return_value=False), + mock.patch.object(executor, "_purge_jobs", return_value=False), + mock.patch.object(executor, "_check_worker_liveness", return_value=False), + ): + executor.sync(session=MagicMock()) + + mock_stats_timer.assert_called_once_with( + "edge_executor.sync.duration", + tags=expected_tags, + ) + @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.running_state") @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.success") @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.fail")