Simple sending of return value to XCom (apache#56481)

This right now is very simple and will need to evolve over time -- it doens't
support XCom backends for instance.

We also don't handle any of the "serialization" format for more advanced
types, such as datetime/time.Time etc. That will come later

Author: ashb

go-sdk/bundle/bundlev1/registry.go

@@ -28,7 +28,6 @@ import (
 )
 
 type (
-	// task is an interface of an task implementation.
 	Task   = worker.Task
 	Bundle = worker.Bundle
 

go-sdk/bundle/bundlev1/task.go

@@ -24,6 +24,8 @@ import (
 	"reflect"
 	"runtime"
 
+	"github.com/apache/airflow/go-sdk/pkg/api"
+	"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
 	"github.com/apache/airflow/go-sdk/sdk"
 )
 
@@ -43,6 +45,7 @@ func NewTaskFunction(fn any) (Task, error) {
 
 func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error {
 	fnType := f.fn.Type()
+	sdkClient := sdk.NewClient()
 
 	reflectArgs := make([]reflect.Value, fnType.NumIn())
 	for i := range reflectArgs {
@@ -54,7 +57,7 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error {
 		case isLogger(in):
 			reflectArgs[i] = reflect.ValueOf(logger)
 		case isClient(in):
-			reflectArgs[i] = reflect.ValueOf(sdk.NewClient())
+			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)
@@ -74,15 +77,26 @@ 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
-	var res any
 	if len(retValues) > 1 && (retValues[0].Kind() != reflect.Ptr || !retValues[0].IsNil()) {
-		res = retValues[0].Interface()
+		res := retValues[0].Interface()
+		f.sendXcom(ctx, res, sdkClient, logger)
 	}
-	// TODO: send the result to XCom
-	_ = res
 	return err
 }
 
+func (f *taskFunction) sendXcom(
+	ctx context.Context,
+	value any,
+	c sdk.XComClient,
+	logger *slog.Logger,
+) {
+	workload := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload)
+	err := c.PushXCom(ctx, workload.TI, api.XComReturnValueKey, value)
+	if err != nil {
+		logger.ErrorContext(ctx, "Unable to set XCom", "err", err)
+	}
+}
+
 func (f *taskFunction) validateFn(fnType reflect.Type) error {
 	if fnType.Kind() != reflect.Func {
 		return fmt.Errorf("expected a func as input but was %s", fnType.Kind())

go-sdk/example/bundle/main.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"runtime"
 	"time"
 
 	v1 "github.com/apache/airflow/go-sdk/bundle/bundlev1"
@@ -56,7 +57,7 @@ func main() {
 	bundlev1server.Serve(&myBundle{})
 }
 
-func extract(ctx context.Context, client sdk.Client, log *slog.Logger) error {
+func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, error) {
 	log.Info("Hello from task")
 	conn, err := client.GetConnection(ctx, "test_http")
 	if err != nil {
@@ -69,15 +70,19 @@ func extract(ctx context.Context, client sdk.Client, log *slog.Logger) error {
 		// Once per loop,.check if we've been asked to cancel!
 		select {
 		case <-ctx.Done():
-			return ctx.Err()
+			return nil, 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")
 
-	return nil
+	ret := map[string]any{
+		"go_version": runtime.Version(),
+	}
+
+	return ret, nil
 }
 
 func transform(ctx context.Context, client sdk.VariableClient, log *slog.Logger) error {