Go SDK: fail a task whose return-value XCom push fails

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.

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(

go-sdk/internal/bundle/task_test.go

@@ -236,19 +236,19 @@ 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
 	failOn  string
 	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,32 +385,47 @@ 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
 		wantRun   bool
 		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)
+}