Vamsi-klu opened a new pull request, #71999:
URL: https://github.com/apache/airflow/pull/71999
# 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.
```
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?
- [x] Yes (Grok 4.6)
Generated-by: Grok 4.6 following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
