This is an automated email from the ASF dual-hosted git repository.
ashb 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 2a9d508e0b4 Optionally emit Dag tags as metric tags (#68568)
2a9d508e0b4 is described below
commit 2a9d508e0b4d55f7511edd2fec32b0f1c77e3f4c
Author: Sebastián Ortega <[email protected]>
AuthorDate: Mon Jul 20 23:52:00 2026 +0900
Optionally emit Dag tags as metric tags (#68568)
* Emit Dag tags as metric tags
Add a [metrics] dag_tags_in_metrics option (default False). When enabled,
each
Dag tag becomes a metric tag on Dag-run and task-instance metrics: tags
with a
colon (e.g. env:prod) split into a key/value pair; plain tags (e.g.
production)
become standalone DogStatsd tags, or tag=true in InfluxDB line protocol.
Built-in
keys (dag_id, run_type, task_id, team_name) win on collision.
Tags come from the DagRun's dag_model.tags. The scheduler hot loop
(get_running_dag_runs_to_examine) and the executor-event failure path
eager-load
them (gated on the flag) to avoid per-DagRun queries; other in-session
emission
paths fall back to a lazy load, so the tags reach all Dag-run and
task-instance
metrics. DagRun.dag_tags_for_stats swallows SQLAlchemyError so a
detached/expired
instance degrades to no tags rather than breaking the caller.
TaskInstance.stats_tags reuses DagRun.dag_tags_for_stats and adds task_id
and
run_type; the Task SDK worker reads the in-memory Dag and also adds
run_type, so
ti.* metrics carry a consistent tag set across the worker and scheduler
emitters.
The build_dag_metric_tags helper lives in the shared observability stats
module.
* Address review feedback on Dag tags metric emission
Keep the scheduler hot path free of per-task conf reads. Warm the dag tags
via the existing
TaskInstance -> DagModel join rather than an extra dag_run hop, since the
DagModel is shared in the identity map. Fold the repeated dogstatsd tag-list
guard into a single helper, and consolidate the tag-formatting tests.
* Address further review nits on Dag tags metric emission
Make session a keyword-only argument on get_running_dag_runs_to_examine, in
line with the convention for session parameters. Reduce
TaskInstance.stats_tags
to the dag run's tags plus the task-level ones, keeping the TI's transiently
resolved team_name when present.
* Source task instance team_name metric tag from the dag run
TaskInstance.stats_tags already reuses the dag run's stats tags, so take
team_name from there too rather than from a separate per-task transient
attribute, keeping a single source of truth for the tag.
The scheduling loop now stashes the resolved team on the dag run so
task-instance metrics still carry it. The heartbeat-purge path already tags
team inline, so its now-redundant per-task assignment is dropped.
---
.../src/airflow/config_templates/config.yml | 11 ++
.../src/airflow/jobs/scheduler_job_runner.py | 33 +++--
airflow-core/src/airflow/models/dagrun.py | 52 +++++++-
airflow-core/src/airflow/models/taskinstance.py | 6 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 38 ++++--
airflow-core/tests/unit/models/test_dagrun.py | 147 ++++++++++++++++-----
.../tests/unit/models/test_taskinstance.py | 57 ++++++--
.../observability/metrics/datadog_logger.py | 44 ++----
.../airflow_shared/observability/metrics/stats.py | 19 ++-
.../observability/metrics/statsd_logger.py | 5 +-
.../tests/observability/metrics/test_stats.py | 84 +++++++++---
.../src/airflow/sdk/execution_time/task_runner.py | 19 ++-
.../task_sdk/execution_time/test_task_runner.py | 70 ++++++++--
13 files changed, 450 insertions(+), 135 deletions(-)
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index 7a65cedb2da..b3dc3f133d2 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -1280,6 +1280,17 @@ metrics:
type: boolean
example: ~
default: "False"
+ dag_tags_in_metrics:
+ description: |
+ Set to ``True`` to include Dag tags as metric tags on all Dag-run and
task-instance metrics.
+ Tags that contain a colon (e.g. ``env:prod``) are split into a
key/value pair. Plain tags
+ (e.g. ``production``) are emitted as standalone DogStatsd tags, or as
``tag=true`` in InfluxDB
+ line-protocol format. Disabled by default to avoid unexpected
cardinality increases, since Dag
+ tags are free-form, user-defined strings.
+ version_added: 3.4.0
+ type: boolean
+ example: ~
+ default: "False"
otel_on:
description: |
Enables sending metrics to OpenTelemetry.
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index c2f3ee05d73..9b9e85c07e3 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -334,6 +334,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
self._scheduler_use_job_schedule = conf.getboolean("scheduler",
"use_job_schedule", fallback=True)
self._parallelism = conf.getint("core", "parallelism")
self._multi_team = conf.getboolean("core", "multi_team")
+ self._dag_tags_in_metrics = conf.getboolean("metrics",
"dag_tags_in_metrics", fallback=False)
self._max_partition_dag_runs_per_loop = MAX_PARTITION_DAG_RUNS_PER_LOOP
self._dag_id_to_team_name: dict[str, str | None] = {}
@@ -741,10 +742,10 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
list(unique_dag_ids),
)
for ti in task_instances_to_examine:
- # Set team as a transient attribute; team lives on the
Bundle, not
- # on the TI/DagRun schema, so we resolve it at scheduling
time.
+ # Team lives on the Bundle, not the TI/DagRun schema, so
resolve it at scheduling
+ # time and stash it on the dag run, where stats_tags reads
it for metric tagging.
if team := dag_id_to_team_name.get(ti.dag_id):
- ti._team_name = team
+ ti.dag_run._team_name = team
executor_slots_available: dict[ExecutorName, int] = {}
# First get a mapping of executor names to slots they have
available
@@ -1265,6 +1266,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
job_id=self.job.id,
scheduler_dag_bag=self.scheduler_dag_bag,
session=session,
+ eagerly_load_dag_tags=self._dag_tags_in_metrics,
)
except Exception as exc:
stats.incr("scheduler.executor_events.failed",
tags={"exception_class": type(exc).__name__})
@@ -1277,7 +1279,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
@classmethod
def process_executor_events(
- cls, executor: BaseExecutor, job_id: int | None, scheduler_dag_bag:
DBDagBag, session: Session
+ cls,
+ executor: BaseExecutor,
+ job_id: int | None,
+ scheduler_dag_bag: DBDagBag,
+ session: Session,
+ eagerly_load_dag_tags: bool = False,
) -> int:
"""
Process task completion events from the executor and update task
instance states.
@@ -1300,6 +1307,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
:param job_id: The scheduler job ID, used to detect task requeuing by
other schedulers
:param scheduler_dag_bag: Serialized DAG bag for retrieving task
definitions
:param session: Database session for task instance updates
+ :param eagerly_load_dag_tags: When True, eager-load dag_model.tags so
the per-finished-task
+ metrics carry Dag tags without a per-TI lazy load. The scheduler
passes its cached flag so
+ the hot path never reads conf; other callers (e.g. ``dag.test()``)
leave it at the default.
:return: Number of events processed from the executor event buffer
@@ -1391,6 +1401,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
.options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
.options(joinedload(TI.dag_version))
)
+ # When emitting Dag tags as metric tags, eager-load dag_model.tags so
the per-finished-task
+ # ti_failures / operator_failures / task.*_duration metrics carry them
without a per-TI lazy load.
+ # TI already joins DagModel by dag_id, so warm tags off that
relationship directly rather than
+ # via the dag_run hop; the DagModel is shared in the identity map, so
dag_run.dag_model.tags is free.
+ if eagerly_load_dag_tags:
+ query =
query.options(selectinload(TI.dag_model).selectinload(DagModel.tags))
# row lock this entire set of taskinstances to make sure the scheduler
doesn't fail when we have
# multi-schedulers
locked_query = with_row_locks(query, of=TI, session=session,
skip_locked=True)
@@ -1927,7 +1943,11 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# examining, rather than making one query per DagRun.
# Materialize into a list because the multi-team block below
iterates
# the result and ScalarResult is a one-pass iterator.
- dag_runs =
list(DagRun.get_running_dag_runs_to_examine(session=session))
+ dag_runs = list(
+ DagRun.get_running_dag_runs_to_examine(
+ session=session,
eagerly_load_dag_tags=self._dag_tags_in_metrics
+ )
+ )
if self._multi_team and dag_runs:
unique_dag_ids = {dr.dag_id for dr in dag_runs}
@@ -3558,9 +3578,6 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if self._multi_team:
unique_dag_ids = {ti.dag_id for ti in
task_instances_without_heartbeats}
dag_id_to_team_name =
self._get_team_names_for_dag_ids(unique_dag_ids, session)
- for ti in task_instances_without_heartbeats:
- if team := dag_id_to_team_name.get(ti.dag_id):
- ti._team_name = team
else:
dag_id_to_team_name = {}
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index a51f16e51f6..d6610b3663e 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -53,16 +53,25 @@ from sqlalchemy import (
update,
)
from sqlalchemy.dialects import postgresql
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.ext.mutable import MutableDict
-from sqlalchemy.orm import Mapped, declared_attr, joinedload, mapped_column,
relationship, synonym, validates
+from sqlalchemy.orm import (
+ Mapped,
+ declared_attr,
+ joinedload,
+ mapped_column,
+ relationship,
+ synonym,
+ validates,
+)
from sqlalchemy.orm.exc import StaleDataError
from sqlalchemy.sql.expression import false, select
from sqlalchemy.sql.functions import coalesce
from airflow._shared.observability.metrics import stats
+from airflow._shared.observability.metrics.stats import build_dag_metric_tags
from airflow._shared.observability.traces import (
DAGRUN_PARENT_TRACE_CONTEXT_KEY,
TASK_SPAN_DETAIL_LEVEL_KEY,
@@ -576,11 +585,36 @@ class DagRun(Base, LoggingMixin):
)
return session.scalar(select_stmt)
+ def dag_tags_for_stats(self) -> dict[str, str]:
+ """Convert dag tags to metric tags. Tags with ':' become key:value;
others are standalone (empty value)."""
+ if not airflow_conf.getboolean("metrics", "dag_tags_in_metrics",
fallback=False):
+ return {}
+ try:
+ # Lazy-loads dag_model.tags if not already loaded. The scheduler
hot loop
+ # (get_running_dag_runs_to_examine) eager-loads these to avoid a
per-DagRun query; other,
+ # low-frequency emission paths fall back to this lazy load. On a
detached/expired DagRun
+ # the load raises — swallow it so metric tagging never breaks the
caller.
+ if not self.dag_model or not self.dag_model.tags:
+ return {}
+ return build_dag_metric_tags(tag.name for tag in
self.dag_model.tags)
+ except SQLAlchemyError:
+ return {}
+
@property
def stats_tags(self) -> dict[str, str]:
- return prune_dict(
- {"dag_id": self.dag_id, "run_type": self.run_type, "team_name":
getattr(self, "_team_name", None)}
+ # prune_dict strips falsy values, so merge dag tags after it runs so
standalone
+ # tags (empty value) are preserved for DogStatsD emission.
+ base = prune_dict(
+ {
+ "dag_id": self.dag_id,
+ # bare value so it serializes as e.g. "scheduled", not
"dagruntype.scheduled"
+ "run_type": getattr(self.run_type, "value", self.run_type),
+ "team_name": getattr(self, "_team_name", None),
+ }
)
+ dag_tags = self.dag_tags_for_stats()
+ # Built-in keys win on collision; dag tags fill in everything else.
+ return {**dag_tags, **base}
def get_state(self):
return self._state
@@ -702,7 +736,9 @@ class DagRun(Base, LoggingMixin):
@classmethod
@retry_db_transaction
- def get_running_dag_runs_to_examine(cls, session: Session) ->
ScalarResult[DagRun]:
+ def get_running_dag_runs_to_examine(
+ cls, *, session: Session, eagerly_load_dag_tags: bool
+ ) -> ScalarResult[DagRun]:
"""
Return the next DagRuns that the scheduler should attempt to schedule.
@@ -733,6 +769,12 @@ class DagRun(Base, LoggingMixin):
.limit(cls.DEFAULT_DAGRUNS_TO_EXAMINE)
)
+ # When dag tags are emitted as metric tags, eagerly load
dag_model.tags so stats_tags does not
+ # fire a per-DagRun N+1 lazy load in the scheduler loop. The caller
owns the feature decision;
+ # the scheduler passes its cached flag so the loop never reads conf.
+ if eagerly_load_dag_tags:
+ query =
query.options(joinedload(cls.dag_model).selectinload(DagModel.tags))
+
query = query.where(DagRun.run_after <= func.now())
result = session.scalars(with_row_locks(query, of=cls,
session=session, skip_locked=True)).unique()
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 1c7de2185e6..0dd10507fc8 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -100,7 +100,6 @@ from airflow.task.priority_strategy import
validate_and_load_priority_weight_str
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS
from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
-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.platform import getuser
@@ -757,9 +756,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
@property
def stats_tags(self) -> dict[str, str]:
"""Returns task instance tags."""
- return prune_dict(
- {"dag_id": self.dag_id, "task_id": self.task_id, "team_name":
getattr(self, "_team_name", None)}
- )
+ # Reuse the dag run's tags and add the task-level ones.
+ return {**self.dag_run.stats_tags, "task_id": self.task_id}
@staticmethod
def insert_mapping(
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 68f4ec09150..dff3adea48b 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -518,8 +518,13 @@ class TestSchedulerJob:
"scheduler.tasks.killed_externally",
tags={"dag_id": dag_id, "task_id": ti1.task_id},
),
- mock.call("operator_failures_EmptyOperator", tags={"dag_id":
dag_id, "task_id": ti1.task_id}),
- mock.call("ti_failures", tags={"dag_id": dag_id, "task_id":
ti1.task_id}),
+ mock.call(
+ "operator_failures_EmptyOperator",
+ tags={"dag_id": dag_id, "task_id": ti1.task_id,
"run_type": "manual"},
+ ),
+ mock.call(
+ "ti_failures", tags={"dag_id": dag_id, "task_id":
ti1.task_id, "run_type": "manual"}
+ ),
],
any_order=True,
)
@@ -645,8 +650,11 @@ class TestSchedulerJob:
"scheduler.tasks.killed_externally",
tags={"dag_id": dag_id, "task_id": task_id},
),
- mock.call("operator_failures_EmptyOperator", tags={"dag_id":
dag_id, "task_id": task_id}),
- mock.call("ti_failures", tags={"dag_id": dag_id, "task_id":
task_id}),
+ mock.call(
+ "operator_failures_EmptyOperator",
+ tags={"dag_id": dag_id, "task_id": task_id, "run_type":
"manual"},
+ ),
+ mock.call("ti_failures", tags={"dag_id": dag_id, "task_id":
task_id, "run_type": "manual"}),
],
any_order=True,
)
@@ -9828,8 +9836,10 @@ class TestSchedulerJob:
assert result2 == mock_executors[1] # Matched by executor name
@conf_vars({("core", "multi_team"): "true"})
- def test_multi_team_sets_team_name_on_task_instances(self, dag_maker,
mock_executors, session):
- """Test that _team_name is set on TaskInstance objects during the
scheduling loop."""
+ @mock.patch("airflow._shared.observability.metrics.stats.timing")
+ def test_multi_team_sets_team_name_on_task_instances(self, mock_timing,
dag_maker, session):
+ """The scheduling loop resolves the bundle's team onto the dag run, so
the QUEUED
+ state-change metric (emitted via TaskInstance.stats_tags) carries
team_name."""
clear_db_teams()
clear_db_dag_bundles()
@@ -9848,19 +9858,21 @@ class TestSchedulerJob:
dr = dag_maker.create_dagrun()
ti = dr.get_task_instance("task_a", session=session)
ti.state = State.SCHEDULED
+ ti.scheduled_dttm = timezone.utcnow()
+ session.merge(ti)
session.flush()
scheduler_job = Job()
self.job_runner = SchedulerJobRunner(job=scheduler_job)
- self.job_runner._multi_team = True
- # Simulate what _executable_task_instances_to_queued does
- dag_id_to_team_name =
self.job_runner._get_team_names_for_dag_ids(["dag_a"], session)
- if team_name := dag_id_to_team_name.get(ti.dag_id):
- ti._team_name = team_name
+ queued_tis =
self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
- assert ti._team_name == "team_a"
- assert ti.stats_tags == {"dag_id": "dag_a", "task_id": "task_a",
"team_name": "team_a"}
+ assert {t.key for t in queued_tis} == {ti.key}
+ scheduled_calls = [
+ c for c in mock_timing.call_args_list if c.args and c.args[0] ==
"task.scheduled_duration"
+ ]
+ assert scheduled_calls, "expected a task.scheduled_duration metric on
QUEUED transition"
+ assert scheduled_calls[0].kwargs["tags"]["team_name"] == "team_a"
@conf_vars({("core", "multi_team"): "true"})
def test_do_scheduling_multi_team_schedules_task_instances(self,
dag_maker, session):
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index 64b4947ec4c..c46c30d1194 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -21,7 +21,7 @@ import datetime
from collections import defaultdict
from collections.abc import Mapping
from contextlib import contextmanager
-from functools import reduce
+from functools import partial, reduce
from typing import TYPE_CHECKING
from unittest import mock
from unittest.mock import ANY, call
@@ -1086,10 +1086,12 @@ class TestDagRun:
)
if state == DagRunState.RUNNING:
- func = DagRun.get_running_dag_runs_to_examine
+ fetch = partial(
+ DagRun.get_running_dag_runs_to_examine, session=session,
eagerly_load_dag_tags=False
+ )
else:
- func = DagRun.get_queued_dag_runs_to_set_running
- runs = func(session).all()
+ fetch = partial(DagRun.get_queued_dag_runs_to_set_running, session)
+ runs = fetch().all()
assert runs == [dr]
@@ -1097,7 +1099,7 @@ class TestDagRun:
session.merge(orm_dag)
session.commit()
- runs = func(session).all()
+ runs = fetch().all()
assert runs == []
@mock.patch("airflow._shared.observability.metrics.stats.timing")
@@ -4563,34 +4565,115 @@ class TestDagRunTracing:
assert trace_id != _EXTERNAL_TRACE_ID
-class TestDagRunStatsTagsTeamName:
- def test_stats_tags_without_team_name(self, dag_maker):
- """stats_tags should not include team_name when _team_name is not
set."""
- with dag_maker("test_dag"):
- EmptyOperator(task_id="t1")
- dr = dag_maker.create_dagrun()
- tags = dr.stats_tags
- assert "team_name" not in tags
- assert tags == {"dag_id": "test_dag", "run_type": "manual"}
+def test_stats_tags_without_team_name(dag_maker):
+ """stats_tags omits team_name when _team_name is not set."""
+ with dag_maker("test_dag"):
+ EmptyOperator(task_id="t1")
+ dr = dag_maker.create_dagrun()
+ assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type}
- def test_stats_tags_with_team_name(self, dag_maker):
- """stats_tags should include team_name when _team_name is set."""
- with dag_maker("test_dag"):
- EmptyOperator(task_id="t1")
- dr = dag_maker.create_dagrun()
- dr._team_name = "my_team"
- tags = dr.stats_tags
- assert tags["team_name"] == "my_team"
- assert tags == {"dag_id": "test_dag", "run_type": "manual",
"team_name": "my_team"}
-
- def test_stats_tags_with_none_team_name(self, dag_maker):
- """stats_tags should not include team_name when _team_name is None."""
- with dag_maker("test_dag"):
- EmptyOperator(task_id="t1")
- dr = dag_maker.create_dagrun()
- dr._team_name = None
- tags = dr.stats_tags
- assert "team_name" not in tags
+
+def test_stats_tags_with_team_name(dag_maker):
+ """stats_tags includes team_name when _team_name is set."""
+ with dag_maker("test_dag"):
+ EmptyOperator(task_id="t1")
+ dr = dag_maker.create_dagrun()
+ dr._team_name = "my_team"
+ assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type,
"team_name": "my_team"}
+
+
+def test_stats_tags_with_none_team_name(dag_maker):
+ """stats_tags omits team_name when _team_name is None."""
+ with dag_maker("test_dag"):
+ EmptyOperator(task_id="t1")
+ dr = dag_maker.create_dagrun()
+ dr._team_name = None
+ assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type}
+
+
+def test_stats_tags_dag_tags_disabled_by_default(dag_maker, session):
+ """With the flag off (the default), dag tags must not leak into metrics."""
+ with dag_maker("disabled_tag_dag", tags=["production", "env:prod"],
session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ _ = dr.dag_model.tags
+ assert dr.stats_tags == {"dag_id": "disabled_tag_dag", "run_type":
dr.run_type}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_without_dag_tags(dag_maker, session):
+ with dag_maker("no_tags_dag", session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ _ = dr.dag_model.tags
+ assert dr.stats_tags == {"dag_id": "no_tags_dag", "run_type": dr.run_type}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_with_standalone_dag_tag(dag_maker, session):
+ with dag_maker("standalone_tag_dag", tags=["production"], session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ _ = dr.dag_model.tags # eager-load so _dag_tags_for_stats sees the tags
+ tags = dr.stats_tags
+ assert tags == {"dag_id": "standalone_tag_dag", "run_type": "manual",
"production": ""}
+ # run_type is the bare value, not a DagRunType enum (serializes as
"manual", not "dagruntype.manual")
+ assert type(tags["run_type"]) is str
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_with_key_value_dag_tag(dag_maker, session):
+ with dag_maker("kv_tag_dag", tags=["env:staging"], session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ _ = dr.dag_model.tags
+ assert dr.stats_tags == {"dag_id": "kv_tag_dag", "run_type": dr.run_type,
"env": "staging"}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_builtin_keys_win_on_collision(dag_maker, session):
+ with dag_maker("collision_dag", tags=["dag_id:sneaky"], session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ _ = dr.dag_model.tags
+ # built-in dag_id wins over the colliding "dag_id:sneaky" tag
+ assert dr.stats_tags == {"dag_id": "collision_dag", "run_type":
dr.run_type}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_lazy_loads_dag_tags_when_not_eager_loaded(dag_maker,
session):
+ """When dag_model is not eager-loaded, stats_tags lazy-loads it in-session
so tags still appear."""
+ from sqlalchemy import inspect as sa_inspect
+
+ with dag_maker("lazy_tag_dag", tags=["env:prod"], session=session):
+ pass
+ dr = dag_maker.create_dagrun()
+ session.expire(dr, ["dag_model"]) # not eager-loaded
+ assert "dag_model" in sa_inspect(dr).unloaded
+
+ # lazy fallback loads dag_model.tags in-session, so the tags are still
emitted
+ assert dr.stats_tags == {"dag_id": "lazy_tag_dag", "run_type":
dr.run_type, "env": "prod"}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_get_running_dag_runs_to_examine_eager_loads_dag_tags(dag_maker,
session):
+ """With the flag on, the scheduler query eager-loads dag_model.tags so
stats_tags fires no lazy load."""
+ from sqlalchemy import inspect as sa_inspect
+
+ with dag_maker("eager_tag_dag", tags=["env:prod"], session=session):
+ pass
+ dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ session.commit()
+
+ dr = next(
+ r
+ for r in DagRun.get_running_dag_runs_to_examine(session=session,
eagerly_load_dag_tags=True)
+ if r.dag_id == "eager_tag_dag"
+ )
+ # dag_model and its tags are already populated — no lazy load needed at
metric-emission time.
+ assert "dag_model" not in sa_inspect(dr).unloaded
+ assert "tags" not in sa_inspect(dr.dag_model).unloaded
+ assert dr.stats_tags == {"dag_id": "eager_tag_dag", "run_type":
dr.run_type, "env": "prod"}
class TestClearPartitionRuns:
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 7f1a193d963..e0d4a1063e9 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -2431,7 +2431,7 @@ class TestTaskInstance:
ti.task = None
ti.state = State.QUEUED
session.flush()
- expected_stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}
+ expected_stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id,
"run_type": dr.run_type}
assert ti.task is None, "Check critical pre-condition"
@@ -4374,47 +4374,82 @@ def
test_task_instance_repr_does_not_raise_for_deferred_columns(dag_maker, sessi
class TestTaskInstanceStatsTagsTeamName:
def test_stats_tags_without_team_name(self, dag_maker, session):
- """stats_tags should not include team_name when _team_name is not
set."""
+ """stats_tags should not include team_name when the dag run has no
team_name."""
with dag_maker("test_dag"):
EmptyOperator(task_id="my_task")
dr = dag_maker.create_dagrun()
ti = dr.get_task_instance("my_task", session=session)
tags = ti.stats_tags
assert "team_name" not in tags
- assert tags == {"dag_id": "test_dag", "task_id": "my_task"}
+ assert tags == {"dag_id": "test_dag", "task_id": "my_task",
"run_type": dr.run_type}
def test_stats_tags_with_team_name(self, dag_maker, session):
- """stats_tags should include team_name when _team_name is set."""
+ """stats_tags takes team_name from the dag run's stats_tags."""
with dag_maker("test_dag"):
EmptyOperator(task_id="my_task")
dr = dag_maker.create_dagrun()
+ dr._team_name = "my_team"
ti = dr.get_task_instance("my_task", session=session)
- ti._team_name = "my_team"
tags = ti.stats_tags
assert tags["team_name"] == "my_team"
- assert tags == {"dag_id": "test_dag", "task_id": "my_task",
"team_name": "my_team"}
+ assert tags == {
+ "dag_id": "test_dag",
+ "task_id": "my_task",
+ "team_name": "my_team",
+ "run_type": dr.run_type,
+ }
def test_stats_tags_with_none_team_name(self, dag_maker, session):
- """stats_tags should not include team_name when _team_name is None."""
+ """stats_tags should not include team_name when the dag run's
team_name is None."""
with dag_maker("test_dag"):
EmptyOperator(task_id="my_task")
dr = dag_maker.create_dagrun()
+ dr._team_name = None
ti = dr.get_task_instance("my_task", session=session)
- ti._team_name = None
tags = ti.stats_tags
assert "team_name" not in tags
+ @conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+ def test_stats_tags_match_worker_tag_set(self, dag_maker, session):
+ """ti_failures (and other ti.* metrics) are emitted from both the
worker
+ (RuntimeTaskInstance.stats_tags) and the scheduler
(TaskInstance.stats_tags); both must produce
+ the same tag set. The worker side is asserted in test_task_runner.py
+ (test_stats_tags_with_standalone_and_key_value_tags); this guards the
scheduler side against drift:
+ dag tags + dag_id + task_id + run_type.
+ """
+ with dag_maker("parity_dag", tags=["env:prod", "production"],
session=session):
+ EmptyOperator(task_id="t1")
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("t1", session=session)
+ tags = ti.stats_tags
+ assert tags == {
+ "dag_id": "parity_dag",
+ "task_id": "t1",
+ "run_type": "manual",
+ "env": "prod",
+ "production": "",
+ }
+ # run_type must be the bare value (not a DagRunType enum) so it
serializes identically to
+ # the worker side, e.g. "manual" not "dagruntype.manual".
+ assert type(tags["run_type"]) is str
+
@pytest.mark.parametrize(
("team_name", "expected_tags"),
[
pytest.param(
"my_team",
- {"dag_id": "test_dag", "task_id": "my_task", "team_name":
"my_team", "queue": "default"},
+ {
+ "dag_id": "test_dag",
+ "task_id": "my_task",
+ "team_name": "my_team",
+ "queue": "default",
+ "run_type": "manual",
+ },
id="with_team",
),
pytest.param(
None,
- {"dag_id": "test_dag", "task_id": "my_task", "queue":
"default"},
+ {"dag_id": "test_dag", "task_id": "my_task", "queue":
"default", "run_type": "manual"},
id="without_team",
),
],
@@ -4430,7 +4465,7 @@ class TestTaskInstanceStatsTagsTeamName:
ti.state = TaskInstanceState.SCHEDULED
ti.scheduled_dttm = timezone.utcnow()
if team_name:
- ti._team_name = team_name
+ dr._team_name = team_name
session.merge(ti)
session.flush()
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py
b/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py
index 129354c3a43..b1fda8208aa 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py
+++
b/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py
@@ -58,6 +58,15 @@ class SafeDogStatsdLogger:
self.stat_name_handler = stat_name_handler
self.statsd_influxdb_enabled = statsd_influxdb_enabled
+ def _build_tags_list(self, tags: dict[str, str] | None) -> list[str]:
+ if not (self.metrics_tags and isinstance(tags, dict)):
+ return []
+ return [
+ (f"{key}:{value}" if value != "" else key)
+ for key, value in tags.items()
+ if self.metric_tags_validator.test(key)
+ ]
+
@validate_stat
def incr(
self,
@@ -68,12 +77,7 @@ class SafeDogStatsdLogger:
tags: dict[str, str] | None = None,
) -> None:
"""Increment stat."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = self._build_tags_list(tags)
if self.metrics_validator.test(stat):
return self.dogstatsd.increment(metric=stat, value=count,
tags=tags_list, sample_rate=rate)
return None
@@ -88,12 +92,7 @@ class SafeDogStatsdLogger:
tags: dict[str, str] | None = None,
) -> None:
"""Decrement stat."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = self._build_tags_list(tags)
if self.metrics_validator.test(stat):
return self.dogstatsd.decrement(metric=stat, value=count,
tags=tags_list, sample_rate=rate)
return None
@@ -109,12 +108,7 @@ class SafeDogStatsdLogger:
tags: dict[str, str] | None = None,
) -> None:
"""Gauge stat."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = self._build_tags_list(tags)
if self.metrics_validator.test(stat):
return self.dogstatsd.gauge(metric=stat, value=value,
tags=tags_list, sample_rate=rate)
return None
@@ -128,12 +122,7 @@ class SafeDogStatsdLogger:
tags: dict[str, str] | None = None,
) -> None:
"""Stats timing."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = self._build_tags_list(tags)
if self.metrics_validator.test(stat):
if isinstance(dt, datetime.timedelta):
dt = dt.total_seconds() * 1000.0
@@ -148,12 +137,7 @@ class SafeDogStatsdLogger:
**kwargs,
) -> Timer:
"""Timer metric that can be cancelled."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = self._build_tags_list(tags)
if stat and self.metrics_validator.test(stat):
return Timer(self.dogstatsd.timed(stat, tags=tags_list, **kwargs))
return Timer()
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/stats.py
b/shared/observability/src/airflow_shared/observability/metrics/stats.py
index 5140314922a..7b51e580cd0 100644
--- a/shared/observability/src/airflow_shared/observability/metrics/stats.py
+++ b/shared/observability/src/airflow_shared/observability/metrics/stats.py
@@ -20,7 +20,7 @@ import logging
import os
import re
import socket
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any
from .base_stats_logger import NoStatsLogger
@@ -35,6 +35,23 @@ log = logging.getLogger(__name__)
_VALID_STAT_NAME_CHARS_RE = re.compile(r"^[a-zA-Z0-9_.-]+$")
_INVALID_STAT_NAME_CHARS_RE = re.compile(r"[^a-zA-Z0-9_.-]")
+
+def build_dag_metric_tags(tag_names: Iterable[str]) -> dict[str, str]:
+ """
+ Convert Dag tag strings into metric tags.
+
+ Tags with a non-empty value after a ``:`` (e.g. ``env:prod``) split into a
+ ``key: value`` pair. Plain tags (e.g. ``production``) and tags with no
value
+ after the colon (e.g. ``env:``) map to an empty string, emitted as a
standalone
+ DogStatsd tag or as ``tag=true`` in InfluxDB line protocol.
+ """
+ result: dict[str, str] = {}
+ for name in tag_names:
+ key, _, value = name.partition(":")
+ result[key] = value
+ return result
+
+
# Module-level singleton state.
_factory: Callable[[], StatsLogger | NoStatsLogger] | None = None
_backend: StatsLogger | NoStatsLogger | None = None
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py
b/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py
index b5f39777241..fd78543e202 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py
+++
b/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py
@@ -52,8 +52,9 @@ def prepare_stat_with_tags(fn: T) -> T:
if stat is not None and tags is not None:
for k, v in tags.items():
if self.metric_tags_validator.test(k):
- if all(c not in [",", "="] for c in f"{v}{k}"):
- stat += f",{k}={v}"
+ v_str = "true" if v == "" else v
+ if all(c not in [",", "="] for c in f"{v_str}{k}"):
+ stat += f",{k}={v_str}"
else:
log.error("Dropping invalid tag: %s=%s.", k, v)
return fn(self, stat, *args, tags=tags, **kwargs)
diff --git a/shared/observability/tests/observability/metrics/test_stats.py
b/shared/observability/tests/observability/metrics/test_stats.py
index 04f2731248a..9bdb4c47312 100644
--- a/shared/observability/tests/observability/metrics/test_stats.py
+++ b/shared/observability/tests/observability/metrics/test_stats.py
@@ -34,6 +34,7 @@ from airflow_shared.observability.exceptions import
InvalidStatsNameException
from airflow_shared.observability.metrics import datadog_logger, statsd_logger
from airflow_shared.observability.metrics.base_stats_logger import StatsLogger
from airflow_shared.observability.metrics.datadog_logger import
SafeDogStatsdLogger
+from airflow_shared.observability.metrics.stats import build_dag_metric_tags
from airflow_shared.observability.metrics.statsd_logger import SafeStatsdLogger
from airflow_shared.observability.metrics.validators import (
PatternAllowListValidator,
@@ -266,6 +267,20 @@ class TestDogStats:
metric="empty", sample_rate=1, value=1, tags=[]
)
+ @pytest.mark.parametrize(
+ ("tags", "expected"),
+ [
+ ({"env": "prod"}, {"env:prod"}),
+ ({"production": ""}, {"production"}),
+ ({"production": "", "env": "staging"}, {"production",
"env:staging"}),
+ ],
+ )
+ def test_key_value_and_standalone_tags(self, tags, expected):
+ dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client,
metrics_tags=True)
+ dogstatsd.incr("my_metric", tags=tags)
+ call_kwargs = self.dogstatsd_client.increment.call_args
+ assert set(call_kwargs.kwargs["tags"]) == expected
+
class TestStatsAllowAndBlockLists:
@pytest.mark.parametrize(
@@ -464,23 +479,35 @@ class TestStatsWithInfluxDBEnabled:
)
self.statsd_client.incr.assert_called_once_with("test_stats_run.delay", 1, 1)
- def test_increment_counter_with_tags(self):
- self.stats.incr(
- "test_stats_run.delay",
- tags={"key0": 0, "key1": "val1", "key2": "val2"},
- )
-
self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key0=0,key1=val1",
1, 1)
-
- def test_increment_counter_with_tags_and_forward_slash(self):
- self.stats.incr("test_stats_run.dag", tags={"path":
"/some/path/dag.py"})
-
self.statsd_client.incr.assert_called_once_with("test_stats_run.dag,path=/some/path/dag.py",
1, 1)
-
- def test_does_not_increment_counter_drops_invalid_tags(self):
- self.stats.incr(
- "test_stats_run.delay",
- tags={"key0,": "val0", "key1": "val1", "key2": "val2", "key3":
"val3"},
- )
-
self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key1=val1",
1, 1)
+ @pytest.mark.parametrize(
+ ("stat", "tags", "expected"),
+ [
+ (
+ "test_stats_run.delay",
+ {"key0": 0, "key1": "val1", "key2": "val2"},
+ "test_stats_run.delay,key0=0,key1=val1",
+ ),
+ (
+ "test_stats_run.dag",
+ {"path": "/some/path/dag.py"},
+ "test_stats_run.dag,path=/some/path/dag.py",
+ ),
+ (
+ "test_stats_run.delay",
+ {"key0,": "val0", "key1": "val1", "key2": "val2", "key3":
"val3"},
+ "test_stats_run.delay,key1=val1",
+ ),
+ # Empty value renders as `=true` in influxdb line protocol.
+ (
+ "test_stats_run.delay",
+ {"production": "", "key1": "val1"},
+ "test_stats_run.delay,production=true,key1=val1",
+ ),
+ ],
+ )
+ def test_increment_counter_with_tags(self, stat, tags, expected):
+ self.stats.incr(stat, tags=tags)
+ self.statsd_client.incr.assert_called_once_with(expected, 1, 1)
def always_invalid(stat_name):
@@ -776,3 +803,26 @@ class TestCustomStatsName:
def teardown_method(self) -> None:
# To avoid side-effect
importlib.reload(airflow_shared.observability.metrics.stats)
+
+
[email protected](
+ ("tag_names", "expected"),
+ [
+ pytest.param([], {}, id="empty"),
+ pytest.param(["production"], {"production": ""}, id="standalone"),
+ pytest.param(["env:prod"], {"env": "prod"}, id="key-value"),
+ pytest.param(
+ ["production", "env:prod", "team:data"],
+ {"production": "", "env": "prod", "team": "data"},
+ id="mixed",
+ ),
+ pytest.param(["a:b:c"], {"a": "b:c"}, id="value-with-colon"),
+ pytest.param(["env:"], {"env": ""}, id="trailing-colon-is-standalone"),
+ ],
+)
+def test_build_dag_metric_tags(tag_names: list[str], expected: dict[str, str])
-> None:
+ assert build_dag_metric_tags(tag_names) == expected
+
+
+def test_build_dag_metric_tags_accepts_generator() -> None:
+ assert build_dag_metric_tags(name for name in ["env:prod"]) == {"env":
"prod"}
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 97d702e8125..31d0831bbad 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -45,6 +45,7 @@ from structlog.contextvars import bind_contextvars
from airflow.dag_processing.bundles.base import BaseDagBundle,
BundleVersionLock
from airflow.dag_processing.bundles.manager import DagBundlesManager
from airflow.sdk._shared.observability.metrics import stats
+from airflow.sdk._shared.observability.metrics.stats import
build_dag_metric_tags
from airflow.sdk._shared.observability.traces import get_task_span_detail_level
from airflow.sdk._shared.template_rendering import truncate_rendered_value
from airflow.sdk.api.client import get_hostname, getuser
@@ -259,10 +260,20 @@ class RuntimeTaskInstance(TaskInstance):
@property
def stats_tags(self) -> dict[str, str]:
- """Metric tags for this task instance, including team_name when
available."""
- tags: dict[str, str] = {"dag_id": self.dag_id, "task_id": self.task_id}
- if self._ti_context_from_server and
self._ti_context_from_server.dag_run.team_name:
- tags["team_name"] = self._ti_context_from_server.dag_run.team_name
+ """Metric tags for this task instance, including dag tags and
team_name when available."""
+ tags: dict[str, str] = {}
+ if conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False):
+ tags.update(build_dag_metric_tags(self.task.dag.tags))
+ # Built-in keys always win on collision.
+ tags.update(dag_id=self.dag_id, task_id=self.task_id)
+ if self._ti_context_from_server:
+ # run_type keeps the tag set consistent with the scheduler-side
TaskInstance.stats_tags.
+ # Coerce the DagRunType enum to its bare value so it serializes as
e.g. "scheduled" rather
+ # than "dagruntype.scheduled" (matching the scheduler, which emits
the plain string).
+ run_type = self._ti_context_from_server.dag_run.run_type
+ tags["run_type"] = getattr(run_type, "value", run_type)
+ if self._ti_context_from_server.dag_run.team_name:
+ tags["team_name"] =
self._ti_context_from_server.dag_run.team_name
return tags
def __rich_repr__(self):
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index bee3f67b5a0..0031974b866 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -5529,11 +5529,11 @@ class TestTaskInstanceMetrics:
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
# verify ti.start was called in legacy format
- backend.incr.assert_any_call(f"ti.start.{ti.dag_id}.{ti.task_id}")
+ backend.incr.assert_any_call(f"ti.start.{ti.dag_id}.{ti.task_id}",
tags={"run_type": "manual"})
# verify ti.start was called in tagged format
backend.incr.assert_any_call(
"ti.start",
- tags={"dag_id": ti.dag_id, "task_id": ti.task_id},
+ tags={"dag_id": ti.dag_id, "task_id": ti.task_id, "run_type":
"manual"},
)
@pytest.mark.parametrize(
@@ -5558,11 +5558,18 @@ class TestTaskInstanceMetrics:
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
# verify ti.finish was called in legacy format
-
backend.incr.assert_any_call(f"ti.finish.{ti.dag_id}.{ti.task_id}.{expected_state}")
+ backend.incr.assert_any_call(
+ f"ti.finish.{ti.dag_id}.{ti.task_id}.{expected_state}",
tags={"run_type": "manual"}
+ )
# verify ti.finish was called in tagged format
backend.incr.assert_any_call(
"ti.finish",
- tags={"dag_id": ti.dag_id, "task_id": ti.task_id, "state":
expected_state},
+ tags={
+ "dag_id": ti.dag_id,
+ "task_id": ti.task_id,
+ "run_type": "manual",
+ "state": expected_state,
+ },
)
def test_operator_successes_metrics_emitted(self, create_runtime_ti,
mock_supervisor_comms):
@@ -5575,7 +5582,7 @@ class TestTaskInstanceMetrics:
mock_get_backend.return_value = backend
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
- stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}
+ stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id,
"run_type": "manual"}
# verify operator_successes in legacy format
backend.incr.assert_any_call("operator_successes_PythonOperator",
tags=stats_tags)
@@ -5596,7 +5603,7 @@ class TestTaskInstanceMetrics:
mock_get_backend.return_value = backend
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
- stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}
+ stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id,
"run_type": "manual"}
# verify operator_failures in legacy format
backend.incr.assert_any_call("operator_failures_PythonOperator",
tags=stats_tags)
@@ -5627,7 +5634,12 @@ class TestTaskInstanceMetrics:
mock_get_backend.return_value = backend
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
- expected = {"dag_id": ti.dag_id, "task_id": ti.task_id,
**expected_tags_extra}
+ expected = {
+ "dag_id": ti.dag_id,
+ "task_id": ti.task_id,
+ "run_type": "manual",
+ **expected_tags_extra,
+ }
backend.incr.assert_any_call("ti.start", tags=expected)
@pytest.mark.parametrize(
@@ -5649,7 +5661,12 @@ class TestTaskInstanceMetrics:
mock_get_backend.return_value = backend
run(ti, context=ti.get_template_context(), log=mock.MagicMock())
- stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id,
"team_name": "team_a"}
+ stats_tags = {
+ "dag_id": ti.dag_id,
+ "task_id": ti.task_id,
+ "run_type": "manual",
+ "team_name": "team_a",
+ }
backend.incr.assert_any_call(
operator_metric,
tags={**stats_tags, "operator_name": "PythonOperator"},
@@ -6438,3 +6455,40 @@ class TestRegisterDeserializationAllowedClasses:
with patch("airflow.sdk.execution_time.task_runner.allow_class",
side_effect=ValueError("nope")):
# Must not raise -- the walk swallows per-class registration
errors.
_register_deserialization_allowed_classes(dag,
structlog.get_logger())
+
+
+def _make_dag_tagged_ti(create_runtime_ti, tags):
+ """Build a RuntimeTaskInstance whose in-memory Dag carries the given
tags."""
+ from airflow.sdk import DAG
+ from airflow.sdk.bases.operator import BaseOperator
+
+ with DAG("tagged_dag", tags=tags):
+ task = BaseOperator(task_id="t")
+ return create_runtime_ti(task=task)
+
+
+def test_stats_tags_dag_tags_disabled_by_default(create_runtime_ti):
+ """With the flag off (the default), dag tags must not leak into metrics."""
+ ti = _make_dag_tagged_ti(create_runtime_ti, ["env:prod", "validation"])
+ assert ti.stats_tags == {"dag_id": "tagged_dag", "task_id": "t",
"run_type": "manual"}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_without_dag_tags(create_runtime_ti):
+ tags = _make_dag_tagged_ti(create_runtime_ti, []).stats_tags
+ assert tags == {"dag_id": "tagged_dag", "task_id": "t", "run_type":
"manual"}
+ # run_type must be a plain str, not a DagRunType enum member: DagRunType
is a str-enum, so the
+ # dict equality above passes either way, but the enum serializes as
"dagruntype.manual" on the wire.
+ assert type(tags["run_type"]) is str
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_with_standalone_and_key_value_tags(create_runtime_ti):
+ ti = _make_dag_tagged_ti(create_runtime_ti, ["env:prod", "validation"])
+ assert ti.stats_tags == {
+ "env": "prod",
+ "validation": "",
+ "dag_id": "tagged_dag",
+ "task_id": "t",
+ "run_type": "manual",
+ }
