Skip to content

Fix asset event partition_key breaking older Task SDK workers - #72327

Open
neel-astro wants to merge 1 commit into
apache:mainfrom
neel-astro:fix-asset-event-partition-key-gating
Open

Fix asset event partition_key breaking older Task SDK workers#72327
neel-astro wants to merge 1 commit into
apache:mainfrom
neel-astro:fix-asset-event-partition-key-gating

Conversation

@neel-astro

Copy link
Copy Markdown

Why

AssetEventDagRunReference.partition_key was added in 3.3.0 (#69115) but registered under the already-released 2026-04-06 Execution API version. Cadwyn applies a version's downgrade instructions only for clients older than that version, so a client requesting exactly 2026-04-06 still receives the field.

Task SDK 1.2.x (Airflow 3.2.x) speaks 2026-04-06, and its AssetEventDagRunReference has no partition_key, with extra="forbid". Every asset-triggered task run by a 3.2.x worker against a 3.3.x API server fails before it starts:

airflow/sdk/api/client.py:231 in start
    return TIRunContext.model_validate_json(resp.read())
pydantic_core._pydantic_core.ValidationError: 1 validation error for TIRunContext
dag_run.consumed_asset_events.0.partition_key
  Extra inputs are not permitted [type=extra_forbidden, input_value=None, input_type=NoneType]

The value is null, so this affects every asset-triggered run, not just partitioned assets. The task dies at supervision start.

Reading partition_key out of each released task-sdk wheel's airflow/sdk/api/datamodels/_generated.py:

task-sdk API_VERSION AssetEventDagRunReference
1.1.5 (3.1.x) 2025-11-05 absent
1.2.0 (3.2.0) 2026-04-06 absent
1.2.1 (3.2.1) 2026-04-06 absent
1.2.2 (3.2.2) 2026-04-06 absent
1.3.0 (3.3.0) 2026-06-30 present
1.3.1 (3.3.1) 2026-06-30 present

AssetEventDagRunReference is the only one of the five models in AddPartitionKeyField that did not gain the field at 2026-04-06. DagRun, AssetEventResponse, TriggerDAGRunPayload and DagRunAssetReference all carry it in 1.2.0+, so their gating is correct and is left alone.

What

Move the AssetEventDagRunReference.partition_key schema instruction and its consumed_asset_events response converter out of v2026_04_06.py into a new AddConsumedAssetEventPartitionKeyField in v2026_06_30.py, matching how the sibling AddPartitionDateField handles DagRun.partition_date (also new in 3.3.0). The HEAD datamodel and the generated Task SDK models are unchanged.

Both halves of the gate move together. The schema(...).didnt_exist instruction drives the per-version OpenAPI document and the converter mutates the wire body, so leaving either behind would ship a spec that contradicts the response.

Only TIRunContext gets a converter, where the two siblings in the same file register three each. DagRun.safe_extract_from_orm defaults consumed_asset_events to [] whenever the relationship is not already loaded, and PATCH /task-instances/{id}/run is the only route that loads it, so bare DagRun responses never carry one. Two tests in versions/v2026_06_30/test_dag_runs.py pin that invariant so it fails loudly if a future change adds eager loading there.

Tests

TestConsumedEventPartitionKeyBackwardCompat moves to the v2026_06_30 suite, where old_ver_client is pinned to exactly 2026-04-06 — the version Task SDK 1.2.x sends. In its previous home it ran at 2025-11-05, older than the mis-placed gate, which is why this passed CI.

Three assertions cover the three ways this can go wrong:

  • a 2026-04-06 client gets the event-level partition_key stripped while the Dag-run-level partition_key survives, in one response (under-moving the gate)
  • a client at exactly 2026-06-30 and a HEAD client both still receive it (over-moving the gate would silently break 1.3.x, which sends 2026-06-30 and carries the field)
  • the served OpenAPI for AssetEventDagRunReference matches the wire body at both versions (moving only one half of the gate)

Beyond the suite, the released 1.2.2 TIRunContext — the model that raised the traceback above — was loaded from its wheel and used to validate a real captured 2026-04-06 response body: it accepts the server's current output and still rejects it once partition_key is re-injected. That checks response bytes against the real released client model; it is not a live 1.2.x worker running end to end.

One pre-existing test fixture changed: old_ver_client in versions/v2026_06_30/test_task_instances.py moved from 2026-06-16 to 2026-04-06. This is behaviour-preserving — 2026-06-16 was never a registered version and resolved to the same bundle — and it aligns the module with its siblings in that directory, which already pin 2026-04-06.

Backport

3.3.0 and 3.3.1 are both released with this bug, so a main-only fix does not help anyone currently affected. Requesting the 3.3.2 milestone.

related: #69115

related: #71861 — that PR fixes a different class of gap in the same two files (missing registrations on the compat previous-run route and nested created_dagruns). The two are semantically orthogonal, but both touch remove_partition_key_from_dag_run, so whichever merges second will conflict there. The resolution is to keep the consumed_asset_events loop deleted and keep that PR's added converters; re-adding the loop alongside the new VersionChange is a silent no-op that the test suite cannot detect.


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

Generated-by: Claude Code (Opus 5) following the guidelines

@boring-cyborg boring-cyborg Bot added area:API Airflow's REST/HTTP API area:task-sdk labels Aug 31, 2026
@boring-cyborg

boring-cyborg Bot commented Aug 31, 2026

Copy link
Copy Markdown

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

  • Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
  • In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
  • Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
  • Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
  • Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
  • Be sure to read the Airflow Coding style.
  • Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
    Apache Airflow is a community-driven project and together we are making it better 🚀.
    In case of doubts contact the developers at:
    Mailing List: dev@airflow.apache.org
    Slack: https://s.apache.org/airflow-slack

@neel-astro
neel-astro force-pushed the fix-asset-event-partition-key-gating branch from cdc7ea8 to 2d70c73 Compare August 31, 2026 11:57
@neel-astro
neel-astro marked this pull request as ready for review August 31, 2026 12:02
@kaxil
kaxil force-pushed the fix-asset-event-partition-key-gating branch from 2d70c73 to 6020a2e Compare September 9, 2026 01:32
@kaxil kaxil added this to the Airflow 3.3.2 milestone Sep 9, 2026
response = old_ver_client.get("/execution/dag-runs/test_dag_run_consumed_event/run1")

assert response.status_code == 200
assert all("partition_key" not in event for event in response.json()["consumed_asset_events"])

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.

consumed_asset_events is always [] on this route: the handler runs a plain select(DagRunModel) with no joinedload, SessionDep is a fresh create_session(scoped=False) so the fixture's in-session append is not visible, and safe_extract_from_orm defaults the unloaded relationship. So this all(...) iterates nothing and would pass with the version gate deleted entirely. Same for the /previous test below. Asserting response.json()["consumed_asset_events"] == [] would pin the invariant that actually holds today, and would still fail loudly if a later change adds eager loading here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:task-sdk

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants