This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dfabf3b7193 Add team_name tagging to multi team dag metrics (#69943)
dfabf3b7193 is described below
commit dfabf3b7193c07291bd9217582d2bbc23b005797
Author: Teghveer Singh Ateliey <[email protected]>
AuthorDate: Fri Jul 24 18:08:55 2026 -0500
Add team_name tagging to multi team dag metrics (#69943)
Added team_name tag to the following:
connection_test.reaped
connection_test.success
connection_test.failed
connection_test.hook_duration
scheduler.executor_heartbeat_duration
triggers.blocked_main_thread
Added a team_name field to TestConnection class. This is to share team_name
from
the scheduler to connection test supervisor without changing the execution
API.
---
airflow-core/docs/core-concepts/multi-team.rst | 11 +-
.../src/airflow/executors/base_executor.py | 1 +
.../airflow/executors/workloads/connection_test.py | 3 +
.../src/airflow/jobs/scheduler_job_runner.py | 18 +++-
.../src/airflow/jobs/triggerer_job_runner.py | 5 +-
.../tests/unit/executors/test_base_executor.py | 30 ++++++
airflow-core/tests/unit/jobs/test_scheduler_job.py | 94 ++++++++++++++++-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 12 ++-
.../edge3/worker_api/v2-edge-generated.yaml | 5 +
.../execution_time/connection_test_supervisor.py | 13 ++-
.../test_connection_test_supervisor.py | 112 +++++++++++++++++++++
11 files changed, 291 insertions(+), 13 deletions(-)
diff --git a/airflow-core/docs/core-concepts/multi-team.rst
b/airflow-core/docs/core-concepts/multi-team.rst
index b45b0379b7f..ab5a008cd4d 100644
--- a/airflow-core/docs/core-concepts/multi-team.rst
+++ b/airflow-core/docs/core-concepts/multi-team.rst
@@ -1060,9 +1060,11 @@ Dags, and global components emit the same metrics
without a ``team_name`` tag.
The ``team_name`` tag is applied to metrics across the following components:
-- **Triggerer**: heartbeat, capacity, and trigger-outcome metrics (for
example, ``triggerer_heartbeat``,
- ``triggers.running``, ``triggers.succeeded``).
-- **Executors**: executor slot gauges (for example, ``executor.open_slots``,
``executor.queued_tasks``).
+- **Triggerer**: heartbeat, capacity, blocked-main-thread, and trigger-outcome
metrics (for example,
+ ``triggerer_heartbeat``, ``triggers.running``, ``triggers.succeeded``,
+ ``triggers.blocked_main_thread``).
+- **Executors**: executor slot gauges and scheduler-observed executor
heartbeat timing (for example,
+ ``executor.open_slots``, ``executor.queued_tasks``,
``scheduler.executor_heartbeat_duration``).
- **Scheduler**: pool slot gauges for team-scoped pools plus task- and
asset-scheduling counters (for
example, ``pool.open_slots``, ``scheduler.tasks.killed_externally``,
``asset.triggered_dagruns``).
- **Dag runs**: dag run timing and lifecycle metrics (for example,
``dagrun.duration.<state>``,
@@ -1073,6 +1075,9 @@ The ``team_name`` tag is applied to metrics across the
following components:
``dag_processing.processor_timeouts``,
``dag_processing.callback_only_count``).
- **Callbacks**: callback execution counters (``callback_success`` /
``callback_failure``, optionally
prefixed).
+- **Connection tests**: per-request worker and reaper metrics for team-owned
connection tests (for
+ example, ``connection_test.success``, ``connection_test.failed``,
``connection_test.hook_duration``,
+ ``connection_test.reaped``). Instance-wide connection-test queue gauges
remain untagged.
.. note::
diff --git a/airflow-core/src/airflow/executors/base_executor.py
b/airflow-core/src/airflow/executors/base_executor.py
index ebd850ad944..8030e9c1445 100644
--- a/airflow-core/src/airflow/executors/base_executor.py
+++ b/airflow-core/src/airflow/executors/base_executor.py
@@ -747,6 +747,7 @@ class BaseExecutor(LoggingMixin):
timeout=workload.timeout,
token=workload.token,
server=server,
+ team_name=workload.team_name,
)
raise ValueError(f"Unknown workload type: {type(workload).__name__}")
diff --git a/airflow-core/src/airflow/executors/workloads/connection_test.py
b/airflow-core/src/airflow/executors/workloads/connection_test.py
index 5f71d6327ae..d3b411af076 100644
--- a/airflow-core/src/airflow/executors/workloads/connection_test.py
+++ b/airflow-core/src/airflow/executors/workloads/connection_test.py
@@ -37,6 +37,7 @@ class TestConnection(BaseWorkloadSchema):
connection_id: str
timeout: int
queue: str | None = None
+ team_name: str | None = None
type: Literal["TestConnection"] = Field(init=False,
default="TestConnection")
@@ -70,6 +71,7 @@ class TestConnection(BaseWorkloadSchema):
connection_id: str,
timeout: int,
queue: str | None = None,
+ team_name: str | None = None,
generator: JWTGenerator | None = None,
) -> TestConnection:
return cls(
@@ -77,5 +79,6 @@ class TestConnection(BaseWorkloadSchema):
connection_id=connection_id,
timeout=timeout,
queue=queue,
+ team_name=team_name,
token=cls.generate_token(str(connection_test_id), generator),
)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index cb4c8f64552..a9efa1c03d6 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -1844,7 +1844,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
for executor in self.executors:
with stats.timer(
"scheduler.executor_heartbeat_duration",
- tags={"executor": type(executor).__name__},
+ tags=prune_dict(
+ {
+ "executor": type(executor).__name__,
+ "team_name": executor.team_name,
+ }
+ ),
):
executor.heartbeat()
@@ -4002,6 +4007,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
connection_id=ct.connection_id,
timeout=timeout,
queue=ct.queue,
+ team_name=team_name,
generator=executor.jwt_generator,
)
executor.queue_workload(workload, session=session)
@@ -4047,7 +4053,15 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
prior_state_value,
ct.team_name,
)
- stats.incr("connection_test.reaped", tags={"prior_state":
prior_state_value})
+ stats.incr(
+ "connection_test.reaped",
+ tags=prune_dict(
+ {
+ "prior_state": prior_state_value,
+ "team_name": ct.team_name if self._multi_team else
None,
+ }
+ ),
+ )
key = ConnectionTestKey(id=str(ct.id))
for executor in self.executors:
if executor.supports_connection_test:
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 759735242ab..f851e491a69 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -1600,7 +1600,10 @@ class TriggerRunner:
time_elapsed,
self.blocked_main_thread_warning_threshold,
)
- stats.incr("triggers.blocked_main_thread")
+ stats.incr(
+ "triggers.blocked_main_thread",
+ tags=prune_dict({"team_name": self.team_name}),
+ )
async def run_trigger(
self,
diff --git a/airflow-core/tests/unit/executors/test_base_executor.py
b/airflow-core/tests/unit/executors/test_base_executor.py
index 8bc7cfd73db..8858eac59f3 100644
--- a/airflow-core/tests/unit/executors/test_base_executor.py
+++ b/airflow-core/tests/unit/executors/test_base_executor.py
@@ -497,10 +497,40 @@ def
test_queue_connection_test_workload_accepted_when_supported():
connection_test_id=uuid4(),
connection_id="test_conn",
timeout=60,
+ team_name="team_a",
)
executor.queue_workload(wl, session=mock.MagicMock(spec=Session))
assert len(executor.queued_connection_tests) == 1
assert executor.queued_connection_tests[wl.key] is wl
+ assert wl.team_name == "team_a"
+
+
[email protected](
+
"airflow.sdk.execution_time.connection_test_supervisor.supervise_connection_test",
+ autospec=True,
+)
+def
test_run_workload_passes_team_name_to_connection_test_supervisor(mock_supervise):
+ """BaseExecutor.run_workload forwards TestConnection.team_name to the
supervisor."""
+ mock_supervise.return_value = 0
+ test_id = uuid4()
+ wl = workloads.TestConnection.make(
+ connection_test_id=test_id,
+ connection_id="test_conn",
+ timeout=60,
+ team_name="team_a",
+ )
+ wl.token = "test-token"
+
+ BaseExecutor.run_workload(wl, server="http://localhost:8080/execution/";)
+
+ mock_supervise.assert_called_once_with(
+ connection_test_id=test_id,
+ connection_id="test_conn",
+ timeout=60,
+ token="test-token",
+ server="http://localhost:8080/execution/";,
+ team_name="team_a",
+ )
def test_trigger_connection_tests_skipped_when_not_supported():
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 782bcf81c4a..4add965a6ab 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -1373,6 +1373,8 @@ class TestSchedulerJob:
def test_executor_heartbeat_emits_timer(self, mock_executors,
configure_testing_dag_bundle):
with configure_testing_dag_bundle(os.devnull):
+ mock_executors[0].team_name = "team_a"
+ mock_executors[1].team_name = None
scheduler_job = Job()
self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1)
with patch("airflow.jobs.scheduler_job_runner.stats.timer") as
mock_timer:
@@ -1385,7 +1387,10 @@ class TestSchedulerJob:
]
assert len(heartbeat_calls) == len(self.job_runner.executors)
for executor, timer_call in zip(self.job_runner.executors,
heartbeat_calls):
- assert timer_call.kwargs.get("tags") == {"executor":
type(executor).__name__}
+ expected_tags = {"executor": type(executor).__name__}
+ if executor.team_name:
+ expected_tags["team_name"] = executor.team_name
+ assert timer_call.kwargs.get("tags") == expected_tags
def test_executor_events_processed(self, mock_executors,
configure_testing_dag_bundle):
with configure_testing_dag_bundle(os.devnull):
@@ -12659,6 +12664,48 @@ class TestDispatchConnectionTests:
assert mock_load.call_args.kwargs["team_name"] == "team_a"
+ @pytest.mark.parametrize(
+ ("multi_team", "row_team_name", "expected_workload_team"),
+ [
+ pytest.param(True, "team_a", "team_a", id="multi_team_with_team"),
+ pytest.param(True, None, None, id="multi_team_without_team"),
+ pytest.param(False, "team_a", None,
id="single_team_ignores_row_team"),
+ ],
+ )
+ @mock.patch.dict(
+ os.environ,
+ {
+ "AIRFLOW__CONNECTION_TEST__MAX_CONCURRENCY": "4",
+ "AIRFLOW__CONNECTION_TEST__TIMEOUT": "60",
+ },
+ )
+ def test_dispatch_puts_team_name_on_workload(
+ self,
+ scheduler_job_runner_for_connection_tests,
+ session,
+ multi_team,
+ row_team_name,
+ expected_workload_team,
+ ):
+ """Queued TestConnection workloads carry team_name only in multi-team
mode."""
+ runner = scheduler_job_runner_for_connection_tests
+ runner._multi_team = multi_team
+
+ session.add(
+ ConnectionTestRequest(
+ conn_type="test_type",
+ connection_id="team_conn",
+ team_name=row_team_name,
+ )
+ )
+ session.commit()
+
+ runner._enqueue_connection_tests(session=session)
+
+ queued = list(runner.executor.queued_connection_tests.values())
+ assert len(queued) == 1
+ assert queued[0].team_name == expected_workload_team
+
@mock.patch.dict(
os.environ,
{
@@ -13038,6 +13085,51 @@ class TestReapStaleConnectionTests:
assert session.get(ConnectionTestRequest, ct_success.id).state ==
ConnectionTestState.SUCCESS
assert session.get(ConnectionTestRequest, ct_failed.id).state ==
ConnectionTestState.FAILED
+ @pytest.mark.parametrize(
+ ("multi_team", "team_name", "expected_tags"),
+ [
+ pytest.param(
+ True,
+ "team_alpha",
+ {"prior_state": "queued", "team_name": "team_alpha"},
+ id="with_team",
+ ),
+ pytest.param(True, None, {"prior_state": "queued"},
id="multi_team_no_team_on_row"),
+ pytest.param(False, "team_alpha", {"prior_state": "queued"},
id="single_team_ignores_row_team"),
+ ],
+ )
+ @mock.patch.dict(os.environ, {"AIRFLOW__CONNECTION_TEST__TIMEOUT": "60"})
+ def test_reap_emits_team_name_tag(
+ self,
+ scheduler_job_runner_for_connection_tests,
+ session,
+ multi_team,
+ team_name,
+ expected_tags,
+ ):
+ """Reaper connection_test.reaped metric includes team_name only when
multi-team is on."""
+ runner = scheduler_job_runner_for_connection_tests
+ runner._multi_team = multi_team
+ initial_time = timezone.utcnow()
+
+ with time_machine.travel(initial_time, tick=False):
+ ct = ConnectionTestRequest(
+ conn_type="test_type",
+ connection_id="reap_team_conn",
+ team_name=team_name,
+ )
+ ct.state = ConnectionTestState.QUEUED
+ session.add(ct)
+ session.commit()
+
+ with (
+ time_machine.travel(initial_time + timedelta(seconds=200),
tick=False),
+ mock.patch("airflow.jobs.scheduler_job_runner.stats.incr") as
mock_stats_incr,
+ ):
+ runner._reap_stale_connection_tests(session=session)
+
+ mock_stats_incr.assert_called_once_with("connection_test.reaped",
tags=expected_tags)
+
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index fa4d68f5fff..31a77b50b65 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -1157,10 +1157,18 @@ class TestTriggerRunner:
mock_stats_incr.assert_not_called()
@pytest.mark.asyncio
- async def test_block_watchdog_logs_when_threshold_is_exceeded(self) ->
None:
+ @pytest.mark.parametrize(
+ ("team_name", "expected_tags"),
+ [
+ pytest.param("team_a", {"team_name": "team_a"}, id="with_team"),
+ pytest.param(None, {}, id="without_team"),
+ ],
+ )
+ async def test_block_watchdog_logs_when_threshold_is_exceeded(self,
team_name, expected_tags) -> None:
with conf_vars({("triggerer",
"blocked_main_thread_warning_threshold"): "0.5"}):
trigger_runner = TriggerRunner()
+ trigger_runner.team_name = team_name
trigger_runner.log = AsyncMock()
async def fake_sleep(_):
@@ -1178,7 +1186,7 @@ class TestTriggerRunner:
assert "configured warning threshold" in log_message
assert elapsed == pytest.approx(0.6)
assert threshold == 0.5
- mock_stats_incr.assert_called_once_with("triggers.blocked_main_thread")
+
mock_stats_incr.assert_called_once_with("triggers.blocked_main_thread",
tags=expected_tags)
def test_run_inline_trigger_canceled(self, session) -> None:
trigger_runner = TriggerRunner()
diff --git
a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
index 4bd6a68af43..13cacba97be 100644
---
a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
+++
b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml
@@ -1385,6 +1385,11 @@ components:
- type: string
- type: 'null'
title: Queue
+ team_name:
+ anyOf:
+ - type: string
+ - type: 'null'
+ title: Team Name
type:
type: string
const: TestConnection
diff --git
a/task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py
index aca24556e29..06a7c3735f7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py
@@ -44,6 +44,7 @@ def supervise_connection_test(
timeout: int,
token: str,
server: str,
+ team_name: str | None = None,
) -> int:
"""Execute a connection test on the worker and report the result via the
Execution API."""
client = Client(base_url=server, token=token)
@@ -51,6 +52,7 @@ def supervise_connection_test(
bind_contextvars(connection_test_id=str(connection_test_id),
connection_id=connection_id)
log.info("Starting connection test", timeout=timeout)
start = time.monotonic()
+ tags = {"team_name": team_name} if team_name else {}
try:
r = client.connection_tests.get_connection(connection_test_id)
@@ -75,7 +77,7 @@ def supervise_connection_test(
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client"
try:
with (
- stats.timer("connection_test.hook_duration"),
+ stats.timer("connection_test.hook_duration", tags=tags),
TimeoutPosix(
seconds=timeout,
error_message=f"Connection test timed out after
{timeout}s",
@@ -95,7 +97,10 @@ def supervise_connection_test(
state = ConnectionTestState.SUCCESS if success else
ConnectionTestState.FAILED
client.connection_tests.update_state(connection_test_id, state,
message)
- stats.incr("connection_test.success" if success else
"connection_test.failed")
+ stats.incr(
+ "connection_test.success" if success else "connection_test.failed",
+ tags=tags,
+ )
log.info(
"Connection test finished",
state=state.value,
@@ -107,7 +112,7 @@ def supervise_connection_test(
timeout=timeout,
duration=round(time.monotonic() - start, 3),
)
- stats.incr("connection_test.failed")
+ stats.incr("connection_test.failed", tags=tags)
client.connection_tests.update_state(
connection_test_id,
ConnectionTestState.FAILED,
@@ -118,7 +123,7 @@ def supervise_connection_test(
"Connection test failed unexpectedly",
duration=round(time.monotonic() - start, 3),
)
- stats.incr("connection_test.failed")
+ stats.incr("connection_test.failed", tags=tags)
client.connection_tests.update_state(
connection_test_id,
ConnectionTestState.FAILED,
diff --git
a/task-sdk/tests/task_sdk/execution_time/test_connection_test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_connection_test_supervisor.py
index 8f48cfa8e6a..d53844b6803 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_connection_test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_connection_test_supervisor.py
@@ -262,3 +262,115 @@ class TestSuperviseConnectionTest:
assert "unrelated" not in observed
assert observed.get("unrelated_error") == "AirflowNotFoundException"
+
+ @pytest.mark.parametrize(
+ ("team_name", "hook_result", "expected_metric", "expected_tags"),
+ [
+ pytest.param(
+ "team_alpha",
+ (True, "Connection OK"),
+ "connection_test.success",
+ {"team_name": "team_alpha"},
+ id="success_with_team",
+ ),
+ pytest.param(
+ "team_alpha",
+ (False, "Connection refused"),
+ "connection_test.failed",
+ {"team_name": "team_alpha"},
+ id="failure_with_team",
+ ),
+ pytest.param(
+ None,
+ (True, "Connection OK"),
+ "connection_test.success",
+ {},
+ id="success_without_team",
+ ),
+ ],
+ )
+
@mock.patch("airflow.sdk.execution_time.connection_test_supervisor.stats.incr")
+
@mock.patch("airflow.sdk.execution_time.connection_test_supervisor.stats.timer")
+ def test_emits_team_name_on_completion(
+ self,
+ mock_timer,
+ mock_incr,
+ MockClient,
+ team_name,
+ hook_result,
+ expected_metric,
+ expected_tags,
+ ):
+ mock_client = MockClient.return_value
+ mock_client.connection_tests.get_connection.return_value =
ConnectionTestConnectionResponse(
+ conn_id="test_conn",
+ conn_type="http",
+ host="httpbin.org",
+ )
+ mock_timer.return_value.__enter__ = mock.Mock(return_value=None)
+ mock_timer.return_value.__exit__ = mock.Mock(return_value=False)
+
+ with mock.patch(
+ "airflow.sdk.definitions.connection.Connection.test_connection",
+ autospec=True,
+ return_value=hook_result,
+ ):
+ _call(team_name=team_name)
+
+ mock_timer.assert_called_once_with("connection_test.hook_duration",
tags=expected_tags)
+ mock_incr.assert_called_once_with(expected_metric, tags=expected_tags)
+
+ @pytest.mark.parametrize(
+ ("team_name", "exception", "expected_tags"),
+ [
+ pytest.param(
+ "team_alpha",
+ AirflowTaskTimeout("Connection test timed out"),
+ {"team_name": "team_alpha"},
+ id="timeout_with_team",
+ ),
+ pytest.param(
+ "team_alpha",
+ RuntimeError("Something broke"),
+ {"team_name": "team_alpha"},
+ id="exception_with_team",
+ ),
+ pytest.param(None, RuntimeError("Something broke"), {},
id="exception_without_team"),
+ ],
+ )
+
@mock.patch("airflow.sdk.execution_time.connection_test_supervisor.stats.incr")
+ def test_emits_team_name_on_failure_paths(
+ self,
+ mock_incr,
+ MockClient,
+ team_name,
+ exception,
+ expected_tags,
+ ):
+ mock_client = MockClient.return_value
+ mock_client.connection_tests.get_connection.return_value =
ConnectionTestConnectionResponse(
+ conn_id="test_conn",
+ conn_type="http",
+ )
+
+ with mock.patch(
+ "airflow.sdk.definitions.connection.Connection.test_connection",
+ autospec=True,
+ side_effect=exception,
+ ):
+ _call(team_name=team_name)
+
+ mock_incr.assert_called_once_with("connection_test.failed",
tags=expected_tags)
+
+
@mock.patch("airflow.sdk.execution_time.connection_test_supervisor.stats.incr")
+ def test_emits_team_name_on_pre_fetch_failure(self, mock_incr, MockClient):
+ """Failures before GET /connection still include workload team_name on
the failed metric."""
+ mock_client = MockClient.return_value
+ mock_client.connection_tests.get_connection.side_effect =
RuntimeError("not found")
+
+ _call(team_name="team_alpha")
+
+ mock_incr.assert_called_once_with(
+ "connection_test.failed",
+ tags={"team_name": "team_alpha"},
+ )
