diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index 0a3aca70666e1..231507a99f2cb 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -25,6 +25,7 @@ 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" ) @@ -32,23 +33,29 @@ import ( 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:""` 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 @@ -56,26 +63,14 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { 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)) -} diff --git a/go-sdk/bundle/bundlev1/task_test.go b/go-sdk/bundle/bundlev1/task_test.go index d5f95c8251e72..a2077e89970cc 100644 --- a/go-sdk/bundle/bundlev1/task_test.go +++ b/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) + }) + } +} diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 284439741f37e..25bbaf89af6d5 100644 --- a/go-sdk/example/bundle/main.go +++ b/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,7 +146,7 @@ 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()) @@ -113,29 +154,48 @@ func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, err } 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") } diff --git a/go-sdk/example/bundle/main_test.go b/go-sdk/example/bundle/main_test.go index 54ff86c10440f..93535a631343d 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/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) } diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go new file mode 100644 index 0000000000000..eff526d43c484 --- /dev/null +++ b/go-sdk/pkg/binding/binding.go @@ -0,0 +1,354 @@ +// 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 turns a task function's parameter list into the concrete +// argument values it is called with at execution time. It is shared by both +// execution paths (the coordinator runtime and the Edge worker), so it depends +// only on the SDK surface and not on the bundle registry. +// +// Two kinds of parameter are supported: +// +// - Injectable runtime values: context.Context, *slog.Logger, and +// sdk.Client (or the narrower sdk.VariableClient / sdk.ConnectionClient). +// - XCom-input structs: a struct (or pointer to one) whose exported fields +// each carry an `xcom:"[,key=]"` tag. Each field is pulled +// from the named upstream task in the current dag run and decoded into the +// field's type, so an author receives an upstream's return value as a typed +// parameter without calling the client explicitly. +// +// Analyze inspects a function once at registration and returns a Plan; Resolve +// builds the call arguments for each execution from that Plan. +package binding + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "reflect" + "strings" + + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +// paramKind classifies how a task-function parameter is filled at execution. +type paramKind int + +const ( + paramContext paramKind = iota + paramLogger + paramClient + paramXComInput +) + +// xcomField records how to pull and decode one field of an XCom-input struct. +type xcomField struct { + fieldIndex int + taskID string + key string + fieldType reflect.Type +} + +// paramPlan describes how Resolve fills a single task-function parameter. For +// an XCom-input struct (kind == paramXComInput) the remaining fields describe +// the struct to build and the per-field pulls to perform. +type paramPlan struct { + kind paramKind + isPtr bool + structType reflect.Type + fields []xcomField +} + +// Plan is the precomputed recipe for filling a task function's parameters. It +// is built once by Analyze and reused for every execution of that function. +type Plan struct { + fnName string + params []paramPlan +} + +// Analyze inspects the parameters of a task function type and builds a Plan. +// fnName appears in error messages only. Every parameter must be an injectable +// runtime type (context.Context, *slog.Logger, sdk.Client / VariableClient / +// ConnectionClient) or an XCom-input struct whose exported fields each carry an +// `xcom:"[,key=]"` tag; anything else is a registration error. +func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { + params := make([]paramPlan, fnType.NumIn()) + for i := range fnType.NumIn() { + p, err := classifyParam(fnName, fnType.In(i)) + if err != nil { + return nil, err + } + params[i] = p + } + return &Plan{fnName: fnName, params: params}, nil +} + +// Resolve builds the ordered argument values for one call: injectable +// parameters receive ctx, logger, or client, and XCom-input structs are pulled +// from their upstream tasks (over client) and decoded. An error fails the task +// before its body runs. +func (p *Plan) Resolve( + ctx context.Context, + logger *slog.Logger, + client sdk.Client, +) ([]reflect.Value, error) { + args := make([]reflect.Value, len(p.params)) + for i, plan := range p.params { + switch plan.kind { + case paramContext: + args[i] = reflect.ValueOf(ctx) + case paramLogger: + args[i] = reflect.ValueOf(logger) + case paramClient: + args[i] = reflect.ValueOf(client) + case paramXComInput: + arg, err := p.resolveXComInput(ctx, client, plan) + if err != nil { + return nil, err + } + args[i] = arg + } + } + return args, nil +} + +// classifyParam decides how a single parameter is filled. Injectable runtime +// types map to their paramKind; anything else must be an XCom-input struct (or +// pointer to one) whose exported fields are all `xcom`-tagged. +func classifyParam(fnName string, in reflect.Type) (paramPlan, error) { + switch { + case isContext(in): + return paramPlan{kind: paramContext}, nil + case isLogger(in): + return paramPlan{kind: paramLogger}, nil + case isClient(in): + return paramPlan{kind: paramClient}, nil + } + + structType := in + isPtr := false + if in.Kind() == reflect.Pointer { + structType = in.Elem() + isPtr = true + } + if structType.Kind() != reflect.Struct { + return paramPlan{}, fmt.Errorf( + "task function %s: parameter of type %s is not an injectable type "+ + "(context.Context, *slog.Logger, sdk.Client/VariableClient/ConnectionClient) "+ + "nor an xcom-input struct", + fnName, in, + ) + } + + var fields []xcomField + for fi := range structType.NumField() { + sf := structType.Field(fi) + if !sf.IsExported() { + continue + } + tag, ok := sf.Tag.Lookup("xcom") + if !ok { + return paramPlan{}, fmt.Errorf( + "task function %s: field %s.%s has no `xcom` tag; every exported field "+ + "of an xcom-input struct must be tagged `xcom:\"\"`", + fnName, structType.Name(), sf.Name, + ) + } + taskID, key, err := parseXComTag(tag) + if err != nil { + return paramPlan{}, fmt.Errorf( + "task function %s: field %s.%s: %w", fnName, structType.Name(), sf.Name, err, + ) + } + if !isDecodableType(sf.Type) { + return paramPlan{}, fmt.Errorf( + "task function %s: field %s.%s has type %s which cannot hold an xcom value", + fnName, structType.Name(), sf.Name, sf.Type, + ) + } + fields = append(fields, xcomField{ + fieldIndex: fi, + taskID: taskID, + key: key, + fieldType: sf.Type, + }) + } + + if len(fields) == 0 { + return paramPlan{}, fmt.Errorf( + "task function %s: xcom-input struct %s has no exported `xcom`-tagged fields", + fnName, structType.Name(), + ) + } + + return paramPlan{ + kind: paramXComInput, + isPtr: isPtr, + structType: structType, + fields: fields, + }, nil +} + +// parseXComTag parses an `xcom` struct tag of the form "" or +// ",key=". The key defaults to the return-value XCom key. +func parseXComTag(tag string) (taskID, key string, err error) { + parts := strings.Split(tag, ",") + taskID = strings.TrimSpace(parts[0]) + if taskID == "" { + return "", "", fmt.Errorf("xcom tag is missing the upstream task id") + } + key = api.XComReturnValueKey + for _, opt := range parts[1:] { + opt = strings.TrimSpace(opt) + if opt == "" { + continue + } + val, ok := strings.CutPrefix(opt, "key=") + if !ok { + return "", "", fmt.Errorf("unknown xcom tag option %q", opt) + } + if val == "" { + return "", "", fmt.Errorf("xcom tag option key= is empty") + } + key = val + } + return taskID, key, nil +} + +// resolveXComInput builds the XCom-input struct described by plan, pulling each +// tagged field from its upstream task in the current dag run and decoding it +// into the field's type. +func (p *Plan) resolveXComInput( + ctx context.Context, + c sdk.XComClient, + plan paramPlan, +) (reflect.Value, error) { + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return reflect.Value{}, fmt.Errorf( + "task function %s: no workload in context, cannot resolve xcom inputs", p.fnName, + ) + } + + structPtr := reflect.New(plan.structType) + structVal := structPtr.Elem() + for _, xf := range plan.fields { + // Pull from the upstream's unmapped instance (map_index -1); mapped + // upstream fan-in is out of scope for now. + raw, err := c.GetXCom( + ctx, + workload.TI.DagId, + workload.TI.RunId, + xf.taskID, + nil, + xf.key, + nil, + ) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: pulling xcom from %q (key %q): %w", + p.fnName, + xf.taskID, + xf.key, + err, + ) + } + decoded, err := decodeXCom(raw, xf.fieldType) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: decoding xcom from %q into field %s: %w", + p.fnName, xf.taskID, xf.fieldType, err, + ) + } + structVal.Field(xf.fieldIndex).Set(decoded) + } + + if plan.isPtr { + return structPtr, nil + } + return structVal, nil +} + +// decodeXCom decodes a raw (generically deserialised) XCom value into target. +// Decoding into a struct is strict: unknown/renamed keys fail rather than +// silently leaving fields zero. Decoding into a map / interface accepts any +// shape, so authors opt into loose decoding by typing the field map[string]any +// or any. A null value is allowed only for a nilable target. +func decodeXCom(raw any, target reflect.Type) (reflect.Value, error) { + out := reflect.New(target) + + if raw == nil { + switch target.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface: + return out.Elem(), nil + default: + return reflect.Value{}, fmt.Errorf( + "xcom value is null but field 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 +} + +// isDecodableType reports whether a value can be JSON-decoded into inType. It +// rejects kinds json cannot target (func, chan, unsafe pointer) and non-empty +// interfaces (only the empty interface `any` is a valid decode target). +func isDecodableType(inType reflect.Type) bool { + switch inType.Kind() { + case reflect.Func, reflect.Chan, reflect.UnsafePointer: + return false + case reflect.Interface: + return inType.NumMethod() == 0 + } + return true +} + +var ( + 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]() +) + +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)) +} diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go new file mode 100644 index 0000000000000..5661bbc1d811f --- /dev/null +++ b/go-sdk/pkg/binding/binding_test.go @@ -0,0 +1,276 @@ +// 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" + "errors" + "log/slog" + "reflect" + "testing" + + "github.com/stretchr/testify/suite" + + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +type extractResult struct { + GoVersion string `json:"go_version"` + Timestamp int64 `json:"timestamp"` +} + +type transformInput struct { + Extracted extractResult `xcom:"extract"` +} + +// xcomCall records one GetXCom invocation so tests can assert the runtime +// pulled the right upstream/key for the current dag run. +type xcomCall struct { + dagID, runID, taskID, key string +} + +// fakeClient is a minimal sdk.Client whose GetXCom returns preconfigured values +// keyed by "/"; absent entries yield sdk.XComNotFound. +type fakeClient struct { + xcoms map[string]any + calls []xcomCall +} + +var _ sdk.Client = (*fakeClient)(nil) + +func (c *fakeClient) GetXCom( + _ context.Context, + dagID, runID, taskID string, + _ *int, + key string, + _ any, +) (any, error) { + c.calls = append(c.calls, xcomCall{dagID, runID, taskID, key}) + if v, ok := c.xcoms[taskID+"/"+key]; ok { + return v, nil + } + return nil, sdk.XComNotFound +} + +func (c *fakeClient) PushXCom(context.Context, api.TaskInstance, string, any) error { return nil } + +func (c *fakeClient) GetVariable( + context.Context, + string, +) (string, error) { + return "", nil +} +func (c *fakeClient) UnmarshalJSONVariable(context.Context, string, any) error { return nil } +func (c *fakeClient) GetConnection(context.Context, string) (sdk.Connection, error) { + return sdk.Connection{}, nil +} + +type BindingSuite struct { + suite.Suite +} + +func TestBindingSuite(t *testing.T) { + suite.Run(t, &BindingSuite{}) +} + +func (s *BindingSuite) TestParseXComTag() { + cases := map[string]struct { + tag string + taskID, key string + errContains string + }{ + "task-only": {tag: "extract", taskID: "extract", key: "return_value"}, + "task-and-key": {tag: "extract,key=foo", taskID: "extract", key: "foo"}, + "trims-spaces": {tag: " extract , key=bar ", taskID: "extract", key: "bar"}, + "empty": {tag: "", errContains: "missing the upstream task id"}, + "empty-task": {tag: ",key=foo", errContains: "missing the upstream task id"}, + "unknown-opt": {tag: "extract,bogus", errContains: `unknown xcom tag option "bogus"`}, + "empty-key-opt": {tag: "extract,key=", errContains: "key= is empty"}, + } + for name, tt := range cases { + s.Run(name, func() { + taskID, key, err := parseXComTag(tt.tag) + if tt.errContains != "" { + if s.Error(err) { + s.Contains(err.Error(), tt.errContains) + } + return + } + s.NoError(err) + s.Equal(tt.taskID, taskID) + s.Equal(tt.key, key) + }) + } +} + +func (s *BindingSuite) TestAnalyze() { + cases := map[string]struct { + fn any + errContains string + }{ + "injectables": { + fn: func(context.Context, *slog.Logger, sdk.Client) error { return nil }, + }, + "valid-input-struct": { + fn: func(in transformInput) error { return nil }, + }, + "untagged-field": { + fn: func(in struct{ Extracted extractResult }) error { return nil }, + errContains: "has no `xcom` tag", + }, + "empty-tag": { + fn: func(in struct { + X string `xcom:""` + }, + ) error { + return nil + }, + errContains: "missing the upstream task id", + }, + "undecodable-field": { + fn: func(in struct { + X chan int `xcom:"up"` + }, + ) error { + return nil + }, + errContains: "cannot hold an xcom value", + }, + "non-struct-param": { + fn: func(x int) error { return nil }, + errContains: "is not an injectable type", + }, + "no-tagged-fields": { + fn: func(in struct{ unexported int }) error { return nil }, + errContains: "no exported `xcom`-tagged fields", + }, + } + for name, tt := range cases { + s.Run(name, func() { + _, err := Analyze(reflect.TypeOf(tt.fn), "fn") + if tt.errContains != "" { + if s.Error(err) { + s.Contains(err.Error(), tt.errContains) + } + return + } + s.NoError(err) + }) + } +} + +func (s *BindingSuite) TestDecodeXCom() { + s.Run("strict-struct-success", func() { + v, err := decodeXCom( + map[string]any{"go_version": "go1.24", "timestamp": 7}, + reflect.TypeFor[extractResult](), + ) + s.Require().NoError(err) + got := v.Interface().(extractResult) + s.Equal("go1.24", got.GoVersion) + s.Equal(int64(7), got.Timestamp) + }) + s.Run("strict-struct-unknown-field-fails", func() { + _, err := decodeXCom( + map[string]any{"go_version": "x", "renamed": 1}, + reflect.TypeFor[extractResult](), + ) + s.Error(err) + }) + s.Run("loose-map-accepts-any-shape", func() { + v, err := decodeXCom( + map[string]any{"anything": 1, "else": "ok"}, + reflect.TypeFor[map[string]any](), + ) + s.Require().NoError(err) + s.Len(v.Interface().(map[string]any), 2) + }) + s.Run("string-into-struct-fails", func() { + _, err := decodeXCom("hello", reflect.TypeFor[extractResult]()) + s.Error(err) + }) + s.Run("null-into-struct-fails", func() { + _, err := decodeXCom(nil, reflect.TypeFor[extractResult]()) + if s.Error(err) { + s.Contains(err.Error(), "not nilable") + } + }) + s.Run("null-into-map-ok", func() { + v, err := decodeXCom(nil, reflect.TypeFor[map[string]any]()) + s.Require().NoError(err) + s.Nil(v.Interface()) + }) +} + +// resolve runs Analyze + Resolve against a fake client, with a workload in +// context so xcom pulls can read the current dag/run. +func (s *BindingSuite) resolve(fn any, client *fakeClient) ([]reflect.Value, error) { + plan, err := Analyze(reflect.TypeOf(fn), "fn") + s.Require().NoError(err) + + workload := api.ExecuteTaskWorkload{ + TI: api.TaskInstance{DagId: "dag1", RunId: "run1", TaskId: "transform"}, + } + ctx := context.WithValue(context.Background(), sdkcontext.WorkloadContextKey, workload) + return plan.Resolve(ctx, slog.Default(), client) +} + +func (s *BindingSuite) TestResolveInjectsXComInput() { + client := &fakeClient{xcoms: map[string]any{ + "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": 99}, + }} + + args, err := s.resolve(func(context.Context, transformInput) error { return nil }, client) + + s.Require().NoError(err) + s.Require().Len(args, 2) + got := args[1].Interface().(transformInput) + s.Equal("go1.24", got.Extracted.GoVersion) + s.Equal(int64(99), got.Extracted.Timestamp) + s.Require().Len(client.calls, 1) + s.Equal(xcomCall{"dag1", "run1", "extract", "return_value"}, client.calls[0]) +} + +func (s *BindingSuite) TestResolveCustomKeyAndPointer() { + type pyInput struct { + FromPython string `xcom:"python_task_1,key=out"` + } + client := &fakeClient{xcoms: map[string]any{ + "python_task_1/out": "value_from_python", + }} + + args, err := s.resolve(func(*pyInput) error { return nil }, client) + + s.Require().NoError(err) + s.Require().Len(args, 1) + got := args[0].Interface().(*pyInput) + s.Require().NotNil(got) + s.Equal("value_from_python", got.FromPython) + s.Equal(xcomCall{"dag1", "run1", "python_task_1", "out"}, client.calls[0]) +} + +func (s *BindingSuite) TestResolveMissingUpstreamFails() { + client := &fakeClient{xcoms: map[string]any{}} // extract not present + + _, err := s.resolve(func(transformInput) error { return nil }, client) + + s.Require().Error(err) + s.True(errors.Is(err, sdk.XComNotFound)) +} diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go index ba25677b8a569..1557e4fc6e6e6 100644 --- a/go-sdk/pkg/execution/client.go +++ b/go-sdk/pkg/execution/client.go @@ -206,6 +206,14 @@ func (c *CoordinatorClient) GetXCom( return nil, fmt.Errorf("decoding xcom result: %w", err) } + if result.Value == nil { + // The supervisor returns a null value for an absent XCom (Airflow's + // xcom_pull yields None when the key is missing), so surface it as + // not-found, matching how the HTTP client maps a 404 and how + // GetVariable handles the same condition above. + return nil, fmt.Errorf("%w: %q", sdk.XComNotFound, key) + } + return result.Value, nil } diff --git a/go-sdk/pkg/execution/client_test.go b/go-sdk/pkg/execution/client_test.go index 5ee9244aab50e..ca1c6b30bf05b 100644 --- a/go-sdk/pkg/execution/client_test.go +++ b/go-sdk/pkg/execution/client_test.go @@ -269,6 +269,68 @@ func TestCoordinatorClientPushXComMapIndex(t *testing.T) { } } +// TestCoordinatorClientGetXComNullIsNotFound verifies that a successful +// response carrying a null value (how the supervisor reports an absent XCom, +// mirroring xcom_pull returning None) is surfaced as sdk.XComNotFound rather +// than handed back as a nil value. +func TestCoordinatorClientGetXComNullIsNotFound(t *testing.T) { + responsePayload, err := encodeRequest(0, map[string]any{ + "type": "XComResult", + "key": "return_value", + "value": nil, + }) + require.NoError(t, err) + + var responseBuf bytes.Buffer + require.NoError(t, writeFrame(&responseBuf, responsePayload)) + + var requestBuf bytes.Buffer + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(&responseBuf, &requestBuf, logger) + client := NewCoordinatorClient(comm) + + _, err = client.GetXCom(context.Background(), "d", "r", "extract", nil, "return_value", nil) + require.Error(t, err) + assert.ErrorIs(t, err, sdk.XComNotFound) +} + +// TestCoordinatorClientPushXComStructUsesJSONTags verifies that a struct value +// crosses the wire using its `json` field tags, not Go field names, so a typed +// XCom round-trips into a json-tagged struct on the consuming side. It also +// confirms a large int64 (e.g. a UnixNano timestamp) is preserved exactly +// rather than coerced through a lossy float. +func TestCoordinatorClientPushXComStructUsesJSONTags(t *testing.T) { + responsePayload := encodeResponseFrame(t, 0, map[string]any{"type": "OKResponse"}, nil) + var responseBuf bytes.Buffer + require.NoError(t, writeFrame(&responseBuf, responsePayload)) + + var requestBuf bytes.Buffer + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(&responseBuf, &requestBuf, logger) + client := NewCoordinatorClient(comm) + + type pushable struct { + GoVersion string `json:"go_version"` + Timestamp int64 `json:"timestamp"` + } + ti := api.TaskInstance{DagId: "d", RunId: "r", TaskId: "t"} + require.NoError(t, client.PushXCom( + context.Background(), ti, "return_value", + pushable{GoVersion: "go1.24", Timestamp: 6767000000000000000}, + )) + + sent, err := readFrame(&requestBuf) + require.NoError(t, err) + + value, ok := sent.Body["value"].(map[string]any) + require.True(t, ok, "value should be encoded as a map, got %T", sent.Body["value"]) + assert.Contains(t, value, "go_version") + assert.Contains(t, value, "timestamp") + assert.NotContains(t, value, "GoVersion") + assert.EqualValues(t, "go1.24", value["go_version"]) + assert.EqualValues(t, 6767000000000000000, value["timestamp"]) +} + // assertNoReadReader fails the test on any Read call. type assertNoReadReader struct{ t *testing.T } diff --git a/go-sdk/pkg/execution/frames.go b/go-sdk/pkg/execution/frames.go index 5d7d7ca3649b5..15af8987e7f4d 100644 --- a/go-sdk/pkg/execution/frames.go +++ b/go-sdk/pkg/execution/frames.go @@ -45,6 +45,13 @@ func encodeRequest(id int64, body map[string]any) ([]byte, error) { var buf bytes.Buffer enc := msgpack.NewEncoder(&buf) enc.UseCompactInts(true) + // Honour `json` struct tags when encoding user-provided values (XCom and + // Variable payloads). Without this, msgpack uses Go field names, so a typed + // XCom pushed as a struct would cross the wire as e.g. "GoVersion" and fail + // to decode into the json-tagged "go_version" the value is read back with + // (and that the HTTP-backed client uses). Protocol frames themselves are + // map[string]any, so this only affects nested user values. + enc.SetCustomStructTag("json") if err := enc.EncodeArrayLen(2); err != nil { return nil, err diff --git a/go-sdk/pkg/execution/logger.go b/go-sdk/pkg/execution/logger.go index 05de45c01f7bc..b5ac787460d87 100644 --- a/go-sdk/pkg/execution/logger.go +++ b/go-sdk/pkg/execution/logger.go @@ -110,13 +110,6 @@ func (h *SocketLogHandler) Enabled(_ context.Context, level slog.Level) bool { func (h *SocketLogHandler) Handle(_ context.Context, r slog.Record) error { entry := make(map[string]any) - // Set standard fields. - entry["event"] = r.Message - entry["level"] = strings.ToLower(r.Level.String()) - if !r.Time.IsZero() { - entry["timestamp"] = r.Time.Format(time.RFC3339Nano) - } - // Apply pre-configured attrs. Keys are already qualified with the groups // active at the WithAttrs call site, so the current h.groups is NOT // applied here — only to record-level attrs below. The stored key already @@ -134,6 +127,16 @@ func (h *SocketLogHandler) Handle(_ context.Context, r slog.Record) error { return true }) + // Set the reserved protocol fields last so they are authoritative: a user + // attribute named "timestamp", "event", or "level" must not overwrite them. + // In particular a non-RFC3339 "timestamp" (e.g. an int) would crash the + // supervisor's log reader, which decodes this field as an RFC3339 datetime. + entry["event"] = r.Message + entry["level"] = strings.ToLower(r.Level.String()) + if !r.Time.IsZero() { + entry["timestamp"] = r.Time.Format(time.RFC3339Nano) + } + line, err := json.Marshal(entry) if err != nil { return err diff --git a/go-sdk/pkg/execution/logger_test.go b/go-sdk/pkg/execution/logger_test.go index ee39fef69f071..87ac95ddc21b0 100644 --- a/go-sdk/pkg/execution/logger_test.go +++ b/go-sdk/pkg/execution/logger_test.go @@ -24,6 +24,7 @@ import ( "log/slog" "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -201,6 +202,32 @@ func TestSocketLogHandlerKeyMapping(t *testing.T) { assert.False(t, hasTime) } +// TestSocketLogHandlerReservedKeysAuthoritative verifies that a user attribute +// named like a reserved protocol field (timestamp/event/level) cannot overwrite +// it. A non-RFC3339 "timestamp" (e.g. an int) would otherwise crash the +// supervisor's log reader, which decodes that field as an RFC3339 datetime. +func TestSocketLogHandlerReservedKeysAuthoritative(t *testing.T) { + var buf bytes.Buffer + handler := NewSocketLogHandler(&buf, slog.LevelDebug) + logger := slog.New(handler) + + logger.Info("real message", + "timestamp", int64(1717000000000000000), + "event", "user event", + "level", "user level", + ) + + var entry map[string]any + require.NoError(t, json.Unmarshal([]byte(strings.TrimSpace(buf.String())), &entry)) + + assert.Equal(t, "real message", entry["event"]) + assert.Equal(t, "info", entry["level"]) + ts, ok := entry["timestamp"].(string) + require.True(t, ok, "timestamp must remain an RFC3339 string, got %T", entry["timestamp"]) + _, err := time.Parse(time.RFC3339Nano, ts) + assert.NoError(t, err, "reserved timestamp must stay RFC3339 despite a user timestamp attr") +} + // TestSocketLogHandlerStringifiesErrorAttr verifies that an attribute whose // value is a stock `error` is rendered as its .Error() string. Without value // resolution the underlying struct's unexported fields would be marshaled as