This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7c6113d9c90 Bind TaskFlow stub-task call arguments in the Go SDK
runtime (#70209)
7c6113d9c90 is described below
commit 7c6113d9c9094a6e177386b66df6a975281c1af8
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Thu Aug 20 13:30:54 2026 +0800
Bind TaskFlow stub-task call arguments in the Go SDK runtime (#70209)
* Bind TaskFlow stub-task call arguments in the Go SDK runtime
A Go task could only reach an upstream task's output by hand-writing a
GetXCom
call against a hard-coded task id, duplicating wiring the Dag file already
owns and breaking silently whenever that upstream was renamed.
#69757 ships the Python half: a `@task.stub` TaskFlow call is captured at
Dag
serialization as an ordered arg-binding spec and returned by ti_run.
Consuming
it here lets a Go task function take the Dag's literals and upstream XComs
as
ordinary typed parameters.
A function declares either flat positional parameters or a single struct
whose
fields bind by name -- kwarg-style, so an unmatched field keeps its zero
value
while an argument no field claims fails the task. Signature problems are
caught
once at registration; per-execution arity, type and spec errors fail the
task
before its body runs, replacing a silent reflect.Zero fill.
* Fail Go SDK TaskFlow argument binding loudly instead of silently
Several binding problems stayed quiet until they were expensive or
confusing.
A parameter nesting an undecodable value failed on every execution rather
than
once when the bundle was built. A struct carrying `arg:` tags whose single
argument no tag matched was decoded whole into the struct, so a typo'd tag
surfaced as a decode error naming the Go type instead of the argument that
matched nothing. A value_schema or from_default of the wrong wire shape was
indistinguishable from an absent one, disabling the declared-type check or
turning a captured stub default into an argument the author supposedly
passed.
The XCom whole-value decode and the concurrent multi-pull failure path were
also reachable from a Dag but exercised only in their simplest shape, and
the
package docs re-explained the whole binding model at three sites, burying
the
rules they were meant to state.
* Fix Go SDK TaskFlow argument binding on Edge Workers and at registration
The Edge Worker's execution API carries no argument spec at all, so failing
a
task there for an argument-count mismatch blamed the Dag author for a limit
of
the transport. Keeping data parameters at their Go zero values is how those
tasks behaved before binding existed, and that path is in maintenance rather
than gaining the spec.
Registration rejected struct shapes that decode without complaint -- a
struct
carrying a callback alongside its data never needed the callback filled --
and
because registering a task panics, one such signature took its whole bundle
down at startup rather than the single task.
Adding a defaulted parameter to a stub is backwards compatible in Python,
and
has to stay so for the Go functions already bound to that stub: the captured
default reaches the wire but needs no Go parameter to receive it.
Untagged fields matched a Go field name verbatim, which no idiomatic
snake_case
stub parameter can produce, so tags were mandatory in practice and a
mismatch
quietly fell back to decoding the argument whole. Folding case and
underscores
makes the untagged form usable, and embedded structs now contribute their
fields the way encoding/json has all along.
A type that decodes itself from JSON also passed registration only to be
rejected at run time by a schema check judging it on its Go kind.
---
.../go_sdk_tests/test_go_sdk_taskflow_binding.py | 170 ++++
go-sdk/README.md | 55 +-
.../adr/0003-coordinator-protocol-msgpack-ipc.md | 12 +-
go-sdk/bundle/bundlev1/task.go | 164 ++--
go-sdk/bundle/bundlev1/task_test.go | 96 +-
.../cmd/airflow-go-pack/pack_integration_test.go | 12 +
go-sdk/dags/go_examples.py | 111 ++-
go-sdk/example/bundle/main.go | 27 +-
go-sdk/example/bundle/main_test.go | 3 +-
.../bundle/taskflowbinding/taskflowbinding.go | 266 ++++++
.../bundle/taskflowbinding/taskflowbinding_test.go | 180 ++++
go-sdk/pkg/binding/binding.go | 899 +++++++++++++++++++
go-sdk/pkg/binding/binding_test.go | 976 +++++++++++++++++++++
go-sdk/pkg/execution/frames.go | 2 +
go-sdk/pkg/execution/frames_test.go | 30 +
go-sdk/pkg/execution/genmodels/models.gen.go | 97 +-
go-sdk/pkg/execution/integration_test.go | 379 ++++++--
go-sdk/pkg/execution/task_runner.go | 120 ++-
18 files changed, 3356 insertions(+), 243 deletions(-)
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py
new file mode 100644
index 00000000000..d078261f8b1
--- /dev/null
+++
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""E2E test for the Go SDK ``taskflow_binding_dag`` example.
+
+The stub Dag's single mixed positional/keyword TaskFlow call carries literals
+of every scalar type, an array literal, a defaulted ``None``, and XComs from
+two upstream Go tasks (an object bound onto a strict Go struct and an array
+bound onto ``[]int``). The Go ``via_flat_args`` task verifies every bound
+value and errors on any mismatch, so a green run *is* the binding assertion;
+the tests here check the run outcome and the summary XCom it pushes.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+import pytest
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+# Allow time for coordinator startup.
+_GO_TASK_TIMEOUT = 300
+
+_DAG_ID = "taskflow_binding_dag"
+
+
+@dataclass
+class _CompletedRun:
+ """The single ``taskflow_binding_dag`` run shared across this module's
tests."""
+
+ client: AirflowClient
+ run_id: str
+ state: str
+ ti_states: dict[str, str]
+
+ def xcom(self, task_id: str, key: str = "return_value"):
+ return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id,
run_id=self.run_id, key=key).get(
+ "value"
+ )
+
+
[email protected](scope="module")
+def completed_run() -> _CompletedRun:
+ """Trigger ``taskflow_binding_dag`` once and wait for it to finish."""
+ client = AirflowClient()
+ resp = client.trigger_dag(_DAG_ID, json={"logical_date":
datetime.now(timezone.utc).isoformat()})
+ run_id = resp["dag_run_id"]
+ state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id,
timeout=_GO_TASK_TIMEOUT)
+ ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id)
+ ti_states = {ti["task_id"]: ti.get("state") for ti in
ti_resp.get("task_instances", [])}
+ return _CompletedRun(client=client, run_id=run_id, state=state,
ti_states=ti_states)
+
+
+def test_all_tasks_succeeded(completed_run: _CompletedRun):
+ """The Go ``via_flat_args`` task errors on any mis-bound argument, so
success here
+ proves every literal, XCom, keyword, and defaulted-None binding was
correct."""
+ assert completed_run.state == "success", (
+ f"expected the run to succeed; got {completed_run.state!r}. task
states: {completed_run.ti_states}"
+ )
+ for task_id in (
+ "make_config",
+ "make_numbers",
+ "make_region",
+ "via_flat_args",
+ "via_struct_no_tags",
+ "via_struct_arg_tag",
+ "via_struct_unmatched_arg",
+ "via_flat_map",
+ "via_struct_map",
+ "via_plain_map",
+ ):
+ assert completed_run.ti_states.get(task_id) == "success",
completed_run.ti_states
+
+
+def test_upstream_xcoms_keep_their_shapes(completed_run: _CompletedRun):
+ """The Go struct arrives as an object XCom, the ``[]int`` as an array, the
region as a string."""
+ assert completed_run.xcom("make_config") == {
+ "environment": "production",
+ "region": "eu-west-1",
+ "debug": True,
+ }
+ assert completed_run.xcom("make_numbers") == [1, 1, 2, 3, 5, 8]
+ assert completed_run.xcom("make_region") == "eu-west-1"
+
+
+def test_via_flat_args_summary_reflects_bound_arguments(completed_run:
_CompletedRun):
+ """``via_flat_args`` re-emits every bound value, confirming types survived
the
+ Python literal / XCom -> Go parameter -> XCom round trip."""
+ assert completed_run.xcom("via_flat_args") == {
+ "name": "summary",
+ "count": 3,
+ "ratio": 2.5,
+ "enabled": True,
+ "tags": ["metrics", "hourly"],
+ "environment": "production",
+ "debug": True,
+ "sum": 20,
+ "note_was_null": True,
+ }
+
+
+def test_via_struct_no_tags_reflects_bound_arguments(completed_run:
_CompletedRun):
+ """``via_struct_no_tags`` demonstrates the Go SDK's name-based struct
binding
+ with no field tags at all: each field falls back to its own Go name,
matched
+ case- and underscore-insensitively, so ``RegionCode`` binds the idiomatic
+ ``region_code``. The region is ``make_region``'s XCom, so a struct field
+ binds an XCom-sourced value here."""
+ assert completed_run.xcom("via_struct_no_tags") == {
+ "region_code": "eu-west-1",
+ "threshold": 0.75,
+ }
+
+
+def test_via_struct_arg_tag_reflects_bound_arguments(completed_run:
_CompletedRun):
+ """``via_struct_arg_tag`` demonstrates explicit ``arg:`` tags: ``Region``
is
+ genuinely renamed to ``region_code`` (bound from ``make_region``'s XCom),
and
+ ``Threshold`` is tagged ``threshold`` to pull the snake_case literal its
+ verbatim field name would miss."""
+ assert completed_run.xcom("via_struct_arg_tag") == {
+ "region": "eu-west-1",
+ "threshold": 0.75,
+ }
+
+
+def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run:
_CompletedRun):
+ """``via_struct_unmatched_arg`` demonstrates mismatch tolerance in both
directions:
+ a struct field whose name has no corresponding TaskFlow call argument
stays at its
+ Go zero value instead of failing the task (kwarg-style, an unpassed name
simply
+ isn't bound), and the stub's defaulted ``sample_rate`` -- captured into
the spec as
+ ``from_default`` -- needs no matching struct field. The task succeeding at
all
+ proves the second half."""
+ assert completed_run.xcom("via_struct_unmatched_arg") == {
+ "region": "eu-west-1",
+ "missing_was_empty": True,
+ }
+
+
+def test_via_flat_map_decodes_single_dict_whole(completed_run: _CompletedRun):
+ """``via_flat_map`` passes one dict literal whose argument name matches no
Go
+ struct field, so the whole map is decoded into the struct (flat
binding)."""
+ assert completed_run.xcom("via_flat_map") == {"region": "eu-west-1",
"count": 3}
+
+
+def test_via_struct_map_binds_single_dict_onto_map_field(completed_run:
_CompletedRun):
+ """``via_struct_map`` passes one dict literal whose argument name binds by
name
+ onto a Go struct's ``map`` field (struct-based binding)."""
+ assert completed_run.xcom("via_struct_map") == {
+ "payload": {"region": "eu-west-1", "count": 3},
+ }
+
+
+def test_via_plain_map_decodes_dict_into_a_typed_map(completed_run:
_CompletedRun):
+ """``via_plain_map`` passes one dict literal onto a Go
``map[string]string``
+ parameter, so it decodes with no struct involved at all."""
+ assert completed_run.xcom("via_plain_map") == {"team": "data", "tier":
"gold"}
diff --git a/go-sdk/README.md b/go-sdk/README.md
index a84373bd291..b93e4a14dd6 100644
--- a/go-sdk/README.md
+++ b/go-sdk/README.md
@@ -105,6 +105,20 @@ A task is an ordinary Go function. The runtime inspects
its signature and inject
`sdk.VariableClient`). An optional `(any, error)` return becomes the task's
XCom; an `error` return marks
the task failed.
+Any other parameter is a **data parameter**, filled in declaration order from
the arguments of the
+Python stub Dag's TaskFlow call. A literal in the Dag file (`transform("uk",
...)`) decodes straight
+into the parameter; an upstream task's output (`transform(..., extract())`) is
pulled from that
+task's XCom in the current Dag run. If the argument count doesn't match, or an
argument's declared
+type can't fill the Go type, the task fails before its body runs.
+
+Stub parameters the Dag author left at their Python defaults are the
exception: they reach the wire
+but need no Go parameter, so adding a defaulted parameter to a stub doesn't
break the Go functions
+already bound to it.
+
+> [!NOTE]
+> Argument binding needs the coordinator path. The Edge Worker sends no
arguments at all, so a task
+> running there keeps its data parameters at their Go zero values.
+
```go
func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any,
error) {
conn, err := client.GetConnection(ctx, "test_http")
@@ -112,12 +126,16 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log
*slog.Logger) (any, er
return map[string]any{"go_version": runtime.Version()}, nil
}
-func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log
*slog.Logger) error {
+// The stub's literal and XCom arguments bind to country and extracted.
+func transform(
+ ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger,
+ country string, extracted map[string]any,
+) error {
val, err := client.GetVariable(ctx, "my_variable")
if err != nil {
return err
}
- log.Info("Obtained variable", "my_variable", val)
+ log.Info("Obtained variable", "my_variable", val, "country", country)
return nil
}
```
@@ -126,6 +144,39 @@ Asking for the narrowest interface a task needs (e.g.
`sdk.VariableClient` inste
unit testing easier and documents which Airflow features the task touches.
`RegisterDags` is the single
source of truth for which `dag_id`s and `task_id`s a bundle can run.
+### Name-based struct binding
+
+When a task's **sole data parameter** is a struct, its fields bind **by name**
instead of by
+position — keyword arguments rather than positional ones, and a friendlier
alternative to a long
+flat parameter list. Being the only data parameter is the opt-in; there is no
marker to add.
+
+```go
+type CombineInput struct {
+ Region string `arg:"region_code"` // renamed
+ Threshold float64
+}
+
+// The stub Dag calls combine(region_code="uk", threshold=0.5).
+func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any,
error) {
+ return nil, nil
+}
+```
+
+An exported field binds the argument matching its own Go name, folding case
and underscores — so
+`Threshold` takes `threshold` and `RegionCode` would take `region_code`. Reach
for an `arg:"<name>"`
+tag when the names genuinely differ, as `Region` does above. Declaration order
is irrelevant on both
+sides, and embedded structs contribute their fields just as they do to
`encoding/json`.
+
+A field no argument matches is left at its Go zero value, like an unpassed
keyword argument. The
+reverse is an error: every argument the Dag author explicitly passed must land
in some field, so a
+typo'd tag fails the task instead of silently dropping the value.
+
+A struct that is **not** the sole data parameter is decoded whole from its one
positional argument
+instead, so `arg:` tags only apply to the sole-parameter form; pairing a
tagged struct with other
+data parameters is rejected at registration. A sole struct parameter also
falls back to whole-value
+decoding when it gets exactly one passed argument no field claims, so a task
can still take an
+upstream object as a single argument.
+
### Reading the task runtime context
Declare an `sdk.TIRunContext` parameter on a task to read the identifiers and
scheduling timestamps of the
diff --git a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md
b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md
index a78c7bb92e4..bb0c24fd298 100644
--- a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md
+++ b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md
@@ -217,7 +217,11 @@ Supervisor Bundle binary (Go)
│ │
├── StartupDetails ────────────────────►│
│ (ti, dag_rel_path, bundle_info, │
- │ start_date, ti_context) │
+ │ start_date, ti_context; the │
+ │ ti_context carries arg_bindings, │
+ │ the positional-argument spec │
+ │ captured from the stub Dag's │
+ │ TaskFlow call) │
│ │
│ ├── lookup task:
│ │ bundle.dags[ti.dag_id]
@@ -225,6 +229,12 @@ Supervisor Bundle binary (Go)
│ │ (returns
TaskState{state:"removed"}
│ │ if not found, mirroring Java)
│ │
+ │ ├── bind arg_bindings onto the task
+ │ │ fn's data parameters (literals
+ │ │ decode directly; xcom refs pull
+ │ │ below); arity/type mismatch
+ │ │ fails the task
+ │ │
│ ├── construct sdk.Client whose
│ │ GetConnection / GetVariable /
│ │ GetXCom / SetXCom calls block
on
diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go
index d31fea84b73..64539c858a8 100644
--- a/go-sdk/bundle/bundlev1/task.go
+++ b/go-sdk/bundle/bundlev1/task.go
@@ -25,68 +25,69 @@ import (
"runtime"
"github.com/apache/airflow/go-sdk/pkg/api"
+ "github.com/apache/airflow/go-sdk/pkg/binding"
"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
"github.com/apache/airflow/go-sdk/sdk"
)
+// TaskWithArgs binds TaskFlow arguments supplied by the coordinator.
+type TaskWithArgs interface {
+ Task
+ ExecuteArgs(ctx context.Context, logger *slog.Logger, args
[]binding.Arg) error
+}
+
type taskFunction struct {
fn reflect.Value
fullName string
+ plan *binding.Plan
}
-var _ Task = (*taskFunction)(nil)
+var _ TaskWithArgs = (*taskFunction)(nil)
-// NewTaskFunction wraps a plain Go function as a Task, validating its
signature
-// (injectable parameters, and a return of error or (result, error)). Bundle
-// authors normally use Dag.AddTask, which calls this for them; use it directly
-// only when building a Task outside the registry.
+// NewTaskFunction validates and wraps a Go function as a Task.
func NewTaskFunction(fn any) (Task, error) {
v := reflect.ValueOf(fn)
fullName := runtime.FuncForPC(v.Pointer()).Name()
- f := &taskFunction{v, fullName}
- return f, f.validateFn(v.Type())
+ f := &taskFunction{fn: v, fullName: fullName}
+ if err := f.validateFn(v.Type()); err != nil {
+ return nil, err
+ }
+ return f, nil
}
+// Execute runs without TaskFlow arguments, as required by the Edge Worker.
func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error
{
- fnType := f.fn.Type()
- var sdkClient sdk.Client
- if injected, ok :=
ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client); ok {
- sdkClient = injected
- } else {
- sdkClient = sdk.NewClient()
+ sdkClient := clientFrom(ctx)
+ return f.call(ctx, sdkClient, f.plan.ResolveUnbound(ctx, logger,
sdkClient), logger)
+}
+
+// ExecuteArgs binds the supplied TaskFlow arguments and runs the task.
+func (f *taskFunction) ExecuteArgs(
+ ctx context.Context,
+ logger *slog.Logger,
+ args []binding.Arg,
+) error {
+ sdkClient := clientFrom(ctx)
+ reflectArgs, err := f.plan.Resolve(ctx, logger, sdkClient, args)
+ if err != nil {
+ return err
}
+ return f.call(ctx, sdkClient, reflectArgs, logger)
+}
- reflectArgs := make([]reflect.Value, fnType.NumIn())
- for i := range reflectArgs {
- in := fnType.In(i)
-
- switch {
- case isTIRunContext(in):
- // sdk.TIRunContext embeds context.Context, so it also
satisfies
- // isContext - this case must come first. The runtime
stores the
- // identifiers/timestamps under RuntimeContextKey;
rebuild the
- // value around the live task context here.
- var ti sdk.TaskInstance
- var dagRun sdk.DagRun
- if stored, ok :=
ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext); ok {
- ti, dagRun = stored.TaskInstance(),
stored.DagRun()
- }
- reflectArgs[i] =
reflect.ValueOf(sdk.NewTIRunContext(ctx, ti, dagRun))
- case isContext(in):
- // Plain context.Context injection is retained for the
Edge Worker
- // runtime path, which does not populate the task
runtime context
- // (TI/DagRun) that sdk.TIRunContext carries. New tasks
should
- // declare sdk.TIRunContext instead.
- reflectArgs[i] = reflect.ValueOf(ctx)
- case isLogger(in):
- reflectArgs[i] = reflect.ValueOf(logger)
- case isClient(in):
- reflectArgs[i] = reflect.ValueOf(sdkClient)
- default:
- // TODO: deal with other value types. For now they will
all be Zero values unless it's a context
- reflectArgs[i] = reflect.Zero(in)
- }
+func clientFrom(ctx context.Context) sdk.Client {
+ if injected, ok :=
ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client); ok {
+ return injected
}
+ return sdk.NewClient()
+}
+
+func (f *taskFunction) call(
+ ctx context.Context,
+ sdkClient sdk.Client,
+ reflectArgs []reflect.Value,
+ logger *slog.Logger,
+) error {
slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs)
retValues := f.fn.Call(reflectArgs)
@@ -150,11 +151,11 @@ func (f *taskFunction) validateFn(fnType reflect.Type)
error {
)
}
- for i := range fnType.NumIn() {
- if err := validateParam(fnType.In(i)); err != nil {
- return fmt.Errorf("task function %s parameter %d: %w",
f.fullName, i, err)
- }
+ plan, err := binding.Analyze(fnType, f.fullName)
+ if err != nil {
+ return err
}
+ f.plan = plan
return nil
}
@@ -168,75 +169,8 @@ func isValidResultType(inType reflect.Type) bool {
return true
}
-var (
- errorType = reflect.TypeFor[error]()
- contextType = reflect.TypeFor[context.Context]()
- tiRunContextType = reflect.TypeFor[sdk.TIRunContext]()
- slogLoggerType = reflect.TypeFor[*slog.Logger]()
-
- clientType = reflect.TypeFor[sdk.Client]()
-)
+var errorType = reflect.TypeFor[error]()
func isError(inType reflect.Type) bool {
return inType != nil && inType.Implements(errorType)
}
-
-func isContext(inType reflect.Type) bool {
- return inType != nil && inType.Implements(contextType)
-}
-
-func isTIRunContext(inType reflect.Type) bool {
- return inType == tiRunContextType
-}
-
-func isLogger(inType reflect.Type) bool {
- return inType != nil && inType.AssignableTo(slogLoggerType)
-}
-
-// isClient reports whether inType's method set is a subset of sdk.Client's,
-// keeping new client capabilities injectable without a hand-kept list.
-func isClient(inType reflect.Type) bool {
- return inType != nil && inType.Kind() == reflect.Interface &&
- inType.NumMethod() > 0 && clientType.Implements(inType)
-}
-
-// validateParam rejects interface parameters Execute cannot inject; they
-// would be bound to nil and panic on first use.
-func validateParam(in reflect.Type) error {
- if in.Kind() != reflect.Interface || isTIRunContext(in) || isClient(in)
{
- return nil
- }
- if isContext(in) {
- // The plain task context injected here cannot satisfy extra
methods.
- if contextType.Implements(in) {
- return nil
- }
- return fmt.Errorf(
- "interface %s adds methods on top of context.Context;
declare sdk.TIRunContext or a separate parameter instead",
- in,
- )
- }
- return fmt.Errorf(
- "interface %s is not injectable (want context.Context,
sdk.TIRunContext, or a subset of sdk.Client): %s",
- in,
- explainClientMismatch(in),
- )
-}
-
-// explainClientMismatch returns why in is not a subset of sdk.Client.
-func explainClientMismatch(in reflect.Type) string {
- if in.NumMethod() == 0 {
- return "empty interfaces cannot be injected"
- }
- for i := range in.NumMethod() {
- m := in.Method(i)
- cm, ok := clientType.MethodByName(m.Name)
- if !ok {
- return fmt.Sprintf("sdk.Client has no method %s",
m.Name)
- }
- if cm.Type != m.Type {
- return fmt.Sprintf("method %s is %s on sdk.Client, not
%s", m.Name, cm.Type, m.Type)
- }
- }
- return "its method set is not a subset of sdk.Client"
-}
diff --git a/go-sdk/bundle/bundlev1/task_test.go
b/go-sdk/bundle/bundlev1/task_test.go
index 3f58e930b87..e4b8dfa672e 100644
--- a/go-sdk/bundle/bundlev1/task_test.go
+++ b/go-sdk/bundle/bundlev1/task_test.go
@@ -20,11 +20,11 @@ package bundlev1
import (
"context"
"log/slog"
- "reflect"
"testing"
"github.com/stretchr/testify/suite"
+ "github.com/apache/airflow/go-sdk/pkg/binding"
"github.com/apache/airflow/go-sdk/pkg/logging"
"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
"github.com/apache/airflow/go-sdk/sdk"
@@ -141,21 +141,6 @@ func (s *TaskSuite) TestClientSubsetInjection() {
s.Require().NoError(task.Execute(context.Background(),
slog.New(logging.NewTeeLogger())))
}
-// TestNamedClientInterfacesAreInjectable guards against sdk.Client dropping an
-// embedded interface, which would break tasks declaring it.
-func (s *TaskSuite) TestNamedClientInterfacesAreInjectable() {
- for name, typ := range map[string]reflect.Type{
- "Client": reflect.TypeFor[sdk.Client](),
- "VariableClient": reflect.TypeFor[sdk.VariableClient](),
- "ConnectionClient": reflect.TypeFor[sdk.ConnectionClient](),
- "XComClient": reflect.TypeFor[sdk.XComClient](),
- } {
- s.True(isClient(typ), "sdk.%s must stay injectable", name)
- }
-}
-
-// TestNonInjectableParamsAreRejected checks registration fails fast on
-// interface parameters Execute cannot inject.
func (s *TaskSuite) TestNonInjectableParamsAreRejected() {
cases := map[string]struct {
fn any
@@ -174,9 +159,9 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() {
},
"method GetVariable is func(context.Context, string)
(string, error) on sdk.Client",
},
- "empty-interface": {
- func(x any) error { return nil },
- "empty interfaces cannot be injected",
+ "func-param": {
+ func(cb func()) error { return nil },
+ "cannot receive a task argument",
},
"context-with-extra-methods": {
func(x interface {
@@ -201,6 +186,79 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() {
}
}
+func (s *TaskSuite) TestExecuteArgsBindsDataParameters() {
+ var gotCountry string
+ var gotMeta map[string]any
+ task, err := NewTaskFunction(func(log *slog.Logger, country string,
meta map[string]any) error {
+ gotCountry = country
+ gotMeta = meta
+ return nil
+ })
+ s.Require().NoError(err)
+
+ tw, ok := task.(TaskWithArgs)
+ s.Require().True(ok, "taskFunction must implement TaskWithArgs")
+
+ err = tw.ExecuteArgs(context.Background(),
slog.New(logging.NewTeeLogger()), []binding.Arg{
+ binding.LiteralArg{Value: "uk"},
+ binding.LiteralArg{Value: map[string]any{"k": "v"}},
+ })
+ s.Require().NoError(err)
+ s.Equal("uk", gotCountry)
+ s.Equal(map[string]any{"k": "v"}, gotMeta)
+}
+
+func (s *TaskSuite) TestExecuteWithoutArgsZeroFillsDataParameters() {
+ type settings struct {
+ Region string
+ Threshold float64
+ }
+ var gotCountry string
+ var gotSettings settings
+ called := false
+ task, err := NewTaskFunction(func(country string, input settings) error
{
+ gotCountry, gotSettings, called = country, input, true
+ return nil
+ })
+ s.Require().NoError(err)
+
+ err = task.Execute(context.Background(),
slog.New(logging.NewTeeLogger()))
+ s.Require().NoError(err)
+ s.True(called, "the task body must run")
+ s.Empty(gotCountry)
+ s.Equal(settings{}, gotSettings)
+}
+
+func (s *TaskSuite) TestExecuteArgsWithoutSpecFailsForDataParameters() {
+ task, err := NewTaskFunction(func(country string) error { return nil })
+ s.Require().NoError(err)
+
+ err = task.(TaskWithArgs).ExecuteArgs(
+ context.Background(), slog.New(logging.NewTeeLogger()), nil,
+ )
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "argument count mismatch")
+ }
+}
+
+func (s *TaskSuite) TestExecuteArgsArityMismatch() {
+ task, err := NewTaskFunction(func(country string) error { return nil })
+ s.Require().NoError(err)
+
+ err = task.(TaskWithArgs).ExecuteArgs(
+ context.Background(),
+ slog.New(logging.NewTeeLogger()),
+ []binding.Arg{
+ binding.LiteralArg{Value: "uk"},
+ binding.LiteralArg{Value: "de"},
+ },
+ )
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "argument count mismatch")
+ s.Contains(err.Error(), "passes 2 positional argument(s)")
+ }
+}
+
// probeKey is an unexported context key used to confirm the live task context
// (not a freshly built one) backs the injected sdk.TIRunContext.
type probeKeyType struct{}
diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
index 84e9a1045f4..ace867461e2 100644
--- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
+++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
@@ -154,6 +154,18 @@ dags:
- "extract"
- "transform"
- "load"
+ taskflow_binding_dag:
+ tasks:
+ - "make_config"
+ - "make_numbers"
+ - "make_region"
+ - "via_flat_args"
+ - "via_struct_no_tags"
+ - "via_struct_arg_tag"
+ - "via_struct_unmatched_arg"
+ - "via_flat_map"
+ - "via_struct_map"
+ - "via_plain_map"
`
assert.Equal(t, expectedManifest, string(metadata))
diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py
index 23e02dd5e49..d6e768a2c37 100644
--- a/go-sdk/dags/go_examples.py
+++ b/go-sdk/dags/go_examples.py
@@ -17,9 +17,11 @@
"""
Python stub Dags mirroring the Go SDK example bundle
(``go-sdk/example/bundle``).
-Two Dags, both backed by the same Go bundle: ``simple_dag`` (extract/transform/
-load, below) and ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task
-timing sequential vs goroutine XCom pulls).
+Three Dags, all backed by the same Go bundle: ``simple_dag``
(extract/transform/
+load, below), ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task
+timing sequential vs goroutine XCom pulls), and ``taskflow_binding_dag`` (one
+task per shape of the TaskFlow argument-binding surface; see its Dag function
+below).
``simple_dag`` sandwiches the Go tasks between two native Python tasks so the
run exercises XCom across the language boundary, the same way
@@ -33,6 +35,10 @@ run exercises XCom across the language boundary, the same way
routed to the ``ExecutableCoordinator``, which locates the bundle by dag_id
and
runs the binary in coordinator mode. ``extract`` returns a map (pushed as its
``return_value`` XCom); ``transform`` reads the ``my_variable`` variable.
+* ``transform`` is called TaskFlow-style -- ``transform("uk", extract())`` --
so
+ the stub captures a positional-argument spec (a literal plus an XCom
+ reference) that the Go runtime binds onto the Go function's ``country`` and
+ ``extracted`` parameters, pulling ``extract``'s XCom on demand.
* ``load`` (``retries=1``) returns an error on its first attempt and succeeds
on the retry, exercising the UP_FOR_RETRY path through the Go coordinator. It
is a leaf (not upstream of ``python_task_2``) so its retry is observable
@@ -65,7 +71,7 @@ def extract(): ...
@task.stub(queue="golang")
-def transform(): ...
+def transform(country: str, extracted: dict): ...
# ``load`` fails on its first attempt and succeeds on the retry, exercising the
@@ -86,7 +92,7 @@ def python_task_2(extracted):
@dag(dag_id="simple_dag")
def simple_dag():
extracted = extract()
- transformed = transform()
+ transformed = transform("uk", extracted)
python_task_1() >> extracted >> transformed
# ``load`` fails once then succeeds on retry; keep it a leaf (not upstream
# of python_task_2) so its retry is observable without affecting the Python
@@ -107,3 +113,98 @@ def concurrent_xcom_dag():
concurrent_xcom_dag()
+
+
[email protected](queue="golang")
+def make_config(): ...
+
+
[email protected](queue="golang")
+def make_numbers(): ...
+
+
[email protected](queue="golang")
+def make_region(): ...
+
+
[email protected](queue="golang")
+def via_flat_args(
+ name: str,
+ count: int,
+ ratio: float,
+ enabled: bool,
+ tags: list,
+ config: dict,
+ numbers: list,
+ note: str | None = None,
+): ...
+
+
+# Go fields match these snake_case names without ``arg:`` tags.
[email protected](queue="golang")
+def via_struct_no_tags(region_code: str, threshold: float): ...
+
+
[email protected](queue="golang")
+def via_struct_arg_tag(region_code: str, threshold: float): ...
+
+
[email protected](queue="golang")
+def via_struct_unmatched_arg(region_code: str, sample_rate: float = 0.1): ...
+
+
[email protected](queue="golang")
+def via_flat_map(config: dict): ...
+
+
[email protected](queue="golang")
+def via_struct_map(payload: dict): ...
+
+
[email protected](queue="golang")
+def via_plain_map(labels: dict): ...
+
+
+@dag(dag_id="taskflow_binding_dag")
+def taskflow_binding_dag():
+ """
+ Exercise the Go SDK's TaskFlow argument-binding surface, one shape per
task.
+
+ A Go task declares either flat data parameters, which bind *positionally*
+ (order matters, every one must be filled), or a single struct, whose fields
+ bind by *name* like keyword arguments -- an unmatched field stays at its Go
+ zero value instead of failing the task.
+
+ * ``via_flat_args``: every scalar literal, an array literal, keyword args,
+ an unpassed ``None`` default, and XComs fanned in from two upstream
tasks.
+ * ``via_struct_no_tags``: fields fall back to their own Go names, matched
+ case- and underscore-insensitively.
+ * ``via_struct_arg_tag``: fields bind via explicit ``arg:`` tags.
+ * ``via_struct_unmatched_arg``: a Go field no argument names, and a stub
+ default no Go field claims.
+ * ``via_flat_map`` / ``via_struct_map`` / ``via_plain_map``: one dict bound
+ whole into a struct, onto a struct's map field, and into a plain Go map.
+
+ Each ``via_struct_*`` call mixes a ``threshold`` literal with
``make_region``'s
+ XCom, so struct fields are proven against both argument sources. The Go
+ tasks verify every bound value and fail on any mismatch.
+ """
+ via_flat_args(
+ "summary",
+ 3,
+ 2.5,
+ True,
+ ["metrics", "hourly"],
+ config=make_config(),
+ numbers=make_numbers(),
+ )
+ region = make_region()
+ via_struct_no_tags(region_code=region, threshold=0.75)
+ via_struct_arg_tag(region_code=region, threshold=0.75)
+ via_struct_unmatched_arg(region_code=region)
+ via_flat_map(config={"region": "eu-west-1", "count": 3})
+ via_struct_map(payload={"region": "eu-west-1", "count": 3})
+ via_plain_map(labels={"team": "data", "tier": "gold"})
+
+
+taskflow_binding_dag()
diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go
index 23e60bd1dd4..66c2b249826 100644
--- a/go-sdk/example/bundle/main.go
+++ b/go-sdk/example/bundle/main.go
@@ -27,6 +27,7 @@ import (
v1 "github.com/apache/airflow/go-sdk/bundle/bundlev1"
"github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server"
"github.com/apache/airflow/go-sdk/example/bundle/concurrentxcom"
+ "github.com/apache/airflow/go-sdk/example/bundle/taskflowbinding"
"github.com/apache/airflow/go-sdk/sdk"
)
@@ -55,6 +56,18 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error {
concurrentDag := dagbag.AddDag("concurrent_xcom_dag")
concurrentDag.AddTaskWithName("pull_xcoms_concurrently",
concurrentxcom.PullXComsConcurrently)
+ bindingDag := dagbag.AddDag("taskflow_binding_dag")
+ bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig)
+ bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers)
+ bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion)
+ bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs)
+ bindingDag.AddTaskWithName("via_struct_no_tags",
taskflowbinding.ViaStructNoTags)
+ bindingDag.AddTaskWithName("via_struct_arg_tag",
taskflowbinding.ViaStructArgTag)
+ bindingDag.AddTaskWithName("via_struct_unmatched_arg",
taskflowbinding.ViaStructUnmatchedArg)
+ bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap)
+ bindingDag.AddTaskWithName("via_struct_map",
taskflowbinding.ViaStructMap)
+ bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap)
+
return nil
}
@@ -127,12 +140,24 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log
*slog.Logger) (any, er
return ret, nil
}
-func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log
*slog.Logger) error {
+// transform receives the stub call's literal and XCom arguments.
+func transform(
+ ctx sdk.TIRunContext,
+ client sdk.VariableClient,
+ log *slog.Logger,
+ country string,
+ extracted map[string]any,
+) error {
// This function takes a VariableClient and not a Client to make unit
testing it easier. See
// `./main_test.go` for an example unit of this task fn. Functionally
taking a `sdk.Client` is the same (as
// Client includes VariableClient) but by using the dedicated type it
can be easier to write unit tests.
//
// It also gives a better indication of what features the tasks use
+ log.InfoContext(ctx, "Bound TaskFlow arguments",
+ "country", country,
+ "extracted_go_version", extracted["go_version"],
+ "extracted_timestamp", extracted["timestamp"],
+ )
key := "my_variable"
val, err := client.GetVariable(ctx, key)
if err != nil {
diff --git a/go-sdk/example/bundle/main_test.go
b/go-sdk/example/bundle/main_test.go
index 474a8406d62..16be1e22e92 100644
--- a/go-sdk/example/bundle/main_test.go
+++ b/go-sdk/example/bundle/main_test.go
@@ -52,8 +52,7 @@ var _ sdk.VariableClient = (*mockVars)(nil)
func Test_transform(t *testing.T) {
log := slog.Default()
// This is not the best test, but it is a good proof of concept -- you
can just call the function.
- // sdk.NewTIRunContext wraps any context to build a TIRunContext in a
test.
ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
- err := transform(ctx, &mockVars{}, log)
+ err := transform(ctx, &mockVars{}, log, "uk",
map[string]any{"go_version": "go1.24"})
assert.NoError(t, err)
}
diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go
b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go
new file mode 100644
index 00000000000..2e2ed740cd7
--- /dev/null
+++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go
@@ -0,0 +1,266 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Package taskflowbinding contains TaskFlow argument-binding examples.
+package taskflowbinding
+
+import (
+ "fmt"
+ "log/slog"
+ "reflect"
+
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+// Config is an object passed through XCom.
+type Config struct {
+ Environment string `json:"environment"`
+ Region string `json:"region"`
+ Debug bool `json:"debug"`
+}
+
+// MakeConfig returns a Config XCom.
+func MakeConfig(log *slog.Logger) (any, error) {
+ cfg := Config{Environment: "production", Region: "eu-west-1", Debug:
true}
+ log.Info(
+ "Pushing config",
+ "environment",
+ cfg.Environment,
+ "region",
+ cfg.Region,
+ "debug",
+ cfg.Debug,
+ )
+ return cfg, nil
+}
+
+// MakeNumbers returns an integer-slice XCom.
+func MakeNumbers(log *slog.Logger) (any, error) {
+ numbers := []int{1, 1, 2, 3, 5, 8}
+ log.Info("Pushing numbers", "numbers", fmt.Sprint(numbers))
+ return numbers, nil
+}
+
+// MakeRegion returns a region XCom.
+func MakeRegion(log *slog.Logger) (any, error) {
+ region := "eu-west-1"
+ log.Info("Pushing region", "region", region)
+ return region, nil
+}
+
+// ViaFlatArgs exercises positional binding for literals, defaults, and XComs.
+func ViaFlatArgs(
+ ctx sdk.TIRunContext,
+ log *slog.Logger,
+ name string,
+ count int,
+ ratio float64,
+ enabled bool,
+ tags []string,
+ config Config,
+ numbers []int,
+ note *string,
+) (any, error) {
+ if name != "summary" || count != 3 || ratio != 2.5 || !enabled {
+ return nil, fmt.Errorf(
+ "scalar literals bound incorrectly: name=%q count=%d
ratio=%v enabled=%v",
+ name, count, ratio, enabled,
+ )
+ }
+ if want := []string{"metrics", "hourly"}; !reflect.DeepEqual(tags,
want) {
+ return nil, fmt.Errorf("array literal bound incorrectly:
tags=%v, want %v", tags, want)
+ }
+ if want := (Config{Environment: "production", Region: "eu-west-1",
Debug: true}); config != want {
+ return nil, fmt.Errorf("object XCom bound incorrectly:
config=%+v, want %+v", config, want)
+ }
+ if want := []int{1, 1, 2, 3, 5, 8}; !reflect.DeepEqual(numbers, want) {
+ return nil, fmt.Errorf("array XCom bound incorrectly:
numbers=%v, want %v", numbers, want)
+ }
+ if note != nil {
+ return nil, fmt.Errorf("defaulted None bound incorrectly:
note=%q, want nil", *note)
+ }
+
+ sum := 0
+ for _, n := range numbers {
+ sum += n
+ }
+ log.InfoContext(ctx, "Bound TaskFlow arguments",
+ "name", name,
+ "count", count,
+ "ratio", ratio,
+ "enabled", enabled,
+ "tags", fmt.Sprint(tags),
+ "environment", config.Environment,
+ "sum", sum,
+ )
+ return map[string]any{
+ "name": name,
+ "count": count,
+ "ratio": ratio,
+ "enabled": enabled,
+ "tags": tags,
+ "environment": config.Environment,
+ "debug": config.Debug,
+ "sum": sum,
+ "note_was_null": note == nil,
+ }, nil
+}
+
+// ViaStructNoTagsInput binds fields by folded Go name.
+type ViaStructNoTagsInput struct {
+ RegionCode string
+ Threshold float64
+}
+
+// ViaStructNoTags exercises binding without `arg:` tags.
+func ViaStructNoTags(
+ ctx sdk.TIRunContext,
+ log *slog.Logger,
+ input ViaStructNoTagsInput,
+) (any, error) {
+ if input.RegionCode != "eu-west-1" || input.Threshold != 0.75 {
+ return nil, fmt.Errorf(
+ "struct fields bound incorrectly: region_code=%q
threshold=%v",
+ input.RegionCode,
+ input.Threshold,
+ )
+ }
+
+ log.InfoContext(ctx, "Bound struct (no tags)",
+ "region_code", input.RegionCode,
+ "threshold", input.Threshold,
+ )
+ return map[string]any{
+ "region_code": input.RegionCode,
+ "threshold": input.Threshold,
+ }, nil
+}
+
+// ViaStructArgTagInput binds fields with explicit `arg:` tags.
+type ViaStructArgTagInput struct {
+ Region string `arg:"region_code"`
+ Threshold float64 `arg:"threshold"`
+}
+
+// ViaStructArgTag exercises explicit field-name binding.
+func ViaStructArgTag(
+ ctx sdk.TIRunContext,
+ log *slog.Logger,
+ input ViaStructArgTagInput,
+) (any, error) {
+ if input.Region != "eu-west-1" || input.Threshold != 0.75 {
+ return nil, fmt.Errorf(
+ "struct fields bound incorrectly: region=%q
threshold=%v",
+ input.Region,
+ input.Threshold,
+ )
+ }
+
+ log.InfoContext(ctx, "Bound struct (arg: tag)",
+ "region", input.Region,
+ "threshold", input.Threshold,
+ )
+ return map[string]any{
+ "region": input.Region,
+ "threshold": input.Threshold,
+ }, nil
+}
+
+// ViaStructUnmatchedArgInput includes a field no argument supplies.
+type ViaStructUnmatchedArgInput struct {
+ Region string `arg:"region_code"`
+ Missing string `arg:"does_not_exist"`
+}
+
+// ViaStructUnmatchedArg exercises unmatched fields and captured defaults.
+func ViaStructUnmatchedArg(
+ ctx sdk.TIRunContext, log *slog.Logger, input
ViaStructUnmatchedArgInput,
+) (any, error) {
+ if input.Region != "eu-west-1" {
+ return nil, fmt.Errorf("struct field bound incorrectly:
region=%q", input.Region)
+ }
+ if input.Missing != "" {
+ return nil, fmt.Errorf(
+ "expected the unmatched field to stay at its Go zero
value, got missing=%q",
+ input.Missing,
+ )
+ }
+
+ log.InfoContext(ctx, "Bound struct (unmatched arg)",
+ "region", input.Region,
+ "missing_was_empty", input.Missing == "",
+ )
+ return map[string]any{
+ "region": input.Region,
+ "missing_was_empty": input.Missing == "",
+ }, nil
+}
+
+// FlatMapConfig receives one dict as a whole value.
+type FlatMapConfig struct {
+ Region string `json:"region"`
+ Count int `json:"count"`
+}
+
+// ViaFlatMap exercises whole-value dict decoding into a struct.
+func ViaFlatMap(
+ ctx sdk.TIRunContext, log *slog.Logger, config FlatMapConfig,
+) (any, error) {
+ if config.Region != "eu-west-1" || config.Count != 3 {
+ return nil, fmt.Errorf(
+ "whole-value map bound incorrectly: region=%q count=%d",
+ config.Region,
+ config.Count,
+ )
+ }
+
+ log.InfoContext(ctx, "Bound whole map into struct",
+ "region", config.Region,
+ "count", config.Count,
+ )
+ return map[string]any{"region": config.Region, "count": config.Count},
nil
+}
+
+// StructMapInput binds a dict to a map field.
+type StructMapInput struct {
+ Payload map[string]any `arg:"payload"`
+}
+
+// ViaStructMap exercises dict binding to a struct field.
+func ViaStructMap(
+ ctx sdk.TIRunContext, log *slog.Logger, input StructMapInput,
+) (any, error) {
+ region, _ := input.Payload["region"].(string)
+ if region != "eu-west-1" {
+ return nil, fmt.Errorf("map field bound incorrectly:
payload=%v", input.Payload)
+ }
+
+ log.InfoContext(ctx, "Bound map onto struct field", "payload",
fmt.Sprint(input.Payload))
+ return map[string]any{"payload": input.Payload}, nil
+}
+
+// ViaPlainMap exercises dict decoding into a typed map.
+func ViaPlainMap(
+ ctx sdk.TIRunContext, log *slog.Logger, labels map[string]string,
+) (any, error) {
+ if labels["team"] != "data" || labels["tier"] != "gold" {
+ return nil, fmt.Errorf("plain map bound incorrectly:
labels=%v", labels)
+ }
+
+ log.InfoContext(ctx, "Bound dict into a plain map", "labels",
fmt.Sprint(labels))
+ return map[string]any{"team": labels["team"], "tier": labels["tier"]},
nil
+}
diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go
b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go
new file mode 100644
index 00000000000..02a3d6c75ad
--- /dev/null
+++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go
@@ -0,0 +1,180 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package taskflowbinding
+
+import (
+ "context"
+ "log/slog"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+func TestViaFlatArgs(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaFlatArgs(ctx, slog.Default(),
+ "summary", 3, 2.5, true,
+ []string{"metrics", "hourly"},
+ Config{Environment: "production", Region: "eu-west-1", Debug:
true},
+ []int{1, 1, 2, 3, 5, 8},
+ nil,
+ )
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaFlatArgs should return a map summary, got %T",
got)
+ assert.Equal(t, 20, summary["sum"])
+ assert.Equal(t, true, summary["note_was_null"])
+}
+
+func TestViaFlatArgsRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaFlatArgs(ctx, slog.Default(),
+ "summary", 3, 2.5, true,
+ []string{"metrics", "hourly"},
+ Config{},
+ []int{1, 1, 2, 3, 5, 8},
+ nil,
+ )
+ assert.ErrorContains(t, err, "object XCom bound incorrectly")
+}
+
+func TestViaStructNoTags(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{
+ RegionCode: "eu-west-1",
+ Threshold: 0.75,
+ })
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaStructNoTags should return a map summary, got
%T", got)
+ assert.Equal(t, "eu-west-1", summary["region_code"])
+}
+
+func TestViaStructNoTagsRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{
+ RegionCode: "wrong-region",
+ Threshold: 0.75,
+ })
+ assert.ErrorContains(t, err, "struct fields bound incorrectly")
+}
+
+func TestViaStructArgTag(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{
+ Region: "eu-west-1",
+ Threshold: 0.75,
+ })
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaStructArgTag should return a map summary, got
%T", got)
+ assert.Equal(t, "eu-west-1", summary["region"])
+}
+
+func TestViaStructArgTagRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{
+ Region: "wrong-region",
+ Threshold: 0.75,
+ })
+ assert.ErrorContains(t, err, "struct fields bound incorrectly")
+}
+
+func TestViaStructUnmatchedArg(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaStructUnmatchedArg(ctx, slog.Default(),
ViaStructUnmatchedArgInput{
+ Region: "eu-west-1",
+ Missing: "",
+ })
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaStructUnmatchedArg should return a map summary,
got %T", got)
+ assert.Equal(t, true, summary["missing_was_empty"])
+}
+
+func TestViaStructUnmatchedArgRejectsNonZeroMissingField(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaStructUnmatchedArg(ctx, slog.Default(),
ViaStructUnmatchedArgInput{
+ Region: "eu-west-1",
+ Missing: "unexpected",
+ })
+ assert.ErrorContains(t, err, "expected the unmatched field to stay at
its Go zero value")
+}
+
+func TestViaFlatMap(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaFlatMap(ctx, slog.Default(), FlatMapConfig{Region:
"eu-west-1", Count: 3})
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaFlatMap should return a map summary, got %T",
got)
+ assert.Equal(t, "eu-west-1", summary["region"])
+ assert.Equal(t, 3, summary["count"])
+}
+
+func TestViaFlatMapRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaFlatMap(ctx, slog.Default(), FlatMapConfig{Region:
"wrong-region", Count: 3})
+ assert.ErrorContains(t, err, "whole-value map bound incorrectly")
+}
+
+func TestViaStructMap(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaStructMap(ctx, slog.Default(), StructMapInput{
+ Payload: map[string]any{"region": "eu-west-1", "count": 3},
+ })
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaStructMap should return a map summary, got %T",
got)
+ assert.Equal(t, map[string]any{"region": "eu-west-1", "count": 3},
summary["payload"])
+}
+
+func TestViaStructMapRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaStructMap(ctx, slog.Default(), StructMapInput{
+ Payload: map[string]any{"region": "wrong-region"},
+ })
+ assert.ErrorContains(t, err, "map field bound incorrectly")
+}
+
+func TestViaPlainMap(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ got, err := ViaPlainMap(ctx, slog.Default(), map[string]string{
+ "team": "data", "tier": "gold",
+ })
+ require.NoError(t, err)
+
+ summary, ok := got.(map[string]any)
+ require.True(t, ok, "ViaPlainMap should return a map summary, got %T",
got)
+ assert.Equal(t, "data", summary["team"])
+ assert.Equal(t, "gold", summary["tier"])
+}
+
+func TestViaPlainMapRejectsWrongBinding(t *testing.T) {
+ ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{},
sdk.DagRun{})
+ _, err := ViaPlainMap(ctx, slog.Default(), map[string]string{"team":
"wrong"})
+ assert.ErrorContains(t, err, "plain map bound incorrectly")
+}
diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go
new file mode 100644
index 00000000000..2a1a7d294ee
--- /dev/null
+++ b/go-sdk/pkg/binding/binding.go
@@ -0,0 +1,899 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Package binding resolves TaskFlow arguments for Go task functions.
+//
+// Runtime values are injected by type. Other parameters bind positionally,
+// except a sole struct whose fields bind by `arg:` tag or folded Go name.
+// Captured defaults may go unclaimed, and a sole untagged struct can decode
one
+// unclaimed argument as a whole value.
+//
+// Analyze validates a function once. Resolve binds each execution, while
+// ResolveUnbound zero-fills data parameters for runtimes without argument
specs.
+package binding
+
+import (
+ "bytes"
+ "context"
+ "encoding"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "log/slog"
+ "reflect"
+ "strings"
+ "sync"
+
+ "github.com/apache/airflow/go-sdk/pkg/api"
+ "github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
+ "github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+// Arg is a literal or XCom-backed TaskFlow argument.
+type Arg interface {
+ // ArgName returns the stub parameter name.
+ ArgName() string
+ // Schema returns the argument schema, if declared.
+ Schema() *genmodels.ArgValueSchema
+ sealedArg()
+}
+
+// XComArg reads an upstream task's return-value XCom.
+type XComArg genmodels.XComArgBinding
+
+// LiteralArg carries an inline value from the Dag file.
+type LiteralArg genmodels.LiteralArgBinding
+
+func (a XComArg) ArgName() string { return a.Name }
+func (a LiteralArg) ArgName() string { return a.Name }
+
+func (a XComArg) Schema() *genmodels.ArgValueSchema { return a.ValueSchema }
+func (a LiteralArg) Schema() *genmodels.ArgValueSchema { return a.ValueSchema }
+
+func (XComArg) sealedArg() {}
+func (LiteralArg) sealedArg() {}
+
+type paramKind int
+
+const (
+ paramTIRunContext paramKind = iota
+ paramContext
+ paramLogger
+ paramClient
+ paramData
+ paramLoneStruct
+)
+
+type structField struct {
+ index []int
+ goName string
+ fieldType reflect.Type
+ argName string
+ tagged bool
+}
+
+type paramPlan struct {
+ kind paramKind
+ typ reflect.Type
+ index int
+ fields []structField
+ tagged bool
+}
+
+// Plan describes how to fill a task function's parameters.
+type Plan struct {
+ fnName string
+ params []paramPlan
+ numData int
+ loneStruct bool
+}
+
+// Analyze validates a task function and builds its binding plan.
+func Analyze(fnType reflect.Type, fnName string) (*Plan, error) {
+ p := &Plan{fnName: fnName, params: make([]paramPlan, fnType.NumIn())}
+ var dataIdxs []int
+ for i := range fnType.NumIn() {
+ plan, err := classifyParam(fnName, fnType.In(i), i)
+ if err != nil {
+ return nil, err
+ }
+ if plan.kind == paramData {
+ p.numData++
+ dataIdxs = append(dataIdxs, i)
+ }
+ p.params[i] = plan
+ }
+
+ if p.numData == 1 {
+ i := dataIdxs[0]
+ if st := structParamType(p.params[i].typ); st != nil {
+ fields, err := buildStructFields(fnName, st, i)
+ if err != nil {
+ return nil, err
+ }
+ p.params[i].kind = paramLoneStruct
+ p.params[i].fields = fields
+ p.params[i].tagged = hasArgTag(st)
+ p.numData = 0
+ p.loneStruct = true
+ return p, nil
+ }
+ }
+
+ // Tags on a non-sole struct would be silently ignored during
whole-value decoding.
+ for _, i := range dataIdxs {
+ if st := structParamType(p.params[i].typ); st != nil &&
hasArgTag(st) {
+ return nil, fmt.Errorf(
+ "task function %s: parameter %d: a struct with
`arg:` tags must be the function's "+
+ "only data parameter (its fields bind
TaskFlow arguments by name); it cannot be "+
+ "combined with other data parameters",
+ fnName, i,
+ )
+ }
+ }
+ return p, nil
+}
+
+// Resolve builds the ordered values for one task call.
+func (p *Plan) Resolve(
+ ctx context.Context,
+ logger *slog.Logger,
+ client sdk.Client,
+ args []Arg,
+) ([]reflect.Value, error) {
+ out := p.resolveInjectables(ctx, logger, client)
+ if p.loneStruct {
+ return p.resolveLoneStructParam(ctx, client, args, out)
+ }
+ return p.resolveFlatParams(ctx, client, args, out)
+}
+
+// ResolveUnbound fills injectables and zero-fills data parameters.
+func (p *Plan) ResolveUnbound(
+ ctx context.Context,
+ logger *slog.Logger,
+ client sdk.Client,
+) []reflect.Value {
+ out := p.resolveInjectables(ctx, logger, client)
+ for i, plan := range p.params {
+ switch plan.kind {
+ case paramData, paramLoneStruct:
+ out[i] = reflect.Zero(plan.typ)
+ }
+ }
+ return out
+}
+
+func (p *Plan) resolveInjectables(
+ ctx context.Context,
+ logger *slog.Logger,
+ client sdk.Client,
+) []reflect.Value {
+ out := make([]reflect.Value, len(p.params))
+ for i, plan := range p.params {
+ switch plan.kind {
+ case paramTIRunContext:
+ // Rebuild the stored metadata around the live task
context.
+ var ti sdk.TaskInstance
+ var dagRun sdk.DagRun
+ if stored, ok :=
ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext); ok {
+ ti, dagRun = stored.TaskInstance(),
stored.DagRun()
+ }
+ out[i] = reflect.ValueOf(sdk.NewTIRunContext(ctx, ti,
dagRun))
+ case paramContext:
+ out[i] = reflect.ValueOf(ctx)
+ case paramLogger:
+ out[i] = reflect.ValueOf(logger)
+ case paramClient:
+ out[i] = reflect.ValueOf(client)
+ case paramData, paramLoneStruct:
+ }
+ }
+ return out
+}
+
+func (p *Plan) resolveFlatParams(
+ ctx context.Context,
+ c sdk.XComClient,
+ args []Arg,
+ out []reflect.Value,
+) ([]reflect.Value, error) {
+ // Captured defaults may exceed the Go function's arity.
+ if len(args) != p.numData {
+ args = dropDefaultedArgs(args)
+ }
+ if len(args) != p.numData {
+ return nil, fmt.Errorf(
+ "task function %s: argument count mismatch: the Dag
passes %d positional argument(s) "+
+ "but the Go function declares %d data
parameter(s)",
+ p.fnName, len(args), p.numData,
+ )
+ }
+ raws, err := p.fetchArgValues(ctx, c, args, nil)
+ if err != nil {
+ return nil, err
+ }
+ flatIdx := 0
+ for i, plan := range p.params {
+ if plan.kind != paramData {
+ continue
+ }
+ v, err := p.decodeArg(
+ args[flatIdx], raws[flatIdx], plan.typ,
+ fmt.Sprintf("argument %d (parameter %d)", flatIdx,
plan.index),
+ )
+ if err != nil {
+ return nil, err
+ }
+ out[i] = v
+ flatIdx++
+ }
+ return out, nil
+}
+
+func (p *Plan) resolveLoneStructParam(
+ ctx context.Context,
+ c sdk.XComClient,
+ args []Arg,
+ out []reflect.Value,
+) ([]reflect.Value, error) {
+ var paramIdx int
+ var plan paramPlan
+ for i, pl := range p.params {
+ if pl.kind == paramLoneStruct {
+ paramIdx, plan = i, pl
+ break
+ }
+ }
+
+ byName := make(map[string]int, len(args))
+ // Do not guess when folded names collide.
+ byFolded := make(map[string]int, len(args))
+ ambiguous := make(map[string]bool, len(args))
+ for i, a := range args {
+ if a == nil {
+ continue
+ }
+ byName[a.ArgName()] = i
+ key := foldArgName(a.ArgName())
+ if _, dup := byFolded[key]; dup {
+ ambiguous[key] = true
+ continue
+ }
+ byFolded[key] = i
+ }
+
+ claimed := make([]bool, len(args))
+ type fieldBind struct {
+ field structField
+ argIdx int
+ }
+ binds := make([]fieldBind, 0, len(plan.fields))
+ for _, sf := range plan.fields {
+ idx, ok := byName[sf.argName]
+ if !ok && !sf.tagged {
+ key := foldArgName(sf.argName)
+ if !ambiguous[key] {
+ idx, ok = byFolded[key]
+ }
+ }
+ if !ok {
+ continue
+ }
+ claimed[idx] = true
+ binds = append(binds, fieldBind{field: sf, argIdx: idx})
+ }
+
+ // Tagged structs never fall back to whole-value decoding, which would
hide tag typos.
+ if len(binds) == 0 && !plan.tagged {
+ if explicit, ok := loneWholeValueArgs(args); ok {
+ return p.resolveWholeStructParam(ctx, c, explicit,
plan, paramIdx, out)
+ }
+ }
+
+ if len(args) == 0 && len(plan.fields) > 0 {
+ return nil, fmt.Errorf(
+ "task function %s: no TaskFlow arg bindings arrived but
the struct declares "+
+ "%d bindable field(s); nothing can fill them on
this execution path",
+ p.fnName, len(plan.fields),
+ )
+ }
+
+ var unclaimed []string
+ for i, c := range claimed {
+ if c {
+ continue
+ }
+ if lit, ok := args[i].(LiteralArg); ok && lit.FromDefault {
+ continue
+ }
+ name := "<nil>"
+ if args[i] != nil {
+ name = fmt.Sprintf("%q", args[i].ArgName())
+ }
+ unclaimed = append(unclaimed, name)
+ }
+ if len(unclaimed) > 0 {
+ return nil, fmt.Errorf(
+ "task function %s: %d TaskFlow call argument(s) not
claimed by any struct "+
+ "field: %s",
+ p.fnName, len(unclaimed), strings.Join(unclaimed, ", "),
+ )
+ }
+
+ raws, err := p.fetchArgValues(ctx, c, args, claimed)
+ if err != nil {
+ return nil, err
+ }
+
+ structType := plan.typ
+ isPtr := structType.Kind() == reflect.Pointer
+ if isPtr {
+ structType = structType.Elem()
+ }
+ structVal := reflect.New(structType).Elem()
+ for _, b := range binds {
+ v, err := p.decodeArg(
+ args[b.argIdx], raws[b.argIdx], b.field.fieldType,
+ fmt.Sprintf("struct field %s (parameter %d)",
b.field.goName, plan.index),
+ )
+ if err != nil {
+ return nil, err
+ }
+ structVal.FieldByIndex(b.field.index).Set(v)
+ }
+
+ if isPtr {
+ out[paramIdx] = structVal.Addr()
+ } else {
+ out[paramIdx] = structVal
+ }
+ return out, nil
+}
+
+func dropDefaultedArgs(args []Arg) []Arg {
+ kept := make([]Arg, 0, len(args))
+ for _, a := range args {
+ if lit, ok := a.(LiteralArg); ok && lit.FromDefault {
+ continue
+ }
+ kept = append(kept, a)
+ }
+ return kept
+}
+
+func loneWholeValueArgs(args []Arg) ([]Arg, bool) {
+ explicit := dropDefaultedArgs(args)
+ if len(explicit) != 1 || explicit[0] == nil {
+ return nil, false
+ }
+ return explicit, true
+}
+
+func (p *Plan) resolveWholeStructParam(
+ ctx context.Context,
+ c sdk.XComClient,
+ args []Arg,
+ plan paramPlan,
+ paramIdx int,
+ out []reflect.Value,
+) ([]reflect.Value, error) {
+ raws, err := p.fetchArgValues(ctx, c, args, nil)
+ if err != nil {
+ return nil, err
+ }
+ v, err := p.decodeArg(
+ args[0], raws[0], plan.typ,
+ fmt.Sprintf("argument %q (parameter %d)", args[0].ArgName(),
plan.index),
+ )
+ if err != nil {
+ return nil, err
+ }
+ out[paramIdx] = v
+ return out, nil
+}
+
+func (p *Plan) fetchArgValues(
+ ctx context.Context,
+ c sdk.XComClient,
+ args []Arg,
+ needed []bool,
+) ([]any, error) {
+ raws := make([]any, len(args))
+ var xcomIdxs []int
+ for i, a := range args {
+ if needed != nil && !needed[i] {
+ continue
+ }
+ switch a := a.(type) {
+ case LiteralArg:
+ raws[i] = a.Value
+ case XComArg:
+ xcomIdxs = append(xcomIdxs, i)
+ }
+ }
+ if len(xcomIdxs) == 0 {
+ return raws, nil
+ }
+
+ workload, ok :=
ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload)
+ if !ok {
+ return nil, fmt.Errorf(
+ "task function %s: no workload in context, cannot
resolve xcom arguments", p.fnName,
+ )
+ }
+ // Stop sibling pulls after the first failure.
+ pullCtx, cancel := context.WithCancel(ctx)
+ defer cancel()
+
+ // Stub calls reference only the unmapped upstream return value.
+ pull := func(i int) error {
+ a := args[i].(XComArg)
+ raw, err := c.GetXCom(
+ pullCtx, workload.TI.DagId, workload.TI.RunId,
a.TaskID, nil,
+ api.XComReturnValueKey, nil,
+ )
+ if err != nil {
+ return fmt.Errorf(
+ "task function %s: argument %q: pulling xcom
from task %q: %w",
+ p.fnName, a.Name, a.TaskID, err,
+ )
+ }
+ raws[i] = raw
+ return nil
+ }
+ var wg sync.WaitGroup
+ errs := make([]error, len(xcomIdxs))
+ for j, i := range xcomIdxs {
+ wg.Go(func() {
+ if err := pull(i); err != nil {
+ errs[j] = err
+ cancel()
+ }
+ })
+ }
+ wg.Wait()
+ if err := errors.Join(errs...); err != nil {
+ return nil, err
+ }
+ return raws, nil
+}
+
+func (p *Plan) decodeArg(
+ arg Arg,
+ raw any,
+ targetType reflect.Type,
+ errCtx string,
+) (reflect.Value, error) {
+ if arg == nil {
+ return reflect.Value{}, fmt.Errorf(
+ "task function %s: %s: nil argument binding", p.fnName,
errCtx,
+ )
+ }
+ if err := checkValueType(arg.Schema(), targetType); err != nil {
+ return reflect.Value{}, fmt.Errorf("task function %s: %s: %w",
p.fnName, errCtx, err)
+ }
+ var source string
+ switch a := arg.(type) {
+ case LiteralArg:
+ source = "literal value"
+ case XComArg:
+ source = fmt.Sprintf("xcom from task %q", a.TaskID)
+ default:
+ return reflect.Value{}, fmt.Errorf(
+ "task function %s: %s: unsupported argument binding
%T", p.fnName, errCtx, arg,
+ )
+ }
+ v, err := decodeValue(raw, targetType)
+ if err != nil {
+ return reflect.Value{}, fmt.Errorf(
+ "task function %s: %s: decoding %s into %s: %w",
+ p.fnName, errCtx, source, targetType, err,
+ )
+ }
+ return v, nil
+}
+
+func classifyParam(fnName string, in reflect.Type, index int) (paramPlan,
error) {
+ switch {
+ case isTIRunContext(in):
+ // TIRunContext also satisfies context.Context, so check it
first.
+ return paramPlan{kind: paramTIRunContext, index: index}, nil
+ case isContext(in):
+ if !contextType.Implements(in) {
+ return paramPlan{}, fmt.Errorf(
+ "task function %s: parameter %d: interface %s
adds methods on top of "+
+ "context.Context; declare
sdk.TIRunContext or a separate parameter instead",
+ fnName, index, in,
+ )
+ }
+ return paramPlan{kind: paramContext, index: index}, nil
+ case isLogger(in):
+ return paramPlan{kind: paramLogger, index: index}, nil
+ case isClient(in):
+ return paramPlan{kind: paramClient, index: index}, nil
+ }
+ if in.Kind() == reflect.Interface && in.NumMethod() > 0 {
+ return paramPlan{}, fmt.Errorf(
+ "task function %s: parameter %d: interface %s is not
injectable "+
+ "(want context.Context, sdk.TIRunContext, or a
subset of sdk.Client): %s",
+ fnName, index, in, explainClientMismatch(in),
+ )
+ }
+ if !isDecodableType(in) {
+ return paramPlan{}, fmt.Errorf(
+ "task function %s: parameter %d: type %s cannot receive
a task argument "+
+ "(func, chan and unsafe-pointer values cannot
be decoded)",
+ fnName, index, in,
+ )
+ }
+ return paramPlan{kind: paramData, typ: in, index: index}, nil
+}
+
+func structParamType(in reflect.Type) reflect.Type {
+ t := in
+ if t.Kind() == reflect.Pointer {
+ t = t.Elem()
+ }
+ if t.Kind() != reflect.Struct {
+ return nil
+ }
+ return t
+}
+
+func hasArgTag(structType reflect.Type) bool {
+ for i := range structType.NumField() {
+ f := structType.Field(i)
+ if embedded := embeddedStructType(f); embedded != nil {
+ if hasArgTag(embedded) {
+ return true
+ }
+ continue
+ }
+ if f.IsExported() && f.Tag.Get("arg") != "" {
+ return true
+ }
+ }
+ return false
+}
+
+func embeddedStructType(f reflect.StructField) reflect.Type {
+ if !f.Anonymous || f.Tag.Get("arg") != "" || f.Type.Kind() !=
reflect.Struct {
+ return nil
+ }
+ return structParamType(f.Type)
+}
+
+func buildStructFields(
+ fnName string,
+ structType reflect.Type,
+ paramIndex int,
+) ([]structField, error) {
+ var fields []structField
+ // Normalized argument name to the field that claims it.
+ claimed := map[string]string{}
+ folded := map[string]string{}
+ err := collectStructFields(
+ fnName, structType, paramIndex, nil, &fields, claimed, folded,
+ )
+ if err != nil {
+ return nil, err
+ }
+ return fields, nil
+}
+
+func collectStructFields(
+ fnName string,
+ structType reflect.Type,
+ paramIndex int,
+ prefix []int,
+ fields *[]structField,
+ claimed map[string]string,
+ folded map[string]string,
+) error {
+ for i := range structType.NumField() {
+ f := structType.Field(i)
+ index := append(append([]int{}, prefix...), i)
+
+ // Embedded structs promote exported fields even when their
type is unexported.
+ if embedded := embeddedStructType(f); embedded != nil {
+ err := collectStructFields(
+ fnName, embedded, paramIndex, index, fields,
claimed, folded,
+ )
+ if err != nil {
+ return err
+ }
+ continue
+ }
+ if !f.IsExported() {
+ continue
+ }
+
+ tag := f.Tag.Get("arg")
+ if !isDecodableType(f.Type) {
+ // Only an explicitly tagged, undecodable field is an
error.
+ if tag == "" {
+ continue
+ }
+ return fmt.Errorf(
+ "task function %s: parameter %d: struct field
%s: type %s cannot receive a task "+
+ "argument (func/chan/unsafe-pointer
values cannot be decoded)",
+ fnName, paramIndex, f.Name, f.Type,
+ )
+ }
+
+ sf := structField{
+ index: index,
+ goName: f.Name,
+ fieldType: f.Type,
+ argName: tag,
+ tagged: tag != "",
+ }
+ if sf.argName == "" {
+ sf.argName = f.Name
+ }
+ if existing, ok := claimed[sf.argName]; ok {
+ return fmt.Errorf(
+ "task function %s: parameter %d: struct fields
%s and %s both bind arg name %q",
+ fnName, paramIndex, existing, f.Name,
sf.argName,
+ )
+ }
+ claimed[sf.argName] = f.Name
+ key := foldArgName(sf.argName)
+ if existing, ok := folded[key]; ok {
+ return fmt.Errorf(
+ "task function %s: parameter %d: struct fields
%s and %s bind arg names that "+
+ "differ only in case or underscores
(%q), which the untagged fallback cannot "+
+ "tell apart",
+ fnName, paramIndex, existing, f.Name,
sf.argName,
+ )
+ }
+ folded[key] = f.Name
+ *fields = append(*fields, sf)
+ }
+ return nil
+}
+
+func foldArgName(name string) string {
+ return strings.ToLower(strings.ReplaceAll(name, "_", ""))
+}
+
+type schemaShape struct {
+ jsonType string
+ format string
+ fragment map[string]any
+}
+
+func shapeOf(fragment map[string]any) schemaShape {
+ jsonType, _ := fragment["type"].(string)
+ format, _ := fragment["format"].(string)
+ return schemaShape{jsonType: jsonType, format: format, fragment:
fragment}
+}
+
+// Unknown schema forms are treated as unconstrained.
+func schemaShapes(
+ schema *genmodels.ArgValueSchema,
+) (shapes []schemaShape, nullable bool, ok bool) {
+ if schema == nil {
+ return nil, false, false
+ }
+ if branches, isUnion := (*schema)["anyOf"].([]any); isUnion {
+ for _, branch := range branches {
+ fragment, isMap := branch.(map[string]any)
+ if !isMap {
+ return nil, false, false
+ }
+ jsonType, _ := fragment["type"].(string)
+ switch jsonType {
+ case "null":
+ nullable = true
+ case "":
+ // One unknown branch makes the union
unconstrained.
+ return nil, false, false
+ default:
+ shapes = append(shapes, shapeOf(fragment))
+ }
+ }
+ return shapes, nullable, len(shapes) > 0
+ }
+ if _, isString := (*schema)["type"].(string); !isString {
+ return nil, false, false
+ }
+ // Copy the defined genmodels.JsonValue map into the plain map type.
+ fragment := make(map[string]any, len(*schema))
+ for k, v := range *schema {
+ fragment[k] = v
+ }
+ return []schemaShape{shapeOf(fragment)}, false, true
+}
+
+func checkValueType(schema *genmodels.ArgValueSchema, target reflect.Type)
error {
+ shapes, nullable, ok := schemaShapes(schema)
+ if !ok {
+ return nil
+ }
+ if nullable && !isNilableType(target) {
+ return fmt.Errorf(
+ "the Dag may pass null for this argument, so Go
parameter type %s must be a "+
+ "pointer (or a slice, map or any)",
+ target,
+ )
+ }
+ t := target
+ if t.Kind() == reflect.Pointer {
+ t = t.Elem()
+ }
+ if t.Kind() == reflect.Interface {
+ return nil
+ }
+ for _, shape := range shapes {
+ if isBindableShape(shape, t) {
+ return nil
+ }
+ }
+ return fmt.Errorf(
+ "the Dag declares %s which cannot bind to Go parameter type %s",
+ describeShapes(shapes), target,
+ )
+}
+
+func isBindableShape(shape schemaShape, t reflect.Type) bool {
+ // Self-decoding types define their own wire representation.
+ if implementsUnmarshaler(t) {
+ return true
+ }
+ switch shape.jsonType {
+ case "string":
+ return t.Kind() == reflect.String
+ case "integer":
+ switch t.Kind() {
+ case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
reflect.Int64,
+ reflect.Uint, reflect.Uint8, reflect.Uint16,
reflect.Uint32, reflect.Uint64,
+ // JSON integers can decode into Go floats.
+ reflect.Float32, reflect.Float64:
+ return true
+ }
+ return false
+ case "number":
+ return t.Kind() == reflect.Float32 || t.Kind() ==
reflect.Float64
+ case "boolean":
+ return t.Kind() == reflect.Bool
+ case "object":
+ return t.Kind() == reflect.Struct || t.Kind() == reflect.Map
+ case "array":
+ return t.Kind() == reflect.Slice || t.Kind() == reflect.Array
+ default:
+ // A type keyword this runtime does not recognize; leave it to
the decode.
+ return true
+ }
+}
+
+func isNilableType(t reflect.Type) bool {
+ switch t.Kind() {
+ case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface:
+ return true
+ }
+ return false
+}
+
+func describeShapes(shapes []schemaShape) string {
+ parts := make([]string, len(shapes))
+ for i, shape := range shapes {
+ parts[i] = fmt.Sprintf("type %q", shape.jsonType)
+ if shape.format != "" {
+ parts[i] += fmt.Sprintf(" format %q", shape.format)
+ }
+ }
+ if len(parts) == 1 {
+ return "JSON-schema " + parts[0]
+ }
+ return "JSON-schema alternatives " + strings.Join(parts, " | ")
+}
+
+func decodeValue(raw any, target reflect.Type) (reflect.Value, error) {
+ out := reflect.New(target)
+
+ if raw == nil {
+ if isNilableType(target) {
+ return out.Elem(), nil
+ }
+ return reflect.Value{}, fmt.Errorf(
+ "value is null but the parameter type %s is not
nilable", target,
+ )
+ }
+
+ blob, err := json.Marshal(raw)
+ if err != nil {
+ return reflect.Value{}, err
+ }
+ dec := json.NewDecoder(bytes.NewReader(blob))
+ dec.DisallowUnknownFields()
+ if err := dec.Decode(out.Interface()); err != nil {
+ return reflect.Value{}, err
+ }
+ return out.Elem(), nil
+}
+
+// Struct fields are checked only if the wire value names them.
+func isDecodableType(inType reflect.Type) bool {
+ if implementsUnmarshaler(inType) {
+ return true
+ }
+ switch inType.Kind() {
+ case reflect.Func, reflect.Chan, reflect.UnsafePointer:
+ return false
+ case reflect.Interface:
+ return inType.NumMethod() == 0
+ case reflect.Pointer, reflect.Slice, reflect.Array:
+ return isDecodableType(inType.Elem())
+ case reflect.Map:
+ return isDecodableType(inType.Key()) &&
isDecodableType(inType.Elem())
+ }
+ return true
+}
+
+func implementsUnmarshaler(t reflect.Type) bool {
+ for _, cand := range []reflect.Type{t, reflect.PointerTo(t)} {
+ if cand.Implements(jsonUnmarshalerType) ||
cand.Implements(textUnmarshalerType) {
+ return true
+ }
+ }
+ return false
+}
+
+var (
+ contextType = reflect.TypeFor[context.Context]()
+ tiRunContextType = reflect.TypeFor[sdk.TIRunContext]()
+ slogLoggerType = reflect.TypeFor[*slog.Logger]()
+ clientType = reflect.TypeFor[sdk.Client]()
+
+ jsonUnmarshalerType = reflect.TypeFor[json.Unmarshaler]()
+ textUnmarshalerType = reflect.TypeFor[encoding.TextUnmarshaler]()
+)
+
+func isContext(inType reflect.Type) bool {
+ return inType != nil && inType.Implements(contextType)
+}
+
+func isTIRunContext(inType reflect.Type) bool {
+ return inType == tiRunContextType
+}
+
+func isLogger(inType reflect.Type) bool {
+ return inType != nil && inType.AssignableTo(slogLoggerType)
+}
+
+// isClient reports whether inType is a non-empty subset of sdk.Client.
+func isClient(inType reflect.Type) bool {
+ return inType != nil && inType.Kind() == reflect.Interface &&
+ inType.NumMethod() > 0 && clientType.Implements(inType)
+}
+
+func explainClientMismatch(in reflect.Type) string {
+ for i := range in.NumMethod() {
+ m := in.Method(i)
+ cm, ok := clientType.MethodByName(m.Name)
+ if !ok {
+ return fmt.Sprintf("sdk.Client has no method %s",
m.Name)
+ }
+ if cm.Type != m.Type {
+ return fmt.Sprintf("method %s is %s on sdk.Client, not
%s", m.Name, cm.Type, m.Type)
+ }
+ }
+ return "its method set is not a subset of sdk.Client"
+}
diff --git a/go-sdk/pkg/binding/binding_test.go
b/go-sdk/pkg/binding/binding_test.go
new file mode 100644
index 00000000000..394df3a2d6d
--- /dev/null
+++ b/go-sdk/pkg/binding/binding_test.go
@@ -0,0 +1,976 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package binding
+
+import (
+ "context"
+ "log/slog"
+ "reflect"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/google/uuid"
+ "github.com/stretchr/testify/suite"
+
+ "github.com/apache/airflow/go-sdk/pkg/api"
+ "github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
+ "github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+type BindingSuite struct {
+ suite.Suite
+}
+
+func TestBindingSuite(t *testing.T) {
+ suite.Run(t, &BindingSuite{})
+}
+
+func argSchema(jsonType string) *genmodels.ArgValueSchema {
+ s := genmodels.ArgValueSchema{"type": jsonType}
+ return &s
+}
+
+func anyOfSchema(branches ...map[string]any) *genmodels.ArgValueSchema {
+ alternatives := make([]any, len(branches))
+ for i, b := range branches {
+ alternatives[i] = b
+ }
+ s := genmodels.ArgValueSchema{"anyOf": alternatives}
+ return &s
+}
+
+// fakeXComClient records concurrent GetXCom calls.
+type fakeXComClient struct {
+ sdk.Client
+
+ values map[string]any // "<task_id>/<key>" -> raw value
+ mu sync.Mutex
+ calls []fakeXComCall
+ err error
+}
+
+type fakeXComCall struct {
+ dagID, runID, taskID, key string
+ mapIndex *int
+}
+
+func (f *fakeXComClient) GetXCom(
+ ctx context.Context,
+ dagID, runID, taskID string,
+ mapIndex *int,
+ key string,
+ _ any,
+) (any, error) {
+ f.mu.Lock()
+ f.calls = append(f.calls, fakeXComCall{dagID, runID, taskID, key,
mapIndex})
+ f.mu.Unlock()
+ if f.err != nil {
+ return nil, f.err
+ }
+ return f.values[taskID+"/"+key], nil
+}
+
+func workloadCtx() context.Context {
+ return context.WithValue(
+ context.Background(),
+ sdkcontext.WorkloadContextKey,
+ api.ExecuteTaskWorkload{
+ TI: api.TaskInstance{
+ Id: uuid.New(),
+ DagId: "dag1",
+ RunId: "run1",
+ TaskId: "transform",
+ },
+ },
+ )
+}
+
+func analyze(s *BindingSuite, fn any) *Plan {
+ plan, err := Analyze(reflect.TypeOf(fn), "testFn")
+ s.Require().NoError(err)
+ return plan
+}
+
+func (s *BindingSuite) resolve(fn any, args []Arg, client sdk.Client)
([]reflect.Value, error) {
+ plan := analyze(s, fn)
+ return plan.Resolve(workloadCtx(), slog.Default(), client, args)
+}
+
+func (s *BindingSuite) TestAnalyzeClassification() {
+ plan := analyze(
+ s,
+ func(ctx sdk.TIRunContext, log *slog.Logger, c
sdk.VariableClient, country string, extracted map[string]any) error {
+ return nil
+ },
+ )
+ s.Equal(2, plan.numData)
+
+ s.Zero(analyze(s, func() error { return nil }).numData)
+ s.Equal(
+ 1,
+ analyze(s, func(x any) error { return nil }).numData,
+ "an `any` parameter is a data parameter",
+ )
+}
+
+func (s *BindingSuite) TestAnalyzeRejections() {
+ cases := map[string]struct {
+ fn any
+ errContains string
+ }{
+ "func-param": {
+ func(cb func()) error { return nil },
+ "cannot receive a task argument",
+ },
+ "chan-param": {
+ func(ch chan int) error { return nil },
+ "cannot receive a task argument",
+ },
+ "pointer-to-func-param": {
+ func(cb *func()) error { return nil },
+ "cannot receive a task argument",
+ },
+ "slice-of-func-param": {
+ func(cbs []func()) error { return nil },
+ "cannot receive a task argument",
+ },
+ "map-with-chan-value-param": {
+ func(m map[string]chan int) error { return nil },
+ "cannot receive a task argument",
+ },
+ "non-client-interface": {
+ func(x interface{ NotAClientMethod() }) error { return
nil },
+ "sdk.Client has no method NotAClientMethod",
+ },
+ "context-with-extra-methods": {
+ func(x interface {
+ context.Context
+ TaskInstance() sdk.TaskInstance
+ },
+ ) error {
+ return nil
+ },
+ "adds methods on top of context.Context",
+ },
+ }
+ for name, tt := range cases {
+ s.Run(name, func() {
+ _, err := Analyze(reflect.TypeOf(tt.fn), "testFn")
+ if s.Assert().Error(err) {
+ s.Assert().Contains(err.Error(), tt.errContains)
+ }
+ })
+ }
+}
+
+// selfDecodingNode must bypass recursive field inspection.
+type selfDecodingNode struct {
+ Cb func()
+ Next *selfDecodingNode
+}
+
+func (n *selfDecodingNode) UnmarshalJSON([]byte) error { return nil }
+
+type recursiveNode struct {
+ Name string
+ Next *recursiveNode
+}
+
+func (s *BindingSuite) TestAnalyzeAcceptsSelfDecodingAndRecursiveTypes() {
+ for name, fn := range map[string]any{
+ "time.Time": func(when time.Time) error { return nil },
+ "slice-of-time": func(when []time.Time) error { return nil
},
+ "self-decoding": func(name string, n selfDecodingNode)
error { return nil },
+ "self-decoding-sole": func(n selfDecodingNode) error { return
nil },
+ "recursive-struct": func(n recursiveNode) error { return nil
},
+ } {
+ s.Run(name, func() {
+ _, err := Analyze(reflect.TypeOf(fn), "testFn")
+ s.Assert().NoError(err)
+ })
+ }
+}
+
+func (s *BindingSuite) TestNamedClientInterfacesAreInjectable() {
+ for name, typ := range map[string]reflect.Type{
+ "Client": reflect.TypeFor[sdk.Client](),
+ "VariableClient": reflect.TypeFor[sdk.VariableClient](),
+ "ConnectionClient": reflect.TypeFor[sdk.ConnectionClient](),
+ "XComClient": reflect.TypeFor[sdk.XComClient](),
+ } {
+ s.True(isClient(typ), "sdk.%s must stay injectable", name)
+ }
+}
+
+func (s *BindingSuite) TestResolveArityMismatch() {
+ fn := func(country string) error { return nil }
+ _, err := s.resolve(fn, nil, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "argument count mismatch")
+ s.Contains(err.Error(), "passes 0 positional argument(s)")
+ s.Contains(err.Error(), "declares 1 data parameter(s)")
+ }
+
+ _, err = s.resolve(
+ func() error { return nil },
+ []Arg{LiteralArg{Value: "uk"}},
+ &fakeXComClient{},
+ )
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "argument count mismatch")
+ }
+}
+
+func (s *BindingSuite) TestResolveLiterals() {
+ fn := func(country string, count int, ratio float64, on bool, tags
[]string, meta map[string]any) error {
+ return nil
+ }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Value: "uk", ValueSchema: argSchema("string")},
+ LiteralArg{Value: 3, ValueSchema: argSchema("integer")},
+ LiteralArg{Value: 1.5, ValueSchema: argSchema("number")},
+ LiteralArg{Value: true, ValueSchema: argSchema("boolean")},
+ LiteralArg{Value: []any{"a", "b"}, ValueSchema:
argSchema("array")},
+ LiteralArg{Value: map[string]any{"k": "v"}, ValueSchema:
argSchema("object")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal("uk", got[0].Interface())
+ s.Equal(3, got[1].Interface())
+ s.Equal(1.5, got[2].Interface())
+ s.Equal(true, got[3].Interface())
+ s.Equal([]string{"a", "b"}, got[4].Interface())
+ s.Equal(map[string]any{"k": "v"}, got[5].Interface())
+}
+
+func (s *BindingSuite) TestResolveSelfDecodingLiterals() {
+ fn := func(when time.Time, id uuid.UUID, ratio float64) error { return
nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Value: "2024-01-02T03:04:05Z",
+ ValueSchema: &genmodels.ArgValueSchema{"type":
"string", "format": "date-time"},
+ },
+ LiteralArg{
+ Value: "6ba7b810-9dad-11d1-80b4-00c04fd430c8",
+ ValueSchema: &genmodels.ArgValueSchema{"type":
"string", "format": "uuid"},
+ },
+ LiteralArg{Value: 3, ValueSchema: argSchema("integer")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC), got[0].Interface())
+ s.Equal(uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8"),
got[1].Interface())
+ s.Equal(3.0, got[2].Interface())
+}
+
+func (s *BindingSuite) TestResolveTypedMapParam() {
+ fn := func(labels map[string]string) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "labels",
+ Value: map[string]any{"team": "data", "tier":
"gold"},
+ ValueSchema: argSchema("object"),
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(map[string]string{"team": "data", "tier": "gold"},
got[0].Interface())
+}
+
+func (s *BindingSuite) TestResolveInterleavedInjectables() {
+ fn := func(log *slog.Logger, country string, ctx context.Context, meta
map[string]any) error {
+ return nil
+ }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Value: "uk", ValueSchema: argSchema("string")},
+ LiteralArg{Value: map[string]any{"k": "v"}, ValueSchema:
argSchema("object")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.NotNil(got[0].Interface().(*slog.Logger))
+ s.Equal("uk", got[1].Interface())
+ s.NotNil(got[2].Interface().(context.Context))
+ s.Equal(map[string]any{"k": "v"}, got[3].Interface())
+}
+
+func (s *BindingSuite) TestCheckValueTypeMatrix() {
+ unionType := &genmodels.ArgValueSchema{"type": []any{"string", "null"}}
+ cases := map[string]struct {
+ schema *genmodels.ArgValueSchema
+ target reflect.Type
+ errContains string
+ }{
+ "string-ok": {argSchema("string"),
reflect.TypeFor[string](), ""},
+ "string-ptr-ok": {argSchema("string"),
reflect.TypeFor[*string](), ""},
+ "string-vs-int": {argSchema("string"),
reflect.TypeFor[int](), "cannot bind"},
+ "integer-ok": {argSchema("integer"),
reflect.TypeFor[int64](), ""},
+ "integer-uint-ok": {argSchema("integer"),
reflect.TypeFor[uint32](), ""},
+ "integer-float-ok": {argSchema("integer"),
reflect.TypeFor[float64](), ""},
+ "number-ok": {argSchema("number"),
reflect.TypeFor[float32](), ""},
+ "number-vs-int": {argSchema("number"),
reflect.TypeFor[int](), "cannot bind"},
+ "boolean-ok": {argSchema("boolean"),
reflect.TypeFor[bool](), ""},
+ "boolean-vs-string": {argSchema("boolean"),
reflect.TypeFor[string](), "cannot bind"},
+ "object-map-ok": {argSchema("object"),
reflect.TypeFor[map[string]int](), ""},
+ "object-struct-ok": {argSchema("object"),
reflect.TypeFor[struct{ A int }](), ""},
+ "object-vs-slice": {argSchema("object"),
reflect.TypeFor[[]int](), "cannot bind"},
+ "array-slice-ok": {argSchema("array"),
reflect.TypeFor[[]string](), ""},
+ "array-array-ok": {argSchema("array"),
reflect.TypeFor[[2]int](), ""},
+ "array-vs-map": {argSchema("array"),
reflect.TypeFor[map[string]any](), "cannot bind"},
+ // Schemas without a usable type are unconstrained.
+ "nil-schema-skips": {nil, reflect.TypeFor[chan int](), ""},
+ "no-type-skips": {&genmodels.ArgValueSchema{},
reflect.TypeFor[string](), ""},
+ "union-type-skips": {unionType, reflect.TypeFor[int](), ""},
+ "unknown-type-skips": {argSchema("uuid"),
reflect.TypeFor[string](), ""},
+ "any-target-skips": {argSchema("string"),
reflect.TypeFor[any](), ""},
+
+ // Pydantic encodes bytes as a string.
+ "binary-to-string": {
+ &genmodels.ArgValueSchema{"type": "string", "format":
"binary"},
+ reflect.TypeFor[string](), "",
+ },
+ "binary-vs-bytes": {
+ &genmodels.ArgValueSchema{"type": "string", "format":
"binary"},
+ reflect.TypeFor[[]byte](), "cannot bind",
+ },
+
+ // Unions match any branch but preserve nullability.
+ "anyof-matching-branch": {
+ anyOfSchema(map[string]any{"type": "integer"},
map[string]any{"type": "string"}),
+ reflect.TypeFor[string](), "",
+ },
+ "anyof-no-branch": {
+ anyOfSchema(map[string]any{"type": "integer"},
map[string]any{"type": "string"}),
+ reflect.TypeFor[bool](), "cannot bind",
+ },
+ "anyof-nullable-ptr-ok": {
+ anyOfSchema(
+ map[string]any{"type": "string"},
+ map[string]any{"type": "null"},
+ ),
+ reflect.TypeFor[*string](), "",
+ },
+ "anyof-nullable-needs-pointer": {
+ anyOfSchema(
+ map[string]any{"type": "string"},
+ map[string]any{"type": "null"},
+ ),
+ reflect.TypeFor[string](), "must be a pointer",
+ },
+ "anyof-nullable-wrong-type": {
+ anyOfSchema(
+ map[string]any{"type": "string"},
+ map[string]any{"type": "null"},
+ ),
+ reflect.TypeFor[*int64](), "cannot bind",
+ },
+ "anyof-unreadable-branch-skips": {
+ anyOfSchema(map[string]any{"type": "string"},
map[string]any{"$ref": "#/x"}),
+ reflect.TypeFor[bool](), "",
+ },
+ }
+ for name, tt := range cases {
+ s.Run(name, func() {
+ err := checkValueType(tt.schema, tt.target)
+ if tt.errContains == "" {
+ s.NoError(err)
+ } else if s.Assert().Error(err) {
+ s.Contains(err.Error(), tt.errContains)
+ }
+ })
+ }
+}
+
+func (s *BindingSuite) TestResolveTypeMismatchFailsLoudly() {
+ fn := func(count int) error { return nil }
+ _, err := s.resolve(
+ fn,
+ []Arg{LiteralArg{Value: "uk", ValueSchema:
argSchema("string")}},
+ &fakeXComClient{},
+ )
+ if s.Assert().Error(err) {
+ s.Contains(
+ err.Error(),
+ `the Dag declares JSON-schema type "string" which
cannot bind to Go parameter type int`,
+ )
+ }
+}
+
+func (s *BindingSuite) TestResolveLiteralDecodeFailure() {
+ fn := func(count int) error { return nil }
+ _, err := s.resolve(fn, []Arg{LiteralArg{Value: "uk"}},
&fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "decoding literal value into int")
+ }
+}
+
+type extractResult struct {
+ GoVersion string `json:"go_version"`
+ Timestamp int64 `json:"timestamp"`
+}
+
+type simpleInput struct {
+ Name string
+}
+
+type twoFieldInput struct {
+ Name string
+ Missing string `arg:"missing"`
+}
+
+type wholeConfig struct {
+ Environment string `json:"environment"`
+ Region string `json:"region"`
+}
+
+type combineInput struct {
+ Name string
+ Count int `arg:"count"`
+}
+
+type reportInput struct {
+ Ratio float64
+ Region string `arg:"region"`
+}
+
+func (s *BindingSuite) TestResolveXComArgs() {
+ client := &fakeXComClient{values: map[string]any{
+ "extract/return_value": map[string]any{"go_version": "go1.24",
"timestamp": int64(42)},
+ "probe/return_value": "probe-value",
+ }}
+
+ fn := func(res extractResult, probe string) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ XComArg{TaskID: "extract", ValueSchema: argSchema("object")},
+ XComArg{TaskID: "probe", ValueSchema: argSchema("string")},
+ }, client)
+ s.Require().NoError(err)
+ s.Equal(extractResult{GoVersion: "go1.24", Timestamp: 42},
got[0].Interface())
+ s.Equal("probe-value", got[1].Interface())
+
+ s.Require().Len(client.calls, 2)
+ taskIDs := make([]string, 0, 2)
+ for _, call := range client.calls {
+ // Pull order is nondeterministic.
+ taskIDs = append(taskIDs, call.taskID)
+ s.Equal("dag1", call.dagID)
+ s.Equal("run1", call.runID)
+ s.Equal(
+ api.XComReturnValueKey,
+ call.key,
+ "an XCom argument always pulls the return-value key",
+ )
+ s.Nil(call.mapIndex, "v1 always pulls the unmapped upstream
instance")
+ }
+ s.ElementsMatch([]string{"extract", "probe"}, taskIDs)
+}
+
+func (s *BindingSuite) TestResolveXComStrictStructDecode() {
+ client := &fakeXComClient{values: map[string]any{
+ "extract/return_value": map[string]any{"go_version": "go1.24",
"renamed_field": 1},
+ }}
+ fn := func(res extractResult) error { return nil }
+ _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client)
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), `decoding xcom from task "extract"`)
+ s.Contains(err.Error(), "unknown field")
+ }
+}
+
+func (s *BindingSuite) TestResolveXComPullFailure() {
+ client := &fakeXComClient{err: sdk.XComNotFound}
+ fn := func(res map[string]any) error { return nil }
+ _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client)
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), `pulling xcom from task "extract"`)
+ }
+}
+
+func (s *BindingSuite) TestResolveMultipleXComPullFailures() {
+ client := &fakeXComClient{err: sdk.XComNotFound}
+ fn := func(a map[string]any, b map[string]any, c map[string]any) error
{ return nil }
+ _, err := s.resolve(fn, []Arg{
+ XComArg{TaskID: "extract_a"},
+ XComArg{TaskID: "extract_b"},
+ XComArg{TaskID: "extract_c"},
+ }, client)
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), `pulling xcom from task "extract_a"`)
+ s.Contains(err.Error(), `pulling xcom from task "extract_b"`)
+ s.Contains(err.Error(), `pulling xcom from task "extract_c"`)
+ }
+}
+
+func (s *BindingSuite) TestResolveWholeStructFromXCom() {
+ client := &fakeXComClient{values: map[string]any{
+ "make_config/return_value": map[string]any{
+ "environment": "production",
+ "region": "eu-west-1",
+ },
+ }}
+ fn := func(cfg wholeConfig) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ XComArg{Name: "cfg", TaskID: "make_config", ValueSchema:
argSchema("object")},
+ }, client)
+ s.Require().NoError(err)
+ s.Equal(
+ wholeConfig{Environment: "production", Region: "eu-west-1"},
+ got[0].Interface(),
+ "an XCom argument no field claims decodes whole into the
struct",
+ )
+}
+
+func (s *BindingSuite) TestResolveXComWithoutWorkload() {
+ plan := analyze(s, func(res map[string]any) error { return nil })
+ _, err := plan.Resolve(
+ context.Background(), slog.Default(), &fakeXComClient{},
+ []Arg{XComArg{TaskID: "extract"}},
+ )
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "no workload in context")
+ }
+}
+
+func (s *BindingSuite) TestResolveNullHandling() {
+ fn := func(meta map[string]any) error { return nil }
+ got, err := s.resolve(
+ fn,
+ []Arg{LiteralArg{Value: nil, ValueSchema: argSchema("object")}},
+ &fakeXComClient{},
+ )
+ s.Require().NoError(err)
+ s.Nil(got[0].Interface())
+
+ fnStr := func(country string) error { return nil }
+ _, err = s.resolve(fnStr, []Arg{LiteralArg{Value: nil}},
&fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "not nilable")
+ }
+}
+
+// fakeArg exercises the defensive branch of the sealed interface.
+type fakeArg struct{}
+
+func (fakeArg) ArgName() string { return "fake" }
+func (fakeArg) Schema() *genmodels.ArgValueSchema { return nil }
+func (fakeArg) sealedArg() {}
+
+func (s *BindingSuite) TestResolveUnsupportedVariant() {
+ fn := func(country string) error { return nil }
+ _, err := s.resolve(fn, []Arg{fakeArg{}}, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "unsupported argument binding
binding.fakeArg")
+ }
+}
+
+func (s *BindingSuite) TestResolveNilArg() {
+ fn := func(country string) error { return nil }
+ _, err := s.resolve(fn, []Arg{nil}, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "nil argument binding")
+ }
+}
+
+func (s *BindingSuite) TestResolveTIRunContextRebuild() {
+ ti := sdk.TaskInstance{DagID: "dag1", RunID: "run1", TaskID:
"transform"}
+ dagRun := sdk.DagRun{DagID: "dag1", RunID: "run1"}
+ ctx := context.WithValue(
+ workloadCtx(),
+ sdkcontext.RuntimeContextKey,
+ sdk.NewTIRunContext(context.Background(), ti, dagRun),
+ )
+
+ plan := analyze(s, func(rc sdk.TIRunContext, country string) error {
return nil })
+ got, err := plan.Resolve(ctx, slog.Default(), &fakeXComClient{}, []Arg{
+ LiteralArg{Value: "uk", ValueSchema: argSchema("string")},
+ })
+ s.Require().NoError(err)
+ rc := got[0].Interface().(sdk.TIRunContext)
+ s.Equal(ti, rc.TaskInstance())
+ s.Equal(dagRun, rc.DagRun())
+ s.Equal("uk", got[1].Interface())
+}
+
+func (s *BindingSuite) TestAnalyzeLoneStructClassification() {
+ plan := analyze(s, func(input simpleInput) error { return nil })
+ s.True(plan.loneStruct, "a sole struct data parameter is resolved by
name at execution")
+ s.Zero(plan.numData)
+
+ ptrPlan := analyze(s, func(input *simpleInput) error { return nil })
+ s.True(ptrPlan.loneStruct, "a pointer to a sole struct is detected the
same way")
+ s.Zero(ptrPlan.numData)
+
+ flatPlan := analyze(s, func(prefix string, cfg wholeConfig) error {
return nil })
+ s.False(flatPlan.loneStruct, "a struct alongside another data parameter
is a flat slot")
+ s.Equal(2, flatPlan.numData)
+
+ scalarPlan := analyze(s, func(name string) error { return nil })
+ s.False(scalarPlan.loneStruct, "a sole non-struct data parameter is
plain positional")
+ s.Equal(1, scalarPlan.numData)
+}
+
+func (s *BindingSuite) TestAnalyzeMultipleStructsAreFlat() {
+ plan := analyze(s, func(a wholeConfig, b wholeConfig) error { return
nil })
+ s.False(plan.loneStruct)
+ s.Equal(2, plan.numData)
+}
+
+func (s *BindingSuite) TestAnalyzeStructValidation() {
+ type duplicateArgNames struct {
+ A string
+ B string `arg:"A"`
+ }
+ type taggedNonDecodableField struct {
+ Bad chan int `arg:"bad"`
+ }
+ type foldedDuplicateArgNames struct {
+ RegionCode string
+ Region_code string
+ }
+
+ cases := map[string]struct {
+ fn any
+ errContains string
+ }{
+ "duplicate-arg-names": {
+ func(input duplicateArgNames) error { return nil },
+ `fields A and B both bind arg name "A"`,
+ },
+ "folded-duplicate-arg-names": {
+ func(input foldedDuplicateArgNames) error { return nil
},
+ "differ only in case or underscores",
+ },
+ "tagged-non-decodable-field": {
+ func(input taggedNonDecodableField) error { return nil
},
+ "cannot receive a task argument",
+ },
+ "tagged-struct-not-sole": {
+ func(prefix string, input combineInput) error { return
nil },
+ "must be the function's only data parameter",
+ },
+ "tagged-struct-trailing": {
+ func(input combineInput, suffix string) error { return
nil },
+ "must be the function's only data parameter",
+ },
+ }
+ for name, tt := range cases {
+ s.Run(name, func() {
+ _, err := Analyze(reflect.TypeOf(tt.fn), "testFn")
+ if s.Assert().Error(err) {
+ s.Assert().Contains(err.Error(), tt.errContains)
+ }
+ })
+ }
+}
+
+func (s *BindingSuite) TestResolveStructAllFields() {
+ fn := func(input combineInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ LiteralArg{Name: "count", Value: 7, ValueSchema:
argSchema("integer")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+
+ input := got[0].Interface().(combineInput)
+ s.Equal("widget", input.Name, "the untagged field claims its verbatim
field name")
+ s.Equal(7, input.Count, "the `arg:` tag claims its named entry")
+}
+
+func (s *BindingSuite) TestResolveStructXComArg() {
+ fn := func(log *slog.Logger, input reportInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ XComArg{Name: "region", TaskID: "make_region", ValueSchema:
argSchema("string")},
+ LiteralArg{Name: "Ratio", Value: 0.5, ValueSchema:
argSchema("number")},
+ }, &fakeXComClient{values: map[string]any{"make_region/return_value":
"east"}})
+ s.Require().NoError(err)
+
+ input := got[1].Interface().(reportInput)
+ s.Equal("east", input.Region, "Region resolves by name despite being
declared after Ratio")
+ s.Equal(0.5, input.Ratio)
+}
+
+func (s *BindingSuite) TestResolveStructSingleClaimedArgBindsByName() {
+ fn := func(input simpleInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal("widget", got[0].Interface().(simpleInput).Name)
+}
+
+func (s *BindingSuite) TestResolveStructPointer() {
+ fn := func(input *simpleInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ input := got[0].Interface().(*simpleInput)
+ s.Require().NotNil(input)
+ s.Equal("widget", input.Name)
+}
+
+func (s *BindingSuite) TestResolveLoneStructBothModes() {
+ fn := func(cfg wholeConfig) error { return nil }
+ want := wholeConfig{Environment: "production", Region: "eu-west-1"}
+
+ named, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Environment", Value: "production",
ValueSchema: argSchema("string")},
+ LiteralArg{Name: "Region", Value: "eu-west-1", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(want, named[0].Interface(), "argument names matching the fields
bind field-by-field")
+
+ whole, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "cfg",
+ Value: map[string]any{"environment":
"production", "region": "eu-west-1"},
+ ValueSchema: argSchema("object"),
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(want, whole[0].Interface(), "a single unclaimed argument
decodes whole into the struct")
+}
+
+type taggedRegionInput struct {
+ Region string `arg:"regon_code"` // deliberate typo
+}
+
+func (s *BindingSuite) TestResolveTaggedStructNeverFallsBackToWholeValue() {
+ fn := func(input taggedRegionInput) error { return nil }
+ _, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "region_code", Value: "eu-west-1",
ValueSchema: argSchema("string")},
+ }, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), `not claimed by any struct field:
"region_code"`)
+ }
+}
+
+func (s *BindingSuite) TestResolveStructUnclaimedArgFailsLoudly() {
+ fn := func(input combineInput) error { return nil }
+ _, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ LiteralArg{Name: "typo", Value: "x", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), `not claimed by any struct field:
"typo"`)
+ }
+}
+
+func (s *BindingSuite) TestResolveStructUnclaimedFromDefaultAllowed() {
+ fn := func(input combineInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ LiteralArg{
+ Name: "threshold",
+ Value: 0.75,
+ ValueSchema: argSchema("number"),
+ FromDefault: true,
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal("widget", got[0].Interface().(combineInput).Name)
+}
+
+func (s *BindingSuite) TestResolveStructEmptySpecFailsLoudly() {
+ fn := func(input simpleInput) error { return nil }
+ for name, args := range map[string][]Arg{"nil-spec": nil, "empty-spec":
{}} {
+ s.Run(name, func() {
+ _, err := s.resolve(fn, args, &fakeXComClient{})
+ if s.Assert().Error(err) {
+ s.Contains(err.Error(), "no TaskFlow arg
bindings arrived")
+ }
+ })
+ }
+}
+
+func (s *BindingSuite) TestResolveStructOnlyDefaultsZeroValues() {
+ fn := func(input twoFieldInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "threshold",
+ Value: 0.75,
+ ValueSchema: argSchema("number"),
+ FromDefault: true,
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ input := got[0].Interface().(twoFieldInput)
+ s.Equal("", input.Name, "no explicit entry arrived; fields keep
kwarg-style zero values")
+ s.Equal("", input.Missing)
+}
+
+func (s *BindingSuite) TestResolveStructUnmatchedFieldZeroValued() {
+ fn := func(input twoFieldInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "Name", Value: "widget", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ input := got[0].Interface().(twoFieldInput)
+ s.Equal("widget", input.Name, "the matched field binds normally")
+ s.Equal("", input.Missing, "the unmatched field is left at its Go zero
value, not an error")
+}
+
+func (s *BindingSuite) TestResolveFlatParamsToleratesCapturedDefaults() {
+ fn := func(country string) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "country", Value: "uk", ValueSchema:
argSchema("string")},
+ LiteralArg{
+ Name: "verbose",
+ Value: false,
+ ValueSchema: argSchema("boolean"),
+ FromDefault: true,
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal("uk", got[0].Interface())
+}
+
+func (s *BindingSuite) TestResolveFlatParamsBindsDefaultsWhenDeclared() {
+ fn := func(country string, verbose bool) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "country", Value: "uk", ValueSchema:
argSchema("string")},
+ LiteralArg{
+ Name: "verbose",
+ Value: true,
+ ValueSchema: argSchema("boolean"),
+ FromDefault: true,
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal("uk", got[0].Interface())
+ s.Equal(true, got[1].Interface())
+}
+
+func (s *BindingSuite) TestResolveWholeStructIgnoresCapturedDefaults() {
+ fn := func(config wholeConfig) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "config",
+ Value: map[string]any{"environment": "prod",
"region": "eu-west-1"},
+ ValueSchema: argSchema("object"),
+ },
+ LiteralArg{
+ Name: "verbose",
+ Value: false,
+ ValueSchema: argSchema("boolean"),
+ FromDefault: true,
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(wholeConfig{Environment: "prod", Region: "eu-west-1"},
got[0].Interface())
+}
+
+type snakeCaseInput struct {
+ RegionCode string
+ Threshold float64
+}
+
+func (s *BindingSuite) TestResolveUntaggedFieldsBindSnakeCaseArguments() {
+ fn := func(input snakeCaseInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "region_code", Value: "eu-west-1",
ValueSchema: argSchema("string")},
+ LiteralArg{Name: "threshold", Value: 0.75, ValueSchema:
argSchema("number")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(snakeCaseInput{RegionCode: "eu-west-1", Threshold: 0.75},
got[0].Interface())
+}
+
+type embeddedCommon struct {
+ Region string `arg:"region"`
+}
+
+type embeddedInput struct {
+ embeddedCommon
+ Threshold float64
+}
+
+func (s *BindingSuite) TestResolveEmbeddedStructFields() {
+ fn := func(input embeddedInput) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "region", Value: "eu-west-1", ValueSchema:
argSchema("string")},
+ LiteralArg{Name: "threshold", Value: 0.75, ValueSchema:
argSchema("number")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ input := got[0].Interface().(embeddedInput)
+ s.Equal("eu-west-1", input.Region)
+ s.Equal(0.75, input.Threshold)
+}
+
+type money struct {
+ Amount string
+}
+
+func (m *money) UnmarshalJSON([]byte) error { m.Amount = "decoded"; return nil
}
+
+func (s *BindingSuite) TestResolveSelfDecodingTypeAgainstStringSchema() {
+ fn := func(price money) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{Name: "price", Value: "12.34", ValueSchema:
argSchema("string")},
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(money{Amount: "decoded"}, got[0].Interface())
+}
+
+type callbackConfig struct {
+ Name string `json:"name"`
+ Cb func() `json:"-"`
+}
+
+func (s *BindingSuite) TestResolveStructWithNonBindableField() {
+ fn := func(config callbackConfig) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "config",
+ Value: map[string]any{"name": "widget"},
+ ValueSchema: argSchema("object"),
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ config := got[0].Interface().(callbackConfig)
+ s.Equal("widget", config.Name)
+ s.Nil(config.Cb)
+}
+
+func (s *BindingSuite) TestResolveEmptyInterfaceDataParam() {
+ fn := func(payload any) error { return nil }
+ got, err := s.resolve(fn, []Arg{
+ LiteralArg{
+ Name: "payload",
+ Value: map[string]any{"k": "v"},
+ ValueSchema: argSchema("object"),
+ },
+ }, &fakeXComClient{})
+ s.Require().NoError(err)
+ s.Equal(map[string]any{"k": "v"}, got[0].Interface())
+}
+
+func (s *BindingSuite) TestResolveUnboundZeroFillsDataParameters() {
+ plan := analyze(s, func(
+ ctx context.Context, country string, config wholeConfig, note
*string,
+ ) error {
+ return nil
+ })
+ got := plan.ResolveUnbound(workloadCtx(), slog.Default(),
&fakeXComClient{})
+ s.Require().Len(got, 4)
+ s.Equal("", got[1].Interface())
+ s.Equal(wholeConfig{}, got[2].Interface())
+ s.True(got[3].IsNil())
+
+ soleStruct := analyze(s, func(input combineInput) error { return nil })
+ s.Equal(
+ combineInput{},
+ soleStruct.ResolveUnbound(workloadCtx(), slog.Default(),
&fakeXComClient{})[0].Interface(),
+ )
+}
diff --git a/go-sdk/pkg/execution/frames.go b/go-sdk/pkg/execution/frames.go
index f9a246286ef..e1124c1b80d 100644
--- a/go-sdk/pkg/execution/frames.go
+++ b/go-sdk/pkg/execution/frames.go
@@ -62,6 +62,8 @@ func encodeRequest(id int64, body any) ([]byte, error) {
var buf bytes.Buffer
enc := msgpack.NewEncoder(&buf)
enc.UseCompactInts(true)
+ // Use JSON field names for user values; explicit msgpack tags still
win.
+ enc.SetCustomStructTag("json")
if err := enc.EncodeArrayLen(2); err != nil {
return nil, err
diff --git a/go-sdk/pkg/execution/frames_test.go
b/go-sdk/pkg/execution/frames_test.go
index 4536a11c692..6dd8d2a9a9a 100644
--- a/go-sdk/pkg/execution/frames_test.go
+++ b/go-sdk/pkg/execution/frames_test.go
@@ -178,3 +178,33 @@ func TestRoundTripMultipleFrames(t *testing.T) {
assert.Equal(t, expected["key"], rawToMap(t, frame.Body)["key"])
}
}
+
+func TestEncodeRequestHonoursJSONTags(t *testing.T) {
+ type userValue struct {
+ GoVersion string `json:"go_version"`
+ Ignored string `json:"-"`
+ Untagged int
+ }
+ type protocolValue struct {
+ Field string `msgpack:"wire_name" json:"json_name"`
+ }
+
+ data, err := encodeRequest(1, map[string]any{
+ "user": userValue{GoVersion: "go1.25", Ignored: "dropped",
Untagged: 7},
+ "protocol": protocolValue{Field: "v"},
+ })
+ require.NoError(t, err)
+
+ var frame []any
+ require.NoError(t, msgpack.Unmarshal(data[:], &frame))
+ require.Len(t, frame, 2)
+ body, ok := frame[1].(map[string]any)
+ require.True(t, ok)
+
+ assert.Equal(
+ t,
+ map[string]any{"go_version": "go1.25", "Untagged": int8(7)},
+ body["user"],
+ )
+ assert.Equal(t, map[string]any{"wire_name": "v"}, body["protocol"])
+}
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index e6861d8c8ad..76978294817 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -20,6 +20,10 @@ package genmodels
import "time"
+type ArgBindings []TaskArgBinding
+
+type ArgValueSchema map[string]JsonValue
+
// Schema for AssetAliasModel used in AssetEventDagRunReference.
type AssetAliasReferenceAssetEventDagRun struct {
// Name corresponds to the JSON schema field "name".
@@ -370,6 +374,9 @@ type DagCallbackRequest struct {
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
+
+ // VersionData corresponds to the JSON schema field "version_data".
+ VersionData *VersionData `msgpack:"version_data,omitempty"`
}
// Request for DAG File Parsing.
@@ -749,6 +756,9 @@ type EmailRequest struct {
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
+
+ // VersionData corresponds to the JSON schema field "version_data".
+ VersionData *VersionData `msgpack:"version_data,omitempty"`
}
type EmailRequestEmailType string
@@ -808,12 +818,22 @@ type GetAssetEventByAsset struct {
// Before corresponds to the JSON schema field "before".
Before interface{} `msgpack:"before,omitempty"`
+ // Extra corresponds to the JSON schema field "extra".
+ Extra *Extra `msgpack:"extra,omitempty"`
+
// Limit corresponds to the JSON schema field "limit".
Limit interface{} `msgpack:"limit,omitempty"`
// Name corresponds to the JSON schema field "name".
Name interface{} `msgpack:"name"`
+ // PartitionKey corresponds to the JSON schema field "partition_key".
+ PartitionKey interface{} `msgpack:"partition_key,omitempty"`
+
+ // PartitionKeyRegexpPattern corresponds to the JSON schema field
+ // "partition_key_regexp_pattern".
+ PartitionKeyRegexpPattern interface{}
`msgpack:"partition_key_regexp_pattern,omitempty"`
+
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
@@ -834,9 +854,19 @@ type GetAssetEventByAssetAlias struct {
// Before corresponds to the JSON schema field "before".
Before interface{} `msgpack:"before,omitempty"`
+ // Extra corresponds to the JSON schema field "extra".
+ Extra *Extra `msgpack:"extra,omitempty"`
+
// Limit corresponds to the JSON schema field "limit".
Limit interface{} `msgpack:"limit,omitempty"`
+ // PartitionKey corresponds to the JSON schema field "partition_key".
+ PartitionKey interface{} `msgpack:"partition_key,omitempty"`
+
+ // PartitionKeyRegexpPattern corresponds to the JSON schema field
+ // "partition_key_regexp_pattern".
+ PartitionKeyRegexpPattern interface{}
`msgpack:"partition_key_regexp_pattern,omitempty"`
+
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
}
@@ -1239,6 +1269,24 @@ type LazyDeserializedDAG struct {
LastLoaded interface{} `msgpack:"last_loaded,omitempty"`
}
+// One positional stub-task argument carrying an inline literal from the Dag
file.
+type LiteralArgBinding struct {
+ // FromDefault corresponds to the JSON schema field "from_default".
+ FromDefault bool `msgpack:"from_default,omitempty"`
+
+ // Kind corresponds to the JSON schema field "kind".
+ Kind string `msgpack:"kind"`
+
+ // Name corresponds to the JSON schema field "name".
+ Name string `msgpack:"name"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value,omitempty"`
+
+ // ValueSchema corresponds to the JSON schema field "value_schema".
+ ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"`
+}
+
type LogicalDates []time.Time
// Add a new value to be redacted in task logs.
@@ -1564,6 +1612,9 @@ type TICount struct {
// Response schema for TaskInstance run context.
type TIRunContext struct {
+ // ArgBindings corresponds to the JSON schema field "arg_bindings".
+ ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"`
+
// Connections corresponds to the JSON schema field "connections".
Connections []ConnectionResponse `msgpack:"connections,omitempty"`
@@ -1596,6 +1647,8 @@ type TIRunContext struct {
XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"`
}
+type TaskArgBinding interface{}
+
type TaskBreadcrumbsResult struct {
// Breadcrumbs corresponds to the JSON schema field "breadcrumbs".
Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem
`msgpack:"breadcrumbs"`
@@ -1636,6 +1689,9 @@ type TaskCallbackRequest struct {
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
+
+ // VersionData corresponds to the JSON schema field "version_data".
+ VersionData *VersionData `msgpack:"version_data,omitempty"`
}
type TaskIds []string
@@ -1777,19 +1833,6 @@ type TriggerDagRun struct {
type TriggerKwargs map[string]JsonValue
-type Warnings []interface{}
-
-// Variable schema for responses with fields that are needed for Runtime.
-type VariableResponse struct {
- // Key corresponds to the JSON schema field "key".
- Key string `msgpack:"key"`
-
- // Value corresponds to the JSON schema field "value".
- Value interface{} `msgpack:"value"`
-}
-
-type VersionData map[string]interface{}
-
// Update the response content part of an existing Human-in-the-loop response.
type UpdateHITLDetail struct {
// ChosenOptions corresponds to the JSON schema field "chosen_options".
@@ -1805,6 +1848,17 @@ type UpdateHITLDetail struct {
Type string `msgpack:"type,omitempty"`
}
+// Variable schema for responses with fields that are needed for Runtime.
+type VariableResponse struct {
+ // Key corresponds to the JSON schema field "key".
+ Key string `msgpack:"key"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value"`
+}
+
+type VersionData map[string]interface{}
+
type ValidateInletsAndOutlets struct {
// TIID corresponds to the JSON schema field "ti_id".
TIID string `msgpack:"ti_id"`
@@ -1824,6 +1878,8 @@ type VariableKeysResult struct {
Type string `msgpack:"type,omitempty"`
}
+type Warnings []interface{}
+
type VariableResult struct {
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
@@ -1835,6 +1891,21 @@ type VariableResult struct {
Value interface{} `msgpack:"value,omitempty"`
}
+// One positional stub-task argument pulled from an upstream task's XCom.
+type XComArgBinding struct {
+ // Kind corresponds to the JSON schema field "kind".
+ Kind string `msgpack:"kind"`
+
+ // Name corresponds to the JSON schema field "name".
+ Name string `msgpack:"name"`
+
+ // TaskID corresponds to the JSON schema field "task_id".
+ TaskID string `msgpack:"task_id"`
+
+ // ValueSchema corresponds to the JSON schema field "value_schema".
+ ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"`
+}
+
type XComCountResponse struct {
// Len corresponds to the JSON schema field "len".
Len int `msgpack:"len"`
diff --git a/go-sdk/pkg/execution/integration_test.go
b/go-sdk/pkg/execution/integration_test.go
index 25593594538..8486a005c3e 100644
--- a/go-sdk/pkg/execution/integration_test.go
+++ b/go-sdk/pkg/execution/integration_test.go
@@ -83,23 +83,33 @@ func buildBundle(t *testing.T, register
func(bundlev1.Registry)) bundlev1.Bundle
return reg
}
-// --- Tests ---
-
-func TestTaskRunnerSuccess(t *testing.T) {
- bundle := buildBundle(t, func(r bundlev1.Registry) {
- r.AddDag("test_dag").AddTask(simpleTask)
- })
-
+func newStartupDetails(
+ taskID string,
+ bindings ...genmodels.TaskArgBinding,
+) *genmodels.StartupDetails {
details := &genmodels.StartupDetails{
TI: genmodels.TaskInstance{
ID: "550e8400-e29b-41d4-a716-446655440000",
DagID: "test_dag",
- TaskID: "simpleTask",
+ TaskID: taskID,
RunID: "run1",
MapIndex: ptr(-1),
},
BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
}
+ if bindings != nil {
+ specs := genmodels.ArgBindings(bindings)
+ details.TIContext.ArgBindings = &specs
+ }
+ return details
+}
+
+func TestTaskRunnerSuccess(t *testing.T) {
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTask(simpleTask)
+ })
+
+ details := newStartupDetails("simpleTask")
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
@@ -113,16 +123,7 @@ func TestTaskRunnerFailure(t *testing.T) {
r.AddDag("test_dag").AddTask(failingTask)
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "failingTask",
- RunID: "run1",
- MapIndex: ptr(-1),
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- }
+ details := newStartupDetails("failingTask")
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
@@ -136,20 +137,9 @@ func TestTaskRunnerRetry(t *testing.T) {
r.AddDag("test_dag").AddTask(failingTask)
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "failingTask",
- RunID: "run1",
- MapIndex: ptr(-1),
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- TIContext: genmodels.TIRunContext{
- ShouldRetry: true,
- MaxTries: 3,
- },
- }
+ details := newStartupDetails("failingTask")
+ details.TIContext.ShouldRetry = true
+ details.TIContext.MaxTries = 3
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
@@ -163,15 +153,7 @@ func TestTaskRunnerTaskNotFound(t *testing.T) {
r.AddDag("test_dag").AddTask(simpleTask)
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "nonexistent",
- RunID: "run1",
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- }
+ details := newStartupDetails("nonexistent")
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
@@ -185,16 +167,7 @@ func TestTaskRunnerPanic(t *testing.T) {
r.AddDag("test_dag").AddTask(panicTask)
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "panicTask",
- RunID: "run1",
- MapIndex: ptr(-1),
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- }
+ details := newStartupDetails("panicTask")
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
@@ -208,26 +181,283 @@ func TestTaskRunnerPanicRetry(t *testing.T) {
r.AddDag("test_dag").AddTask(panicTask)
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "panicTask",
- RunID: "run1",
- MapIndex: ptr(-1),
+ details := newStartupDetails("panicTask")
+ details.TIContext.ShouldRetry = true
+ details.TIContext.MaxTries = 3
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertRetryTask(t, result, "panic: something went wrong")
+}
+
+func TestTaskRunnerBindsArgs(t *testing.T) {
+ var gotCountry string
+ var gotMeta map[string]any
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(log *slog.Logger, country string, meta
map[string]any) error {
+ gotCountry = country
+ gotMeta = meta
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{
+ "name": "country",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "string"},
+ "value": "uk",
},
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- TIContext: genmodels.TIRunContext{
- ShouldRetry: true,
- MaxTries: 3,
+ map[string]any{
+ "name": "meta",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "object"},
+ "value": map[string]any{"k": "v"},
+ },
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertSucceedTask(t, result)
+ assert.Equal(t, "uk", gotCountry)
+ assert.Equal(t, map[string]any{"k": "v"}, gotMeta)
+}
+
+func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) {
+ ran := false
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(country string, meta map[string]any) error {
+ ran = true
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{
+ "name": "country",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "string"},
+ "value": "uk",
+ },
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertTaskState(t, result, genmodels.TaskStateStateFailed)
+ assert.False(t, ran, "the task body must not run on an arity mismatch")
+}
+
+type regionInput struct {
+ Region string `arg:"region"`
+}
+
+func TestTaskRunnerBindsStructArgs(t *testing.T) {
+ var got regionInput
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(input regionInput) error {
+ got = input
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{
+ "name": "region",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "string"},
+ "value": "eu-west-1",
},
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertSucceedTask(t, result)
+ assert.Equal(t, "eu-west-1", got.Region)
+}
+
+func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) {
+ var got regionInput
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(input regionInput) error {
+ got = input
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{
+ "name": "region",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "string"},
+ "value": "eu-west-1",
+ },
+ map[string]any{
+ "name": "threshold",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "number"},
+ "value": 0.75,
+ "from_default": true,
+ },
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertSucceedTask(t, result)
+ assert.Equal(t, "eu-west-1", got.Region)
+}
+
+func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) {
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(count int) error { return nil })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{
+ "name": "count",
+ "kind": "literal",
+ "value_schema": map[string]any{"type": "string"},
+ "value": "uk",
+ },
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertTaskState(t, result, genmodels.TaskStateStateFailed)
+}
+
+func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) {
+ ran := false
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(country string) error {
+ ran = true
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{"name": "country", "kind": "template", "value":
"x"},
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertTaskState(t, result, genmodels.TaskStateStateFailed)
+ assert.False(t, ran, "the task body must not run on an unknown binding
kind")
+}
+
+func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) {
+ ran := false
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(country string) error {
+ ran = true
+ return nil
+ })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ "bogus",
+ )
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertTaskState(t, result, genmodels.TaskStateStateFailed)
+ assert.False(t, ran, "the task body must not run on a malformed binding
element")
+}
+
+func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) {
+ cases := []struct {
+ name string
+ spec map[string]any
+ }{
+ {name: "missing name", spec: map[string]any{"kind": "literal",
"value": "x"}},
+ {name: "empty name", spec: map[string]any{"name": "", "kind":
"literal", "value": "x"}},
+ {name: "xcom missing task_id", spec: map[string]any{"name":
"country", "kind": "xcom"}},
+ {
+ name: "xcom empty task_id",
+ spec: map[string]any{"name": "country", "kind": "xcom",
"task_id": ""},
+ },
+ {
+ name: "value_schema not a map",
+ spec: map[string]any{
+ "name": "country", "kind": "literal", "value":
"x", "value_schema": "string",
+ },
+ },
+ {
+ name: "from_default not a bool",
+ spec: map[string]any{
+ "name": "country", "kind": "literal", "value":
"x", "from_default": "true",
+ },
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ ran := false
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+
r.AddDag("test_dag").AddTaskWithName("transform",
+ func(country string) error {
+ ran = true
+ return nil
+ })
+ })
+
+ details := newStartupDetails("transform", tc.spec)
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil),
io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle,
details, comm, logger)
+ assertTaskState(t, result,
genmodels.TaskStateStateFailed)
+ assert.False(t, ran, "the task body must not run on an
incomplete binding spec")
+ })
}
+}
+
+func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) {
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("transform",
+ func(country string) error { return nil })
+ })
+
+ details := newStartupDetails(
+ "transform",
+ map[string]any{"name": "country", "kind": "template", "value":
"x"},
+ )
+ details.TIContext.ShouldRetry = true
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
result := RunTask(context.Background(), bundle, details, comm, logger)
- assertRetryTask(t, result, "panic: something went wrong")
+ assertRetryTask(t, result, `unknown kind "template"`)
}
func TestRunTaskHonorsContextCancellation(t *testing.T) {
@@ -236,16 +466,7 @@ func TestRunTaskHonorsContextCancellation(t *testing.T) {
func(ctx context.Context) error { return ctx.Err() })
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "ctxcheck",
- RunID: "run1",
- MapIndex: ptr(-1),
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- }
+ details := newStartupDetails("ctxcheck")
// A cancelled root context must reach the user task through RunTask's
// threading; the task surfaces ctx.Err(), which RunTask maps to failed.
@@ -333,16 +554,8 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) {
})
})
- details := &genmodels.StartupDetails{
- TI: genmodels.TaskInstance{
- ID: "550e8400-e29b-41d4-a716-446655440000",
- DagID: "test_dag",
- TaskID: "ctxgrab",
- RunID: "run1",
- MapIndex: ptr(5),
- },
- BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- }
+ details := newStartupDetails("ctxgrab")
+ details.TI.MapIndex = ptr(5)
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
diff --git a/go-sdk/pkg/execution/task_runner.go
b/go-sdk/pkg/execution/task_runner.go
index c656b4cfd71..e24601db2f0 100644
--- a/go-sdk/pkg/execution/task_runner.go
+++ b/go-sdk/pkg/execution/task_runner.go
@@ -28,6 +28,7 @@ import (
"github.com/apache/airflow/go-sdk/bundle/bundlev1"
"github.com/apache/airflow/go-sdk/pkg/api"
+ "github.com/apache/airflow/go-sdk/pkg/binding"
"github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
"github.com/apache/airflow/go-sdk/sdk"
@@ -124,7 +125,109 @@ func RunTask(
ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey,
sdk.Client(client))
ctx = context.WithValue(ctx, sdkcontext.RuntimeContextKey,
runtimeContext)
- return executeTask(ctx, task, details.TIContext.ShouldRetry, logger)
+ args, err := convertArgBindings(details.TIContext.ArgBindings)
+ if err != nil {
+ logger.Error("Invalid arg_bindings spec from supervisor",
+ "dag_id", details.TI.DagID,
+ "task_id", details.TI.TaskID,
+ "error", err,
+ )
+ if details.TIContext.ShouldRetry {
+ return genmodels.RetryTask{
+ EndDate: time.Now().UTC(),
+ RetryReason: err.Error(),
+ }
+ }
+ return genmodels.TaskState{
+ State: genmodels.TaskStateStateFailed,
+ EndDate: time.Now().UTC(),
+ }
+ }
+
+ return executeTask(ctx, task, args, details.TIContext.ShouldRetry,
logger)
+}
+
+func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg,
error) {
+ if specsPtr == nil || len(*specsPtr) == 0 {
+ return nil, nil
+ }
+ specs := *specsPtr
+ args := make([]binding.Arg, len(specs))
+ for i, raw := range specs {
+ m, ok := raw.(map[string]any)
+ if !ok {
+ return nil, fmt.Errorf("arg_bindings[%d]: unexpected
wire shape %T", i, raw)
+ }
+ name, ok := m["name"].(string)
+ if !ok || name == "" {
+ return nil, fmt.Errorf("arg_bindings[%d]: missing or
empty name", i)
+ }
+ valueSchema, err := argValueSchema(m["value_schema"])
+ if err != nil {
+ return nil, fmt.Errorf("arg_bindings[%d] (%q): %w", i,
name, err)
+ }
+ switch kind, _ := m["kind"].(string); kind {
+ case "xcom":
+ taskID, ok := m["task_id"].(string)
+ if !ok || taskID == "" {
+ return nil, fmt.Errorf(
+ "arg_bindings[%d] (%q): missing or
empty task_id for xcom kind",
+ i,
+ name,
+ )
+ }
+ args[i] = binding.XComArg{
+ Kind: kind,
+ Name: name,
+ TaskID: taskID,
+ ValueSchema: valueSchema,
+ }
+ case "literal":
+ fromDefault, err := optionalBool(m["from_default"])
+ if err != nil {
+ return nil, fmt.Errorf("arg_bindings[%d] (%q):
%w", i, name, err)
+ }
+ args[i] = binding.LiteralArg{
+ Kind: kind,
+ Name: name,
+ Value: m["value"],
+ ValueSchema: valueSchema,
+ FromDefault: fromDefault,
+ }
+ default:
+ return nil, fmt.Errorf("arg_bindings[%d]: unknown kind
%q", i, kind)
+ }
+ }
+ return args, nil
+}
+
+func argValueSchema(raw any) (*genmodels.ArgValueSchema, error) {
+ if raw == nil {
+ return nil, nil
+ }
+ m, ok := raw.(map[string]any)
+ if !ok {
+ return nil, fmt.Errorf("value_schema has unexpected wire shape
%T", raw)
+ }
+ if len(m) == 0 {
+ return nil, nil
+ }
+ schema := make(genmodels.ArgValueSchema, len(m))
+ for k, v := range m {
+ schema[k] = v
+ }
+ return &schema, nil
+}
+
+func optionalBool(raw any) (bool, error) {
+ if raw == nil {
+ return false, nil
+ }
+ b, ok := raw.(bool)
+ if !ok {
+ return false, fmt.Errorf("from_default has unexpected wire
shape %T", raw)
+ }
+ return b, nil
}
// mapIndexPtr normalizes the supervisor's map_index into the optional form
@@ -145,6 +248,7 @@ func mapIndexPtr(mapIndex *int) *int {
func executeTask(
ctx context.Context,
task bundlev1.Task,
+ args []binding.Arg,
shouldRetry bool,
logger *slog.Logger,
) (result any) {
@@ -168,7 +272,19 @@ func executeTask(
}
}()
- if err := task.Execute(ctx, logger); err != nil {
+ var err error
+ if tw, ok := task.(bundlev1.TaskWithArgs); ok {
+ err = tw.ExecuteArgs(ctx, logger, args)
+ } else if len(args) > 0 {
+ err = fmt.Errorf(
+ "task received %d positional argument(s) from the Dag
but its implementation "+
+ "does not support argument binding (does not
implement TaskWithArgs)",
+ len(args),
+ )
+ } else {
+ err = task.Execute(ctx, logger)
+ }
+ if err != nil {
logger.ErrorContext(ctx, "Task failed", "error", err)
// A task that fails when ti_context.should_retry is set is
reported as
// UP_FOR_RETRY via RetryTask; otherwise it terminates as
FAILED.
