FrankYang0529 commented on code in PR #74075:
URL: https://github.com/apache/airflow/pull/74075#discussion_r4172293663
##########
go-sdk/internal/bundle/task.go:
##########
@@ -202,6 +280,13 @@ func (f *taskFunction) validateFn(
fnType.Out(fnType.NumOut()-1).Kind(),
)
}
+ if f.findSkipped != nil && fnType.NumOut() != 2 {
+ return fmt.Errorf(
+ "task function %s returns only an error, but a task
that skips downstream tasks "+
+ "must return `<result>, error`",
+ f.fullName,
+ )
+ }
Review Comment:
Removed it. Thanks.
##########
go-sdk/internal/bundle/task.go:
##########
@@ -139,9 +155,71 @@ func (f *taskFunction) call(
res := retValues[0].Interface()
f.sendXcom(ctx, res, sdkClient, logger)
}
+ if err == nil && f.findSkipped != nil {
+ return f.skipDownstream(ctx, sdkClient,
f.findSkipped(retValues[0].Interface()), logger)
+ }
return err
}
+// skipMixinXComKey is the key of the XCom that lists the tasks a task
skipped. When one of those
+// tasks is cleared, NotPreviouslySkippedDep in Airflow core reads the XCom
and skips the cleared
+// task again instead of running it. SkipMixin in the standard provider writes
the same key.
+const skipMixinXComKey = "skipmixin_key"
+
+type skipDownstreamTasksKey struct{}
+
+// WithSkipDownstreamTasks returns a copy of ctx that carries skip, the
function that a task from
+// NewPositionalBranchFunction calls to skip tasks downstream of it. The
runtime passes skip in the
+// context and not in sdk.Client, so that task functions cannot call it.
Airflow core reads the
+// skipmixin_key XCom only from a task that has _can_skip_downstream set in
the serialized Dag. A
+// task skipped by a task without that flag would run when someone clears it.
+func WithSkipDownstreamTasks(
+ ctx context.Context,
+ skip func(ctx context.Context, taskIDs []string) error,
+) context.Context {
+ return context.WithValue(ctx, skipDownstreamTasksKey{}, skip)
+}
+
+func (f *taskFunction) skipDownstream(
+ ctx context.Context,
+ client sdk.Client,
+ taskIDs []string,
+ logger *slog.Logger,
+) error {
+ skip, ok := ctx.Value(skipDownstreamTasksKey{}).(func(context.Context,
[]string) error)
Review Comment:
Updated it. `Execute` now checks for the skip function and the task
instance, and clears the XCom, before it resolves the arguments and runs `fn`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
