Bind mapped stub-task arguments in the Go SDK runtime - #70571
Draft
jason810496 wants to merge 55 commits into
Draft
Bind mapped stub-task arguments in the Go SDK runtime#70571jason810496 wants to merge 55 commits into
jason810496 wants to merge 55 commits into
Conversation
Stub Dags could only declare argless tasks, so cross-language dataflow required hand-written GetXCom calls inside each Go task. Capturing the TaskFlow call's argument spec at parse time and delivering it through the Execution API and StartupDetails lets a Go task receive upstream outputs and Dag-file literals as plain typed parameters, with loud arity/type errors instead of silently zero-filled values.
"stub_args" leaked the _StubOperator implementation detail into the wire contract that foreign-language SDKs code-generate against; "arg bindings" names what the data actually is -- the ordered spec a runtime binds onto the task function. Renaming now, before the field ships in a released execution API or supervisor schema version, keeps the contract clean without any compatibility shims.
The single try/except made Airflow 3.0 (whose SDK predates KNOWN_CONTEXT_KEYS) fall back to the Airflow 2 import paths and fail; the arg-capture tests imported airflow.sdk directly, which does not exist on 2.11; and the .expand() rejection relies on the supports_expand opt-out that only ships with Airflow 3.4. Also reword the context-key rejection to stop implying foreign runtimes have no task context -- the lang SDKs inject their own natively; stub signatures just must not declare Airflow context parameters.
Only the Multi-Lang stub-task path needs the serialized-dag machinery and the arg-binding models, so regular task-run requests should not pay for them: the TaskArgBinding datamodels move to a dedicated module and the serialized-dag imports become local to the stub lookup. The OpenAPI schema is unchanged (component names stay the same), which is why no execution API version bump accompanies this commit.
simple_dag only exercises the minimal binding: one literal and one XCom argument. The new taskflow_binding_dag locks in the rest of the surface end to end -- scalar and array literals, keyword arguments, a defaulted None, and XCom fan-in from two upstream Go tasks bound onto a strict struct and a typed slice -- with the Go task verifying every bound value so binding regressions fail the example run loudly.
Naming every stub argument as a separate flat Go parameter gets unwieldy as the argument count grows, and there was no way for a Go task to pull an XCom that the Python TaskFlow call itself never passed. A struct that embeds sdk.TaskInput lets a task bind many arguments by name (or an explicit ad hoc XCom pull) onto one parameter instead, while the existing flat/positional binding keeps working unchanged for functions that don't opt in. This required adding a name to the wire-level TaskArgBinding spec so a struct field can look itself up by the Dag's TaskFlow argument name regardless of declaration order on either side, since Go cannot recover a plain function parameter's name via reflection the way it can for a struct's fields.
These tests built stub tasks with the raw stub(fn)(...) call instead of the @task.stub decorator every real Dag (Go/TS/Java examples) already uses, so a reader comparing the tests to real usage saw a syntax the feature doesn't actually ship.
As a plain Literal type alias, the field's generated model came out under a generic, field-derived name (DataType) in both the task-sdk client model and the Go SDK's generated types, rather than the ArgBindingDataType name declared in the source. A real Enum class carries its own name through codegen, so providers/standard can import it directly instead of re-deriving the same string vocabulary by hand, with a hand-written fallback for Airflow 2 where the execution-API generated models aren't importable.
The combined TaskInput example mixed all three field-binding modes (arg: tag, no tag, xcom: tag) into one struct, so no single task demonstrated any one mode in isolation. Split it into via_struct_no_tags, via_struct_arg_tag, and via_struct_xcom_tag, and renamed combine to via_flat_args to make the positional/keyword-style split between flat and struct binding legible at the call-site naming level. A TaskInput struct field whose name has no matching TaskFlow call argument now stays at its Go zero value instead of failing the task -- keyword-argument semantics (an unpassed name falls back to its default) rather than the strict arity check flat, positional parameters get. via_struct_unmatched_arg exercises this directly.
The flat kind-discriminated shape kept foreign-language codegen simple but left each variant's contract implicit: task_id was nullable even though every xcom binding has one, and value/key were dead weight on the opposite kind. Modelling arg_bindings as a kind-discriminated union makes the contracts explicit on every wire (OpenAPI, supervisor schema, Go, TypeScript) - xcom bindings now require task_id - and lets the Go runtime mirror the split as a sealed sum type instead of branching on a string field, so malformed specs fail the task before its body runs.
An ad hoc `xcom:"<task-id>"` pull baked the upstream task id into the compiled Go binary, hiding a data dependency from the Dag file that owns task wiring on the Python side (the example even needed a manual >> to order the pull's upstream). Fields now bind exclusively by argument name -- an `arg:"<name>"` tag, or the snake_cased field name when the tag is omitted -- so every value a task consumes stays visible in its TaskFlow call, and a task that needs an extra XCom can still ask for it explicitly through the injected client.
The snake_cased fallback silently rewrote Go field names into wire argument names, hiding the cross-language mapping from the reader; and because an unmatched TaskInput field kwarg-style falls back to its zero value, a wrong guess about the conversion never failed loudly. Matching the field name verbatim removes that magic: every snake_case Python parameter a field binds is now spelled out as an explicit `arg:` tag in the Go source. The e2e module also still referenced the via_struct_xcom_tag task removed with the xcom struct tag, which would have failed the suite against the current Dag.
Most workloads are not stub operators, so constructing the discriminated-union adapter at module import made every execution API process pay for it up front. Moving it next to the TaskArgBinding models behind a cached getter defers the cost to the first stub-task run and leaves the _STUB_TASK_TYPE gate as the only stub-specific module-level state in the route.
The XCom key was always return_value for a TaskFlow call, so the key field carried no information; it is removed end to end (datamodel, serialized spec, supervisor schema, generated task-sdk and Go models) and indexing a stub argument by a custom key now fails at parse time instead of being silently representable. Mixing flat positional data parameters with a TaskInput struct in one Go task signature was too ambiguous to reason about, so Analyze now rejects it: a function declares one binding shape or the other. The Go binding sum type and its DataType vocabulary are now defined in terms of the generated supervisor-schema models rather than hand-written mirrors, so they cannot drift from the wire contract, and every via_struct_* example task now binds an XCom-sourced argument (make_region) alongside a literal so struct-field binding is exercised with both sources end to end. The TypeScript supervisor model bump is left out of this PR on purpose.
The 2026-06-30 execution API version already shipped in Airflow 3.3.0, so appending the arg_bindings migration to it would mutate a released version, which the execution API versioning policy forbids; the change now opens version 2026-07-30, matching the supervisor-schema date. The rest addresses a local multi-reviewer audit of the branch: - The airflow-go-pack integration test's expected manifest was missing the make_region task added to the example bundle, failing go test. - The per-field cadwyn didnt_exist instructions on XComArgBinding and LiteralArgBinding name fields could never apply (arg_bindings is stripped wholesale on downgrade) and are dropped on both the execution API and the supervisor schema side; the supervisor-schema change class is renamed so the two same-named migrations cannot be confused. - The stub decorator's hand-rolled version-split imports now go through the common.compat sdk seam (new PlainXComArg and KNOWN_CONTEXT_KEYS exports). - The Go runtime validates required wire-spec fields (name, xcom task_id) instead of silently binding empty strings, populates the carried Kind discriminant, and reports a binding bookkeeping bug as a task error instead of panicking the worker. - The supports_expand opt-out is now covered by task-sdk-level tests, the ti_run serialized-dag scan moved onto LazyDeserializedDAG next to its sibling accessors, and assorted review nits (exception types, stale wording, enum comparisons) are fixed.
A multi-angle review of the branch surfaced gaps at the edges of the new binding contract: - An unrecognized serialized spec escaped ti_run as an opaque 500 on provider/core version skew; it now returns a structured invalid_arg_bindings error, per the route-boundary convention. - The new parse-time signature checks broke previously importable argless stub Dags (e.g. a **kwargs or ti parameter); they now fire only when a TaskFlow call actually passes arguments. - Stubs called with arguments inside a mapped task group serialized a spec with no map-index dimension and failed (or mis-bound) at runtime; they are now rejected at parse time. - A Go TaskInput struct had to mirror every stub parameter -- including defaulted ones the author never passed -- or fail every run, while an empty spec silently zero-filled the whole struct, contradicting the documented fail-loud behavior. Literal entries captured from signature defaults now carry from_default on the wire (inside the still-in-progress 2026-07-30 schema) and may go unclaimed; a spec that never arrives fails loudly when the struct declares bindable fields. - NaN/Infinity literals passed the parse-time JSON check only to fail far away (or silently bind 0.0); json.dumps now rejects them. - A malformed spec bypassed ShouldRetry while an equally permanent arity error retried; both now share retry semantics. - ti_run no longer issues two queries and re-parses the serialized-Dag blob on every stub-task start: single joinedload query plus a per-(dag_version, task) cache of the immutable extracted spec, and XCom pulls in the Go binding path now run concurrently.
Reviewing the Python-side arg-binding contract and the Go runtime that consumes it in one PR ties the core/task-sdk review to Go SDK internals. Scoping this PR to the contract lets it merge on its own; the Go SDK consumption (pkg/binding, task-runner dispatch, example bundle, e2e test) lands stacked on top from feature/go-sdk/taskflow-arg-binding.
The API server already holds a DBDagBag with configurable LRU+TTL caching of deserialized Dags; a route-local raw-blob accessor plus a hand-rolled module-level cache duplicated that machinery with its own eviction story. The serialized _arg_bindings field survives full deserialization onto the task object, so ti_run can read it off the dag_bag-resolved Dag version directly, and the LazyDeserializedDAG.get_task_arg_bindings accessor goes away.
The from_default flag records provenance, not value equality: the Go TaskInput struct mode fails unclaimed explicit arguments but tolerates unclaimed defaults, so an author-passed value that happens to equal the signature default must not be flagged. Pin that boundary at the capture site.
The supervisor schema bump ships in this PR, and the version-pin guards (TestSupervisorSchemaVersionMatchesSnapshot, the ts-sdk supervisor schema hook) rightly insist the pinned constant and the generated TypeScript models follow the schema in the same change. The Go runtime consumption of the new arg_bindings field stays in the stacked feature/go-sdk/taskflow-arg-binding branch.
Review feedback on the arg-binding contract asked to reuse JSON Schema rather than invent an ArgBindingDataType vocabulary, so foreign runtimes can validate bound values with plain JSON-schema semantics. Each binding now ships an optional value_schema fragment carrying the standard type and format keywords: unions and Optionals map to type lists instead of degrading to "any", and int/float/datetime/date/time/timedelta annotations gain int64/double/date-time/date/time/duration formats the bare type name cannot convey. An unconstrained argument omits the field entirely, and unknown keywords from newer providers are tolerated. The kind discriminator stays required: giving it a server-side default drops it from the OpenAPI required list and datamodel-codegen then emits Literal | None, which pydantic rejects as a tagged-union discriminator. Also retargets the execution API and supervisor schema version to 2026-10-30 to match the Airflow 3.4 release train, documents the optional _arg_bindings task property in the serialized-Dag schema (no serializer version bump: optional field, unchanged logic), adds the Dag version id to the ti_run spec-validation failure log, and rewords the stub .expand() rejection rationale: arg types are uniform across map indexes, the blocker is per-index value resolution at runtime.
…d mapper The hand-written annotation-to-fragment mapper duplicated what pydantic already exposes publicly: TypeAdapter(annotation).json_schema(), with a GenerateJsonSchema subclass layering the int64/double numeric formats a foreign runtime needs. Delegating to pydantic removes the bespoke union logic and buys richer, standard fragments for free -- anyOf for unions (per-member formats survive mixed unions now), items/additionalProperties for parameterized containers, and enum for Literal annotations -- while anything pydantic cannot schema (arbitrary classes anywhere in the annotation) still degrades to a decode-only binding. value_schema becomes a free-form JSON object on the wire instead of a typed model: a typed model silently strips every keyword it does not know when the spec is re-serialized along the server-to-supervisor delivery path, which would corrupt exactly the open-vocabulary fragments this contract promises to carry verbatim. The provider declares its now-direct pydantic dependency (edge3/http precedent). Trade-offs accepted: fragment shapes follow the pydantic version active at parse time, and pendulum.DateTime annotations degrade to decode-only since pydantic has no schema for arbitrary datetime subclasses. No new execution-API version: the field's shape changes inside the still-unreleased 2026-10-30 version this PR introduces.
pendulum.DateTime is the most common temporal annotation in Airflow Dag code, but pydantic cannot generate schemas for datetime subclasses, so stub arg specs silently degraded to decode-only checks. Temporal subclasses now normalize to their stdlib bases (recursively through unions and containers) before schema generation. Also drop the standard provider's explicit pydantic dependency (apache-airflow already provides it), trim the verbose comments introduced with the arg-binding spec, and shorten the invalid-arg-bindings error message.
…annotation Airflow 2.x base installs do not ship pydantic (it is an optional extra there), so after dropping the standard provider's explicit pydantic dependency, importing the stub decorator would crash Dag parsing on such installs. Value schemas now degrade to the decode-only fallback the wire contract already supports when pydantic is absent. Annotations pydantic can validate but not schema-ify (e.g. Callable) raise PydanticInvalidForJsonSchema, which escaped the existing handler and crashed Dag parsing; it is now caught alongside PydanticSchemaGenerationError. Temporal normalization now runs only as a retry after direct schema generation fails, so temporal subclasses that carry their own pydantic schema keep it.
The cross-arch pack integration test asserted a hardcoded schema date, so bumping the Go SDK's supervisor schema version (2026-06-16 -> 2026-10-30 on this branch) broke it in CI while the unit tests, which inject the version explicitly, kept passing. The packed binary and the test compile from the same module tree, so referencing execution.SupervisorSchemaVersion asserts the same value the binary embeds and future bumps cannot drift.
Banning .expand() on @task.stub blocked a core dynamic-mapping pattern for foreign-runtime Dags with no workaround. A mapped stub never instantiates at parse time, so instead of a parse-time capture, ti_run now derives the per-map-index arg spec from the serialized expand input, mirroring the task-sdk's DictOfListsExpandInput index decomposition: literal expands resolve to their element server-side, expands over a mapped upstream bind that upstream's XCom row via the new map_index field, and expands over an unmapped upstream's output carry the new element_index field so the runtime picks the right element of the pulled list. Value schemas stay parse-time-only and are omitted for mapped stubs, falling back to the decode-only contract. The derivation grew into enough business logic that it lives in a new execution_api services package (mirroring core_api's services layout) rather than the routes module, and the generic map-index decomposition sits on SchedulerDictOfListsExpandInput beside its map-length helpers, mirroring where the task-sdk twin keeps the same arithmetic. The supports_expand opt-out this branch added to the task-sdk decorator machinery existed only for the stub ban, so it is reverted. Stubs with arguments inside a mapped task group stay rejected: those instances have no expand input of their own to derive bindings from.
On Python 3.10, isinstance(list[X], type) is True and issubclass on the alias silently consults the origin, so the plain-class branch swallowed parametrized generics before the origin/args reconstruction could rewrite their arguments; list[pendulum.DateTime] then degraded to no value schema at all. Python 3.11+ returns False there, which is why the regression only surfaced on the 3.10 CI jobs. Detecting parametrized generics first restores the normalization on every supported version.
The map_index and element_index fields added to XComArgBinding in the supervisor schema must land together with the generated TypeScript output, which the check-ts-sdk-supervisor-schema static check enforces by regenerating and diffing the file.
Split out of the arg-binding contract PR (feature/lang-sdk/ taskflow-stub-dag) so the Python-side wire model can merge on its own review track; this stacked branch carries the Go SDK consumption of TIRunContext.arg_bindings and the end-to-end coverage.
Task authors previously had to embed the zero-size sdk.TaskInput marker to opt a struct into per-field, name-based TaskFlow argument binding. The marker was redundant: a function whose sole data parameter is a struct is unambiguous, and the per-execution argument spec (argument names, arity, and from_default) is enough to choose between binding fields by name and decoding one argument whole. Removing it drops boilerplate every keyword-style task had to carry and that had no analogue in the Python or Java SDKs.
Clearing only the upstream of a queued mapped stub and re-running it to an empty list records a TaskMap length of 0 while the expanded TI still exists; decomposing its map index then divided by zero and surfaced as an opaque catch-all 500. The task-sdk twin of this arithmetic guards mapped lengths below 1, so mirror it and route the failure through the structured invalid_arg_bindings error like every other undeliverable binding.
ti_run derived arg bindings for every stub task regardless of the client's negotiated API version, so a stub Dag using a construct the derivation rejects (e.g. expand_kwargs) went from running with its args ignored to hard-failing with a 500 after a server upgrade -- even for clients whose responses have arg_bindings stripped anyway. The cadwyn migration only pops the field from successful responses; it cannot gate the computation, so consult the negotiated version before deriving.
A mapped stub never instantiates at parse time, so ti_run derived its arg bindings blind to the stub signature: the spec came out in call-site dict order while the wire contract promises declaration order (a positional binder like the Go SDK's flat mode then receives swapped values), and parameters filled from signature defaults were silently dropped where the unmapped path ships from_default entries. Declaration order, defaults, and value schemas can only come from the real function, which exists nowhere but the Dag processor, so expose a classmethod the core serializer can call while serializing the mapped operator (wired up in a follow-up commit). Building the metadata also validates the mapping at parse time, so expand_kwargs() on a parameterful stub, partial() kwargs over a mapped upstream, and mappings that do not bind to the signature fail as Dag import errors instead of per-TI 500s at run time.
The server-side derivation for mapped stubs was blind to the stub signature: it emitted the spec in call-site dict order while the wire contract promises declaration order (a positional binder like the Go SDK's flat mode then receives silently swapped values), dropped parameters filled from signature defaults where the unmapped path ships from_default entries, and could not attach value schemas. The Dag serializer now consults an optional operator-class hook while serializing a mapped operator -- the one point where operator_class and python_callable are still the real objects -- and stores the stub's per-parameter metadata under _mapped_arg_binding_params. ti_run walks that metadata in declaration order, fills expanded, partial, and defaulted parameters alike, and carries each parameter's value schema, closing the mapped/unmapped contract gap (apache#70523). Dags serialized without the metadata (an older provider) keep the legacy ignored-args behavior instead of receiving order-uncertain bindings, and the old derivation's rejections stay as backstops for such Dags.
TypeAdapter construction is one of pydantic's most expensive operations and ran fresh for every annotated stub parameter on every Dag file re-parse, which the Dag processor repeats continuously. Annotations are static, so cache the generated fragment per annotation for the process lifetime, deep-copying on the way out so embedded specs never alias the cache, and falling back to uncached generation for unhashable annotations.
Both binding validators duplicated the canonical XCom return-value key as a string literal, evading constant-based refactors and diverging from the neighboring code (serialization's xcom_arg already compares against XCOM_RETURN_KEY). The constant is importable in both contexts: common.compat.sdk re-exports it for the provider, airflow.models.xcom for core.
Reviewing unmapped TaskFlow delivery and per-map-index derivation together made the PR hard to land, so this PR narrows to the unmapped contract: mapped (.expand()) stubs keep the released ignored-args behavior (they capture no parse-time spec, so ti_run naturally delivers no bindings), documented on the stub decorator. The mapped derivation -- the serializer capture hook, per-parameter metadata, map-index decomposition, and their tests -- moves wholesale to the stacked follow-up branch feature/lang-sdk/taskflow-stub-dag-mapped.
XComArgBinding carried map_index and element_index for the per-map-index delivery that now lands in the stacked follow-up branch; the unmapped path never sets either, so this PR ships the contract without them. The task-sdk client models, supervisor schema snapshot, and ts-sdk types are regenerated accordingly; the follow-up re-adds the fields with its derivation.
Restore the mapped (.expand()) stub arg-binding support split out of the unmapped PR: the Dag serializer captures per-parameter metadata (declaration order, defaults, value schemas) from the stub signature via the get_mapped_serialized_fields hook, ti_run derives per-map-index bindings from it with the map-index decomposition on SchedulerDictOfListsExpandInput, and XComArgBinding regains the map_index/element_index delivery fields across the task-sdk and ts-sdk generated models.
The unmapped PR narrows to delivering plain TaskFlow call arguments; consuming per-map-index bindings (map_index row selection, element_index extraction), the via_expand example and e2e coverage, and the has_mapped_dependants/mapped_length supervisor recording move wholesale to the stacked follow-up branch feature/go-sdk/taskflow-arg-binding-mapped.
The supervisor schema no longer carries map_index/element_index on XComArgBinding (they move to the mapped follow-up), so the generated models and the binding package's scope note follow suit.
…/go-sdk/taskflow-arg-binding-mapped
Restore the mapped delivery split out of the unmapped PR: the binding package pulls the upstream row a map_index selects and extracts the element an element_index addresses, the via_expand example and e2e coverage exercise .expand() over Go TaskFlow arguments end to end, and the supervisor records mapped_length for Go stub tasks feeding dynamic task mapping (has_mapped_dependants on TIRunContext). ti_run keeps the negotiated-version gate so pre-arg-bindings clients skip the derivation.
boring-cyborg
Bot
added
area:API
This was referenced Jul 28, 2026
This was referenced Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
The mapped lang-SDK PR makes the execution API derive a per-map-index arg-binding spec for a mapped
@task.stub. This PR makes the Go SDK actually consume it: a stub expanded with.expand()delivers its element to the native Go function, and a Go stub's output can feed a downstream.expand().Supported dynamic-mapping forms
A mapped
@task.stubproduces one arg-binding per map index (declaration order), and the Go runtime binds each onto the native function's parameters. Each form and how the runtime delivers it:Expanded arguments —
.expand(param=…), value differs per index:via.expand(country=["uk", "fr", "de"])via.expand(extracted=extract())element_indexof the listvia.expand(extracted=seed.expand(n=[1, 2]))map_indexdirectlyMultiple expanded arguments — cross product:
.expand(a=…, b=…)combine.expand(a=["x", "y"], b=[1, 2, 3])→ 6 instancesPartial arguments —
.partial(param=…), constant across every index:via.partial(region="uk").expand(…)via.partial(config=load_config()).expand(…)Defaulted arguments:
retries: int = 3left unpassedfrom_default; keyword-stylesdk.TaskInputfields may leave it unclaimedOne DAG exercising every form, and the Go function receiving it:
Rejected loudly (server-side, mirroring the provider's parse-time checks):
.expand_kwargs(); apartial()kwarg over a mapped upstream's aggregated output;.map()/.zip()/concator custom-key XCom; non-JSON literals; a mapped stub TI still atmap_index=-1. A mapped stub in an older-provider Dag delivers no bindings and keeps the legacy ignored-args behavior.How
pkg/binding):XComArggainsMapIndex/ElementIndex(from regeneratedgenmodels).Resolvenow pulls the specific upstream rowMapIndexselects (expand over a mapped upstream); whenElementIndexis set it takes that element of the pulled sequence (expand over an unmapped upstream's list output), with typed/out-of-range errors. An unmapped argument carries neither and takes the whole value; literal expands are resolved to their element server-side and arrive as plain literals — so the flat-vs-struct binding surface from Bind TaskFlow stub-task call arguments in the Go SDK runtime #70209 is unchanged.mapped_length): a foreign runtime can't inspect the Dag to learn its return value feeds a downstream.expand(). New server-derivedTIRunContext.has_mapped_dependantsflag (computed fromiter_mapped_dependants); when set, the supervisor recordsmapped_length = len(value)on the return-valueSetXComon the task's behalf — the foreign-runtime analogue of the Python task runner's_push_xcom_if_neededlogic — so the scheduler can expand the Go stub's mapped dependants.(arg_bindings, has_mapped_dependants); both fields are version-gated to2026-10-30so pre-arg-bindings clients skip the derivation entirely.Was generative AI tooling used to co-author this PR?