jason810496 commented on code in PR #74321:
URL: https://github.com/apache/airflow/pull/74321#discussion_r4192338102
##########
go-sdk/internal/bundle/task.go:
##########
@@ -157,19 +163,24 @@ func (f *taskFunction) call(
)
}
}
+ 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
+ }
+ value, skipped, err := f.decide(retValues[0].Interface())
+ if err != nil {
+ return err
+ }
+ f.sendXcom(ctx, value, sdkClient, logger)
Review Comment:
Unlike the non-branch path below (which guards against pushing a nil-pointer
result), this push is unconditional. Latent today since neither `If` nor
`Switch` returns a pointer, but `DecideFunc` is exported and general-purpose, a
future decider returning a pointer would silently push a non-nil interface
wrapping a nil pointer.
```suggestion
rv := reflect.ValueOf(value)
if rv.Kind() != reflect.Ptr || !rv.IsNil() {
f.sendXcom(ctx, value, sdkClient, logger)
}
```
##########
go-sdk/airflow/dag.go:
##########
@@ -116,9 +117,9 @@ type TaskRef struct {
// It is nil for a task that runs a Go function. A task from
TriggerDagRun runs no Go
// function, so its resultType, inputs and task are nil.
triggerDagRun *TriggerDagRunSpec
- // ifRef is the IfRef that DagRef.If or TaskGroupRef.If returned for
the task. It is nil for a
- // task from DagRef.Task or TaskGroupRef.Task.
- ifRef *IfRef
+ // decider is the IfRef or the SwitchRef that an If or a Switch method
returned for the task. It
+ // is nil for a task from DagRef.Task or TaskGroupRef.Task.
+ decider decider
Review Comment:
Heads up for the merge with #74319 (left a matching note there): #74319 adds
`go-sdk/airflow/serialize.go` with `if task.ifRef != nil` at line 462,
referencing the field this PR renames to `decider`. Whoever merges second needs
to update that check to `task.decider != nil` with the same type-switch
`markRegistered` uses below, otherwise it won't compile, or
`_can_skip_downstream` silently breaks for Switch tasks.
##########
go-sdk/airflow/switch.go:
##########
@@ -0,0 +1,253 @@
+// 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 airflow
+
+import (
+ "fmt"
+ "reflect"
+ "slices"
+ "strconv"
+ "strings"
+
+ "github.com/apache/airflow/go-sdk/internal/bundle"
+)
+
+// SwitchRef is a switch that [DagRef.Switch] or [TaskGroupRef.Switch] added
to a Dag. Each call to
+// [SwitchRef.Case] names a task that the switch can choose.
+type SwitchRef struct {
+ // task is the task that runs the decider function. The cases run after
task, but none of them
+ // takes the result of task.
+ task *TaskRef
+ // cases holds the tasks that Case named, in the order Case named them.
+ cases []*TaskRef
+}
+
+// Switch adds a task that runs the decider function fn, and returns the
switch. Use
+// [SwitchRef.Case] to name each task that fn can choose. fn returns
(*airflow.TaskRef, error), and
+// the TaskRef that it returns is the case that runs. So fn needs the TaskRefs
of the cases, for
+// example through the receiver of a method value:
+//
+// type paths struct{ long, short *airflow.TaskRef }
+//
+// func (p *paths) pickPath(actx airflow.Context, rows []string)
(*airflow.TaskRef, error) {
+// if len(rows) > 1000 {
+// return p.long, nil
+// }
+// return p.short, nil
+// }
+//
+// Where the Dag is built:
+//
+// extracted := dag.Task(extract)
+// p := &paths{long: dag.Task(handleLong), short: dag.Task(handleShort)}
+//
+// dag.Switch(p.pickPath,
airflow.Inputs(extracted)).Case(p.long).Case(p.short)
+//
+// The task that Switch adds then gets the task_id pickPath, the name of the
method. fn can also be
+// a function literal that closes over the TaskRefs. A function literal has no
name to take the
+// task_id from, so pass Switch a [TaskSpec] that sets TaskID. A function that
reads the TaskRefs
+// from package-level variables works only when the code that builds the Dag
runs once: a second
+// Dag built by the same code would overwrite the variables.
+//
+// In every other way fn follows the rules of a function passed to
[DagRef.Task]: it takes a
+// [Context] first, and [Inputs] fills the parameters after the Context. The
task that Switch adds
+// gets its task_id in the same way as a task from DagRef.Task, and a TaskSpec
sets its attributes.
+//
+// When the task runs, the result of fn decides what it skips:
+// - a case skips every other case
+// - a TaskRef that is not a case, or a nil one, fails the task with an
error that names the
+// TaskRef and the cases
+// - an error fails the task
+//
+// When the task fails, it pushes no XCom and skips nothing. Otherwise its
return_value XCom holds
+// the task_id of the case that fn chose.
+//
+// fn chooses exactly one case. A Python branch can choose several tasks by
returning a list of
+// task_ids, but a switch cannot. To run several tasks together, order them
after one task and make
+// that task a case, or give each of them a condition of its own with
[DagRef.If].
+//
+// A switch has no default case, because a Python branch operator has none. To
run a task when no
+// other case fits, make it a case and have fn return it.
+//
+// A case runs after the switch, so naming it records an edge from the switch
to it, as
+// [TaskRef.Before] would. When fn chooses a case, the switch skips only the
cases that fn did not
+// choose. Whether a task after the cases runs is then up to its trigger rule.
With the default
+// trigger rule, a task is skipped when one of its upstream tasks is skipped.
So a task after both
+// cases of the example above would never run. Give that task a trigger rule
like
+// [TriggerRuleNoneFailedMinOneSuccess], which runs a task when at least one
of its upstream tasks
+// succeeds and none fails:
+//
+// report := dag.Task(writeReport, airflow.TaskSpec{
+// TriggerRule: airflow.TriggerRuleNoneFailedMinOneSuccess,
+// })
+// report.After(p.long, p.short)
+//
+// Do not make report a case as well. The switch skips every case that fn does
not choose, so when
+// fn chooses p.short, the switch would skip report too. A Python branch works
differently here: it
+// does not skip a task that runs after the task it chooses. So the Python
pattern
+// branch >> [optional, join] with optional >> join needs a change in Go. Make
join a task like
+// report, and give the path without optional a case of its own, such as a
task that does nothing.
+//
+// [BundleRef.Register] panics if a switch has no case.
+//
+// Switch panics for the same reasons as DagRef.Task does for a Go function.
It also panics if fn
+// comes from [TriggerDagRun] or does not return (*airflow.TaskRef, error).
+func (d *DagRef) Switch(fn any, opts ...TaskOption) *SwitchRef {
+ return d.addSwitch("airflow.DagRef.Switch", nil, fn, opts)
+}
+
+// addSwitch adds a switch for DagRef.Switch and TaskGroupRef.Switch. group is
the task group that
+// the switch is added through, and nil for DagRef.Switch.
+func (d *DagRef) addSwitch(
+ method string, group *TaskGroupRef, fn any, opts []TaskOption,
+) *SwitchRef {
+ if _, ok := fn.(TriggerDagRunTask); ok {
+ panic(fmt.Sprintf(
+ "%s: Dag %q: fn comes from airflow.TriggerDagRun, "+
+ "but a decider function is a Go function that
returns (*airflow.TaskRef, error)",
+ method, d.dagID,
+ ))
+ }
+ switchRef := &SwitchRef{}
+ d.addTask(method, group, fn, opts, switchRef)
+ return switchRef
+}
+
+// Case names task as a task that the switch s can choose. When the decider
function of s chooses
+// another case, the task is skipped. Case returns s, so that a switch and all
of its cases fit in
+// one statement, as in the example of [DagRef.Switch]:
+//
+// dag.Switch(p.pickPath,
airflow.Inputs(extracted)).Case(p.long).Case(p.short)
+//
+// The task runs after the decider, but its function does not take the result
of the decider as a
+// parameter.
+//
+// Case panics if:
+// - s is not the SwitchRef that DagRef.Switch or TaskGroupRef.Switch
returned, for example a
+// copy of it
+// - task is nil, or neither DagRef.Task nor TaskGroupRef.Task added task to
the Dag of the
+// switch
+// - task is already a case of s
+// - the Dag is already registered
+func (s *SwitchRef) Case(task *TaskRef) *SwitchRef {
+ const method = "airflow.SwitchRef.Case"
+ // When the decider runs, it reads the cases from the SwitchRef that
Switch returned. A case
+ // given to a copy of that SwitchRef would never reach the decider.
+ if s == nil || s.task == nil || s.task.decider != s {
Review Comment:
`Case` doesn't reject the decider's own task as a case: `pick :=
dag.Switch(fn); pick.Case(pick.task)` passes this check (the decider is a
registered task, so the ownership check holds) and only surfaces later as a
confusing "contains a cycle" panic at `Register()`, instead of a clear, named
error like every other `Case` misuse gets.
```suggestion
if s == nil || s.task == nil || s.task.decider != s || task == s.task {
```
--
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]
