Go-SDK: Support TaskFlow

go-sdk/bundle/bundlev1/task.go

@@ -25,57 +25,52 @@ 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"
 )
 
 type taskFunction struct {
 	fn       reflect.Value
 	fullName string
+	// plan describes how each parameter is filled at execution (injected
+	// runtime value or XCom pull). It is built once at registration by
+	// validateFn via the binding package.
+	plan *binding.Plan
 }
 
 var _ Task = (*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 wraps a plain Go function as a Task, validating its signature:
+// each parameter must be an injectable runtime value (context.Context,
+// *slog.Logger, sdk.Client / VariableClient / ConnectionClient) or an
+// XCom-input struct whose exported fields carry `xcom:"<task_id>"` tags, and the
+// function must return 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.
 func NewTaskFunction(fn any) (Task, error) {
 	v := reflect.ValueOf(fn)
 	fullName := runtime.FuncForPC(v.Pointer()).Name()
-	f := &taskFunction{v, fullName}
+	f := &taskFunction{fn: v, fullName: fullName}
 	return f, f.validateFn(v.Type())
 }
 
 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()
 	}
 
-	reflectArgs := make([]reflect.Value, fnType.NumIn())
-	for i := range reflectArgs {
-		in := fnType.In(i)
-
-		switch {
-		case isContext(in):
-			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)
-		}
+	reflectArgs, err := f.plan.Resolve(ctx, logger, sdkClient)
+	if err != nil {
+		return err
 	}
+
 	slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs)
 	retValues := f.fn.Call(reflectArgs)
 
-	var err error
 	if errResult := retValues[len(retValues)-1].Interface(); errResult != nil {
 		var ok bool
 		if err, ok = errResult.(error); !ok {
@@ -86,7 +81,7 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error {
 		}
 	}
 	// If there are two results, convert the first only if it's not a nil pointer
-	if len(retValues) > 1 && (retValues[0].Kind() != reflect.Ptr || !retValues[0].IsNil()) {
+	if len(retValues) > 1 && (retValues[0].Kind() != reflect.Pointer || !retValues[0].IsNil()) {
 		res := retValues[0].Interface()
 		f.sendXcom(ctx, res, sdkClient, logger)
 	}
@@ -134,6 +129,15 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error {
 			fnType.Out(fnType.NumOut()-1).Kind(),
 		)
 	}
+
+	// Parameters: build the injection / XCom-binding plan once, so any wiring
+	// mistake (an unrecognised parameter type, an untagged input-struct field,
+	// a malformed tag) fails at registration rather than at execution.
+	plan, err := binding.Analyze(fnType, f.fullName)
+	if err != nil {
+		return err
+	}
+	f.plan = plan
 	return nil
 }
 
@@ -147,30 +151,8 @@ func isValidResultType(inType reflect.Type) bool {
 	return true
 }
 
-var (
-	errorType      = reflect.TypeFor[error]()
-	contextType    = reflect.TypeFor[context.Context]()
-	slogLoggerType = reflect.TypeFor[*slog.Logger]()
-
-	connClientType = reflect.TypeFor[sdk.ConnectionClient]()
-	varClientType  = reflect.TypeFor[sdk.VariableClient]()
-	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 isLogger(inType reflect.Type) bool {
-	return inType != nil && inType.AssignableTo(slogLoggerType)
-}
-
-func isClient(inType reflect.Type) bool {
-	return inType != nil && (inType.AssignableTo(clientType) ||
-		inType.AssignableTo(connClientType) ||
-		inType.AssignableTo(varClientType))
-}

go-sdk/bundle/bundlev1/task_test.go

@@ -117,3 +117,45 @@ func (s *TaskSuite) TestArgumentBinding() {
 		})
 	}
 }
+
+// TestXComInputRegistration checks that AddTask/NewTaskFunction surfaces the
+// binding package's signature analysis: a valid xcom-input struct registers,
+// and a wiring mistake fails at registration. The decoding/pull behaviour
+// itself is covered by the binding package's own tests.
+func (s *TaskSuite) TestXComInputRegistration() {
+	cases := map[string]struct {
+		fn          any
+		errContains string
+	}{
+		"valid-input-struct": {
+			fn: func(in struct {
+				Extracted string `xcom:"extract"`
+			},
+			) error {
+				return nil
+			},
+		},
+		"untagged-field": {
+			fn: func(in struct {
+				Extracted string
+			},
+			) error {
+				return nil
+			},
+			errContains: "has no `xcom` tag",
+		},
+	}
+
+	for name, tt := range cases {
+		s.Run(name, func() {
+			_, err := NewTaskFunction(tt.fn)
+			if tt.errContains != "" {
+				if s.Error(err) {
+					s.Contains(err.Error(), tt.errContains)
+				}
+				return
+			}
+			s.NoError(err)
+		})
+	}
+}

go-sdk/example/bundle/main.go

@@ -60,7 +60,48 @@ func main() {
 	}
 }
 
-func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, error) {
+// ExtractResult is extract's return value. Returning it from the task pushes it
+// as the task's return_value XCom, ready for a downstream task to pull.
+type ExtractResult struct {
+	GoVersion string `json:"go_version"`
+	Timestamp int64  `json:"timestamp"`
+}
+
+// ExtractInput declares extract's XCom inputs. The `xcom:"python_task_1"` tag
+// binds this field to the return_value of the upstream Python task
+// `python_task_1`, so the Go task receives a Python task's output without
+// calling client.GetXCom itself. The runtime pulls and decodes it before
+// extract runs.
+type ExtractInput struct {
+	FromPython string `xcom:"python_task_1"`
+}
+
+// TransformInput binds transform's only XCom input to extract's return_value.
+// Because the field is a dedicated struct, decoding is strict: a renamed or
+// unexpected key fails the task instead of silently leaving fields zero. To
+// decode loosely (e.g. for an evolving or cross-language producer), type the
+// field map[string]any instead.
+type TransformInput struct {
+	Extracted ExtractResult `xcom:"extract"`
+}
+
+// TransformResult is transform's return value.
+type TransformResult struct {
+	Variable  string        `json:"variable"`
+	Extracted ExtractResult `json:"extracted"`
+}
+
+// LoadInput binds load's parameter to transform's return_value.
+type LoadInput struct {
+	Transformed TransformResult `xcom:"transform"`
+}
+
+func extract(
+	ctx context.Context,
+	client sdk.Client,
+	log *slog.Logger,
+	in ExtractInput,
+) (ExtractResult, error) {
 	log.Info("Hello from task")
 
 	// Log every field the runtime context exposes. The fields are namespaced
@@ -105,37 +146,56 @@ func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, err
 		// Once per loop,.check if we've been asked to cancel!
 		select {
 		case <-ctx.Done():
-			return nil, ctx.Err()
+			return ExtractResult{}, ctx.Err()
 		default:
 		}
 		log.Info("After the beep the time will be", "time", time.Now())
 		time.Sleep(2 * time.Second)
 	}
 	log.Info("Goodbye from task")
 
-	ret := map[string]any{
-		"go_version": runtime.Version(),
-		"timestamp":  time.Now().UnixNano(),
-	}
-
-	return ret, nil
+	return ExtractResult{
+		GoVersion: runtime.Version(),
+		Timestamp: time.Now().UnixNano(),
+	}, nil
 }
 
-func transform(ctx context.Context, client sdk.VariableClient, log *slog.Logger) 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
+func transform(
+	ctx context.Context,
+	client sdk.VariableClient,
+	log *slog.Logger,
+	in TransformInput,
+) (TransformResult, error) {
+	// `in.Extracted` is the Taskflow-injected return value of the upstream
+	// `extract` task. The explicit client-pull pattern still works alongside
+	// it: here we also read a Variable directly. transform takes a
+	// VariableClient (not the full sdk.Client) to make unit testing easier --
+	// see `./main_test.go`.
+	// Note: avoid an attribute literally named "timestamp" — it is a reserved
+	// field in the structured log records the coordinator sends to the
+	// supervisor (which parses it as an RFC3339 datetime).
+	log.Info("Got upstream XCom from 'extract'",
+		"go_version", in.Extracted.GoVersion,
+		"extracted_timestamp", in.Extracted.Timestamp,
+	)
+
 	key := "my_variable"
 	val, err := client.GetVariable(ctx, key)
 	if err != nil {
-		return err
+		return TransformResult{}, err
 	}
 	log.Info("Obtained variable", key, val)
-	return nil
+
+	return TransformResult{Variable: val, Extracted: in.Extracted}, nil
 }
 
-func load() error {
+func load(log *slog.Logger, in LoadInput) error {
+	log.Info(
+		"Got upstream XCom from 'transform'",
+		"variable",
+		in.Transformed.Variable,
+		"extracted",
+		in.Transformed.Extracted,
+	)
 	return fmt.Errorf("Please fail")
 }

go-sdk/example/bundle/main_test.go

@@ -52,6 +52,10 @@ 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.
-	err := transform(context.Background(), &mockVars{}, log)
+	// The Taskflow-injected input is supplied directly, so no XCom client is needed in the unit test.
+	in := TransformInput{Extracted: ExtractResult{GoVersion: "go1.24", Timestamp: 1}}
+	res, err := transform(context.Background(), &mockVars{}, log, in)
 	assert.NoError(t, err)
+	assert.Equal(t, "value1", res.Variable)
+	assert.Equal(t, "go1.24", res.Extracted.GoVersion)
 }