Skip to content

Support TaskFlow call syntax on stub tasks for the Go SDK#69757

Draft
jason810496 wants to merge 3 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag
Draft

Support TaskFlow call syntax on stub tasks for the Go SDK#69757
jason810496 wants to merge 3 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag

Conversation

@jason810496

Copy link
Copy Markdown
Member

Why

@task.stub tasks can only be declared argless today, so a Go task that needs an upstream's output has to hand-write GetXCom calls (with the upstream task_id hard-coded in Go, duplicating the wiring the Dag file already expresses). This PR makes the natural TaskFlow call work across the language boundary:

@task.stub(queue="golang")
def transform(country: str, extracted: dict): ...

with DAG(...):
    transform("uk", extract())   # extract() is a normal Python @task
// The runtime binds "uk" onto country and pulls extract's XCom into extracted.
func transform(ctx sdk.TIRunContext, log *slog.Logger, country string, extracted map[string]any) error

What

  • Parse time (providers/standard): _StubOperator binds the TaskFlow call to the stub's signature and serializes an ordered positional-arg spec with the Dag (_arg_bindings): each entry is an XCom reference or a JSON literal, with data_type derived from the parameter annotation.
  • Server (execution API): ti_run returns the spec as a new optional TIRunContext.arg_bindings field — extracted from the serialized-dag blob only for _StubOperator tasks, and stripped for older API versions (Cadwyn change in the in-progress 2026-06-30 block). The supervisor needs zero changes because TIRunContext is embedded verbatim in StartupDetails.
  • Supervisor schema: new in-progress version 2026-07-30 (AddArgBindingsToTIRunContext); the downgrade path strips the field, so bundles pinned to 2026-06-16 keep working unchanged.
  • Go SDK: new pkg/binding package. Analyze classifies parameters once at registration (injectables like context.Context/*slog.Logger vs JSON-decodable data parameters); Resolve binds the spec onto the data parameters in order — literals decode directly, XCom entries pull on demand. Arity or declared-type mismatches fail the task loudly before the body runs, replacing the old silent reflect.Zero fill.
  • Fail-loud v1 exclusions at parse time: .expand()/.partial() (via a new generic supports_expand opt-out), map/zip/concat XComArgs, *args/**kwargs, context-key parameter names, and non-JSON-serializable literals.

Was generative AI tooling used to co-author this PR?

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant