jason810496 commented on code in PR #70760:
URL: https://github.com/apache/airflow/pull/70760#discussion_r3713003266
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -458,6 +458,25 @@ def _get_team_names_for_dag_ids(
# Ensure all requested dag_ids are in the result (with None for those
not found)
return {dag_id: self._dag_id_to_team_name.get(dag_id) for dag_id in
dag_ids}
+ def _stamp_team_names(self, dag_runs: Iterable[DagRun], session: Session)
-> None:
Review Comment:
```suggestion
def _stamp_team_names(self, dag_runs: Collection[DagRun], session:
Session) -> None:
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -458,6 +458,25 @@ def _get_team_names_for_dag_ids(
# Ensure all requested dag_ids are in the result (with None for those
not found)
return {dag_id: self._dag_id_to_team_name.get(dag_id) for dag_id in
dag_ids}
+ def _stamp_team_names(self, dag_runs: Iterable[DagRun], session: Session)
-> None:
+ """
+ Stamp ``_team_name`` on each DagRun.
+
+ Team names are resolved via ``_get_team_names_for_dag_ids``, which
caches results in
+ ``self._dag_id_to_team_name`` for the duration of the current
scheduler loop. In
+ practice this means the first call per loop issues one batched query;
subsequent calls
+ for the same dag_ids are pure dict reads with no DB round-trip.
+ """
+ if not self._multi_team:
+ return
+ runs = list(dag_runs)
+ if not runs:
+ return
+ team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in
runs}, session)
+ for dr in runs:
Review Comment:
With the `Collection[DagRun]` annotation, the `list()` materialisation is no
longer needed.
```suggestion
if not dag_runs:
return
team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in
dag_runs}, session)
for dr in dag_runs:
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2981,6 +2995,9 @@ def _schedule_dag_run(
execute=False,
)
+ # dag_run was reloaded from DB above, so _team_name set on the
original object is lost.
+ # Team name should be added before listeners are called in
notify_dagrun_state_changed()
Review Comment:
The first line does not hold: `_team_name` is not a mapped attribute
(`DagRun.stats_tags` reads it with `getattr(self, "_team_name", None)`), and
the reload above is a `select()` by primary key on the *same* session, so
SQLAlchemy's identity map returns the object that is already loaded, i.e.
`dag_run_reloaded is dag_run`. Even if it were expired or refreshed, that only
touches mapped columns, never a plain instance attribute. So nothing is lost by
the reload, and in the production path (`_do_scheduling` stamps, then
`_schedule_all_dag_runs` -> `_schedule_dag_run`) the run arrives here already
stamped.
Suggest dropping the inaccurate sentence and keeping the same wording as the
other three call sites:
```suggestion
# Team name should be added before listeners are called in
notify_dagrun_state_changed()
```
---
Drafted-by: Claude Code (Opus 5); reviewed by @jason810496 before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
