Skip to content

Add partition_key to Metadata so yield matches outlet_events - #71993

Closed
Vamsi-klu wants to merge 3 commits into
apache:mainfrom
Vamsi-klu:pr3/69117-metadata-partition-parity
Closed

Add partition_key to Metadata so yield matches outlet_events#71993
Vamsi-klu wants to merge 3 commits into
apache:mainfrom
Vamsi-klu:pr3/69117-metadata-partition-parity

Conversation

@Vamsi-klu

@Vamsi-klu Vamsi-klu commented Aug 23, 2026

Copy link
Copy Markdown
Contributor

Add partition_key to Metadata so yield matches outlet_events

closes: #69117

What I did

Added optional Metadata.partition_key. Sync and async runners call add_partitions when it is set. Extra still merges and never becomes a key. Docs say the two APIs agree for extra, and that a partitioned consumer needs an explicit key. Newsfragment 71993.bugfix.rst.

yield Metadata(asset, extra={"row_count": 1}, partition_key="us")
# same as
outlet_events[asset].extra = {"row_count": 1}
outlet_events[asset].add_partitions("us")

Why I did

Docs said the two APIs were the same. For extra they are. For partitions they were not. outlet_events[ASSET].extra = {"section": "XY"} does not select a partition, so a partitioned consumer gets no APDR. Reporter also hit outlet_events=None as a positional parameter. Injection already worked; the test keeps it that way.

How I did

Metadata.partition_key: str | None = None is the last field, so Metadata(asset, extra) still works. After extra.update, outlet_events[asset].add_partitions(partition_key). Invalid keys fail the TI. Alias emit does not fan out keys; partition_key is recorded on the concrete asset accessor only. Alias-only emit still uses the producing Dag run's key.

Same partition key still batches onto one APDR. Distinct keys still make distinct runs. No scheduler change.

What's the impact

Users can yield a partition key the same way they call add_partitions. Extra-only emit to a partitioned consumer still warns and does not create an empty Dag run.

What's the testing

task-sdk/tests/task_sdk/execution_time/test_callback_runner.py

  • test_two_yields_merge_extra_and_collect_partition_keys
  • test_add_partitions_list_matches_two_metadata_yields
  • test_extra_does_not_imply_partition_key
  • test_invalid_partition_key_raises
  • test_alias_and_partition_key_records_key_on_asset
  • test_alias_as_metadata_asset_with_partition_key_raises_type_error

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

  • test_emits_shared_extra_on_each_partition_key

airflow-core/tests/unit/models/test_taskinstance.py

  • test_mapped_outlet_events_extra_injection
  • test_mapped_metadata_partition_key_and_extra_per_ti
  • test_metadata_invalid_partition_key_fails_task_without_event

airflow-core/tests/unit/jobs/test_scheduler_job.py

  • test_mapped_producer_partition_keys_match_between_apis
  • test_extra_only_mapped_emit_does_not_create_apdr
  • test_runtime_partition_key_without_partition_date_still_queues
  • test_partitioned_event_does_not_queue_non_partitioned_consumer
  • test_multi_partition_payloads_share_merged_extra
uv run --project task-sdk pytest task-sdk/tests/task_sdk/execution_time/test_callback_runner.py \
  task-sdk/tests/task_sdk/execution_time/test_task_runner.py -k 'serialize_outlet or partition or Metadata or add_partitions' -q

uv run --project airflow-core pytest airflow-core/tests/unit/models/test_taskinstance.py \
  airflow-core/tests/unit/jobs/test_scheduler_job.py \
  -k 'outlet_events or partition_key or Metadata or mapped_producer_partition' -q

Was generative AI tooling used to co-author this PR?
  • Yes (Grok 4.6)

Generated-by: Grok 4.6 following the guidelines

Docs treated yield Metadata and outlet_events as the same API. Extra
already merged onto the same accessor, but partitioned consumers need
an explicit key and extra was never a substitute.

closes: apache#69117
The upgrade note pointed at an unmerged batching PR, alias docs did not
say where partition_key is recorded, and one ADRQ assertion counted the
whole table.
@boring-cyborg boring-cyborg Bot added area:Scheduler including HA (high availability) scheduler area:task-sdk kind:documentation labels Aug 23, 2026
@Vamsi-klu

Vamsi-klu commented Aug 23, 2026

Copy link
Copy Markdown
Contributor Author

cc @dstandish @jroachgolf84 @ashb @kaxil @Lee-W @uranusjr — Metadata.partition_key so yield matches outlet_events.add_partitions.


Drafted-by: Cursor Grok 4.6

The check rejects issue-numbered fragments so the changelog entry
links to the change that actually merged.

Co-authored-by: Cursor <cursoragent@cursor.com>
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 area:task-sdk kind:documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Inconsistency between yielding Metadata and outlet_events extras for partitioned Assets

1 participant