Skip to content

Fix missing _team_name on DagRun before some listener calls - #70760

Open
kacpermuda wants to merge 4 commits into
apache:mainfrom
kacpermuda:stamp-team-name-on-dagrun-before-listeners
Open

Fix missing _team_name on DagRun before some listener calls#70760
kacpermuda wants to merge 4 commits into
apache:mainfrom
kacpermuda:stamp-team-name-on-dagrun-before-listeners

Conversation

@kacpermuda

@kacpermuda kacpermuda commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

In multi-team deployments the scheduler fires on_dag_run_running / success / failed listener hooks from several code paths. Before this change, DagRun._team_name was not guaranteed to be set when those hooks fired, so OpenLineage (and any other plugin reading _team_name) would miss the team on some events.

_team_name is a private attribute — it is not part of any public API and is not guaranteed to be present on every DagRun object a listener receives. Since the value is already cached in self._dag_id_to_team_name for metrics purposes, setting it consistently before each hook fires adds no cost. Listeners that wish to use it should treat it as best-effort and guard accordingly.

This PR ensures _team_name is stamped on every ORM DagRun object before notify_dagrun_state_changed() is called, covering all five listener code paths. The stamping is cheap: _get_team_names_for_dag_ids caches results in self._dag_id_to_team_name for the lifetime of a scheduler loop, so most calls are a dict read with no DB query. I've also added comments to the existing stamping paths, just to make sure it's clear why it's done before listener call.


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Sonnet 4.6)

Generated-by: Claude Code (Sonnet 4.6) following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@kacpermuda
kacpermuda requested review from XD-DENG and ashb as code owners July 30, 2026 12:02
@boring-cyborg boring-cyborg Bot added the area:Scheduler including HA (high availability) scheduler label Jul 30, 2026
@kacpermuda

Copy link
Copy Markdown
Collaborator Author

cc @ferruzzi , FYI as I see that you worked on adding this before

@kacpermuda
kacpermuda force-pushed the stamp-team-name-on-dagrun-before-listeners branch from a0393d4 to 45cf7c3 Compare July 30, 2026 13:08
@kacpermuda kacpermuda changed the title fix: Add _team_name on DagRun consistently before every listener call Fix missing _team_name on DagRun before some listener calls Jul 30, 2026
@kacpermuda
kacpermuda force-pushed the stamp-team-name-on-dagrun-before-listeners branch from 45cf7c3 to e83184a Compare July 30, 2026 14:27
@eladkal eladkal added this to the Airflow 3.3.1 milestone Jul 30, 2026
@eladkal eladkal added type:bug-fix Changelog: Bug Fixes backport-to-v3-3-test Backport to v3-3-test labels Jul 30, 2026
@uranusjr

Copy link
Copy Markdown
Member

I would introduce a helper (say _stamp_team_names) to do the stamping instead of repeating the same logic multiple times. It would also help clean up the verbose comment added on every stamping block.

@kacpermuda
kacpermuda force-pushed the stamp-team-name-on-dagrun-before-listeners branch from e83184a to 9fe1f70 Compare July 31, 2026 10:46
@kacpermuda

Copy link
Copy Markdown
Collaborator Author

I would introduce a helper (say _stamp_team_names) to do the stamping instead of repeating the same logic multiple times. It would also help clean up the verbose comment added on every stamping block.

@uranusjr Extracted the logic to a separate helper, ptal

Comment on lines +9645 to +9646
clear_db_teams()
clear_db_dag_bundles()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clearing the db at the beginning of the tests should be a code smell. I believe this issue was fixed in #69093, if not, then we need to figure out what tests are not cleaning up after themselves and fix them.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct, I should have just used the fixtures currently present. Adjusted the code.

bundle = DagBundleModel(name="test_bundle")
bundle.teams.append(team)
session.add(bundle)
session.flush()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These tests all have a pretty big repetitive block at the top, maybe move the setup into a shared fixture if none of the existing fixtures do what you need? look at the testing_team fixture though, I think that covers most/all of this boilerplate.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct, used testing_team and testing_dag_bundle fixtures instead, and linked them together in a new small fixture that I reused in all tests. Should be clean now, thanks for spotting that.

@vatsrahul1001

Copy link
Copy Markdown
Contributor

@kacpermuda can you resolve comments?

@kacpermuda
kacpermuda force-pushed the stamp-team-name-on-dagrun-before-listeners branch 2 times, most recently from 9fdc186 to 21524ac Compare August 3, 2026 12:27
@kacpermuda

Copy link
Copy Markdown
Collaborator Author

@vatsrahul1001 addressed all the comments

@vatsrahul1001

Copy link
Copy Markdown
Contributor

LGTM!, can be merged after code owners review

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, LGTM overall.

Comment thread airflow-core/src/airflow/jobs/scheduler_job_runner.py Outdated
Comment thread airflow-core/src/airflow/jobs/scheduler_job_runner.py Outdated
Comment thread airflow-core/src/airflow/jobs/scheduler_job_runner.py Outdated
@kacpermuda
kacpermuda force-pushed the stamp-team-name-on-dagrun-before-listeners branch from 21524ac to 560334c Compare August 4, 2026 14:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:Scheduler including HA (high availability) scheduler backport-to-v3-3-test Backport to v3-3-test type:bug-fix Changelog: Bug Fixes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants