From 1ec4567acdf27be62380cc182c2d15fd8f92f774 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Thu, 30 Jul 2026 15:08:05 +0200 Subject: [PATCH 1/4] Fix missing _team_name on DagRun before some listener calls --- .../src/airflow/jobs/scheduler_job_runner.py | 25 +++ .../tests/unit/jobs/test_scheduler_job.py | 199 ++++++++++++++++++ 2 files changed, 224 insertions(+) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index b6a29995d5937..af41a7800f392 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1734,6 +1734,10 @@ def _update_dag_run_state_for_paused_dags(self, *, session: Session = NEW_SESSIO .group_by(DagRun) ) ) + # Stamp _team_name before update_state() calls notify_dagrun_state_changed(), + # which fires on_dag_run_success/failed listeners. + # _get_team_names_for_dag_ids caches results in self._dag_id_to_team_name, so + # this is usually a dict read with no DB query. if self._multi_team and paused_runs: paused_dag_ids = {dr.dag_id for dr in paused_runs} paused_team_mapping = self._get_team_names_for_dag_ids(paused_dag_ids, session) @@ -1996,6 +2000,12 @@ def _do_scheduling(self, session: Session) -> int: ) ) + # Stamp _team_name before _schedule_all_dag_runs calls notify_dagrun_state_changed() + # (on_dag_run_success/failed listeners). Runs that just transitioned QUEUED→RUNNING + # in _start_queued_dagruns may already have _team_name set (SQLAlchemy identity map + # returns the same Python objects); overwriting with the same value is harmless. + # _get_team_names_for_dag_ids caches results in self._dag_id_to_team_name, so this + # is usually a dict read with no DB query. if self._multi_team and dag_runs: unique_dag_ids = {dr.dag_id for dr in dag_runs} dr_team_mapping = self._get_team_names_for_dag_ids(unique_dag_ids, session) @@ -2844,6 +2854,15 @@ def _update_state(dag: SerializedDAG, dag_run: DagRun): partial(self.scheduler_dag_bag.get_dag_for_run, session=session) ) + # Stamp _team_name on each queued DagRun before the listener fires. + # Results are usually already cached in self._dag_id_to_team_name from the + # current or previous scheduler loop, so this is typically a dict read with no DB query. + if self._multi_team and dag_runs: + queued_team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in dag_runs}, session) + for dr in dag_runs: + if team := queued_team_map.get(dr.dag_id): + dr._team_name = team + for dag_run in dag_runs: dag_id = dag_run.dag_id run_id = dag_run.run_id @@ -2981,6 +3000,12 @@ def _schedule_dag_run( execute=False, ) + # dag_run was reloaded from DB above, so _team_name set on the original object is lost. + # Re-stamp it before the listener fires; the result is almost certainly cached in + # self._dag_id_to_team_name already, so this is a dict read with no DB query. + if self._multi_team: + if team := self._get_team_names_for_dag_ids([dag_run.dag_id], session).get(dag_run.dag_id): + dag_run._team_name = team dag_run.notify_dagrun_state_changed(msg="timed_out") if dag_run.end_date and dag_run.start_date: duration = dag_run.end_date - dag_run.start_date diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 53ea5bc933f3f..9382d7f58951b 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9635,6 +9635,205 @@ def test_dag_timeout_notifies_with_timed_out_msg(self, mock_get_listener_manager assert call_args.kwargs["msg"] == "timed_out" assert call_args.kwargs["dag_run"] == dag_run + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.models.dagrun.get_listener_manager") + def test_dag_start_notifies_listener_with_team_name(self, mock_get_listener_manager, dag_maker, session): + """Test that on_dag_run_running receives dag_run with _team_name set.""" + mock_listener_manager = MagicMock() + mock_get_listener_manager.return_value = mock_listener_manager + + clear_db_teams() + clear_db_dag_bundles() + + team = Team(name="test_team") + session.add(team) + session.flush() + + bundle = DagBundleModel(name="test_bundle") + bundle.teams.append(team) + session.add(bundle) + session.flush() + + with dag_maker(dag_id="test_dag_start_team", bundle_name="test_bundle", session=session): + EmptyOperator(task_id="test_task") + + dag_maker.create_dagrun(run_id="test_run", state=DagRunState.QUEUED) + session.commit() + + mock_executor = MagicMock() + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) + + self.job_runner._start_queued_dagruns(session) + + mock_listener_manager.hook.on_dag_run_running.assert_called_once() + call_args = mock_listener_manager.hook.on_dag_run_running.call_args + assert call_args.kwargs["dag_run"]._team_name == "test_team" + + @conf_vars({("core", "multi_team"): "true"}) + @time_machine.travel(DEFAULT_DATE, tick=False) + @mock.patch("airflow.models.dagrun.get_listener_manager") + def test_dag_timeout_notifies_listener_with_team_name( + self, mock_get_listener_manager, dag_maker, session + ): + """Test that on_dag_run_failed receives dag_run with _team_name set when a DAG times out.""" + mock_listener_manager = MagicMock() + mock_get_listener_manager.return_value = mock_listener_manager + + clear_db_teams() + clear_db_dag_bundles() + + team = Team(name="test_team") + session.add(team) + session.flush() + + bundle = DagBundleModel(name="test_bundle") + bundle.teams.append(team) + session.add(bundle) + session.flush() + + with dag_maker( + dag_id="test_dag_timeout_team", + bundle_name="test_bundle", + session=session, + dagrun_timeout=timedelta(seconds=60), + ): + EmptyOperator(task_id="test_task") + + dag_run = dag_maker.create_dagrun(run_id="test_run", state=DagRunState.RUNNING) + # We set it to double dagrun timeout so the timeout path is taken. + dag_run.start_date = DEFAULT_DATE - timedelta(seconds=120) + session.merge(dag_run) + session.commit() + + mock_executor = MagicMock() + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) + + self.job_runner._schedule_dag_run(dag_run, session) + + mock_listener_manager.hook.on_dag_run_failed.assert_called_once() + call_args = mock_listener_manager.hook.on_dag_run_failed.call_args + assert call_args.kwargs["msg"] == "timed_out" + assert call_args.kwargs["dag_run"]._team_name == "test_team" + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.models.dagrun.get_listener_manager") + def test_dag_success_notifies_listener_with_team_name( + self, mock_get_listener_manager, dag_maker, session + ): + """Test that on_dag_run_success receives dag_run with _team_name set.""" + mock_listener_manager = MagicMock() + mock_get_listener_manager.return_value = mock_listener_manager + + clear_db_teams() + clear_db_dag_bundles() + + team = Team(name="test_team") + session.add(team) + session.flush() + + bundle = DagBundleModel(name="test_bundle") + bundle.teams.append(team) + session.add(bundle) + session.flush() + + with dag_maker(dag_id="test_dag_success_team", bundle_name="test_bundle", session=session): + EmptyOperator(task_id="test_task") + + dag_run = dag_maker.create_dagrun() + ti = dag_run.get_task_instance("test_task") + ti.set_state(TaskInstanceState.SUCCESS, session=session) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[MockExecutor(do_update=False)]) + + self.job_runner._do_scheduling(session) + + mock_listener_manager.hook.on_dag_run_success.assert_called_once() + call_args = mock_listener_manager.hook.on_dag_run_success.call_args + assert call_args.kwargs["dag_run"]._team_name == "test_team" + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch("airflow.models.dagrun.get_listener_manager") + def test_dag_failure_notifies_listener_with_team_name( + self, mock_get_listener_manager, dag_maker, session + ): + """Test that on_dag_run_failed receives dag_run with _team_name set.""" + mock_listener_manager = MagicMock() + mock_get_listener_manager.return_value = mock_listener_manager + + clear_db_teams() + clear_db_dag_bundles() + + team = Team(name="test_team") + session.add(team) + session.flush() + + bundle = DagBundleModel(name="test_bundle") + bundle.teams.append(team) + session.add(bundle) + session.flush() + + with dag_maker(dag_id="test_dag_failure_team", bundle_name="test_bundle", session=session): + EmptyOperator(task_id="test_task") + + dag_run = dag_maker.create_dagrun() + ti = dag_run.get_task_instance("test_task") + ti.set_state(TaskInstanceState.FAILED, session=session) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[MockExecutor(do_update=False)]) + + self.job_runner._do_scheduling(session) + + mock_listener_manager.hook.on_dag_run_failed.assert_called_once() + call_args = mock_listener_manager.hook.on_dag_run_failed.call_args + assert call_args.kwargs["dag_run"]._team_name == "test_team" + + @conf_vars({("core", "multi_team"): "true"}) + @time_machine.travel(DEFAULT_DATE, tick=False) + @mock.patch("airflow.models.dagrun.get_listener_manager") + def test_dag_paused_success_notifies_listener_with_team_name( + self, mock_get_listener_manager, dag_maker, session + ): + """Test that on_dag_run_success receives dag_run with _team_name set for paused DAGs.""" + mock_listener_manager = MagicMock() + mock_get_listener_manager.return_value = mock_listener_manager + + clear_db_teams() + clear_db_dag_bundles() + + team = Team(name="test_team") + session.add(team) + session.flush() + + bundle = DagBundleModel(name="test_bundle") + bundle.teams.append(team) + session.add(bundle) + session.flush() + + with dag_maker(dag_id="test_dag_paused_team", bundle_name="test_bundle", session=session) as dag: + EmptyOperator(task_id="test_task") + + dag_run = dag_maker.create_dagrun() + dag_run.last_scheduling_decision = DEFAULT_DATE - timedelta(minutes=1) + ti = dag_run.get_task_instance("test_task") + ti.set_state(TaskInstanceState.SUCCESS, session=session) + dm = DagModel.get_dagmodel(dag.dag_id, session=session) + dm.is_paused = True + session.flush() + + mock_executor = MagicMock() + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) + + self.job_runner._update_dag_run_state_for_paused_dags(session=session) + + mock_listener_manager.hook.on_dag_run_success.assert_called_once() + call_args = mock_listener_manager.hook.on_dag_run_success.call_args + assert call_args.kwargs["dag_run"]._team_name == "test_team" + @mock.patch("airflow.models.Deadline.handle_miss") def test_process_expired_deadlines(self, mock_handle_miss, session, dag_maker): """Verify all expired and unhandled deadlines (and only those) are processed by the scheduler.""" From cf3986508e2f1638ac598876f229bf92df179bcc Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Fri, 31 Jul 2026 12:46:08 +0200 Subject: [PATCH 2/4] Extract logic to helper _stamp_team_names method --- .../src/airflow/jobs/scheduler_job_runner.py | 62 ++++++++----------- 1 file changed, 27 insertions(+), 35 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index af41a7800f392..2513b0f485cd9 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/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: + if team := team_map.get(dr.dag_id): + dr._team_name = team + def _get_workload_team_name(self, workload: SchedulerWorkload, session: Session) -> str | None: """ Resolve team name for a workload using the DAG > Bundle > Team relationship chain. @@ -1734,16 +1753,8 @@ def _update_dag_run_state_for_paused_dags(self, *, session: Session = NEW_SESSIO .group_by(DagRun) ) ) - # Stamp _team_name before update_state() calls notify_dagrun_state_changed(), - # which fires on_dag_run_success/failed listeners. - # _get_team_names_for_dag_ids caches results in self._dag_id_to_team_name, so - # this is usually a dict read with no DB query. - if self._multi_team and paused_runs: - paused_dag_ids = {dr.dag_id for dr in paused_runs} - paused_team_mapping = self._get_team_names_for_dag_ids(paused_dag_ids, session) - for dr in paused_runs: - if team := paused_team_mapping.get(dr.dag_id): - dr._team_name = team + # Team name should be added before listeners are called in update_state() + self._stamp_team_names(paused_runs, session) for dag_run in paused_runs: dag = self.scheduler_dag_bag.get_dag_for_run(dag_run=dag_run, session=session) if dag is not None: @@ -2000,18 +2011,8 @@ def _do_scheduling(self, session: Session) -> int: ) ) - # Stamp _team_name before _schedule_all_dag_runs calls notify_dagrun_state_changed() - # (on_dag_run_success/failed listeners). Runs that just transitioned QUEUED→RUNNING - # in _start_queued_dagruns may already have _team_name set (SQLAlchemy identity map - # returns the same Python objects); overwriting with the same value is harmless. - # _get_team_names_for_dag_ids caches results in self._dag_id_to_team_name, so this - # is usually a dict read with no DB query. - if self._multi_team and dag_runs: - unique_dag_ids = {dr.dag_id for dr in dag_runs} - dr_team_mapping = self._get_team_names_for_dag_ids(unique_dag_ids, session) - for dr in dag_runs: - if team := dr_team_mapping.get(dr.dag_id): - dr._team_name = team + # Team name should be added before listeners are called in _schedule_all_dag_runs() + self._stamp_team_names(dag_runs, session) callback_tuples = self._schedule_all_dag_runs(guard, dag_runs, session) @@ -2854,14 +2855,8 @@ def _update_state(dag: SerializedDAG, dag_run: DagRun): partial(self.scheduler_dag_bag.get_dag_for_run, session=session) ) - # Stamp _team_name on each queued DagRun before the listener fires. - # Results are usually already cached in self._dag_id_to_team_name from the - # current or previous scheduler loop, so this is typically a dict read with no DB query. - if self._multi_team and dag_runs: - queued_team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in dag_runs}, session) - for dr in dag_runs: - if team := queued_team_map.get(dr.dag_id): - dr._team_name = team + # Team name should be added before listeners are called in notify_dagrun_state_changed() + self._stamp_team_names(dag_runs, session) for dag_run in dag_runs: dag_id = dag_run.dag_id @@ -3001,11 +2996,8 @@ def _schedule_dag_run( ) # dag_run was reloaded from DB above, so _team_name set on the original object is lost. - # Re-stamp it before the listener fires; the result is almost certainly cached in - # self._dag_id_to_team_name already, so this is a dict read with no DB query. - if self._multi_team: - if team := self._get_team_names_for_dag_ids([dag_run.dag_id], session).get(dag_run.dag_id): - dag_run._team_name = team + # Team name should be added before listeners are called in notify_dagrun_state_changed() + self._stamp_team_names([dag_run], session) dag_run.notify_dagrun_state_changed(msg="timed_out") if dag_run.end_date and dag_run.start_date: duration = dag_run.end_date - dag_run.start_date From 6d276b6fa9a88613726111618835b9e4c38547aa Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 3 Aug 2026 08:27:24 -0400 Subject: [PATCH 3/4] Address review comments --- .../tests/unit/jobs/test_scheduler_job.py | 100 +++++------------- 1 file changed, 25 insertions(+), 75 deletions(-) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 9382d7f58951b..824074057a352 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -342,6 +342,14 @@ def set_instance_attrs(self) -> Generator: yield self.null_exec = None + @pytest.fixture + def team_bundle(self, testing_team, testing_dag_bundle, session): + team = session.merge(testing_team) + bundle = session.scalar(select(DagBundleModel).where(DagBundleModel.name == "testing")) + bundle.teams.append(team) + session.flush() + return bundle + @pytest.fixture def mock_executors(self): mock_jwt_generator = MagicMock(spec=JWTGenerator) @@ -9637,24 +9645,14 @@ def test_dag_timeout_notifies_with_timed_out_msg(self, mock_get_listener_manager @conf_vars({("core", "multi_team"): "true"}) @mock.patch("airflow.models.dagrun.get_listener_manager") - def test_dag_start_notifies_listener_with_team_name(self, mock_get_listener_manager, dag_maker, session): + def test_dag_start_notifies_listener_with_team_name( + self, mock_get_listener_manager, dag_maker, session, team_bundle + ): """Test that on_dag_run_running receives dag_run with _team_name set.""" mock_listener_manager = MagicMock() mock_get_listener_manager.return_value = mock_listener_manager - clear_db_teams() - clear_db_dag_bundles() - - team = Team(name="test_team") - session.add(team) - session.flush() - - bundle = DagBundleModel(name="test_bundle") - bundle.teams.append(team) - session.add(bundle) - session.flush() - - with dag_maker(dag_id="test_dag_start_team", bundle_name="test_bundle", session=session): + with dag_maker(dag_id="test_dag_start_team", bundle_name="testing", session=session): EmptyOperator(task_id="test_task") dag_maker.create_dagrun(run_id="test_run", state=DagRunState.QUEUED) @@ -9668,33 +9666,21 @@ def test_dag_start_notifies_listener_with_team_name(self, mock_get_listener_mana mock_listener_manager.hook.on_dag_run_running.assert_called_once() call_args = mock_listener_manager.hook.on_dag_run_running.call_args - assert call_args.kwargs["dag_run"]._team_name == "test_team" + assert call_args.kwargs["dag_run"]._team_name == "testing" @conf_vars({("core", "multi_team"): "true"}) @time_machine.travel(DEFAULT_DATE, tick=False) @mock.patch("airflow.models.dagrun.get_listener_manager") def test_dag_timeout_notifies_listener_with_team_name( - self, mock_get_listener_manager, dag_maker, session + self, mock_get_listener_manager, dag_maker, session, team_bundle ): """Test that on_dag_run_failed receives dag_run with _team_name set when a DAG times out.""" mock_listener_manager = MagicMock() mock_get_listener_manager.return_value = mock_listener_manager - clear_db_teams() - clear_db_dag_bundles() - - team = Team(name="test_team") - session.add(team) - session.flush() - - bundle = DagBundleModel(name="test_bundle") - bundle.teams.append(team) - session.add(bundle) - session.flush() - with dag_maker( dag_id="test_dag_timeout_team", - bundle_name="test_bundle", + bundle_name="testing", session=session, dagrun_timeout=timedelta(seconds=60), ): @@ -9715,30 +9701,18 @@ def test_dag_timeout_notifies_listener_with_team_name( mock_listener_manager.hook.on_dag_run_failed.assert_called_once() call_args = mock_listener_manager.hook.on_dag_run_failed.call_args assert call_args.kwargs["msg"] == "timed_out" - assert call_args.kwargs["dag_run"]._team_name == "test_team" + assert call_args.kwargs["dag_run"]._team_name == "testing" @conf_vars({("core", "multi_team"): "true"}) @mock.patch("airflow.models.dagrun.get_listener_manager") def test_dag_success_notifies_listener_with_team_name( - self, mock_get_listener_manager, dag_maker, session + self, mock_get_listener_manager, dag_maker, session, team_bundle ): """Test that on_dag_run_success receives dag_run with _team_name set.""" mock_listener_manager = MagicMock() mock_get_listener_manager.return_value = mock_listener_manager - clear_db_teams() - clear_db_dag_bundles() - - team = Team(name="test_team") - session.add(team) - session.flush() - - bundle = DagBundleModel(name="test_bundle") - bundle.teams.append(team) - session.add(bundle) - session.flush() - - with dag_maker(dag_id="test_dag_success_team", bundle_name="test_bundle", session=session): + with dag_maker(dag_id="test_dag_success_team", bundle_name="testing", session=session): EmptyOperator(task_id="test_task") dag_run = dag_maker.create_dagrun() @@ -9752,30 +9726,18 @@ def test_dag_success_notifies_listener_with_team_name( mock_listener_manager.hook.on_dag_run_success.assert_called_once() call_args = mock_listener_manager.hook.on_dag_run_success.call_args - assert call_args.kwargs["dag_run"]._team_name == "test_team" + assert call_args.kwargs["dag_run"]._team_name == "testing" @conf_vars({("core", "multi_team"): "true"}) @mock.patch("airflow.models.dagrun.get_listener_manager") def test_dag_failure_notifies_listener_with_team_name( - self, mock_get_listener_manager, dag_maker, session + self, mock_get_listener_manager, dag_maker, session, team_bundle ): """Test that on_dag_run_failed receives dag_run with _team_name set.""" mock_listener_manager = MagicMock() mock_get_listener_manager.return_value = mock_listener_manager - clear_db_teams() - clear_db_dag_bundles() - - team = Team(name="test_team") - session.add(team) - session.flush() - - bundle = DagBundleModel(name="test_bundle") - bundle.teams.append(team) - session.add(bundle) - session.flush() - - with dag_maker(dag_id="test_dag_failure_team", bundle_name="test_bundle", session=session): + with dag_maker(dag_id="test_dag_failure_team", bundle_name="testing", session=session): EmptyOperator(task_id="test_task") dag_run = dag_maker.create_dagrun() @@ -9789,31 +9751,19 @@ def test_dag_failure_notifies_listener_with_team_name( mock_listener_manager.hook.on_dag_run_failed.assert_called_once() call_args = mock_listener_manager.hook.on_dag_run_failed.call_args - assert call_args.kwargs["dag_run"]._team_name == "test_team" + assert call_args.kwargs["dag_run"]._team_name == "testing" @conf_vars({("core", "multi_team"): "true"}) @time_machine.travel(DEFAULT_DATE, tick=False) @mock.patch("airflow.models.dagrun.get_listener_manager") def test_dag_paused_success_notifies_listener_with_team_name( - self, mock_get_listener_manager, dag_maker, session + self, mock_get_listener_manager, dag_maker, session, team_bundle ): """Test that on_dag_run_success receives dag_run with _team_name set for paused DAGs.""" mock_listener_manager = MagicMock() mock_get_listener_manager.return_value = mock_listener_manager - clear_db_teams() - clear_db_dag_bundles() - - team = Team(name="test_team") - session.add(team) - session.flush() - - bundle = DagBundleModel(name="test_bundle") - bundle.teams.append(team) - session.add(bundle) - session.flush() - - with dag_maker(dag_id="test_dag_paused_team", bundle_name="test_bundle", session=session) as dag: + with dag_maker(dag_id="test_dag_paused_team", bundle_name="testing", session=session) as dag: EmptyOperator(task_id="test_task") dag_run = dag_maker.create_dagrun() @@ -9832,7 +9782,7 @@ def test_dag_paused_success_notifies_listener_with_team_name( mock_listener_manager.hook.on_dag_run_success.assert_called_once() call_args = mock_listener_manager.hook.on_dag_run_success.call_args - assert call_args.kwargs["dag_run"]._team_name == "test_team" + assert call_args.kwargs["dag_run"]._team_name == "testing" @mock.patch("airflow.models.Deadline.handle_miss") def test_process_expired_deadlines(self, mock_handle_miss, session, dag_maker): From 560334c51bc434b43faf455680511262bda57e47 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 4 Aug 2026 10:11:13 -0400 Subject: [PATCH 4/4] Address more review comments --- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 2513b0f485cd9..6e6d5dd78e7d3 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -458,7 +458,7 @@ 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: + def _stamp_team_names(self, dag_runs: Collection[DagRun], session: Session) -> None: """ Stamp ``_team_name`` on each DagRun. @@ -469,11 +469,10 @@ def _stamp_team_names(self, dag_runs: Iterable[DagRun], session: Session) -> Non """ if not self._multi_team: return - runs = list(dag_runs) - if not runs: + if not dag_runs: return - team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in runs}, session) - for dr in runs: + team_map = self._get_team_names_for_dag_ids({dr.dag_id for dr in dag_runs}, session) + for dr in dag_runs: if team := team_map.get(dr.dag_id): dr._team_name = team @@ -2995,7 +2994,6 @@ 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() self._stamp_team_names([dag_run], session) dag_run.notify_dagrun_state_changed(msg="timed_out")