Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions task-sdk/src/airflow/sdk/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,14 @@ class DagRunProtocol(Protocol):
partition_key: str | None
partition_date: AwareDatetime | None
note: str | None
# Semantically ``team_name`` is a property of the Dag Bundle (Dag, DagRun, etc resolve it via the
# owning bundle) — every run of a given Dag has the same
# ``team_name``. It is denormalized onto the run here because the
# Execution API already includes it in the ``DagRun`` payload delivered
# as part of ``TIRunContext`` at task start, so exposing it on the run
# protocol keeps the SDK in sync with the server schema and saves the
# worker an extra ``get_dag()`` round trip.
team_name: str | None

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.

This feels like the wrong place to expose this. team_name is not a property of the dag run, but of the DAG.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@ashb You're right that team_name is semantically a property of the Dag (resolved through its bundle), not of the run - every DagRun of the same Dag has the same team_name, and there's no team_name column on the DagRun ORM either.

Where it does already live is the server-side Execution API DagRun payload - it was added there (see AddTeamNameField in execution_api/versions/v2026_06_16.py) so it can be delivered as part of TIRunContext at task start, saving the worker an extra get_dag() round trip. This PR is just keeping DagRunProtocol in sync with that server schema so SDK consumers reading dag_run.team_name don't hit a type-checker failure.

I've pushed a comment on the field that spells this out - the denormalization is a deliberate choice at the server-schema layer, not something we're claiming is the "true" home of the field. If you'd rather see this promoted to DagResponse (and thus DagResult on the SDK side) with team_name removed from DagRun, happy to open a follow-up that does that server-side refactor - but I think it's worth keeping the SDK protocol in sync with the server as long as the server exposes it here.

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.

Yeah, basically this was the cleanest way to plumb the team_name to the context on the worker side.

This feels like the wrong place to expose this. team_name is not a property of the dag run, but of the DAG.

This isn't even true, I wish it was, but the community would not allow it to be a property of DAG or Task (or indeed DagRun) when designing the multi team feature. It is only a property of the Dag Bundle, everything else (Dag, Dag Run, task, etc) all just infer from Bundle.

dag_run sub model on the TIRunContext already has conf, notes, triggering_user_name, etc and it was the closest thing that maps to team_name. And adding a whole new model just for team_name is way over kill. And there is already code that iterates over dag_run and exports the AIRFLOW_CTX_* env vars.

tl;dr: It was the closest match without adding hundreds of lines of code and new mechanisms just for team, this PR is just updating typing.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@ashb Given the design context @o-nikolas laid out above - team lives on the bundle and DagRun sub-model on TIRunContext is the pragmatic carrier that already ships from the server - this PR is purely a typing sync with the already-shipped Execution API schema. If moving team_name off DagRun server-side is desirable, happy to file a follow-up issue. If you're OK with landing this typing sync as-is, could you resolve the thread?



class RuntimeTaskInstanceProtocol(Protocol):
Expand Down