Deliver value_schema on arg bindings for mapped stub tasks
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
### Background
https://github.com/apache/airflow/pull/69757 lets `@task.stub` tasks be called TaskFlow-style: a parse-time capture builds the arg-binding spec (parameter `name`, `kind`, `value`, and a pydantic-generated `value_schema` JSON-schema fragment inferred from the stub function's annotations) and ships it to lang-SDK runtimes via `TIRunContext.arg_bindings`.
Mapped stubs (`.expand()`) never instantiate `_StubOperator` at parse time — `_TaskDecorator._expand()` builds a generic `DecoratedMappedOperator` — so no spec lands in the serialized Dag. Instead ti_run derives per-map-index bindings server-side in `airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py` (`_resolve_mapped_stub_arg_bindings`). The stub function's annotations are not available on the API server, so derived bindings omit `value_schema` and foreign runtimes fall back to decode-only checks.
### What needs to happen
1. Add a capture point at `.expand()` time in the task-sdk decorator machinery (operator-class hook in `_TaskDecorator._expand()`, or at mapped-operator serialization time in the Dag processor, where the callable is importable) so the stub's per-parameter schemas (`_infer_value_schema` in `providers/standard/src/airflow/providers/standard/decorators/stub.py`) ride with the mapped operator.
2. Serialize the captured schema map on mapped operators (new serialized field + `airflow-core/src/airflow/serialization/schema.json` entry).
3. Merge the schemas into the server-derived bindings in `_resolve_mapped_stub_arg_bindings` so mapped bindings carry `value_schema` like unmapped ones.
4. Degrade gracefully on Airflow 2.x installs of the standard provider (no hook/pydantic → bindings omit `value_schema`, today's behavior).
5. Tests: task-sdk expand capture, Dag serialization round trip, execution API ti_run response.
### Acceptance criteria
- A mapped `@task.stub` with annotated parameters receives `value_schema` on each derived binding in `TIRunContext.arg_bindings`.
- The limitation note + tracking link in `_resolve_mapped_stub_arg_bindings` are removed.
- The omission contract (absent key = unconstrained) and unmapped behavior are unchanged.
### Context
- Originating PR: https://github.com/apache/airflow/pull/69757
- Review thread: https://github.com/apache/airflow/pull/69757#discussion_r3654919186
- Parent effort: #66937
Contributor guide
Assessment
This issue has not been assessed yet.