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 1a50a18a3c6 Go SDK: fail a task whose return-value XCom push fails
(#75443)
1a50a18a3c6 is described below
commit 1a50a18a3c604e75c6f5f059e90e4a5e8777fcf0
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Fri Oct 23 21:40:00 2026 +0800
Go SDK: fail a task whose return-value XCom push fails (#75443)
* Go SDK: fail a task whose return-value XCom push fails, fix the Label doc
A failed return-value push was logged and the task was left successful, so
the
downstream tasks that read it failed later with XComNotFound. Fail the task
instead, and stop pushing a result the task never produced when its function
already returned an error, as the branch path does.
The guide's labelled-edge example did not compile: Label takes only the
text,
and On names the tasks its edges reach.
* Go SDK: address review on the return-value XCom push
Push the task instance from the runtime context as is, share its lookup
with startBranch, and cover the failed push on the branch path. Merge the
copied test helper into runTask.
* Go SDK: test that a task pushes no nil-pointer return value
Only the branch path covered this check, and a nil pointer that slipped
past it would be pushed as null now that a push can fail the task.
---------
---
.../authoring-and-scheduling/language-sdks/go.rst | 9 +-
go-sdk/internal/bundle/task.go | 58 +++++------
go-sdk/internal/bundle/task_test.go | 116 +++++++++++++++++----
3 files changed, 127 insertions(+), 56 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
index 8e87a094aa4..3aaf906a3d8 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
@@ -204,12 +204,13 @@ When a task only has to run after another, without
receiving its value, draw the
cleaned := dag.Task(cleanup)
- loaded.Before(cleaned) // loaded >> cleaned
- cleaned.After(loaded, summarized) // [loaded, summarized] >>
cleaned
- loaded.Before(airflow.Label(cleaned, "done")) // a labelled edge
+ loaded.Before(cleaned) // loaded >> cleaned
+ cleaned.After(loaded, summarized) // [loaded, summarized]
>> cleaned
+ loaded.Before(airflow.Label("done").On(cleaned)) // a labelled edge
Both take any number of tasks or task groups, and drawing a dependency that
already exists changes nothing.
-``airflow.Label`` puts a label on the edge, which the Airflow UI shows on the
graph.
+``airflow.Label`` takes the text, and ``On`` names the tasks the label's edges
reach, which is Python's
+``loaded >> Label("done") >> cleaned``. The Airflow UI shows the label on the
graph.
Pass a value through ``airflow.Inputs`` when the downstream task needs it, and
use ``Before`` or ``After`` when it
only needs to run in order.
diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go
index a8aff19e9af..6a8e2e67d82 100644
--- a/go-sdk/internal/bundle/task.go
+++ b/go-sdk/internal/bundle/task.go
@@ -183,28 +183,28 @@ func (f *taskFunction) call(
)
}
}
+ if err != nil {
+ return err
+ }
if branch != nil {
- // The task pushes the value that decide returns, and decide
runs only on a result that fn
- // returned without an error. So the task pushes no XCom when
fn or decide fails.
- if err != nil {
- return err
- }
+ // The task pushes the value that decide returns.
value, skipped, err := f.decide(retValues[0].Interface())
if err != nil {
return err
}
rv := reflect.ValueOf(value)
if rv.Kind() != reflect.Ptr || !rv.IsNil() {
- f.sendXcom(ctx, value, sdkClient, logger)
+ if err := pushReturnValue(ctx, value, sdkClient); err
!= nil {
+ return err
+ }
}
return branch.skipDownstream(ctx, sdkClient, skipped, logger)
}
// 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()) {
- res := retValues[0].Interface()
- f.sendXcom(ctx, res, sdkClient, logger)
+ return pushReturnValue(ctx, retValues[0].Interface(), sdkClient)
}
- return err
+ return nil
}
// skipMixinXComKey is the key of the XCom that lists the tasks a task
skipped. When one of those
@@ -240,11 +240,19 @@ func startBranch(ctx context.Context) (*branchRun, error)
{
if !ok {
return nil, errors.New("the task runtime cannot skip downstream
tasks")
}
+ ti, err := taskInstanceFrom(ctx)
+ if err != nil {
+ return nil, err
+ }
+ return &branchRun{skip: skip, ti: ti}, nil
+}
+
+func taskInstanceFrom(ctx context.Context) (sdk.TaskInstance, error) {
runtimeContext, ok :=
ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext)
if !ok {
- return nil, errors.New("task runtime context is missing")
+ return sdk.TaskInstance{}, errors.New("task runtime context is
missing")
}
- return &branchRun{skip: skip, ti: runtimeContext.TaskInstance()}, nil
+ return runtimeContext.TaskInstance(), nil
}
// skipDownstream records taskIDs in the skipmixin_key XCom, and then skips
those tasks.
@@ -268,27 +276,17 @@ func (b *branchRun) skipDownstream(
return nil
}
-func (f *taskFunction) sendXcom(
- ctx context.Context,
- value any,
- c sdk.XComClient,
- logger *slog.Logger,
-) {
- runtimeContext, ok :=
ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext)
- if !ok {
- logger.ErrorContext(ctx, "Unable to set XCom", "err", "task
runtime context is missing")
- return
- }
- ti := runtimeContext.TaskInstance()
- err := c.PushXCom(ctx, sdk.TaskInstance{
- DagID: ti.DagID,
- RunID: ti.RunID,
- TaskID: ti.TaskID,
- MapIndex: ti.MapIndex,
- }, sdk.XComReturnValueKey, value)
+// pushReturnValue pushes the task's return value. A failure fails the task:
leaving it successful
+// would only move the error to the downstream tasks, which fail later on the
XCom that is not there.
+func pushReturnValue(ctx context.Context, value any, c sdk.XComClient) error {
+ ti, err := taskInstanceFrom(ctx)
if err != nil {
- logger.ErrorContext(ctx, "Unable to set XCom", "err", err)
+ return fmt.Errorf("setting the return value XCom: %w", err)
}
+ if err := c.PushXCom(ctx, ti, sdk.XComReturnValueKey, value); err !=
nil {
+ return fmt.Errorf("setting the return value XCom: %w", err)
+ }
+ return nil
}
func (f *taskFunction) validateFn(
diff --git a/go-sdk/internal/bundle/task_test.go
b/go-sdk/internal/bundle/task_test.go
index 4d3aeadf307..b75fc9ea8de 100644
--- a/go-sdk/internal/bundle/task_test.go
+++ b/go-sdk/internal/bundle/task_test.go
@@ -236,11 +236,11 @@ func (s *TaskSuite)
TestExecuteRequiresCoordinatorClient() {
}
}
-// branchClient records each XCom that a task pushes and keeps the last value
of each key.
+// recordingClient records each XCom that a task pushes and keeps the last
value of each key.
// PushXCom fails for the call whose record equals failOn. Its skip method
stands in for the
// function that the runtime passes through WithSkipDownstreamTasks. skip
records each call in the
// same list as the XComs, so that the tests see the order of the calls. It
returns skipErr.
-type branchClient struct {
+type recordingClient struct {
sdk.Client
calls []string
values map[string]any
@@ -248,7 +248,7 @@ type branchClient struct {
skipErr error
}
-func (c *branchClient) PushXCom(
+func (c *recordingClient) PushXCom(
_ context.Context,
ti sdk.TaskInstance,
key string,
@@ -266,15 +266,15 @@ func (c *branchClient) PushXCom(
return nil
}
-func (c *branchClient) skip(_ context.Context, taskIDs []string) error {
+func (c *recordingClient) skip(_ context.Context, taskIDs []string) error {
c.calls = append(c.calls, fmt.Sprintf("SkipDownstreamTasks %v",
taskIDs))
return c.skipErr
}
-// runBranch runs task as the task instance decide of dag1 with client. When
ti is false, the
+// runTask runs task as the task instance decide of dag1 with client. When ti
is false, the
// context has no task instance. When canSkip is false, the context has no
function to skip
// tasks with.
-func runBranch(task Task, client *branchClient, ti, canSkip bool) error {
+func runTask(task Task, client *recordingClient, ti, canSkip bool) error {
ctx := context.WithValue(
context.Background(),
sdkcontext.SdkClientContextKey,
@@ -305,8 +305,8 @@ func (s *TaskSuite)
TestBranchFunctionPushesTheValueAndSkipsTheTasksThatDecideRe
)
s.Require().NoError(err)
- client := &branchClient{}
- s.Require().NoError(runBranch(task, client, true, true))
+ client := &recordingClient{}
+ s.Require().NoError(runTask(task, client, true, true))
s.Equal(true, got)
s.Equal([]string{
@@ -323,8 +323,8 @@ func (s *TaskSuite) TestBranchFunctionPushesNoNilPointer() {
)
s.Require().NoError(err)
- client := &branchClient{}
- s.Require().NoError(runBranch(task, client, true, true))
+ client := &recordingClient{}
+ s.Require().NoError(runTask(task, client, true, true))
s.Equal([]string{
"PushXCom decide skipmixin_key map[skipped:[load]]",
@@ -341,8 +341,8 @@ func (s *TaskSuite) TestBranchFunctionWithNothingToSkip() {
)
s.Require().NoError(err)
- client := &branchClient{}
- s.Require().NoError(runBranch(task, client, true, true))
+ client := &recordingClient{}
+ s.Require().NoError(runTask(task, client, true, true))
s.Equal([]string{"PushXCom decide return_value true"},
client.calls)
})
@@ -364,10 +364,10 @@ func (s *TaskSuite)
TestBranchFunctionThatFailsPushesAndSkipsNothing() {
})
s.Require().NoError(err)
- client := &branchClient{}
+ client := &recordingClient{}
func() {
defer func() { _ = recover() }()
- s.Error(runBranch(task, client, true, true))
+ s.Error(runTask(task, client, true, true))
}()
s.False(called)
@@ -385,14 +385,29 @@ func (s *TaskSuite)
TestBranchFunctionFailsWhenDecideReturnsAnError() {
)
s.Require().NoError(err)
- client := &branchClient{}
- s.EqualError(runBranch(task, client, true, true), "not one of the
cases")
+ client := &recordingClient{}
+ s.EqualError(runTask(task, client, true, true), "not one of the cases")
s.Empty(client.calls)
}
+func (s *TaskSuite) TestBranchFunctionFailsWhenItCannotPushTheValue() {
+ task, err := NewPositionalBranchFunction(
+ func(contexttest.Context) (bool, error) { return true, nil },
+ func(any) (any, []string, error) { return "chosen",
[]string{"load"}, nil },
+ )
+ s.Require().NoError(err)
+
+ client := &recordingClient{failOn: "PushXCom decide return_value
chosen"}
+ s.EqualError(
+ runTask(task, client, true, true),
+ "setting the return value XCom: xcom refused",
+ )
+ s.Equal([]string{"PushXCom decide return_value chosen"}, client.calls)
+}
+
func (s *TaskSuite) TestBranchFunctionFailsWhenItCannotSkip() {
cases := map[string]struct {
- client *branchClient
+ client *recordingClient
withTI bool
canSkip bool
wantErr string
@@ -400,17 +415,17 @@ func (s *TaskSuite)
TestBranchFunctionFailsWhenItCannotSkip() {
wantCalls []string
}{
"runtime cannot skip": {
- client: &branchClient{},
+ client: &recordingClient{},
withTI: true,
wantErr: "the task runtime cannot skip downstream
tasks",
},
"no task instance": {
- client: &branchClient{},
+ client: &recordingClient{},
canSkip: true,
wantErr: "task runtime context is missing",
},
"recording the skipped tasks fails": {
- client: &branchClient{
+ client: &recordingClient{
failOn: "PushXCom decide skipmixin_key
map[skipped:[load]]",
},
withTI: true,
@@ -423,7 +438,7 @@ func (s *TaskSuite)
TestBranchFunctionFailsWhenItCannotSkip() {
},
},
"skip fails": {
- client: &branchClient{skipErr: errors.New("supervisor
refused")},
+ client: &recordingClient{skipErr:
errors.New("supervisor refused")},
withTI: true,
canSkip: true,
wantErr: `skipping the downstream tasks ["load"]:
supervisor refused`,
@@ -447,9 +462,66 @@ func (s *TaskSuite)
TestBranchFunctionFailsWhenItCannotSkip() {
)
s.Require().NoError(err)
- s.EqualError(runBranch(task, tt.client, tt.withTI,
tt.canSkip), tt.wantErr)
+ s.EqualError(runTask(task, tt.client, tt.withTI,
tt.canSkip), tt.wantErr)
s.Equal(tt.wantRun, ran, "whether fn ran")
s.Equal(tt.wantCalls, tt.client.calls)
})
}
}
+
+func (s *TaskSuite) TestPushesTheReturnValue() {
+ task, err := NewTaskFunction(func(contexttest.Context) (int, error) {
return 7, nil })
+ s.Require().NoError(err)
+
+ client := &recordingClient{}
+ s.Require().NoError(runTask(task, client, true, false))
+
+ s.Equal([]string{"PushXCom decide return_value 7"}, client.calls)
+}
+
+func (s *TaskSuite) TestPushesNoNilPointer() {
+ task, err := NewTaskFunction(func(contexttest.Context) (*int, error) {
return nil, nil })
+ s.Require().NoError(err)
+
+ client := &recordingClient{}
+ s.Require().NoError(runTask(task, client, true, false))
+
+ s.Empty(client.calls)
+}
+
+func (s *TaskSuite) TestFailsWhenItCannotPushTheReturnValue() {
+ cases := map[string]struct {
+ client *recordingClient
+ withTI bool
+ wantErr string
+ }{
+ "push fails": {
+ client: &recordingClient{failOn: "PushXCom decide
return_value 7"},
+ withTI: true,
+ wantErr: "setting the return value XCom: xcom refused",
+ },
+ "no task instance": {
+ client: &recordingClient{},
+ wantErr: "setting the return value XCom: task runtime
context is missing",
+ },
+ }
+ for name, tt := range cases {
+ s.Run(name, func() {
+ task, err := NewTaskFunction(func(contexttest.Context)
(int, error) { return 7, nil })
+ s.Require().NoError(err)
+
+ s.EqualError(runTask(task, tt.client, tt.withTI,
false), tt.wantErr)
+ })
+ }
+}
+
+func (s *TaskSuite) TestFailedTaskPushesNoReturnValue() {
+ task, err := NewTaskFunction(
+ func(contexttest.Context) (int, error) { return 0,
errors.New("no table") },
+ )
+ s.Require().NoError(err)
+
+ client := &recordingClient{}
+ s.EqualError(runTask(task, client, true, false), "no table")
+ s.Empty(client.calls)
+}
