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 6330b4de9c5 Go SDK: reject a second TaskSpec or Inputs when dag.Task
applies it (#74057)
6330b4de9c5 is described below
commit 6330b4de9c50f20542736b8f0ab2901c28ec49e3
Author: PoAn Yang <[email protected]>
AuthorDate: Fri Oct 2 14:20:36 2026 +0900
Go SDK: reject a second TaskSpec or Inputs when dag.Task applies it (#74057)
Signed-off-by: PoAn Yang <[email protected]>
---
go-sdk/adr/0008-native-dag-interface.md | 2 +-
go-sdk/airflow/dag.go | 42 +++++++++++++++----------
go-sdk/airflow/dag_test.go | 17 ++++++++--
go-sdk/airflow/inputs.go | 31 +++++++++----------
go-sdk/airflow/inputs_test.go | 55 ++++++++++++++++++++++++++++++---
go-sdk/airflow/task_option.go | 25 +++++++++++----
go-sdk/airflow/trigger_dag_run_test.go | 4 +++
7 files changed, 128 insertions(+), 48 deletions(-)
diff --git a/go-sdk/adr/0008-native-dag-interface.md
b/go-sdk/adr/0008-native-dag-interface.md
index cd34e90f174..aebb92f5514 100644
--- a/go-sdk/adr/0008-native-dag-interface.md
+++ b/go-sdk/adr/0008-native-dag-interface.md
@@ -135,7 +135,7 @@ type TaskSpec struct {
// TaskOption is sealed: its only method is unexported, so a task takes
SDK-defined options and
// nothing else. TaskSpec and the value Inputs returns both implement it.
-type TaskOption interface{ applyTask(*taskConfig) }
+type TaskOption interface{ applyTask(*taskConfig) error }
func Inputs(refs ...*TaskRef) TaskOption
diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go
index e0b48277909..757ee52f817 100644
--- a/go-sdk/airflow/dag.go
+++ b/go-sdk/airflow/dag.go
@@ -173,21 +173,20 @@ func (d *DagRef) Task(fn any, opts ...TaskOption)
*TaskRef {
d.dagID, i, opt,
))
}
- opt.applyTask(&cfg)
- }
- if len(cfg.specs) > 1 {
- panic(fmt.Sprintf(
- "airflow.DagRef.Task: task %q of Dag %q got %d
airflow.TaskSpec values; "+
- "set all of the task's attributes in one
TaskSpec",
- findTaskName(fn, cfg.specs), d.dagID, len(cfg.specs),
- ))
+ if err := opt.applyTask(&cfg); err != nil {
+ // A task from TriggerDagRun takes no Inputs at all,
and the check after this loop
+ // reports that. err says to merge the Inputs into one,
which would not fix a task from
+ // TriggerDagRun.
+ if _, ok := opt.(inputs); ok && isTrigger {
+ continue
+ }
+ panic(fmt.Sprintf(
+ "airflow.DagRef.Task: task %q of Dag %q: %v",
findTaskName(fn, opts), d.dagID, err,
+ ))
+ }
}
- var spec TaskSpec
- if len(cfg.specs) == 1 {
- spec = cfg.specs[0]
- }
- taskID := spec.TaskID
+ taskID := cfg.spec.TaskID
if taskID == "" && isTrigger {
panic(fmt.Sprintf(
"airflow.DagRef.Task: Dag %q: a task from
airflow.TriggerDagRun with DagID %q has no "+
@@ -220,7 +219,7 @@ func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef {
var upstreams []*TaskRef
var resultType reflect.Type
if isTrigger {
- if len(cfg.inputs) > 0 {
+ if cfg.hasInputs {
panic(fmt.Sprintf(
"airflow.DagRef.Task: task %q of Dag %q comes
from airflow.TriggerDagRun and "+
"takes no airflow.Inputs, because it
has no Go function to pass the results to",
@@ -239,7 +238,7 @@ func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef {
task := &TaskRef{
dag: d,
taskID: taskID,
- spec: copySpec(spec),
+ spec: copySpec(cfg.spec),
resultType: resultType,
inputs: upstreams,
task: wrapped,
@@ -263,8 +262,17 @@ func (d *DagRef) markRegistered() {
func funcName(fn any) string { return
runtime.FuncForPC(reflect.ValueOf(fn).Pointer()).Name() }
// findTaskName names a task in an error that Task raises before it settles
the task_id.
-func findTaskName(fn any, specs []TaskSpec) string {
- for _, spec := range specs {
+func findTaskName(fn any, opts []TaskOption) string {
+ for _, opt := range opts {
+ var spec TaskSpec
+ switch opt := opt.(type) {
+ case TaskSpec:
+ spec = opt
+ case *TaskSpec:
+ if opt != nil {
+ spec = *opt
+ }
+ }
if spec.TaskID != "" {
return spec.TaskID
}
diff --git a/go-sdk/airflow/dag_test.go b/go-sdk/airflow/dag_test.go
index 5991283180b..ada37091ee3 100644
--- a/go-sdk/airflow/dag_test.go
+++ b/go-sdk/airflow/dag_test.go
@@ -218,6 +218,7 @@ func TestTaskNeedsTaskIDWhenFnHasNoName(t *testing.T) {
func TestTaskRejectsASecondTaskSpec(t *testing.T) {
literal := func(Context) error { return nil }
+ var nilSpec *TaskSpec
tests := []struct {
name string
@@ -238,12 +239,24 @@ func TestTaskRejectsASecondTaskSpec(t *testing.T) {
opts: []TaskOption{TaskSpec{}, TaskSpec{TaskID:
"given"}},
task: `given`,
},
+ {
+ name: "TaskID in a *TaskSpec",
+ fn: literal,
+ opts: []TaskOption{TaskSpec{}, &TaskSpec{TaskID:
"given"}},
+ task: `given`,
+ },
{
name: "no TaskID",
fn: ExtractRows,
opts: []TaskOption{TaskSpec{}, TaskSpec{}},
task: `ExtractRows`,
},
+ {
+ name: "nil *TaskSpec after the second spec",
+ fn: ExtractRows,
+ opts: []TaskOption{TaskSpec{}, TaskSpec{}, nilSpec},
+ task: `ExtractRows`,
+ },
{
name: "no TaskID and no function name",
fn: literal,
@@ -262,8 +275,8 @@ func TestTaskRejectsASecondTaskSpec(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
msg := panicMessage(t, func() { Dag("etl").Task(tt.fn,
tt.opts...) })
assert.Regexp(t,
- `^airflow\.DagRef\.Task: task "`+tt.task+`" of
Dag "etl" got 2 airflow\.TaskSpec `+
- `values; set all of the task's
attributes in one TaskSpec$`,
+ `^airflow\.DagRef\.Task: task "`+tt.task+`" of
Dag "etl": got more than one `+
+ `airflow\.TaskSpec; set all of the
task's attributes in one TaskSpec$`,
msg,
)
})
diff --git a/go-sdk/airflow/inputs.go b/go-sdk/airflow/inputs.go
index 3e3ad7f0a0d..62430dba3c8 100644
--- a/go-sdk/airflow/inputs.go
+++ b/go-sdk/airflow/inputs.go
@@ -18,6 +18,7 @@
package airflow
import (
+ "errors"
"fmt"
"reflect"
"slices"
@@ -27,7 +28,15 @@ import (
type inputs []*TaskRef
-func (in inputs) applyTask(c *taskConfig) { c.inputs = append(c.inputs, in) }
+func (in inputs) applyTask(c *taskConfig) error {
+ if c.hasInputs {
+ return errors.New(
+ "got more than one airflow.Inputs; pass all of the
task's inputs to one airflow.Inputs",
+ )
+ }
+ c.inputs, c.hasInputs = in, true
+ return nil
+}
// Inputs passes the results of tasks to the task that [DagRef.Task] adds, and
makes each of
// those tasks an upstream task of the new one. It is the Go form of a Python
TaskFlow call such
@@ -54,22 +63,10 @@ func (in inputs) applyTask(c *taskConfig) { c.inputs =
append(c.inputs, in) }
// parameter type, as in a Go function call. Pass at most one Inputs to a task.
func Inputs(refs ...*TaskRef) TaskOption { return inputs(refs) }
-// checkInputs returns the tasks that task taskID got through Inputs. given
has one entry for
-// each Inputs passed to the task. checkInputs panics if there is more than
one entry, or if the
-// tasks do not match the parameters that a function of type fnType takes
after the Context.
-func (d *DagRef) checkInputs(taskID string, fnType reflect.Type, given
[][]*TaskRef) []*TaskRef {
- if len(given) > 1 {
- panic(fmt.Sprintf(
- "airflow.DagRef.Task: task %q of Dag %q got %d
airflow.Inputs values; "+
- "pass all of the task's inputs to one
airflow.Inputs",
- taskID, d.dagID, len(given),
- ))
- }
- var tasks []*TaskRef
- if len(given) == 1 {
- tasks = given[0]
- }
-
+// checkInputs returns the tasks that task taskID got through Inputs, in a new
slice. It panics if
+// any of those tasks was not added to d by DagRef.Task, or if the tasks do
not match the
+// parameters that a function of type fnType takes after the Context.
+func (d *DagRef) checkInputs(taskID string, fnType reflect.Type, tasks
[]*TaskRef) []*TaskRef {
for i, upstream := range tasks {
switch {
case upstream == nil:
diff --git a/go-sdk/airflow/inputs_test.go b/go-sdk/airflow/inputs_test.go
index b4dd5dc4ec2..af0c86c7372 100644
--- a/go-sdk/airflow/inputs_test.go
+++ b/go-sdk/airflow/inputs_test.go
@@ -383,10 +383,55 @@ func TestTaskPanicsOnAnInputFromOutsideTheDag(t
*testing.T) {
func TestTaskRejectsASecondInputs(t *testing.T) {
dag := Dag("etl")
read := dag.Task(readRows)
+ literal := func(Context, rowSet) (int, error) { return 0, nil }
- assert.PanicsWithValue(t,
- `airflow.DagRef.Task: task "countRows" of Dag "etl" got 2
airflow.Inputs values; `+
- `pass all of the task's inputs to one airflow.Inputs`,
- func() { dag.Task(countRows, Inputs(read), Inputs()) },
- )
+ tests := []struct {
+ name string
+ fn any
+ opts []TaskOption
+ // task is a regexp for how the panic message names the task.
+ task string
+ }{
+ {
+ name: "second Inputs is empty",
+ fn: countRows,
+ opts: []TaskOption{Inputs(read), Inputs()},
+ task: `countRows`,
+ },
+ {
+ name: "first Inputs is empty",
+ fn: countRows,
+ opts: []TaskOption{Inputs(), Inputs(read)},
+ task: `countRows`,
+ },
+ {
+ name: "TaskID in a TaskSpec after the Inputs",
+ fn: countRows,
+ opts: []TaskOption{Inputs(read), Inputs(read),
TaskSpec{TaskID: "count_rows"}},
+ task: `count_rows`,
+ },
+ {
+ name: "nil option after the second Inputs",
+ fn: countRows,
+ opts: []TaskOption{Inputs(read), Inputs(read), nil},
+ task: `countRows`,
+ },
+ {
+ name: "no TaskID and no function name",
+ fn: literal,
+ opts: []TaskOption{Inputs(read), Inputs(read)},
+ task: `github\.com/apache/airflow/go-sdk/airflow\.` +
+ `TestTaskRejectsASecondInputs\.func\d+`,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ msg := panicMessage(t, func() { dag.Task(tt.fn,
tt.opts...) })
+ assert.Regexp(t,
+ `^airflow\.DagRef\.Task: task "`+tt.task+`" of
Dag "etl": got more than one `+
+ `airflow\.Inputs; pass all of the
task's inputs to one airflow\.Inputs$`,
+ msg,
+ )
+ })
+ }
}
diff --git a/go-sdk/airflow/task_option.go b/go-sdk/airflow/task_option.go
index d4d78d733dd..91efff4ff0d 100644
--- a/go-sdk/airflow/task_option.go
+++ b/go-sdk/airflow/task_option.go
@@ -17,6 +17,8 @@
package airflow
+import "errors"
+
// TaskOption is an option to [DagRef.Task]. There are two kinds: a [TaskSpec]
sets the
// attributes of the task that DagRef.Task adds, and [Inputs] passes the
results of other tasks
// to that task.
@@ -24,13 +26,24 @@ package airflow
// Its only method is unexported, so a type outside this package cannot
declare it.
// A struct that embeds a TaskSpec or a TaskOption still satisfies the
interface, and Task
// panics when it is given one.
-type TaskOption interface{ applyTask(*taskConfig) }
+type TaskOption interface{ applyTask(*taskConfig) error }
type taskConfig struct {
- // specs and inputs keep every TaskSpec and every Inputs passed to
DagRef.Task, so that Task
- // can reject a second TaskSpec or a second Inputs instead of merging
it into the first.
- specs []TaskSpec
- inputs [][]*TaskRef
+ spec TaskSpec
+ inputs []*TaskRef
+ // hasSpec is true once DagRef.Task has applied a TaskSpec, and
hasInputs is true once it has
+ // applied an Inputs. The spec and inputs fields cannot show that an
option was applied,
+ // because TaskSpec{} leaves spec at its zero value and Inputs() leaves
inputs nil.
+ hasSpec bool
+ hasInputs bool
}
-func (s TaskSpec) applyTask(c *taskConfig) { c.specs = append(c.specs, s) }
+func (s TaskSpec) applyTask(c *taskConfig) error {
+ if c.hasSpec {
+ return errors.New(
+ "got more than one airflow.TaskSpec; set all of the
task's attributes in one TaskSpec",
+ )
+ }
+ c.spec, c.hasSpec = s, true
+ return nil
+}
diff --git a/go-sdk/airflow/trigger_dag_run_test.go
b/go-sdk/airflow/trigger_dag_run_test.go
index 7cc825073f2..b812f146cec 100644
--- a/go-sdk/airflow/trigger_dag_run_test.go
+++ b/go-sdk/airflow/trigger_dag_run_test.go
@@ -87,6 +87,7 @@ func TestTriggerDagRunNeedsATaskID(t *testing.T) {
assert.PanicsWithValue(t, want, func() { Dag("etl").Task(trigger) })
assert.PanicsWithValue(t, want, func() { Dag("etl").Task(trigger,
TaskSpec{}) })
+ assert.PanicsWithValue(t, want, func() { Dag("etl").Task(trigger,
Inputs(), Inputs()) })
}
// buildNestedConf returns a conf in which maps nest depth levels deep.
encoding/json decodes at
@@ -250,6 +251,9 @@ func TestTriggerDagRunTakesNoInputs(t *testing.T) {
assert.PanicsWithValue(t, want, func() {
dag.Task(trigger, TaskSpec{TaskID: "trigger_downstream"},
Inputs())
})
+ assert.PanicsWithValue(t, want, func() {
+ dag.Task(trigger, TaskSpec{TaskID: "trigger_downstream"},
Inputs(extracted), Inputs())
+ })
assert.NotPanics(t,
func() { dag.Task(trigger, TaskSpec{TaskID:
"trigger_downstream"}) },
"a rejected task does not take its task_id",
