From 944df9236cd66313525c17e6755e61df82ef6b0e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 5 May 2026 11:16:55 +0800 Subject: [PATCH 1/9] Implement execution server and task runner for coordinator protocol - Add `messages_test.go` to test message decoding and encoding functionalities. - Introduce `serde.go` for serialization of various data types to Airflow's format. - Create `serde_test.go` to validate serialization logic and ensure correctness. - Implement `server.go` to handle communication with the supervisor and manage task execution. - Add `task_runner.go` to execute tasks based on received startup details and handle success/failure. --- go-sdk/bundle/bundlev1/registry.go | 58 ++++-- go-sdk/pkg/execution/dag_parser.go | 53 ++++++ go-sdk/pkg/execution/serde.go | 279 +++++++++++++++++++++++++++++ go-sdk/pkg/execution/serde_test.go | 226 +++++++++++++++++++++++ 4 files changed, 598 insertions(+), 18 deletions(-) create mode 100644 go-sdk/pkg/execution/dag_parser.go create mode 100644 go-sdk/pkg/execution/serde.go create mode 100644 go-sdk/pkg/execution/serde_test.go diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 80253fb36653b..a83ffbe1c6b85 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -75,9 +75,16 @@ type ( AddDag(dagId string) Dag } - // TaskInfo describes a registered task by its user-visible id. + // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses + // it to render the per-task block of a DagFileParsingResult. TaskInfo struct { + // ID is the user-visible task id (the function name unless overridden + // via AddTaskWithName). ID string + // TypeName is the unqualified Go function name (e.g. "extract"). + TypeName string + // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). + PkgPath string } // DagInfo describes a registered dag together with its tasks in @@ -87,9 +94,9 @@ type ( Tasks []TaskInfo } - // EnumerableBundle exposes the dag/task identity recorded by RegisterDags. - // The default registry implements it; airflow-go-pack relies on it to read - // a bundle's dag/task ids without executing any task. + // EnumerableBundle exposes the dag/task identity recorded by + // RegisterDags. The default registry implements it; the coordinator-mode + // runtime relies on it for the DAG-parse one-shot. EnumerableBundle interface { OrderedDags() []DagInfo } @@ -97,6 +104,7 @@ type ( registry struct { sync.RWMutex taskFuncMap map[string]map[string]Task + taskInfo map[string]map[string]TaskInfo dagOrder []string taskOrder map[string][]string } @@ -122,26 +130,35 @@ func (d dagShim) AddTaskWithName(taskId string, fn any) { func New() Registry { return ®istry{ taskFuncMap: make(map[string]map[string]Task), + taskInfo: make(map[string]map[string]TaskInfo), taskOrder: make(map[string][]string), } } +func splitFullName(fullName string) (typeName, pkgPath string) { + // fullName looks like "main.extract" or "github.com/x/y.MyTask"; method + // values get a "-fm" suffix. + lastDot := strings.LastIndex(fullName, ".") + if lastDot < 0 { + return strings.TrimSuffix(fullName, "-fm"), "" + } + return strings.TrimSuffix(fullName[lastDot+1:], "-fm"), fullName[:lastDot] +} + func getFnName(fn reflect.Value) string { fullName := runtime.FuncForPC(fn.Pointer()).Name() - parts := strings.Split(fullName, ".") - fnName := parts[len(parts)-1] - // Go adds `-fm` suffix to a method names - return strings.TrimSuffix(fnName, "-fm") + name, _ := splitFullName(fullName) + return name } func (r *registry) AddDag(dagId string) Dag { - r.Lock() - defer r.Unlock() - + r.RWMutex.Lock() + defer r.RWMutex.Unlock() if _, exists := r.taskFuncMap[dagId]; exists { panic(fmt.Errorf("Dag %q already exists in bundle", dagId)) } r.taskFuncMap[dagId] = make(map[string]Task) + r.taskInfo[dagId] = make(map[string]TaskInfo) r.dagOrder = append(r.dagOrder, dagId) return dagShim{dagId, r} } @@ -164,22 +181,27 @@ func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) } + val := reflect.ValueOf(fn) + fullName := runtime.FuncForPC(val.Pointer()).Name() + typeName, pkgPath := splitFullName(fullName) + r.RWMutex.Lock() defer r.RWMutex.Unlock() dagTasks, exists := r.taskFuncMap[dagId] - if !exists { dagTasks = make(map[string]Task) r.taskFuncMap[dagId] = dagTasks + r.taskInfo[dagId] = make(map[string]TaskInfo) r.dagOrder = append(r.dagOrder, dagId) } - _, exists = dagTasks[taskId] - if exists { + if _, exists := dagTasks[taskId]; exists { panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId)) } + dagTasks[taskId] = task + r.taskInfo[dagId][taskId] = TaskInfo{ID: taskId, TypeName: typeName, PkgPath: pkgPath} r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) } @@ -195,9 +217,9 @@ func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { return task, exists } -// OrderedDags returns the registered dags in AddDag order, each with its tasks -// in registration order. The returned slice is freshly allocated; callers may -// mutate it freely. +// OrderedDags returns the registered dags in the order AddDag was called, +// each with its tasks in the order AddTask / AddTaskWithName was called. The +// returned slice is freshly allocated; callers may mutate it freely. func (r *registry) OrderedDags() []DagInfo { r.RLock() defer r.RUnlock() @@ -207,7 +229,7 @@ func (r *registry) OrderedDags() []DagInfo { taskIDs := r.taskOrder[dagID] tasks := make([]TaskInfo, 0, len(taskIDs)) for _, tid := range taskIDs { - tasks = append(tasks, TaskInfo{ID: tid}) + tasks = append(tasks, r.taskInfo[dagID][tid]) } out = append(out, DagInfo{DagID: dagID, Tasks: tasks}) } diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go new file mode 100644 index 0000000000000..734a5ded77f72 --- /dev/null +++ b/go-sdk/pkg/execution/dag_parser.go @@ -0,0 +1,53 @@ +// 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 execution + +import ( + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// ParseDags processes a DagFileParseRequest by serialising every dag +// registered on bundle to DagSerialization v3 and returning the result as a +// DagFileParsingResult body. bundle is the materialised registry produced by +// running BundleProvider.RegisterDags. +func ParseDags(bundle bundlev1.Bundle, req *DagFileParseRequest) map[string]any { + fileloc := req.File + bundlePath := req.BundlePath + relativeFileloc := computeRelativeFileloc(fileloc, bundlePath) + + var dags []bundlev1.DagInfo + if enum, ok := bundle.(bundlev1.EnumerableBundle); ok { + dags = enum.OrderedDags() + } + + serializedDags := make([]any, 0, len(dags)) + for _, d := range dags { + serializedDags = append(serializedDags, map[string]any{ + "data": map[string]any{ + "__version": 3, + "dag": SerializeDag(d, fileloc, relativeFileloc), + }, + }) + } + + return map[string]any{ + "type": "DagFileParsingResult", + "fileloc": fileloc, + "serialized_dags": serializedDags, + } +} diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go new file mode 100644 index 0000000000000..02ec641c89dcb --- /dev/null +++ b/go-sdk/pkg/execution/serde.go @@ -0,0 +1,279 @@ +// 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 execution + +import ( + "fmt" + "path/filepath" + "reflect" + "sort" + "time" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// serializeValue recursively serializes a value with Airflow's type/var encoding. +// This matches Python's BaseSerialization.serialize() output: +// - primitives (string, bool, int, float) pass through unchanged +// - time.Time -> {"__type": "datetime", "__var": epoch_seconds_float} +// - time.Duration -> {"__type": "timedelta", "__var": total_seconds_float} +// - map[string]any -> {"__type": "dict", "__var": {k: serialize(v), ...}} +// - []any -> direct array with each element serialized +func serializeValue(value any) any { + if value == nil { + return nil + } + switch v := value.(type) { + case string, bool: + return v + case int: + return v + case int8: + return int(v) + case int16: + return int(v) + case int32: + return int(v) + case int64: + return v + case float32: + return float64(v) + case float64: + return v + case time.Time: + epochSec := float64(v.Unix()) + float64(v.Nanosecond())/1e9 + return map[string]any{ + "__type": "datetime", + "__var": epochSec, + } + case time.Duration: + return map[string]any{ + "__type": "timedelta", + "__var": v.Seconds(), + } + case map[string]any: + serialized := make(map[string]any, len(v)) + for k, val := range v { + serialized[k] = serializeValue(val) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case []string: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + case []any: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + default: + // Use reflection to handle typed maps and slices that don't match + // the concrete types above (e.g., map[string]map[string][]string). + rv := reflect.ValueOf(v) + switch rv.Kind() { + case reflect.Map: + serialized := make(map[string]any, rv.Len()) + for _, key := range rv.MapKeys() { + serialized[fmt.Sprint(key.Interface())] = serializeValue(rv.MapIndex(key).Interface()) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case reflect.Slice, reflect.Array: + result := make([]any, rv.Len()) + for i := range result { + result[i] = serializeValue(rv.Index(i).Interface()) + } + return result + default: + return v + } + } +} + +// unwrapTypeEncoding extracts the "__var" part from a type-encoded value. +// In Python's serialize_to_json, non-decorated fields are serialized then unwrapped. +func unwrapTypeEncoding(value any) any { + m, ok := value.(map[string]any) + if !ok { + return value + } + if _, hasType := m["__type"]; !hasType { + return value + } + if v, hasVar := m["__var"]; hasVar { + return v + } + return value +} + +// serializeTimetable converts a schedule string to the Airflow timetable format. +func serializeTimetable(schedule *string) map[string]any { + if schedule == nil { + return map[string]any{ + "__type": "airflow.timetables.simple.NullTimetable", + "__var": map[string]any{}, + } + } + switch *schedule { + case "@once": + return map[string]any{ + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": map[string]any{}, + } + case "@continuous": + return map[string]any{ + "__type": "airflow.timetables.simple.ContinuousTimetable", + "__var": map[string]any{}, + } + default: + return map[string]any{ + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": map[string]any{ + "expression": *schedule, + "timezone": "UTC", + "interval": 0.0, + "run_immediately": false, + }, + } + } +} + +// serializeTask converts a task to the Airflow serialization format. +func serializeTask(taskID, typeName, pkgPath string, downstream []string) map[string]any { + if typeName == "" { + typeName = taskID + } + if pkgPath == "" { + pkgPath = "main" + } + data := map[string]any{ + "task_id": taskID, + "task_type": typeName, + "_task_module": pkgPath, + "language": "go", + } + if len(downstream) > 0 { + sorted := make([]string, len(downstream)) + copy(sorted, downstream) + sort.Strings(sorted) + data["downstream_task_ids"] = sorted + } + return map[string]any{ + "__type": "operator", + "__var": data, + } +} + +// serializeTaskGroup creates a flat root task group containing all task IDs. +func serializeTaskGroup(taskIDs []string) map[string]any { + children := make(map[string]any, len(taskIDs)) + for _, id := range taskIDs { + children[id] = []any{"operator", id} + } + return map[string]any{ + "_group_id": nil, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": children, + "upstream_group_ids": []any{}, + "downstream_group_ids": []any{}, + "upstream_task_ids": []any{}, + "downstream_task_ids": []any{}, + } +} + +// serializeParams converts DAG params to Airflow's serialization format. +func serializeParams(params map[string]any) []any { + if len(params) == 0 { + return []any{} + } + result := make([]any, 0, len(params)) + for k, v := range params { + result = append(result, []any{ + k, + map[string]any{ + "__class": "airflow.sdk.definitions.param.Param", + "default": serializeValue(v), + "description": nil, + "schema": serializeValue(map[string]any{}), + "source": nil, + }, + }) + } + return result +} + +// SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 +// format. The Go SDK's bundlev1.Dag interface does not (yet) carry per-DAG +// metadata like schedule, start_date, tags, etc., so the encoding emits +// schema defaults for those fields. The optional-field handling below is +// kept (gated on nil checks) so the encoder can grow naturally as the +// bundle metadata surface expands. +func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { + taskIDs := make([]string, len(info.Tasks)) + tasks := make([]any, len(info.Tasks)) + for i, t := range info.Tasks { + taskIDs[i] = t.ID + tasks[i] = serializeTask(t.ID, t.TypeName, t.PkgPath, nil) + } + + return map[string]any{ + // Required fields (always present) + "dag_id": info.DagID, + "fileloc": fileloc, + "relative_fileloc": relativeFileloc, + "timezone": "UTC", + "timetable": serializeTimetable(nil), + "tasks": tasks, + "dag_dependencies": []any{}, + "task_group": serializeTaskGroup(taskIDs), + "edge_info": map[string]any{}, + "params": serializeParams(nil), + "deadline": nil, + "allowed_run_types": nil, + } +} + +// computeRelativeFileloc computes the relative file location from the bundle path. +func computeRelativeFileloc(fileloc, bundlePath string) string { + if fileloc == "" { + return "" + } + if bundlePath == "" { + return "." + } + rel, err := filepath.Rel(bundlePath, fileloc) + if err != nil { + return "." + } + if rel == "" { + return "." + } + return rel +} diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go new file mode 100644 index 0000000000000..afddaf5e700b1 --- /dev/null +++ b/go-sdk/pkg/execution/serde_test.go @@ -0,0 +1,226 @@ +// 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 execution + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +func TestSerializeValuePrimitives(t *testing.T) { + assert.Nil(t, serializeValue(nil)) + assert.Equal(t, "hello", serializeValue("hello")) + assert.Equal(t, true, serializeValue(true)) + assert.Equal(t, 42, serializeValue(42)) + assert.Equal(t, float64(3.14), serializeValue(3.14)) +} + +func TestSerializeValueDatetime(t *testing.T) { + ts := time.Date(2024, 1, 15, 10, 30, 0, 500000000, time.UTC) + result := serializeValue(ts) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "datetime", m["__type"]) + epochSec := m["__var"].(float64) + expected := float64(ts.Unix()) + 0.5 + assert.InDelta(t, expected, epochSec, 0.001) +} + +func TestSerializeValueTimedelta(t *testing.T) { + dur := 90 * time.Second + result := serializeValue(dur) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "timedelta", m["__type"]) + assert.Equal(t, 90.0, m["__var"]) +} + +func TestSerializeValueMap(t *testing.T) { + input := map[string]any{ + "key1": "val1", + "key2": 42, + } + result := serializeValue(input) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "dict", m["__type"]) + inner := m["__var"].(map[string]any) + assert.Equal(t, "val1", inner["key1"]) + assert.Equal(t, 42, inner["key2"]) +} + +func TestSerializeValueSlice(t *testing.T) { + input := []any{"a", 1, true} + result := serializeValue(input) + arr, ok := result.([]any) + require.True(t, ok) + assert.Len(t, arr, 3) + assert.Equal(t, "a", arr[0]) +} + +func TestUnwrapTypeEncoding(t *testing.T) { + wrapped := map[string]any{ + "__type": "datetime", + "__var": 1705313400.5, + } + assert.Equal(t, 1705313400.5, unwrapTypeEncoding(wrapped)) + + assert.Equal(t, "hello", unwrapTypeEncoding("hello")) + assert.Equal(t, 42, unwrapTypeEncoding(42)) +} + +func TestSerializeTimetable(t *testing.T) { + t.Run("nil schedule", func(t *testing.T) { + result := serializeTimetable(nil) + assert.Equal(t, "airflow.timetables.simple.NullTimetable", result["__type"]) + }) + + t.Run("@once", func(t *testing.T) { + s := "@once" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.simple.OnceTimetable", result["__type"]) + }) + + t.Run("@continuous", func(t *testing.T) { + s := "@continuous" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.simple.ContinuousTimetable", result["__type"]) + }) + + t.Run("cron expression", func(t *testing.T) { + s := "0 12 * * *" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.trigger.CronTriggerTimetable", result["__type"]) + v := result["__var"].(map[string]any) + assert.Equal(t, "0 12 * * *", v["expression"]) + assert.Equal(t, "UTC", v["timezone"]) + assert.Equal(t, 0.0, v["interval"]) + assert.Equal(t, false, v["run_immediately"]) + }) +} + +func TestSerializeTask(t *testing.T) { + result := serializeTask("extract", "extract", "main", []string{"transform"}) + assert.Equal(t, "operator", result["__type"]) + data := result["__var"].(map[string]any) + assert.Equal(t, "extract", data["task_id"]) + assert.Equal(t, "extract", data["task_type"]) + assert.Equal(t, "main", data["_task_module"]) + assert.Equal(t, "go", data["language"]) + assert.Equal(t, []string{"transform"}, data["downstream_task_ids"]) +} + +func TestSerializeTaskNoDownstream(t *testing.T) { + result := serializeTask("load", "load", "main", nil) + data := result["__var"].(map[string]any) + _, hasDownstream := data["downstream_task_ids"] + assert.False(t, hasDownstream) +} + +func TestSerializeTaskGroup(t *testing.T) { + result := serializeTaskGroup([]string{"t1", "t2"}) + assert.Nil(t, result["_group_id"]) + assert.Equal(t, true, result["prefix_group_id"]) + assert.Equal(t, "CornflowerBlue", result["ui_color"]) + + children := result["children"].(map[string]any) + assert.Equal(t, []any{"operator", "t1"}, children["t1"]) + assert.Equal(t, []any{"operator", "t2"}, children["t2"]) +} + +func TestSerializeParams(t *testing.T) { + t.Run("empty", func(t *testing.T) { + result := serializeParams(nil) + assert.Equal(t, []any{}, result) + }) + + t.Run("with values", func(t *testing.T) { + result := serializeParams(map[string]any{"key1": "default_val"}) + assert.Len(t, result, 1) + pair := result[0].([]any) + assert.Equal(t, "key1", pair[0]) + paramMap := pair[1].(map[string]any) + assert.Equal(t, "airflow.sdk.definitions.param.Param", paramMap["__class"]) + assert.Equal(t, "default_val", paramMap["default"]) + }) +} + +func TestSerializeDagMinimal(t *testing.T) { + info := bundlev1.DagInfo{DagID: "test_dag"} + result := SerializeDag(info, "/path/to/bundle", ".") + + assert.Equal(t, "test_dag", result["dag_id"]) + assert.Equal(t, "/path/to/bundle", result["fileloc"]) + assert.Equal(t, ".", result["relative_fileloc"]) + assert.Equal(t, "UTC", result["timezone"]) + + tt := result["timetable"].(map[string]any) + assert.Equal(t, "airflow.timetables.simple.NullTimetable", tt["__type"]) + + _, hasDesc := result["description"] + assert.False(t, hasDesc) + _, hasCatchup := result["catchup"] + assert.False(t, hasCatchup) +} + +func TestSerializeDagWithTasks(t *testing.T) { + info := bundlev1.DagInfo{ + DagID: "etl", + Tasks: []bundlev1.TaskInfo{ + {ID: "extract", TypeName: "extract", PkgPath: "main"}, + {ID: "load", TypeName: "load", PkgPath: "main"}, + }, + } + result := SerializeDag(info, "/bundle/main.go", "main.go") + + tasks := result["tasks"].([]any) + require.Len(t, tasks, 2) + first := tasks[0].(map[string]any) + v := first["__var"].(map[string]any) + assert.Equal(t, "extract", v["task_id"]) + assert.Equal(t, "extract", v["task_type"]) + assert.Equal(t, "main", v["_task_module"]) + assert.Equal(t, "go", v["language"]) + + tg := result["task_group"].(map[string]any) + children := tg["children"].(map[string]any) + assert.Contains(t, children, "extract") + assert.Contains(t, children, "load") +} + +func TestComputeRelativeFileloc(t *testing.T) { + tests := []struct { + fileloc string + bundlePath string + want string + }{ + {"", "", ""}, + {"/a/b/c.go", "", "."}, + {"/bundles/my/dags.go", "/bundles/my", "dags.go"}, + {"/bundles/my/sub/dags.go", "/bundles/my", "sub/dags.go"}, + } + for _, tt := range tests { + result := computeRelativeFileloc(tt.fileloc, tt.bundlePath) + assert.Equal(t, tt.want, result, "fileloc=%q bundlePath=%q", tt.fileloc, tt.bundlePath) + } +} From ae15c12629604e38b558a1317165e5446e789733 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 11 May 2026 17:46:30 +0800 Subject: [PATCH 2/9] Enhance task and DAG registration with optional specifications for improved configuration flexibility --- go-sdk/bundle/bundlev1/registry.go | 131 ++++++++++++++--- go-sdk/bundle/bundlev1/registry_test.go | 75 +++++++++- go-sdk/example/bundle/main.go | 12 +- go-sdk/pkg/execution/serde.go | 167 +++++++++++++++++++-- go-sdk/pkg/execution/serde_test.go | 184 +++++++++++++++++++++++- 5 files changed, 533 insertions(+), 36 deletions(-) diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index a83ffbe1c6b85..2f03f6f8a88a2 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,6 +23,7 @@ import ( "runtime" "strings" "sync" + "time" "github.com/apache/airflow/go-sdk/pkg/worker" ) @@ -42,7 +43,9 @@ type ( // functions that implement the dag's tasks. Dag interface { // AddTask registers fn as a task, deriving the task_id from fn's own - // name (so it must match the @task.stub name in the Python dag). + // name (so it must match the @task.stub name in the Python dag). An + // optional TaskSpec configures scheduling attributes such as queue, + // pool, and retries; passing more than one spec panics. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -55,13 +58,13 @@ type ( // the task, and a non-nil first result is pushed as the task's // return-value XCom. Passing a non-function, or a function whose return // signature does not match, panics at registration time. - AddTask(fn any) + AddTask(fn any, spec ...TaskSpec) // AddTaskWithName is like AddTask but sets task_id explicitly instead of // deriving it from the function name. Use it when the Go function name // cannot match the Python @task.stub id, for example for an anonymous // function or a differing name. - AddTaskWithName(taskId string, fn any) + AddTaskWithName(taskId string, fn any, spec ...TaskSpec) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -70,9 +73,73 @@ type ( Registry interface { Bundle // AddDag registers a dag by its dag_id (matching the Python stub dag) - // and returns a Dag handle for attaching tasks. Registering the same - // dag_id twice panics. - AddDag(dagId string) Dag + // and returns a Dag handle for attaching tasks. An optional DagSpec + // configures dag-level attributes such as schedule and tags; passing + // more than one spec panics. Registering the same dag_id twice panics. + AddDag(dagId string, spec ...DagSpec) Dag + } + + // TaskSpec is the optional configuration applied to a task at registration + // time. Every field is optional: a zero value means "unset" and the + // scheduler falls back to its serialization-schema default. The field + // names mirror the keys defined under "operator" in + // airflow-core/src/airflow/serialization/schema.json. + TaskSpec struct { + Queue string + Pool string + PoolSlots int + Retries int + RetryDelay time.Duration + MaxRetryDelay time.Duration + RetryExponentialBackoff float64 + PriorityWeight int + WeightRule string + TriggerRule string + Owner string + ExecutionTimeout time.Duration + Executor string + StartDate time.Time + EndDate time.Time + DependsOnPast bool + WaitForDownstream bool + // DoXComPush, EmailOnFailure, and EmailOnRetry default to true in the + // scheduler. A nil pointer means "unset" so the field is omitted from + // the serialized payload; pass Bool(false) to explicitly opt out. + DoXComPush *bool + EmailOnFailure *bool + EmailOnRetry *bool + DocMD string + MapIndexTemplate string + MaxActiveTisPerDag int + MaxActiveTisPerDagrun int + } + + // DagSpec is the optional configuration applied to a DAG at registration + // time. Every field is optional: a zero value means "unset" and the + // scheduler falls back to its serialization-schema default. The field + // names mirror the keys defined under "dag" in + // airflow-core/src/airflow/serialization/schema.json. + DagSpec struct { + // Schedule is "@once", "@continuous", a cron expression, or "" for + // NullTimetable (no schedule). + Schedule string + Description string + StartDate time.Time + EndDate time.Time + Tags []string + DagDisplayName string + DocMD string + MaxActiveTasks int + MaxActiveRuns int + MaxConsecutiveFailedDagRuns int + DagrunTimeout time.Duration + Catchup bool + FailFast bool + RenderTemplateAsNativeObj bool + DisableBundleVersioning bool + // IsPausedUponCreation has no schema default. nil means "unset"; pass + // Bool(true) or Bool(false) to set it explicitly. + IsPausedUponCreation *bool } // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses @@ -85,12 +152,18 @@ type ( TypeName string // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). PkgPath string + // Spec carries the optional per-task configuration supplied at + // registration. The zero value means "no overrides". + Spec TaskSpec } // DagInfo describes a registered dag together with its tasks in // registration order. DagInfo struct { DagID string + // Spec carries the optional per-dag configuration supplied at + // registration. The zero value means "no overrides". + Spec DagSpec Tasks []TaskInfo } @@ -105,6 +178,7 @@ type ( sync.RWMutex taskFuncMap map[string]map[string]Task taskInfo map[string]map[string]TaskInfo + dagSpec map[string]DagSpec dagOrder []string taskOrder map[string][]string } @@ -115,12 +189,32 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any) { - d.registry.registerTask(d.dagId, fn) +func (d dagShim) AddTask(fn any, spec ...TaskSpec) { + d.registry.registerTask(d.dagId, fn, optionalSpec(spec, "AddTask")) } -func (d dagShim) AddTaskWithName(taskId string, fn any) { - d.registry.registerTaskWithName(d.dagId, taskId, fn) +func (d dagShim) AddTaskWithName(taskId string, fn any, spec ...TaskSpec) { + d.registry.registerTaskWithName(d.dagId, taskId, fn, optionalSpec(spec, "AddTaskWithName")) +} + +// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / +// DagSpec where nil means "leave at schema default": +// +// v1.TaskSpec{DoXComPush: v1.Bool(false)} +func Bool(b bool) *bool { + return &b +} + +func optionalSpec[T any](specs []T, caller string) T { + switch len(specs) { + case 0: + var zero T + return zero + case 1: + return specs[0] + default: + panic(fmt.Errorf("%s accepts at most one spec, got %d", caller, len(specs))) + } } // New returns an empty Registry on which dags and tasks can be registered. The @@ -131,6 +225,7 @@ func New() Registry { return ®istry{ taskFuncMap: make(map[string]map[string]Task), taskInfo: make(map[string]map[string]TaskInfo), + dagSpec: make(map[string]DagSpec), taskOrder: make(map[string][]string), } } @@ -151,7 +246,8 @@ func getFnName(fn reflect.Value) string { return name } -func (r *registry) AddDag(dagId string) Dag { +func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { + dagSpec := optionalSpec(spec, "AddDag") r.RWMutex.Lock() defer r.RWMutex.Unlock() if _, exists := r.taskFuncMap[dagId]; exists { @@ -159,11 +255,12 @@ func (r *registry) AddDag(dagId string) Dag { } r.taskFuncMap[dagId] = make(map[string]Task) r.taskInfo[dagId] = make(map[string]TaskInfo) + r.dagSpec[dagId] = dagSpec r.dagOrder = append(r.dagOrder, dagId) return dagShim{dagId, r} } -func (r *registry) registerTask(dagId string, fn any) { +func (r *registry) registerTask(dagId string, fn any, spec TaskSpec) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -172,10 +269,10 @@ func (r *registry) registerTask(dagId string, fn any) { fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn) + r.registerTaskWithName(dagId, fnName, fn, spec) } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { +func (r *registry) registerTaskWithName(dagId, taskId string, fn any, spec TaskSpec) { task, err := NewTaskFunction(fn) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) @@ -185,6 +282,8 @@ func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { fullName := runtime.FuncForPC(val.Pointer()).Name() typeName, pkgPath := splitFullName(fullName) + info := TaskInfo{ID: taskId, TypeName: typeName, PkgPath: pkgPath, Spec: spec} + r.RWMutex.Lock() defer r.RWMutex.Unlock() @@ -201,7 +300,7 @@ func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { } dagTasks[taskId] = task - r.taskInfo[dagId][taskId] = TaskInfo{ID: taskId, TypeName: typeName, PkgPath: pkgPath} + r.taskInfo[dagId][taskId] = info r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) } @@ -231,7 +330,7 @@ func (r *registry) OrderedDags() []DagInfo { for _, tid := range taskIDs { tasks = append(tasks, r.taskInfo[dagID][tid]) } - out = append(out, DagInfo{DagID: dagID, Tasks: tasks}) + out = append(out, DagInfo{DagID: dagID, Spec: r.dagSpec[dagID], Tasks: tasks}) } return out } diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index f405c8f531578..6856d487e2311 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -144,12 +144,83 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { got := enum.OrderedDags() s.Require().Len(got, 3) + taskIDs := func(tasks []TaskInfo) []string { + ids := make([]string, 0, len(tasks)) + for _, t := range tasks { + ids = append(ids, t.ID) + } + return ids + } + s.Equal("zeta", got[0].DagID) - s.Equal([]TaskInfo{{ID: "z1"}, {ID: "z2"}}, got[0].Tasks) + s.Equal([]string{"z1", "z2"}, taskIDs(got[0].Tasks)) s.Equal("alpha", got[1].DagID) - s.Equal([]TaskInfo{{ID: "a1"}}, got[1].Tasks) + s.Equal([]string{"a1"}, taskIDs(got[1].Tasks)) s.Equal("mid", got[2].DagID) s.Empty(got[2].Tasks) } + +func (s *RegistrySuite) TestAddTask_WithSpec() { + s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 1) + s.Require().Len(dags[0].Tasks, 1) + got := dags[0].Tasks[0] + s.Equal("myTask", got.ID) + s.Equal("high_mem", got.Spec.Queue) + s.Equal(3, got.Spec.Retries) + s.Require().NotNil(got.Spec.DoXComPush) + s.False(*got.Spec.DoXComPush) +} + +func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { + s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 1) + s.Require().Len(dags[0].Tasks, 1) + got := dags[0].Tasks[0] + s.Equal("special", got.ID) + s.Equal("gpu", got.Spec.Queue) + s.Equal("gpu_pool", got.Spec.Pool) +} + +func (s *RegistrySuite) TestAddTask_TooManySpecsPanics() { + s.PanicsWithError("AddTask accepts at most one spec, got 2", func() { + s.dag.AddTask(myTask, TaskSpec{}, TaskSpec{}) + }) +} + +func (s *RegistrySuite) TestAddDag_WithSpec() { + dag2 := s.reg.AddDag( + "dag2", + DagSpec{Schedule: "@daily", Tags: []string{"etl"}, MaxActiveRuns: 4}, + ) + s.NotNil(dag2) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 2) + var got DagInfo + for _, d := range dags { + if d.DagID == "dag2" { + got = d + break + } + } + s.Equal("dag2", got.DagID) + s.Equal("@daily", got.Spec.Schedule) + s.Equal([]string{"etl"}, got.Spec.Tags) + s.Equal(4, got.Spec.MaxActiveRuns) +} + +func (s *RegistrySuite) TestAddDag_TooManySpecsPanics() { + s.PanicsWithError("AddDag accepts at most one spec, got 2", func() { + s.reg.AddDag("dag3", DagSpec{}, DagSpec{}) + }) +} diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..32d55d53d753d 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -46,10 +46,14 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { } func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) - simpleDag.AddTask(load) + simpleDag := dagbag.AddDag("simple_dag", v1.DagSpec{ + Schedule: "@daily", + Description: "Example Go-authored Dag", + Tags: []string{"example", "go-sdk"}, + }) + simpleDag.AddTask(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}) + simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}) + simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}) // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go index 02ec641c89dcb..5b42b9b7a352c 100644 --- a/go-sdk/pkg/execution/serde.go +++ b/go-sdk/pkg/execution/serde.go @@ -162,19 +162,22 @@ func serializeTimetable(schedule *string) map[string]any { } // serializeTask converts a task to the Airflow serialization format. -func serializeTask(taskID, typeName, pkgPath string, downstream []string) map[string]any { +func serializeTask(info bundlev1.TaskInfo, downstream []string) map[string]any { + typeName := info.TypeName if typeName == "" { - typeName = taskID + typeName = info.ID } + pkgPath := info.PkgPath if pkgPath == "" { pkgPath = "main" } data := map[string]any{ - "task_id": taskID, + "task_id": info.ID, "task_type": typeName, "_task_module": pkgPath, "language": "go", } + applyTaskSpec(data, info.Spec) if len(downstream) > 0 { sorted := make([]string, len(downstream)) copy(sorted, downstream) @@ -187,6 +190,142 @@ func serializeTask(taskID, typeName, pkgPath string, downstream []string) map[st } } +// applyTaskSpec mirrors Python BaseSerialization's "omit hard-coded default" +// behavior: each TaskSpec field is written into data only when it differs +// from the schema default declared in +// airflow-core/src/airflow/serialization/schema.json. A zero-valued field is +// always considered "unset" and is skipped. +func applyTaskSpec(data map[string]any, s bundlev1.TaskSpec) { + if s.Queue != "" && s.Queue != "default" { + data["queue"] = s.Queue + } + if s.Pool != "" && s.Pool != "default_pool" { + data["pool"] = s.Pool + } + if s.PoolSlots != 0 && s.PoolSlots != 1 { + data["pool_slots"] = s.PoolSlots + } + if s.Retries != 0 { + data["retries"] = s.Retries + } + if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { + data["retry_delay"] = unwrapTypeEncoding(serializeValue(s.RetryDelay)) + } + if s.MaxRetryDelay != 0 { + data["max_retry_delay"] = unwrapTypeEncoding(serializeValue(s.MaxRetryDelay)) + } + if s.RetryExponentialBackoff != 0 { + data["retry_exponential_backoff"] = s.RetryExponentialBackoff + } + if s.PriorityWeight != 0 && s.PriorityWeight != 1 { + data["priority_weight"] = s.PriorityWeight + } + if s.WeightRule != "" && s.WeightRule != "downstream" { + data["weight_rule"] = s.WeightRule + } + if s.TriggerRule != "" && s.TriggerRule != "all_success" { + data["trigger_rule"] = s.TriggerRule + } + if s.Owner != "" && s.Owner != "airflow" { + data["owner"] = s.Owner + } + if s.ExecutionTimeout != 0 { + data["execution_timeout"] = unwrapTypeEncoding(serializeValue(s.ExecutionTimeout)) + } + if s.Executor != "" { + data["executor"] = s.Executor + } + if !s.StartDate.IsZero() { + data["start_date"] = unwrapTypeEncoding(serializeValue(s.StartDate)) + } + if !s.EndDate.IsZero() { + data["end_date"] = unwrapTypeEncoding(serializeValue(s.EndDate)) + } + if s.DependsOnPast { + data["depends_on_past"] = true + } + if s.WaitForDownstream { + data["wait_for_downstream"] = true + } + // do_xcom_push / email_on_failure / email_on_retry default to true; only + // emit when an explicit false overrides the default. + if s.DoXComPush != nil && !*s.DoXComPush { + data["do_xcom_push"] = false + } + if s.EmailOnFailure != nil && !*s.EmailOnFailure { + data["email_on_failure"] = false + } + if s.EmailOnRetry != nil && !*s.EmailOnRetry { + data["email_on_retry"] = false + } + if s.DocMD != "" { + data["doc_md"] = s.DocMD + } + if s.MapIndexTemplate != "" { + data["map_index_template"] = s.MapIndexTemplate + } + if s.MaxActiveTisPerDag != 0 { + data["max_active_tis_per_dag"] = s.MaxActiveTisPerDag + } + if s.MaxActiveTisPerDagrun != 0 { + data["max_active_tis_per_dagrun"] = s.MaxActiveTisPerDagrun + } +} + +// applyDagSpec writes optional DAG-level fields onto data, omitting any +// field equal to its schema default. See applyTaskSpec for the convention. +func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { + if s.Description != "" { + data["description"] = s.Description + } + if !s.StartDate.IsZero() { + data["start_date"] = unwrapTypeEncoding(serializeValue(s.StartDate)) + } + if !s.EndDate.IsZero() { + data["end_date"] = unwrapTypeEncoding(serializeValue(s.EndDate)) + } + if len(s.Tags) > 0 { + tags := make([]any, len(s.Tags)) + for i, t := range s.Tags { + tags[i] = t + } + data["tags"] = tags + } + if s.DagDisplayName != "" { + data["dag_display_name"] = s.DagDisplayName + } + if s.DocMD != "" { + data["doc_md"] = s.DocMD + } + if s.MaxActiveTasks != 0 && s.MaxActiveTasks != 16 { + data["max_active_tasks"] = s.MaxActiveTasks + } + if s.MaxActiveRuns != 0 && s.MaxActiveRuns != 16 { + data["max_active_runs"] = s.MaxActiveRuns + } + if s.MaxConsecutiveFailedDagRuns != 0 { + data["max_consecutive_failed_dag_runs"] = s.MaxConsecutiveFailedDagRuns + } + if s.DagrunTimeout != 0 { + data["dagrun_timeout"] = unwrapTypeEncoding(serializeValue(s.DagrunTimeout)) + } + if s.Catchup { + data["catchup"] = true + } + if s.FailFast { + data["fail_fast"] = true + } + if s.RenderTemplateAsNativeObj { + data["render_template_as_native_obj"] = true + } + if s.DisableBundleVersioning { + data["disable_bundle_versioning"] = true + } + if s.IsPausedUponCreation != nil { + data["is_paused_upon_creation"] = *s.IsPausedUponCreation + } +} + // serializeTaskGroup creates a flat root task group containing all task IDs. func serializeTaskGroup(taskIDs []string) map[string]any { children := make(map[string]any, len(taskIDs)) @@ -230,26 +369,30 @@ func serializeParams(params map[string]any) []any { } // SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 -// format. The Go SDK's bundlev1.Dag interface does not (yet) carry per-DAG -// metadata like schedule, start_date, tags, etc., so the encoding emits -// schema defaults for those fields. The optional-field handling below is -// kept (gated on nil checks) so the encoder can grow naturally as the -// bundle metadata surface expands. +// format. Required fields are always present; optional fields from +// info.Spec are emitted only when they differ from their schema default +// (see applyDagSpec). func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { taskIDs := make([]string, len(info.Tasks)) tasks := make([]any, len(info.Tasks)) for i, t := range info.Tasks { taskIDs[i] = t.ID - tasks[i] = serializeTask(t.ID, t.TypeName, t.PkgPath, nil) + tasks[i] = serializeTask(t, nil) } - return map[string]any{ + var schedule *string + if info.Spec.Schedule != "" { + s := info.Spec.Schedule + schedule = &s + } + + result := map[string]any{ // Required fields (always present) "dag_id": info.DagID, "fileloc": fileloc, "relative_fileloc": relativeFileloc, "timezone": "UTC", - "timetable": serializeTimetable(nil), + "timetable": serializeTimetable(schedule), "tasks": tasks, "dag_dependencies": []any{}, "task_group": serializeTaskGroup(taskIDs), @@ -258,6 +401,8 @@ func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[st "deadline": nil, "allowed_run_types": nil, } + applyDagSpec(result, info.Spec) + return result } // computeRelativeFileloc computes the relative file location from the bundle path. diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index afddaf5e700b1..6121d24e285e8 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -120,7 +120,10 @@ func TestSerializeTimetable(t *testing.T) { } func TestSerializeTask(t *testing.T) { - result := serializeTask("extract", "extract", "main", []string{"transform"}) + result := serializeTask( + bundlev1.TaskInfo{ID: "extract", TypeName: "extract", PkgPath: "main"}, + []string{"transform"}, + ) assert.Equal(t, "operator", result["__type"]) data := result["__var"].(map[string]any) assert.Equal(t, "extract", data["task_id"]) @@ -128,15 +131,124 @@ func TestSerializeTask(t *testing.T) { assert.Equal(t, "main", data["_task_module"]) assert.Equal(t, "go", data["language"]) assert.Equal(t, []string{"transform"}, data["downstream_task_ids"]) + _, hasQueue := data["queue"] + assert.False(t, hasQueue, "queue should be omitted when unset") } func TestSerializeTaskNoDownstream(t *testing.T) { - result := serializeTask("load", "load", "main", nil) + result := serializeTask( + bundlev1.TaskInfo{ID: "load", TypeName: "load", PkgPath: "main"}, + nil, + ) data := result["__var"].(map[string]any) _, hasDownstream := data["downstream_task_ids"] assert.False(t, hasDownstream) } +func TestSerializeTaskCustomQueue(t *testing.T) { + result := serializeTask( + bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }, + nil, + ) + data := result["__var"].(map[string]any) + assert.Equal(t, "high_mem", data["queue"]) +} + +func TestSerializeTaskDefaultQueueOmitted(t *testing.T) { + result := serializeTask( + bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "default"}, + }, + nil, + ) + data := result["__var"].(map[string]any) + _, hasQueue := data["queue"] + assert.False(t, hasQueue, "queue=\"default\" matches the schema default and should be omitted") +} + +func TestApplyTaskSpec_EmitsAndOmits(t *testing.T) { + spec := bundlev1.TaskSpec{ + Queue: "gpu", + Pool: "gpu_pool", + PoolSlots: 4, + Retries: 3, + RetryDelay: 60 * time.Second, + MaxRetryDelay: 10 * time.Minute, + RetryExponentialBackoff: 2.0, + PriorityWeight: 5, + WeightRule: "upstream", + TriggerRule: "all_done", + Owner: "data-eng", + ExecutionTimeout: 45 * time.Second, + Executor: "KubernetesExecutor", + DependsOnPast: true, + WaitForDownstream: true, + DoXComPush: bundlev1.Bool(false), + EmailOnFailure: bundlev1.Bool(false), + EmailOnRetry: bundlev1.Bool(false), + DocMD: "## task", + MapIndexTemplate: "{{ task.task_id }}", + MaxActiveTisPerDag: 2, + MaxActiveTisPerDagrun: 1, + } + data := map[string]any{} + applyTaskSpec(data, spec) + + assert.Equal(t, "gpu", data["queue"]) + assert.Equal(t, "gpu_pool", data["pool"]) + assert.Equal(t, 4, data["pool_slots"]) + assert.Equal(t, 3, data["retries"]) + assert.Equal(t, 60.0, data["retry_delay"]) + assert.Equal(t, 600.0, data["max_retry_delay"]) + assert.Equal(t, 2.0, data["retry_exponential_backoff"]) + assert.Equal(t, 5, data["priority_weight"]) + assert.Equal(t, "upstream", data["weight_rule"]) + assert.Equal(t, "all_done", data["trigger_rule"]) + assert.Equal(t, "data-eng", data["owner"]) + assert.Equal(t, 45.0, data["execution_timeout"]) + assert.Equal(t, "KubernetesExecutor", data["executor"]) + assert.Equal(t, true, data["depends_on_past"]) + assert.Equal(t, true, data["wait_for_downstream"]) + assert.Equal(t, false, data["do_xcom_push"]) + assert.Equal(t, false, data["email_on_failure"]) + assert.Equal(t, false, data["email_on_retry"]) + assert.Equal(t, "## task", data["doc_md"]) + assert.Equal(t, "{{ task.task_id }}", data["map_index_template"]) + assert.Equal(t, 2, data["max_active_tis_per_dag"]) + assert.Equal(t, 1, data["max_active_tis_per_dagrun"]) +} + +func TestApplyTaskSpec_OmitsSchemaDefaults(t *testing.T) { + // Values equal to schema defaults must be dropped. + spec := bundlev1.TaskSpec{ + Queue: "default", + Pool: "default_pool", + PoolSlots: 1, + Retries: 0, + RetryDelay: 300 * time.Second, + PriorityWeight: 1, + WeightRule: "downstream", + TriggerRule: "all_success", + Owner: "airflow", + DoXComPush: bundlev1.Bool(true), + EmailOnFailure: bundlev1.Bool(true), + EmailOnRetry: bundlev1.Bool(true), + } + data := map[string]any{} + applyTaskSpec(data, spec) + assert.Empty(t, data, "all fields equal schema defaults; nothing should be emitted") +} + +func TestApplyTaskSpec_EmptySpecNoOp(t *testing.T) { + data := map[string]any{} + applyTaskSpec(data, bundlev1.TaskSpec{}) + assert.Empty(t, data) +} + func TestSerializeTaskGroup(t *testing.T) { result := serializeTaskGroup([]string{"t1", "t2"}) assert.Nil(t, result["_group_id"]) @@ -188,7 +300,10 @@ func TestSerializeDagWithTasks(t *testing.T) { DagID: "etl", Tasks: []bundlev1.TaskInfo{ {ID: "extract", TypeName: "extract", PkgPath: "main"}, - {ID: "load", TypeName: "load", PkgPath: "main"}, + { + ID: "load", TypeName: "load", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }, }, } result := SerializeDag(info, "/bundle/main.go", "main.go") @@ -201,6 +316,11 @@ func TestSerializeDagWithTasks(t *testing.T) { assert.Equal(t, "extract", v["task_type"]) assert.Equal(t, "main", v["_task_module"]) assert.Equal(t, "go", v["language"]) + _, hasQueue := v["queue"] + assert.False(t, hasQueue, "extract has no queue set; field should be omitted") + + second := tasks[1].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "high_mem", second["queue"]) tg := result["task_group"].(map[string]any) children := tg["children"].(map[string]any) @@ -208,6 +328,64 @@ func TestSerializeDagWithTasks(t *testing.T) { assert.Contains(t, children, "load") } +func TestSerializeDagWithSpec(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + info := bundlev1.DagInfo{ + DagID: "etl", + Spec: bundlev1.DagSpec{ + Schedule: "@daily", + Description: "Extract, transform, load", + StartDate: start, + Tags: []string{"prod", "etl"}, + DagDisplayName: "ETL Pipeline", + DocMD: "## ETL", + MaxActiveTasks: 32, + MaxActiveRuns: 4, + MaxConsecutiveFailedDagRuns: 3, + DagrunTimeout: 2 * time.Hour, + Catchup: true, + FailFast: true, + RenderTemplateAsNativeObj: true, + DisableBundleVersioning: true, + IsPausedUponCreation: bundlev1.Bool(true), + }, + } + result := SerializeDag(info, "/bundle/main.go", "main.go") + + tt := result["timetable"].(map[string]any) + assert.Equal(t, "airflow.timetables.trigger.CronTriggerTimetable", tt["__type"]) + v := tt["__var"].(map[string]any) + assert.Equal(t, "@daily", v["expression"]) + + assert.Equal(t, "Extract, transform, load", result["description"]) + assert.Equal(t, []any{"prod", "etl"}, result["tags"]) + assert.Equal(t, "ETL Pipeline", result["dag_display_name"]) + assert.Equal(t, "## ETL", result["doc_md"]) + assert.Equal(t, 32, result["max_active_tasks"]) + assert.Equal(t, 4, result["max_active_runs"]) + assert.Equal(t, 3, result["max_consecutive_failed_dag_runs"]) + assert.Equal(t, (2 * time.Hour).Seconds(), result["dagrun_timeout"]) + assert.Equal(t, true, result["catchup"]) + assert.Equal(t, true, result["fail_fast"]) + assert.Equal(t, true, result["render_template_as_native_obj"]) + assert.Equal(t, true, result["disable_bundle_versioning"]) + assert.Equal(t, true, result["is_paused_upon_creation"]) + + // start_date is a raw epoch number, not the type-wrapped form. + startDate := result["start_date"].(float64) + assert.InDelta(t, float64(start.Unix()), startDate, 0.001) +} + +func TestApplyDagSpec_OmitsSchemaDefaults(t *testing.T) { + spec := bundlev1.DagSpec{ + MaxActiveTasks: 16, + MaxActiveRuns: 16, + } + data := map[string]any{} + applyDagSpec(data, spec) + assert.Empty(t, data, "values equal to schema defaults must be omitted") +} + func TestComputeRelativeFileloc(t *testing.T) { tests := []struct { fileloc string From e86d11f010d6f4b40e39abedfe0561d02270c087 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 12 May 2026 10:38:10 +0800 Subject: [PATCH 3/9] Support setting downstream at AddTask method --- go-sdk/bundle/bundlev1/registry.go | 68 +++++++++++++----- go-sdk/bundle/bundlev1/registry_test.go | 87 +++++++++++++++++++----- go-sdk/example/bundle/main.go | 6 +- go-sdk/pkg/execution/integration_test.go | 10 +-- go-sdk/pkg/execution/serde.go | 14 ++-- go-sdk/pkg/execution/serde_test.go | 52 +++++++------- go-sdk/pkg/worker/runner_test.go | 4 +- 7 files changed, 169 insertions(+), 72 deletions(-) diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 2f03f6f8a88a2..f6a933cb50096 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -42,10 +42,12 @@ type ( // Dag is the handle returned by Registry.AddDag. Use it to attach the Go // functions that implement the dag's tasks. Dag interface { - // AddTask registers fn as a task, deriving the task_id from fn's own - // name (so it must match the @task.stub name in the Python dag). An - // optional TaskSpec configures scheduling attributes such as queue, - // pool, and retries; passing more than one spec panics. + // AddTask registers fn as a task in this Dag using fn's Go name as + // the task id (so it must match the @task.stub name in the Python + // dag). spec carries optional per-task configuration (pass TaskSpec{} + // for defaults). depends lists task ids in the same Dag that must run + // before this one; each must already be registered. Pass nil for no + // dependencies. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -58,13 +60,13 @@ type ( // the task, and a non-nil first result is pushed as the task's // return-value XCom. Passing a non-function, or a function whose return // signature does not match, panics at registration time. - AddTask(fn any, spec ...TaskSpec) + AddTask(fn any, spec TaskSpec, depends []string) - // AddTaskWithName is like AddTask but sets task_id explicitly instead of - // deriving it from the function name. Use it when the Go function name - // cannot match the Python @task.stub id, for example for an anonymous - // function or a differing name. - AddTaskWithName(taskId string, fn any, spec ...TaskSpec) + // AddTaskWithName is like AddTask but sets task_id explicitly instead + // of deriving it from the function name. Use it when the Go function + // name cannot match the Python @task.stub id, for example for an + // anonymous function or a differing name. + AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -155,6 +157,10 @@ type ( // Spec carries the optional per-task configuration supplied at // registration. The zero value means "no overrides". Spec TaskSpec + // Downstream lists task ids that depend on this task, populated as + // later tasks declare this id in their AddTask `depends` argument. + // Order is registration order; the serializer sorts before emit. + Downstream []string } // DagInfo describes a registered dag together with its tasks in @@ -189,12 +195,12 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any, spec ...TaskSpec) { - d.registry.registerTask(d.dagId, fn, optionalSpec(spec, "AddTask")) +func (d dagShim) AddTask(fn any, spec TaskSpec, depends []string) { + d.registry.registerTask(d.dagId, fn, spec, depends) } -func (d dagShim) AddTaskWithName(taskId string, fn any, spec ...TaskSpec) { - d.registry.registerTaskWithName(d.dagId, taskId, fn, optionalSpec(spec, "AddTaskWithName")) +func (d dagShim) AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) { + d.registry.registerTaskWithName(d.dagId, taskId, fn, spec, depends) } // Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / @@ -260,7 +266,7 @@ func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { return dagShim{dagId, r} } -func (r *registry) registerTask(dagId string, fn any, spec TaskSpec) { +func (r *registry) registerTask(dagId string, fn any, spec TaskSpec, depends []string) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -269,10 +275,15 @@ func (r *registry) registerTask(dagId string, fn any, spec TaskSpec) { fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn, spec) + r.registerTaskWithName(dagId, fnName, fn, spec, depends) } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any, spec TaskSpec) { +func (r *registry) registerTaskWithName( + dagId, taskId string, + fn any, + spec TaskSpec, + depends []string, +) { task, err := NewTaskFunction(fn) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) @@ -299,6 +310,29 @@ func (r *registry) registerTaskWithName(dagId, taskId string, fn any, spec TaskS panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId)) } + // Resolve depends to upstream TaskInfo entries, validating each exists. + // We dedupe so a repeated id in `depends` only records one downstream + // edge on the parent. + seen := make(map[string]bool, len(depends)) + for _, dep := range depends { + if dep == taskId { + panic(fmt.Errorf("task %q cannot depend on itself in DAG %q", taskId, dagId)) + } + if seen[dep] { + continue + } + seen[dep] = true + parent, ok := r.taskInfo[dagId][dep] + if !ok { + panic(fmt.Errorf( + "task %q depends on unknown task %q in DAG %q; register upstream tasks first", + taskId, dep, dagId, + )) + } + parent.Downstream = append(parent.Downstream, taskId) + r.taskInfo[dagId][dep] = parent + } + dagTasks[taskId] = task r.taskInfo[dagId][taskId] = info r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index 6856d487e2311..364b7b8a8dc37 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -67,14 +67,14 @@ func (s *RegistrySuite) TestAddDag_DuplicatePanics() { } func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask) + s.dag.AddTask(myTask, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "myTask") s.True(exists) s.NotNil(task) } func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "special") s.True(exists) s.NotNil(task) @@ -85,20 +85,20 @@ func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { } func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) }) } func (s *RegistrySuite) TestAddTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func") + s.dag.AddTask("not a func", TaskSpec{}, nil) }) } func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs) + s.dag.AddTask(myTaskWithArgs, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) @@ -108,13 +108,13 @@ func (s *RegistrySuite) TestAddTask_InvalidReturnType() { s.PanicsWithError( "error registering task \"NotErrorRet\" for DAG \"dag1\": expected task function github.com/apache/airflow/go-sdk/bundle/bundlev1.NotErrorRet last return value to return error but found int", func() { - s.dag.AddTask(NotErrorRet) + s.dag.AddTask(NotErrorRet, TaskSpec{}, nil) }, ) } func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask) + s.dag.AddTask(errorTask, TaskSpec{}, nil) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } @@ -130,11 +130,11 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { // Register dags out of alphabetical order so the test fails if OrderedDags // ever sorts instead of preserving AddDag order. zeta := reg.AddDag("zeta") - zeta.AddTaskWithName("z1", myTask) - zeta.AddTaskWithName("z2", myTask) + zeta.AddTaskWithName("z1", myTask, TaskSpec{}, nil) + zeta.AddTaskWithName("z2", myTask, TaskSpec{}, nil) alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask) + alpha.AddTaskWithName("a1", myTask, TaskSpec{}, nil) reg.AddDag("mid") // dag with no tasks @@ -163,7 +163,7 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { } func (s *RegistrySuite) TestAddTask_WithSpec() { - s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) + s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}, nil) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -178,7 +178,7 @@ func (s *RegistrySuite) TestAddTask_WithSpec() { } func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}) + s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}, nil) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -190,9 +190,64 @@ func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { s.Equal("gpu_pool", got.Spec.Pool) } -func (s *RegistrySuite) TestAddTask_TooManySpecsPanics() { - s.PanicsWithError("AddTask accepts at most one spec, got 2", func() { - s.dag.AddTask(myTask, TaskSpec{}, TaskSpec{}) +func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("transform", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.Equal([]string{"transform"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform"].Downstream) + s.Nil(byID["load"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_FanOutFanIn() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("transform_a", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("transform_b", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform_a", "transform_b"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.ElementsMatch([]string{"transform_a", "transform_b"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform_a"].Downstream) + s.Equal([]string{"load"}, byID["transform_b"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_DependsDuplicatesIgnored() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract", "extract"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.Equal([]string{"load"}, byID["extract"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_DependsUnknownPanics() { + s.PanicsWithError( + `task "load" depends on unknown task "extract" in DAG "dag1"; register upstream tasks first`, + func() { + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract"}) + }, + ) +} + +func (s *RegistrySuite) TestAddTask_DependsOnSelfPanics() { + s.PanicsWithError(`task "self" cannot depend on itself in DAG "dag1"`, func() { + s.dag.AddTaskWithName("self", myTask, TaskSpec{}, []string{"self"}) }) } diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 32d55d53d753d..5d13af315bc9c 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -51,9 +51,9 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { Description: "Example Go-authored Dag", Tags: []string{"example", "go-sdk"}, }) - simpleDag.AddTask(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}) - simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}) - simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}) + simpleDag.AddTask(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}, nil) + simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}, []string{"extract"}) + simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}, []string{"transform"}) // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2559359453866..21faf488974d3 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -87,7 +87,7 @@ func buildBundle(t *testing.T, register func(bundlev1.Registry)) bundlev1.Bundle func TestTaskRunnerSuccess(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask) + r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -110,7 +110,7 @@ func TestTaskRunnerSuccess(t *testing.T) { func TestTaskRunnerFailure(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask) + r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -160,7 +160,7 @@ func TestTaskRunnerRetry(t *testing.T) { func TestTaskRunnerTaskNotFound(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask) + r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -182,7 +182,7 @@ func TestTaskRunnerTaskNotFound(t *testing.T) { func TestTaskRunnerPanic(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask) + r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -412,7 +412,7 @@ func TestServeStartupDetailsEndToEnd(t *testing.T) { provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTask(simpleTask) + r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) return nil }, } diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go index 5b42b9b7a352c..0d9618c3990eb 100644 --- a/go-sdk/pkg/execution/serde.go +++ b/go-sdk/pkg/execution/serde.go @@ -161,8 +161,10 @@ func serializeTimetable(schedule *string) map[string]any { } } -// serializeTask converts a task to the Airflow serialization format. -func serializeTask(info bundlev1.TaskInfo, downstream []string) map[string]any { +// serializeTask converts a task to the Airflow serialization format. The +// downstream_task_ids slice is read from info.Downstream (populated by the +// registry from each task's `depends` argument) and sorted for stable JSON. +func serializeTask(info bundlev1.TaskInfo) map[string]any { typeName := info.TypeName if typeName == "" { typeName = info.ID @@ -178,9 +180,9 @@ func serializeTask(info bundlev1.TaskInfo, downstream []string) map[string]any { "language": "go", } applyTaskSpec(data, info.Spec) - if len(downstream) > 0 { - sorted := make([]string, len(downstream)) - copy(sorted, downstream) + if len(info.Downstream) > 0 { + sorted := make([]string, len(info.Downstream)) + copy(sorted, info.Downstream) sort.Strings(sorted) data["downstream_task_ids"] = sorted } @@ -377,7 +379,7 @@ func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[st tasks := make([]any, len(info.Tasks)) for i, t := range info.Tasks { taskIDs[i] = t.ID - tasks[i] = serializeTask(t, nil) + tasks[i] = serializeTask(t) } var schedule *string diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index 6121d24e285e8..6ad434e539b10 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -120,10 +120,10 @@ func TestSerializeTimetable(t *testing.T) { } func TestSerializeTask(t *testing.T) { - result := serializeTask( - bundlev1.TaskInfo{ID: "extract", TypeName: "extract", PkgPath: "main"}, - []string{"transform"}, - ) + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"transform"}, + }) assert.Equal(t, "operator", result["__type"]) data := result["__var"].(map[string]any) assert.Equal(t, "extract", data["task_id"]) @@ -135,36 +135,36 @@ func TestSerializeTask(t *testing.T) { assert.False(t, hasQueue, "queue should be omitted when unset") } +func TestSerializeTaskDownstreamSorted(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"transform", "audit", "load"}, + }) + data := result["__var"].(map[string]any) + assert.Equal(t, []string{"audit", "load", "transform"}, data["downstream_task_ids"]) +} + func TestSerializeTaskNoDownstream(t *testing.T) { - result := serializeTask( - bundlev1.TaskInfo{ID: "load", TypeName: "load", PkgPath: "main"}, - nil, - ) + result := serializeTask(bundlev1.TaskInfo{ID: "load", TypeName: "load", PkgPath: "main"}) data := result["__var"].(map[string]any) _, hasDownstream := data["downstream_task_ids"] assert.False(t, hasDownstream) } func TestSerializeTaskCustomQueue(t *testing.T) { - result := serializeTask( - bundlev1.TaskInfo{ - ID: "extract", TypeName: "extract", PkgPath: "main", - Spec: bundlev1.TaskSpec{Queue: "high_mem"}, - }, - nil, - ) + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }) data := result["__var"].(map[string]any) assert.Equal(t, "high_mem", data["queue"]) } func TestSerializeTaskDefaultQueueOmitted(t *testing.T) { - result := serializeTask( - bundlev1.TaskInfo{ - ID: "extract", TypeName: "extract", PkgPath: "main", - Spec: bundlev1.TaskSpec{Queue: "default"}, - }, - nil, - ) + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "default"}, + }) data := result["__var"].(map[string]any) _, hasQueue := data["queue"] assert.False(t, hasQueue, "queue=\"default\" matches the schema default and should be omitted") @@ -299,7 +299,10 @@ func TestSerializeDagWithTasks(t *testing.T) { info := bundlev1.DagInfo{ DagID: "etl", Tasks: []bundlev1.TaskInfo{ - {ID: "extract", TypeName: "extract", PkgPath: "main"}, + { + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"load"}, + }, { ID: "load", TypeName: "load", PkgPath: "main", Spec: bundlev1.TaskSpec{Queue: "high_mem"}, @@ -318,9 +321,12 @@ func TestSerializeDagWithTasks(t *testing.T) { assert.Equal(t, "go", v["language"]) _, hasQueue := v["queue"] assert.False(t, hasQueue, "extract has no queue set; field should be omitted") + assert.Equal(t, []string{"load"}, v["downstream_task_ids"]) second := tasks[1].(map[string]any)["__var"].(map[string]any) assert.Equal(t, "high_mem", second["queue"]) + _, hasDownstream := second["downstream_task_ids"] + assert.False(t, hasDownstream, "leaf task has no downstream") tg := result["task_group"].(map[string]any) children := tg["children"].(map[string]any) diff --git a/go-sdk/pkg/worker/runner_test.go b/go-sdk/pkg/worker/runner_test.go index 01ab4ba62b29a..201a963ac01af 100644 --- a/go-sdk/pkg/worker/runner_test.go +++ b/go-sdk/pkg/worker/runner_test.go @@ -163,7 +163,7 @@ func (s *WorkerSuite) TestStartContextErrorTaskDoesntStart() { s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { wasCalled = true return nil - }) + }, bundlev1.TaskSpec{}, nil) // Setup the mock s.ti.EXPECT(). @@ -196,7 +196,7 @@ func (s *WorkerSuite) TestTaskHeartbeatsWhileRunning() { s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { time.Sleep(time.Second) return nil - }) + }, bundlev1.TaskSpec{}, nil) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateSuccess) From 8ba040c73e2c43ae71ab12d06122d920190184fb Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 3 Jun 2026 11:41:39 +0800 Subject: [PATCH 4/9] Go-SDK: Restore DagFileParseRequest coordinator parse path Re-add the DagFileParseRequest message, its decoder wiring in decodeIncomingBody, and the server dispatch case so the supervisor can request DAG parsing and receive a DagFileParsingResult. Also fix two AddTaskWithName test calls broken by the registration signature change. --- go-sdk/example/bundle/main.go | 2 +- go-sdk/pkg/execution/dag_parser.go | 3 +- go-sdk/pkg/execution/integration_test.go | 124 +++++++++++++++++++++-- go-sdk/pkg/execution/server.go | 17 +++- go-sdk/pkg/worker/runner_test.go | 2 +- 5 files changed, 137 insertions(+), 11 deletions(-) diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 5d13af315bc9c..59dde6b28ec0a 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -57,7 +57,7 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently, v1.TaskSpec{}, nil) return nil } diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go index 734a5ded77f72..4b6836ab4d4b4 100644 --- a/go-sdk/pkg/execution/dag_parser.go +++ b/go-sdk/pkg/execution/dag_parser.go @@ -19,13 +19,14 @@ package execution import ( "github.com/apache/airflow/go-sdk/bundle/bundlev1" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" ) // ParseDags processes a DagFileParseRequest by serialising every dag // registered on bundle to DagSerialization v3 and returning the result as a // DagFileParsingResult body. bundle is the materialised registry produced by // running BundleProvider.RegisterDags. -func ParseDags(bundle bundlev1.Bundle, req *DagFileParseRequest) map[string]any { +func ParseDags(bundle bundlev1.Bundle, req *genmodels.DagFileParseRequest) map[string]any { fileloc := req.File bundlePath := req.BundlePath relativeFileloc := computeRelativeFileloc(fileloc, bundlePath) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 21faf488974d3..6abb963f391dd 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -29,6 +29,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/vmihailenco/msgpack/v5" "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" @@ -85,6 +86,64 @@ func buildBundle(t *testing.T, register func(bundlev1.Registry)) bundlev1.Bundle // --- Tests --- +func TestDagParsing(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + d := r.AddDag("test_dag") + d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + }) + + req := &genmodels.DagFileParseRequest{ + File: "/bundles/test/main.go", + BundlePath: "/bundles/test", + } + + result := ParseDags(bundle, req) + + assert.Equal(t, "DagFileParsingResult", result["type"]) + assert.Equal(t, "/bundles/test/main.go", result["fileloc"]) + + serializedDags, ok := result["serialized_dags"].([]any) + require.True(t, ok) + require.Len(t, serializedDags, 1) + + dagEntry := serializedDags[0].(map[string]any) + data := dagEntry["data"].(map[string]any) + assert.Equal(t, 3, data["__version"]) + + dagMap := data["dag"].(map[string]any) + assert.Equal(t, "test_dag", dagMap["dag_id"]) + + tt := dagMap["timetable"].(map[string]any) + assert.Equal(t, "airflow.timetables.simple.NullTimetable", tt["__type"]) + + tasks := dagMap["tasks"].([]any) + require.Len(t, tasks, 1) + taskMap := tasks[0].(map[string]any) + assert.Equal(t, "operator", taskMap["__type"]) + taskData := taskMap["__var"].(map[string]any) + assert.Equal(t, "simpleTask", taskData["task_id"]) + assert.Equal(t, "go", taskData["language"]) +} + +func TestDagParsingMultipleDagsPreservesOrder(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag("dag2").AddTask(failingTask, bundlev1.TaskSpec{}, nil) + }) + + req := &genmodels.DagFileParseRequest{File: "/bundle/main.go", BundlePath: "/bundle"} + result := ParseDags(bundle, req) + + serializedDags := result["serialized_dags"].([]any) + require.Len(t, serializedDags, 2) + + dag1Data := serializedDags[0].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "dag1", dag1Data["dag_id"]) + + dag2Data := serializedDags[1].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "dag2", dag2Data["dag_id"]) +} + func TestTaskRunnerSuccess(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) @@ -133,7 +192,7 @@ func TestTaskRunnerFailure(t *testing.T) { func TestTaskRunnerRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask) + r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -205,7 +264,7 @@ func TestTaskRunnerPanic(t *testing.T) { func TestTaskRunnerPanicRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask) + r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -233,7 +292,7 @@ func TestTaskRunnerPanicRetry(t *testing.T) { func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("ctxcheck", - func(ctx context.Context) error { return ctx.Err() }) + func(ctx context.Context) error { return ctx.Err() }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -270,7 +329,7 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -330,7 +389,7 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -406,6 +465,59 @@ func startSupervisor( return commLn.Addr().String(), logsLn.Addr().String(), commCh, logsCh, cleanup } +func TestServeDagFileParseEndToEnd(t *testing.T) { + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + defer cleanup() + + provider := &fakeProvider{ + register: func(r bundlev1.Registry) error { + d := r.AddDag("simple_dag") + d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + return nil + }, + } + + done := make(chan error, 1) + go func() { done <- Serve(provider, commAddr, logsAddr) }() + + commConn := <-commCh + require.NotNil(t, commConn) + defer commConn.Close() + logsConn := <-logsCh + require.NotNil(t, logsConn) + defer logsConn.Close() + + // Send DagFileParseRequest as a request frame. + payload, err := encodeRequest(0, map[string]any{ + "type": "DagFileParseRequest", + "file": "/bundle/main.go", + "bundle_path": "/bundle", + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, payload)) + + frame, err := readFrame(commConn) + require.NoError(t, err) + assert.Equal(t, int64(0), frame.ID) + require.True(t, isNilRaw(frame.Err)) + assert.Equal(t, "DagFileParsingResult", peekBodyType(frame.Body)) + + var body map[string]any + require.NoError(t, msgpack.Unmarshal(frame.Body, &body)) + + dags := body["serialized_dags"].([]any) + require.Len(t, dags, 1) + dag := dags[0].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "simple_dag", dag["dag_id"]) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after parse result") + } +} + func TestServeStartupDetailsEndToEnd(t *testing.T) { commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) defer cleanup() @@ -478,7 +590,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }) + }, bundlev1.TaskSpec{}, nil) return nil }, } diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 7e6e05ef94dc9..71ec3fc82616c 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,11 @@ // the bundle binary is launched with --comm/--logs by the Airflow supervisor // (Python ExecutableCoordinator), bundlev1server.Serve dispatches here. // -// The first inbound frame on the comm socket is a StartupDetails message -// that drives multi-round task execution. +// The first inbound frame on the comm socket selects between two +// sub-protocols: +// +// - DagFileParseRequest: one-shot, returns DagFileParsingResult and exits. +// - StartupDetails: multi-round task execution. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -163,6 +166,16 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { } switch msg := body.(type) { + case *genmodels.DagFileParseRequest: + logger.Debug("DAG parsing mode", "file", msg.File) + result := ParseDags(bundle, msg) + // Bound the terminal write so a wedged socket cannot hang shutdown. + _ = commConn.SetWriteDeadline(time.Now().Add(terminalSendTimeout)) + if err := comm.SendRequest(frame.ID, result); err != nil { + return fmt.Errorf("sending parse result: %w", err) + } + logger.Debug("DAG parsing complete") + case *genmodels.StartupDetails: logger.Debug("Task execution mode", "dag_id", msg.TI.DagID, diff --git a/go-sdk/pkg/worker/runner_test.go b/go-sdk/pkg/worker/runner_test.go index 201a963ac01af..90cdc2e9ffc77 100644 --- a/go-sdk/pkg/worker/runner_test.go +++ b/go-sdk/pkg/worker/runner_test.go @@ -232,7 +232,7 @@ func (s *WorkerSuite) TestTaskHeartbeatConflictStopsTask() { case <-time.After(2 * time.Second): return fmt.Errorf("task context was not cancelled") } - }) + }, bundlev1.TaskSpec{}, nil) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateFailed) From 9e22d199862b29045bea5d39c3698939188b9a5d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 3 Jun 2026 14:19:28 +0800 Subject: [PATCH 5/9] Go-SDK: Match Python DagSerialization output in bundle serde Align the Go bundle serializer with Python's serialize_dag so a Go-authored Dag produces the same serialized shape Python would for the fields the SDK models, verified field-by-field against a reference dump: - Always emit the DAG fields that have no JSON-schema default and that Python therefore never omits: catchup, disable_bundle_versioning, max_active_tasks, max_active_runs (default 16 from [core]) and max_consecutive_failed_dag_runs. - Sort tags, matching Python's set-backed serialization (stable dag_hash). - Always emit template_fields on each task. Cron schedules still serialize to CronTriggerTimetable, which matches only the default [scheduler] create_cron_data_intervals=False; honoring the non-default value needs the supervisor to pass scheduler config over the coordinator protocol, tracked in https://github.com/apache/airflow/issues/67938. --- go-sdk/pkg/execution/serde.go | 81 ++++++++++++++++++++++-------- go-sdk/pkg/execution/serde_test.go | 55 ++++++++++++++++---- 2 files changed, 105 insertions(+), 31 deletions(-) diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go index 0d9618c3990eb..0aca23993d99e 100644 --- a/go-sdk/pkg/execution/serde.go +++ b/go-sdk/pkg/execution/serde.go @@ -27,6 +27,14 @@ import ( "github.com/apache/airflow/go-sdk/bundle/bundlev1" ) +// Per-DAG defaults that Python resolves from [core] config when the DAG does +// not override them. The serializer always emits these fields (they have no +// JSON-schema default to omit against), so we fall back to the same values. +const ( + defaultMaxActiveTasksPerDag = 16 // [core] max_active_tasks_per_dag + defaultMaxActiveRunsPerDag = 16 // [core] max_active_runs_per_dag +) + // serializeValue recursively serializes a value with Airflow's type/var encoding. // This matches Python's BaseSerialization.serialize() output: // - primitives (string, bool, int, float) pass through unchanged @@ -130,6 +138,17 @@ func unwrapTypeEncoding(value any) any { } // serializeTimetable converts a schedule string to the Airflow timetable format. +// +// TODO: respect [scheduler] create_cron_data_intervals (and, once timedelta +// schedules are supported, create_delta_data_intervals). Python's +// _create_timetable selects CronDataIntervalTimetable when +// create_cron_data_intervals is True and CronTriggerTimetable when False; we +// hardcode CronTriggerTimetable, which matches only the default (False). The +// Go bundle binary cannot read airflow.cfg, so the supervisor (Python +// ExecutableCoordinator) must send these scheduler flags to the lang-SDK over +// the coordinator protocol (e.g. on DagFileParseRequest) before we can honor +// non-default deployments. Until that channel exists this stays default-only. +// Tracked at https://github.com/apache/airflow/issues/67938 func serializeTimetable(schedule *string) map[string]any { if schedule == nil { return map[string]any{ @@ -178,6 +197,10 @@ func serializeTask(info bundlev1.TaskInfo) map[string]any { "task_type": typeName, "_task_module": pkgPath, "language": "go", + // Python's operator serializer always emits template_fields (its + // list value never matches the tuple default it is compared against), + // so it is unconditional here too. Go tasks have no template fields. + "template_fields": []any{}, } applyTaskSpec(data, info.Spec) if len(info.Downstream) > 0 { @@ -274,8 +297,12 @@ func applyTaskSpec(data map[string]any, s bundlev1.TaskSpec) { } } -// applyDagSpec writes optional DAG-level fields onto data, omitting any -// field equal to its schema default. See applyTaskSpec for the convention. +// applyDagSpec writes DAG-level fields onto data. Fields with a JSON-schema +// default (description, dates, tags, fail_fast, …) are omitted when unset, as +// in applyTaskSpec. Fields with no schema default (catchup, +// disable_bundle_versioning, max_active_tasks, max_active_runs, +// max_consecutive_failed_dag_runs) are always emitted, because Python's +// serializer never omits them. func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { if s.Description != "" { data["description"] = s.Description @@ -287,8 +314,13 @@ func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { data["end_date"] = unwrapTypeEncoding(serializeValue(s.EndDate)) } if len(s.Tags) > 0 { - tags := make([]any, len(s.Tags)) - for i, t := range s.Tags { + // Python stores tags in a set and serializes them sorted (for stable + // dag_hash); mirror that here regardless of registration order. + sorted := make([]string, len(s.Tags)) + copy(sorted, s.Tags) + sort.Strings(sorted) + tags := make([]any, len(sorted)) + for i, t := range sorted { tags[i] = t } data["tags"] = tags @@ -299,30 +331,36 @@ func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { if s.DocMD != "" { data["doc_md"] = s.DocMD } - if s.MaxActiveTasks != 0 && s.MaxActiveTasks != 16 { - data["max_active_tasks"] = s.MaxActiveTasks - } - if s.MaxActiveRuns != 0 && s.MaxActiveRuns != 16 { - data["max_active_runs"] = s.MaxActiveRuns - } - if s.MaxConsecutiveFailedDagRuns != 0 { - data["max_consecutive_failed_dag_runs"] = s.MaxConsecutiveFailedDagRuns - } + // max_active_tasks / max_active_runs / max_consecutive_failed_dag_runs and + // the catchup / disable_bundle_versioning booleans have no schema default, + // so Python's serializer never omits them — it always writes the resolved + // value (the per-DAG default falls back to the matching [core] config). + // Emit them unconditionally to match, using the config defaults for unset + // (zero) fields. + maxActiveTasks := s.MaxActiveTasks + if maxActiveTasks == 0 { + maxActiveTasks = defaultMaxActiveTasksPerDag + } + data["max_active_tasks"] = maxActiveTasks + maxActiveRuns := s.MaxActiveRuns + if maxActiveRuns == 0 { + maxActiveRuns = defaultMaxActiveRunsPerDag + } + data["max_active_runs"] = maxActiveRuns + data["max_consecutive_failed_dag_runs"] = s.MaxConsecutiveFailedDagRuns + data["catchup"] = s.Catchup + data["disable_bundle_versioning"] = s.DisableBundleVersioning if s.DagrunTimeout != 0 { data["dagrun_timeout"] = unwrapTypeEncoding(serializeValue(s.DagrunTimeout)) } - if s.Catchup { - data["catchup"] = true - } + // fail_fast and render_template_as_native_obj both have schema default + // false, so Python omits them when false; keep that behavior. if s.FailFast { data["fail_fast"] = true } if s.RenderTemplateAsNativeObj { data["render_template_as_native_obj"] = true } - if s.DisableBundleVersioning { - data["disable_bundle_versioning"] = true - } if s.IsPausedUponCreation != nil { data["is_paused_upon_creation"] = *s.IsPausedUponCreation } @@ -371,9 +409,8 @@ func serializeParams(params map[string]any) []any { } // SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 -// format. Required fields are always present; optional fields from -// info.Spec are emitted only when they differ from their schema default -// (see applyDagSpec). +// format. Required fields are always present; spec-driven fields are emitted +// per the rules in applyDagSpec (some always, some only when set). func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { taskIDs := make([]string, len(info.Tasks)) tasks := make([]any, len(info.Tasks)) diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index 6ad434e539b10..e7aa19c70ae66 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -131,6 +131,8 @@ func TestSerializeTask(t *testing.T) { assert.Equal(t, "main", data["_task_module"]) assert.Equal(t, "go", data["language"]) assert.Equal(t, []string{"transform"}, data["downstream_task_ids"]) + // template_fields is always present (matches Python), empty for Go tasks. + assert.Equal(t, []any{}, data["template_fields"]) _, hasQueue := data["queue"] assert.False(t, hasQueue, "queue should be omitted when unset") } @@ -289,10 +291,19 @@ func TestSerializeDagMinimal(t *testing.T) { tt := result["timetable"].(map[string]any) assert.Equal(t, "airflow.timetables.simple.NullTimetable", tt["__type"]) + // Optional spec fields stay omitted when unset. _, hasDesc := result["description"] assert.False(t, hasDesc) - _, hasCatchup := result["catchup"] - assert.False(t, hasCatchup) + _, hasTags := result["tags"] + assert.False(t, hasTags) + + // Python always serializes these (no schema default to omit against), so + // a minimal Dag still carries them at their resolved [core] defaults. + assert.Equal(t, false, result["catchup"]) + assert.Equal(t, false, result["disable_bundle_versioning"]) + assert.Equal(t, defaultMaxActiveTasksPerDag, result["max_active_tasks"]) + assert.Equal(t, defaultMaxActiveRunsPerDag, result["max_active_runs"]) + assert.Equal(t, 0, result["max_consecutive_failed_dag_runs"]) } func TestSerializeDagWithTasks(t *testing.T) { @@ -364,7 +375,8 @@ func TestSerializeDagWithSpec(t *testing.T) { assert.Equal(t, "@daily", v["expression"]) assert.Equal(t, "Extract, transform, load", result["description"]) - assert.Equal(t, []any{"prod", "etl"}, result["tags"]) + // Tags are emitted sorted, matching Python's set-backed serialization. + assert.Equal(t, []any{"etl", "prod"}, result["tags"]) assert.Equal(t, "ETL Pipeline", result["dag_display_name"]) assert.Equal(t, "## ETL", result["doc_md"]) assert.Equal(t, 32, result["max_active_tasks"]) @@ -382,14 +394,39 @@ func TestSerializeDagWithSpec(t *testing.T) { assert.InDelta(t, float64(start.Unix()), startDate, 0.001) } -func TestApplyDagSpec_OmitsSchemaDefaults(t *testing.T) { - spec := bundlev1.DagSpec{ - MaxActiveTasks: 16, - MaxActiveRuns: 16, +func TestApplyDagSpec_AlwaysEmitsNonSchemaDefaultFields(t *testing.T) { + // catchup, disable_bundle_versioning and the three max_* fields have no + // JSON-schema default, so Python always serializes them. An empty spec + // must still emit them at their resolved [core] defaults. + data := map[string]any{} + applyDagSpec(data, bundlev1.DagSpec{}) + + assert.Equal(t, false, data["catchup"]) + assert.Equal(t, false, data["disable_bundle_versioning"]) + assert.Equal(t, defaultMaxActiveTasksPerDag, data["max_active_tasks"]) + assert.Equal(t, defaultMaxActiveRunsPerDag, data["max_active_runs"]) + assert.Equal(t, 0, data["max_consecutive_failed_dag_runs"]) + + // Fields with a false/empty schema default stay omitted when unset. + for _, k := range []string{ + "description", "start_date", "end_date", "tags", "dag_display_name", + "doc_md", "dagrun_timeout", "fail_fast", "render_template_as_native_obj", + "is_paused_upon_creation", + } { + _, ok := data[k] + assert.Falsef(t, ok, "%q should be omitted when unset", k) } +} + +func TestApplyDagSpec_OmitsFalseSchemaDefaultBooleans(t *testing.T) { + // fail_fast and render_template_as_native_obj default to false in the + // schema, so Python omits them unless explicitly true. data := map[string]any{} - applyDagSpec(data, spec) - assert.Empty(t, data, "values equal to schema defaults must be omitted") + applyDagSpec(data, bundlev1.DagSpec{FailFast: false, RenderTemplateAsNativeObj: false}) + _, hasFailFast := data["fail_fast"] + assert.False(t, hasFailFast) + _, hasNative := data["render_template_as_native_obj"] + assert.False(t, hasNative) } func TestComputeRelativeFileloc(t *testing.T) { From 8b8f7f345e68dd6fe39a08a797492f465344df01 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 16:36:57 +0000 Subject: [PATCH 6/9] Go-SDK: Generate TaskSpec from the dag serialization schema Replace the hand-written TaskSpec struct and its hand-mapped serializer omit-if-default rules with spec.gen.go, generated from the "operator" definition in airflow-core/src/airflow/serialization/schema.json by a local gen tool (same pattern as pkg/execution/genmodels). Field types and schema defaults can no longer drift from what the scheduler deserializes; regenerate with `just generate-specs`. DagSpec and the registration Info structs move to the hand-written spec.go next to it: Schedule is an SDK-level concept the schema has no scalar for, and several dag keys are always-emitted with [core]-config fallbacks the schema cannot express. --- go-sdk/Justfile | 4 + go-sdk/bundle/bundlev1/gen.go | 29 +++ go-sdk/bundle/bundlev1/gen/main.go | 349 +++++++++++++++++++++++++++++ go-sdk/bundle/bundlev1/registry.go | 101 --------- go-sdk/bundle/bundlev1/spec.gen.go | 158 +++++++++++++ go-sdk/bundle/bundlev1/spec.go | 92 ++++++++ go-sdk/pkg/execution/serde.go | 92 +------- go-sdk/pkg/execution/serde_test.go | 18 +- 8 files changed, 653 insertions(+), 190 deletions(-) create mode 100644 go-sdk/bundle/bundlev1/gen.go create mode 100644 go-sdk/bundle/bundlev1/gen/main.go create mode 100644 go-sdk/bundle/bundlev1/spec.gen.go create mode 100644 go-sdk/bundle/bundlev1/spec.go diff --git a/go-sdk/Justfile b/go-sdk/Justfile index a88599969634b..8312d483335ee 100644 --- a/go-sdk/Justfile +++ b/go-sdk/Justfile @@ -50,3 +50,7 @@ protoc: # Regenerate the coordinator-protocol data models from the task-sdk supervisor schema generate-models: go generate ./pkg/execution/genmodels/... + +# Regenerate the bundle spec structs from the Airflow dag-serialization schema +generate-specs: + go generate ./bundle/bundlev1/... diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go new file mode 100644 index 0000000000000..ab3c3718de63c --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen.go @@ -0,0 +1,29 @@ +// 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. + +// spec.gen.go is generated from the Airflow dag-serialization schema by the +// local gen tool: the TaskSpec struct and its SchemaFields omit-if-default +// rules come from the "operator" definition, so they cannot drift from what +// the scheduler deserializes. Don't edit spec.gen.go by hand; change the +// schema (or the field allowlist in ./gen) and re-run `just generate-specs` +// (go generate). DagSpec and the registration Info structs stay hand-written +// in spec.go: Schedule is an SDK-level concept the schema has no scalar for, +// and several dag keys are always-emitted with [core]-config fallbacks the +// schema cannot express. +package bundlev1 + +//go:generate go run ./gen -schema ../../../airflow-core/src/airflow/serialization/schema.json -out spec.gen.go diff --git a/go-sdk/bundle/bundlev1/gen/main.go b/go-sdk/bundle/bundlev1/gen/main.go new file mode 100644 index 0000000000000..085110676dfae --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -0,0 +1,349 @@ +// 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. + +// Command gen generates spec.gen.go from the Airflow dag-serialization schema +// (airflow-core/src/airflow/serialization/schema.json). +// +// It reads the "operator" definition and emits the TaskSpec struct plus its +// SchemaFields method, so the Go field types and the omit-if-default rules the +// serializer relies on cannot drift from the schema. Only the keys listed in +// taskSpecFields are emitted: they are the task attributes a bundle author may +// set from Go, a deliberate subset of the schema (the rest of the "operator" +// keys are serializer internals such as task_type or template_fields, or +// Python-only concerns such as templating and callbacks). +// +// Field defaults are read from the schema, never hard-coded here: a field is +// serialized only when it is set (non-zero) and differs from its schema +// default, mirroring Python BaseSerialization's behaviour of omitting values +// the scheduler will re-derive. Booleans whose schema default is true become +// *bool so nil can mean "unset" while an explicit false still serializes. +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "go/format" + "log" + "math" + "os" + "strings" +) + +// fieldDef selects one schema property for the generated spec struct. goType +// overrides the mechanical schema→Go type mapping for the few keys whose +// schema type is too loose (e.g. "number" that is semantically a float). +type fieldDef struct { + key string + goType string +} + +// taskSpecFields lists the "operator" keys exposed on TaskSpec, in struct +// order. Adding a task attribute means adding its key here and regenerating. +var taskSpecFields = []fieldDef{ + {key: "queue"}, + {key: "pool"}, + {key: "pool_slots"}, + {key: "retries"}, + {key: "retry_delay"}, + {key: "max_retry_delay"}, + {key: "retry_exponential_backoff", goType: "float64"}, + {key: "priority_weight"}, + {key: "weight_rule"}, + {key: "trigger_rule"}, + {key: "owner"}, + {key: "execution_timeout"}, + {key: "executor"}, + {key: "start_date"}, + {key: "end_date"}, + {key: "depends_on_past"}, + {key: "wait_for_downstream"}, + {key: "do_xcom_push"}, + {key: "email_on_failure"}, + {key: "email_on_retry"}, + {key: "doc_md"}, + {key: "map_index_template"}, + {key: "max_active_tis_per_dag"}, + {key: "max_active_tis_per_dagrun"}, +} + +// initialisms maps snake_case segments that must keep non-Title capitalization +// in Go names (do_xcom_push -> DoXComPush, doc_md -> DocMD). +var initialisms = map[string]string{ + "id": "ID", + "md": "MD", + "ui": "UI", + "uri": "URI", + "xcom": "XCom", +} + +type propSchema struct { + Type any `json:"type"` + Ref string `json:"$ref"` + Default any `json:"default"` +} + +type schemaDoc struct { + Definitions map[string]struct { + Properties map[string]propSchema `json:"properties"` + } `json:"definitions"` +} + +// field is one resolved spec field: schema key + default joined with the Go +// name and type the generated struct uses. +type field struct { + key string + goName string + goType string + def any // schema default (nil when absent) + hasDef bool // schema declares a default + refName string +} + +func main() { + schemaPath := flag.String("schema", "", "path to the dag-serialization schema.json") + outPath := flag.String("out", "spec.gen.go", "path to write the generated spec file") + pkg := flag.String("package", "bundlev1", "package name for the generated file") + flag.Parse() + + if *schemaPath == "" { + log.Fatal("gen: -schema is required") + } + raw, err := os.ReadFile(*schemaPath) + if err != nil { + log.Fatalf("gen: reading schema: %v", err) + } + var doc schemaDoc + if err := json.Unmarshal(raw, &doc); err != nil { + log.Fatalf("gen: parsing schema: %v", err) + } + operator, ok := doc.Definitions["operator"] + if !ok { + log.Fatal(`gen: schema has no "operator" definition`) + } + + fields := make([]field, 0, len(taskSpecFields)) + for _, fd := range taskSpecFields { + prop, ok := operator.Properties[fd.key] + if !ok { + log.Fatalf("gen: schema operator definition has no property %q", fd.key) + } + f, err := resolveField(fd, prop) + if err != nil { + log.Fatalf("gen: property %q: %v", fd.key, err) + } + fields = append(fields, f) + } + + src, err := render(*pkg, fields) + if err != nil { + log.Fatalf("gen: rendering: %v", err) + } + if err := os.WriteFile(*outPath, src, 0o644); err != nil { + log.Fatalf("gen: writing %s: %v", *outPath, err) + } +} + +func resolveField(fd fieldDef, prop propSchema) (field, error) { + f := field{ + key: fd.key, + goName: goName(fd.key), + def: prop.Default, + hasDef: prop.Default != nil, + } + if ref := strings.TrimPrefix(prop.Ref, "#/definitions/"); ref != prop.Ref { + f.refName = ref + } + + switch { + case fd.goType != "": + f.goType = fd.goType + case f.refName == "timedelta": + f.goType = "time.Duration" + case f.refName == "datetime": + f.goType = "time.Time" + default: + typ, _ := prop.Type.(string) + switch typ { + case "string": + f.goType = "string" + case "integer": + f.goType = "int" + case "number": + f.goType = "int" + if n, ok := prop.Default.(float64); ok && n != math.Trunc(n) { + f.goType = "float64" + } + case "boolean": + f.goType = "bool" + if d, ok := prop.Default.(bool); ok && d { + // A plain bool cannot distinguish "unset" from an explicit + // false when the schema default is true. + f.goType = "*bool" + } + default: + return field{}, fmt.Errorf("unsupported schema type %v (ref %q)", prop.Type, prop.Ref) + } + } + return f, nil +} + +func goName(key string) string { + parts := strings.Split(key, "_") + for i, p := range parts { + if init, ok := initialisms[p]; ok { + parts[i] = init + continue + } + parts[i] = strings.ToUpper(p[:1]) + p[1:] + } + return strings.Join(parts, "") +} + +const licenseHeader = `// 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. +` + +func render(pkg string, fields []field) ([]byte, error) { + var b bytes.Buffer + b.WriteString( + "// Code generated by bundlev1/gen from airflow-core/src/airflow/serialization/schema.json, DO NOT EDIT.\n", + ) + b.WriteString(licenseHeader) + fmt.Fprintf(&b, "\npackage %s\n\n", pkg) + b.WriteString("import \"time\"\n\n") + + b.WriteString(`// TaskSpec is the optional configuration applied to a task at registration +// time. Every field is optional: the zero value (nil for the *bool fields) +// means "unset" and the scheduler falls back to the schema default. Each +// field maps to the same-named key of the "operator" definition in +// airflow-core/src/airflow/serialization/schema.json. +type TaskSpec struct { +`) + for _, f := range fields { + fmt.Fprintf(&b, "\t// %s maps to the schema key %q", f.goName, f.key) + if f.hasDef { + fmt.Fprintf(&b, " (schema default %s)", defaultDoc(f)) + } + b.WriteString(".\n") + fmt.Fprintf(&b, "\t%s %s\n", f.goName, f.goType) + } + b.WriteString("}\n\n") + + b.WriteString(`// SchemaFields returns the schema-keyed value of every field that is set and +// differs from its schema default, mirroring Python BaseSerialization's +// omission of values the scheduler re-derives. time.Duration and time.Time +// values are returned as-is; the serializer owns the wire encoding. +func (s TaskSpec) SchemaFields() map[string]any { + m := map[string]any{} +`) + for _, f := range fields { + emitCondition(&b, f) + } + b.WriteString("\treturn m\n}\n") + + return format.Source(b.Bytes()) +} + +// defaultDoc renders a field's schema default for its doc comment. +func defaultDoc(f field) string { + if f.refName == "timedelta" { + return fmt.Sprintf("%v seconds", f.def) + } + if s, ok := f.def.(string); ok { + return fmt.Sprintf("%q", s) + } + return fmt.Sprintf("%v", f.def) +} + +// emitCondition writes the "set and not schema default" guard for one field. +func emitCondition(b *bytes.Buffer, f field) { + name := "s." + f.goName + key := f.key + switch { + case f.goType == "time.Time": + fmt.Fprintf(b, "\tif !%s.IsZero() {\n\t\tm[%q] = %s\n\t}\n", name, key, name) + case f.goType == "time.Duration": + cond := fmt.Sprintf("%s != 0", name) + if secs, ok := f.def.(float64); ok && secs != 0 { + cond += fmt.Sprintf(" && %s != %s", name, durationLit(secs)) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + case f.goType == "*bool": + notDefault := "*" + name + if def, _ := f.def.(bool); def { + notDefault = "!*" + name + } + fmt.Fprintf( + b, + "\tif %s != nil && %s {\n\t\tm[%q] = *%s\n\t}\n", + name, + notDefault, + key, + name, + ) + case f.goType == "bool": + notDefault := name + if def, _ := f.def.(bool); def { + notDefault = "!" + name + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", notDefault, key, name) + case f.goType == "string": + cond := fmt.Sprintf("%s != \"\"", name) + if def, ok := f.def.(string); ok && def != "" { + cond += fmt.Sprintf(" && %s != %q", name, def) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + default: // int / float64 + cond := fmt.Sprintf("%s != 0", name) + if def, ok := f.def.(float64); ok && def != 0 { + cond += fmt.Sprintf(" && %s != %s", name, numLit(def, f.goType)) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + } +} + +// durationLit renders a schema timedelta default (seconds) as a Go duration +// expression, e.g. 300 -> "300 * time.Second". +func durationLit(secs float64) string { + if secs == math.Trunc(secs) { + return fmt.Sprintf("%d * time.Second", int64(secs)) + } + return fmt.Sprintf("time.Duration(%v * float64(time.Second))", secs) +} + +func numLit(v float64, goType string) string { + if goType == "int" { + return fmt.Sprintf("%d", int64(v)) + } + return fmt.Sprintf("%v", v) +} diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index f6a933cb50096..4407bb01f1516 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,7 +23,6 @@ import ( "runtime" "strings" "sync" - "time" "github.com/apache/airflow/go-sdk/pkg/worker" ) @@ -81,98 +80,6 @@ type ( AddDag(dagId string, spec ...DagSpec) Dag } - // TaskSpec is the optional configuration applied to a task at registration - // time. Every field is optional: a zero value means "unset" and the - // scheduler falls back to its serialization-schema default. The field - // names mirror the keys defined under "operator" in - // airflow-core/src/airflow/serialization/schema.json. - TaskSpec struct { - Queue string - Pool string - PoolSlots int - Retries int - RetryDelay time.Duration - MaxRetryDelay time.Duration - RetryExponentialBackoff float64 - PriorityWeight int - WeightRule string - TriggerRule string - Owner string - ExecutionTimeout time.Duration - Executor string - StartDate time.Time - EndDate time.Time - DependsOnPast bool - WaitForDownstream bool - // DoXComPush, EmailOnFailure, and EmailOnRetry default to true in the - // scheduler. A nil pointer means "unset" so the field is omitted from - // the serialized payload; pass Bool(false) to explicitly opt out. - DoXComPush *bool - EmailOnFailure *bool - EmailOnRetry *bool - DocMD string - MapIndexTemplate string - MaxActiveTisPerDag int - MaxActiveTisPerDagrun int - } - - // DagSpec is the optional configuration applied to a DAG at registration - // time. Every field is optional: a zero value means "unset" and the - // scheduler falls back to its serialization-schema default. The field - // names mirror the keys defined under "dag" in - // airflow-core/src/airflow/serialization/schema.json. - DagSpec struct { - // Schedule is "@once", "@continuous", a cron expression, or "" for - // NullTimetable (no schedule). - Schedule string - Description string - StartDate time.Time - EndDate time.Time - Tags []string - DagDisplayName string - DocMD string - MaxActiveTasks int - MaxActiveRuns int - MaxConsecutiveFailedDagRuns int - DagrunTimeout time.Duration - Catchup bool - FailFast bool - RenderTemplateAsNativeObj bool - DisableBundleVersioning bool - // IsPausedUponCreation has no schema default. nil means "unset"; pass - // Bool(true) or Bool(false) to set it explicitly. - IsPausedUponCreation *bool - } - - // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses - // it to render the per-task block of a DagFileParsingResult. - TaskInfo struct { - // ID is the user-visible task id (the function name unless overridden - // via AddTaskWithName). - ID string - // TypeName is the unqualified Go function name (e.g. "extract"). - TypeName string - // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). - PkgPath string - // Spec carries the optional per-task configuration supplied at - // registration. The zero value means "no overrides". - Spec TaskSpec - // Downstream lists task ids that depend on this task, populated as - // later tasks declare this id in their AddTask `depends` argument. - // Order is registration order; the serializer sorts before emit. - Downstream []string - } - - // DagInfo describes a registered dag together with its tasks in - // registration order. - DagInfo struct { - DagID string - // Spec carries the optional per-dag configuration supplied at - // registration. The zero value means "no overrides". - Spec DagSpec - Tasks []TaskInfo - } - // EnumerableBundle exposes the dag/task identity recorded by // RegisterDags. The default registry implements it; the coordinator-mode // runtime relies on it for the DAG-parse one-shot. @@ -203,14 +110,6 @@ func (d dagShim) AddTaskWithName(taskId string, fn any, spec TaskSpec, depends [ d.registry.registerTaskWithName(d.dagId, taskId, fn, spec, depends) } -// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / -// DagSpec where nil means "leave at schema default": -// -// v1.TaskSpec{DoXComPush: v1.Bool(false)} -func Bool(b bool) *bool { - return &b -} - func optionalSpec[T any](specs []T, caller string) T { switch len(specs) { case 0: diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go new file mode 100644 index 0000000000000..a3557ae841570 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -0,0 +1,158 @@ +// Code generated by bundlev1/gen from airflow-core/src/airflow/serialization/schema.json, DO NOT EDIT. +// 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 bundlev1 + +import "time" + +// TaskSpec is the optional configuration applied to a task at registration +// time. Every field is optional: the zero value (nil for the *bool fields) +// means "unset" and the scheduler falls back to the schema default. Each +// field maps to the same-named key of the "operator" definition in +// airflow-core/src/airflow/serialization/schema.json. +type TaskSpec struct { + // Queue maps to the schema key "queue" (schema default "default"). + Queue string + // Pool maps to the schema key "pool" (schema default "default_pool"). + Pool string + // PoolSlots maps to the schema key "pool_slots" (schema default 1). + PoolSlots int + // Retries maps to the schema key "retries" (schema default 0). + Retries int + // RetryDelay maps to the schema key "retry_delay" (schema default 300 seconds). + RetryDelay time.Duration + // MaxRetryDelay maps to the schema key "max_retry_delay". + MaxRetryDelay time.Duration + // RetryExponentialBackoff maps to the schema key "retry_exponential_backoff" (schema default 0). + RetryExponentialBackoff float64 + // PriorityWeight maps to the schema key "priority_weight" (schema default 1). + PriorityWeight int + // WeightRule maps to the schema key "weight_rule" (schema default "downstream"). + WeightRule string + // TriggerRule maps to the schema key "trigger_rule" (schema default "all_success"). + TriggerRule string + // Owner maps to the schema key "owner" (schema default "airflow"). + Owner string + // ExecutionTimeout maps to the schema key "execution_timeout". + ExecutionTimeout time.Duration + // Executor maps to the schema key "executor". + Executor string + // StartDate maps to the schema key "start_date". + StartDate time.Time + // EndDate maps to the schema key "end_date". + EndDate time.Time + // DependsOnPast maps to the schema key "depends_on_past" (schema default false). + DependsOnPast bool + // WaitForDownstream maps to the schema key "wait_for_downstream" (schema default false). + WaitForDownstream bool + // DoXComPush maps to the schema key "do_xcom_push" (schema default true). + DoXComPush *bool + // EmailOnFailure maps to the schema key "email_on_failure" (schema default true). + EmailOnFailure *bool + // EmailOnRetry maps to the schema key "email_on_retry" (schema default true). + EmailOnRetry *bool + // DocMD maps to the schema key "doc_md". + DocMD string + // MapIndexTemplate maps to the schema key "map_index_template". + MapIndexTemplate string + // MaxActiveTisPerDag maps to the schema key "max_active_tis_per_dag". + MaxActiveTisPerDag int + // MaxActiveTisPerDagrun maps to the schema key "max_active_tis_per_dagrun". + MaxActiveTisPerDagrun int +} + +// SchemaFields returns the schema-keyed value of every field that is set and +// differs from its schema default, mirroring Python BaseSerialization's +// omission of values the scheduler re-derives. time.Duration and time.Time +// values are returned as-is; the serializer owns the wire encoding. +func (s TaskSpec) SchemaFields() map[string]any { + m := map[string]any{} + if s.Queue != "" && s.Queue != "default" { + m["queue"] = s.Queue + } + if s.Pool != "" && s.Pool != "default_pool" { + m["pool"] = s.Pool + } + if s.PoolSlots != 0 && s.PoolSlots != 1 { + m["pool_slots"] = s.PoolSlots + } + if s.Retries != 0 { + m["retries"] = s.Retries + } + if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { + m["retry_delay"] = s.RetryDelay + } + if s.MaxRetryDelay != 0 { + m["max_retry_delay"] = s.MaxRetryDelay + } + if s.RetryExponentialBackoff != 0 { + m["retry_exponential_backoff"] = s.RetryExponentialBackoff + } + if s.PriorityWeight != 0 && s.PriorityWeight != 1 { + m["priority_weight"] = s.PriorityWeight + } + if s.WeightRule != "" && s.WeightRule != "downstream" { + m["weight_rule"] = s.WeightRule + } + if s.TriggerRule != "" && s.TriggerRule != "all_success" { + m["trigger_rule"] = s.TriggerRule + } + if s.Owner != "" && s.Owner != "airflow" { + m["owner"] = s.Owner + } + if s.ExecutionTimeout != 0 { + m["execution_timeout"] = s.ExecutionTimeout + } + if s.Executor != "" { + m["executor"] = s.Executor + } + if !s.StartDate.IsZero() { + m["start_date"] = s.StartDate + } + if !s.EndDate.IsZero() { + m["end_date"] = s.EndDate + } + if s.DependsOnPast { + m["depends_on_past"] = s.DependsOnPast + } + if s.WaitForDownstream { + m["wait_for_downstream"] = s.WaitForDownstream + } + if s.DoXComPush != nil && !*s.DoXComPush { + m["do_xcom_push"] = *s.DoXComPush + } + if s.EmailOnFailure != nil && !*s.EmailOnFailure { + m["email_on_failure"] = *s.EmailOnFailure + } + if s.EmailOnRetry != nil && !*s.EmailOnRetry { + m["email_on_retry"] = *s.EmailOnRetry + } + if s.DocMD != "" { + m["doc_md"] = s.DocMD + } + if s.MapIndexTemplate != "" { + m["map_index_template"] = s.MapIndexTemplate + } + if s.MaxActiveTisPerDag != 0 { + m["max_active_tis_per_dag"] = s.MaxActiveTisPerDag + } + if s.MaxActiveTisPerDagrun != 0 { + m["max_active_tis_per_dagrun"] = s.MaxActiveTisPerDagrun + } + return m +} diff --git a/go-sdk/bundle/bundlev1/spec.go b/go-sdk/bundle/bundlev1/spec.go new file mode 100644 index 0000000000000..d9af9e1316f72 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.go @@ -0,0 +1,92 @@ +// 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 bundlev1 + +import "time" + +// This file holds the hand-written companions of the generated spec.gen.go: +// DagSpec (whose Schedule field and always-emit/config-fallback keys the +// schema cannot express mechanically) and the registration Info structs that +// carry Go-side identity the schema knows nothing about. + +type ( + // DagSpec is the optional configuration applied to a DAG at registration + // time. Every field is optional: a zero value means "unset" and the + // scheduler falls back to its serialization-schema default. The field + // names mirror the keys defined under "dag" in + // airflow-core/src/airflow/serialization/schema.json. + DagSpec struct { + // Schedule is "@once", "@continuous", a cron expression, or "" for + // NullTimetable (no schedule). + Schedule string + Description string + StartDate time.Time + EndDate time.Time + Tags []string + DagDisplayName string + DocMD string + MaxActiveTasks int + MaxActiveRuns int + MaxConsecutiveFailedDagRuns int + DagrunTimeout time.Duration + Catchup bool + FailFast bool + RenderTemplateAsNativeObj bool + DisableBundleVersioning bool + // IsPausedUponCreation has no schema default. nil means "unset"; pass + // Bool(true) or Bool(false) to set it explicitly. + IsPausedUponCreation *bool + } + + // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses + // it to render the per-task block of a DagFileParsingResult. + TaskInfo struct { + // ID is the user-visible task id (the function name unless overridden + // via AddTaskWithName). + ID string + // TypeName is the unqualified Go function name (e.g. "extract"). + TypeName string + // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). + PkgPath string + // Spec carries the optional per-task configuration supplied at + // registration. The zero value means "no overrides". + Spec TaskSpec + // Downstream lists task ids that depend on this task, populated as + // later tasks declare this id in their AddTask `depends` argument. + // Order is registration order; the serializer sorts before emit. + Downstream []string + } + + // DagInfo describes a registered dag together with its tasks in + // registration order. + DagInfo struct { + DagID string + // Spec carries the optional per-dag configuration supplied at + // registration. The zero value means "no overrides". + Spec DagSpec + Tasks []TaskInfo + } +) + +// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / +// DagSpec where nil means "leave at schema default": +// +// v1.TaskSpec{DoXComPush: v1.Bool(false)} +func Bool(b bool) *bool { + return &b +} diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go index 0aca23993d99e..b323688a68ab2 100644 --- a/go-sdk/pkg/execution/serde.go +++ b/go-sdk/pkg/execution/serde.go @@ -202,7 +202,15 @@ func serializeTask(info bundlev1.TaskInfo) map[string]any { // so it is unconditional here too. Go tasks have no template fields. "template_fields": []any{}, } - applyTaskSpec(data, info.Spec) + // TaskSpec.SchemaFields (generated from schema.json) returns only the + // fields that are set and differ from their schema default, so this + // mirrors Python BaseSerialization's "omit hard-coded default" behavior. + // serializeValue converts time.Duration / time.Time to their wire floats; + // unwrapTypeEncoding strips the __type wrapper because operator fields are + // stored unwrapped. + for key, value := range info.Spec.SchemaFields() { + data[key] = unwrapTypeEncoding(serializeValue(value)) + } if len(info.Downstream) > 0 { sorted := make([]string, len(info.Downstream)) copy(sorted, info.Downstream) @@ -215,88 +223,6 @@ func serializeTask(info bundlev1.TaskInfo) map[string]any { } } -// applyTaskSpec mirrors Python BaseSerialization's "omit hard-coded default" -// behavior: each TaskSpec field is written into data only when it differs -// from the schema default declared in -// airflow-core/src/airflow/serialization/schema.json. A zero-valued field is -// always considered "unset" and is skipped. -func applyTaskSpec(data map[string]any, s bundlev1.TaskSpec) { - if s.Queue != "" && s.Queue != "default" { - data["queue"] = s.Queue - } - if s.Pool != "" && s.Pool != "default_pool" { - data["pool"] = s.Pool - } - if s.PoolSlots != 0 && s.PoolSlots != 1 { - data["pool_slots"] = s.PoolSlots - } - if s.Retries != 0 { - data["retries"] = s.Retries - } - if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { - data["retry_delay"] = unwrapTypeEncoding(serializeValue(s.RetryDelay)) - } - if s.MaxRetryDelay != 0 { - data["max_retry_delay"] = unwrapTypeEncoding(serializeValue(s.MaxRetryDelay)) - } - if s.RetryExponentialBackoff != 0 { - data["retry_exponential_backoff"] = s.RetryExponentialBackoff - } - if s.PriorityWeight != 0 && s.PriorityWeight != 1 { - data["priority_weight"] = s.PriorityWeight - } - if s.WeightRule != "" && s.WeightRule != "downstream" { - data["weight_rule"] = s.WeightRule - } - if s.TriggerRule != "" && s.TriggerRule != "all_success" { - data["trigger_rule"] = s.TriggerRule - } - if s.Owner != "" && s.Owner != "airflow" { - data["owner"] = s.Owner - } - if s.ExecutionTimeout != 0 { - data["execution_timeout"] = unwrapTypeEncoding(serializeValue(s.ExecutionTimeout)) - } - if s.Executor != "" { - data["executor"] = s.Executor - } - if !s.StartDate.IsZero() { - data["start_date"] = unwrapTypeEncoding(serializeValue(s.StartDate)) - } - if !s.EndDate.IsZero() { - data["end_date"] = unwrapTypeEncoding(serializeValue(s.EndDate)) - } - if s.DependsOnPast { - data["depends_on_past"] = true - } - if s.WaitForDownstream { - data["wait_for_downstream"] = true - } - // do_xcom_push / email_on_failure / email_on_retry default to true; only - // emit when an explicit false overrides the default. - if s.DoXComPush != nil && !*s.DoXComPush { - data["do_xcom_push"] = false - } - if s.EmailOnFailure != nil && !*s.EmailOnFailure { - data["email_on_failure"] = false - } - if s.EmailOnRetry != nil && !*s.EmailOnRetry { - data["email_on_retry"] = false - } - if s.DocMD != "" { - data["doc_md"] = s.DocMD - } - if s.MapIndexTemplate != "" { - data["map_index_template"] = s.MapIndexTemplate - } - if s.MaxActiveTisPerDag != 0 { - data["max_active_tis_per_dag"] = s.MaxActiveTisPerDag - } - if s.MaxActiveTisPerDagrun != 0 { - data["max_active_tis_per_dagrun"] = s.MaxActiveTisPerDagrun - } -} - // applyDagSpec writes DAG-level fields onto data. Fields with a JSON-schema // default (description, dates, tags, fail_fast, …) are omitted when unset, as // in applyTaskSpec. Fields with no schema default (catchup, diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index e7aa19c70ae66..a883626b3cac5 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -172,7 +172,13 @@ func TestSerializeTaskDefaultQueueOmitted(t *testing.T) { assert.False(t, hasQueue, "queue=\"default\" matches the schema default and should be omitted") } -func TestApplyTaskSpec_EmitsAndOmits(t *testing.T) { +func applySpecFields(data map[string]any, s bundlev1.TaskSpec) { + for k, v := range s.SchemaFields() { + data[k] = unwrapTypeEncoding(serializeValue(v)) + } +} + +func TestTaskSpecSchemaFields_EmitsAndOmits(t *testing.T) { spec := bundlev1.TaskSpec{ Queue: "gpu", Pool: "gpu_pool", @@ -198,7 +204,7 @@ func TestApplyTaskSpec_EmitsAndOmits(t *testing.T) { MaxActiveTisPerDagrun: 1, } data := map[string]any{} - applyTaskSpec(data, spec) + applySpecFields(data, spec) assert.Equal(t, "gpu", data["queue"]) assert.Equal(t, "gpu_pool", data["pool"]) @@ -224,7 +230,7 @@ func TestApplyTaskSpec_EmitsAndOmits(t *testing.T) { assert.Equal(t, 1, data["max_active_tis_per_dagrun"]) } -func TestApplyTaskSpec_OmitsSchemaDefaults(t *testing.T) { +func TestTaskSpecSchemaFields_OmitsSchemaDefaults(t *testing.T) { // Values equal to schema defaults must be dropped. spec := bundlev1.TaskSpec{ Queue: "default", @@ -241,13 +247,13 @@ func TestApplyTaskSpec_OmitsSchemaDefaults(t *testing.T) { EmailOnRetry: bundlev1.Bool(true), } data := map[string]any{} - applyTaskSpec(data, spec) + applySpecFields(data, spec) assert.Empty(t, data, "all fields equal schema defaults; nothing should be emitted") } -func TestApplyTaskSpec_EmptySpecNoOp(t *testing.T) { +func TestTaskSpecSchemaFields_EmptySpecNoOp(t *testing.T) { data := map[string]any{} - applyTaskSpec(data, bundlev1.TaskSpec{}) + applySpecFields(data, bundlev1.TaskSpec{}) assert.Empty(t, data) } From 2676244ff2e62a1becfae5d47178e068ae99a51a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 04:45:02 +0000 Subject: [PATCH 7/9] Go-SDK: Derive TaskSpec field selection from the serialization schema The generator read field types and defaults from schema.json, but the field list itself was a hand-written allowlist, so a new author-settable operator key would go silently missing from TaskSpec. Selecting every eligible scalar key from the schema - with serializer-owned and Python-only keys excluded by documented rules - inverts the drift direction: new schema keys surface in the regenerated spec.gen.go diff, and stale exclusions fail generation. This also surfaces ignore_first_depends_on_past and wait_for_past_depends_before_skipping, which the hand list had missed. --- go-sdk/bundle/bundlev1/gen.go | 11 +- go-sdk/bundle/bundlev1/gen/main.go | 200 ++++++++++++++++++++--------- go-sdk/bundle/bundlev1/spec.gen.go | 98 +++++++------- go-sdk/pkg/execution/serde_test.go | 48 +++---- 4 files changed, 226 insertions(+), 131 deletions(-) diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go index ab3c3718de63c..59846166bd759 100644 --- a/go-sdk/bundle/bundlev1/gen.go +++ b/go-sdk/bundle/bundlev1/gen.go @@ -16,11 +16,12 @@ // under the License. // spec.gen.go is generated from the Airflow dag-serialization schema by the -// local gen tool: the TaskSpec struct and its SchemaFields omit-if-default -// rules come from the "operator" definition, so they cannot drift from what -// the scheduler deserializes. Don't edit spec.gen.go by hand; change the -// schema (or the field allowlist in ./gen) and re-run `just generate-specs` -// (go generate). DagSpec and the registration Info structs stay hand-written +// local gen tool: the TaskSpec field set, its Go types, and its SchemaFields +// omit-if-default rules all derive from the "operator" definition, so they +// cannot drift from what the scheduler deserializes. Don't edit spec.gen.go +// by hand; change the schema (or the exclusion rules in ./gen) and re-run +// `just generate-specs` (go generate). DagSpec and the registration Info +// structs stay hand-written // in spec.go: Schedule is an SDK-level concept the schema has no scalar for, // and several dag keys are always-emitted with [core]-config fallbacks the // schema cannot express. diff --git a/go-sdk/bundle/bundlev1/gen/main.go b/go-sdk/bundle/bundlev1/gen/main.go index 085110676dfae..c1714dc0c3348 100644 --- a/go-sdk/bundle/bundlev1/gen/main.go +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -19,12 +19,24 @@ // (airflow-core/src/airflow/serialization/schema.json). // // It reads the "operator" definition and emits the TaskSpec struct plus its -// SchemaFields method, so the Go field types and the omit-if-default rules the -// serializer relies on cannot drift from the schema. Only the keys listed in -// taskSpecFields are emitted: they are the task attributes a bundle author may -// set from Go, a deliberate subset of the schema (the rest of the "operator" -// keys are serializer internals such as task_type or template_fields, or -// Python-only concerns such as templating and callbacks). +// SchemaFields method, so the exposed fields, their Go types, and the +// omit-if-default rules the serializer relies on cannot drift from the schema. +// The field set is derived from the schema rather than hand-listed: every +// scalar property (string/integer/number/boolean, or a timedelta/datetime ref) +// becomes a TaskSpec field, in schema order, unless one of these rules skips +// it: +// +// - "_"-prefixed keys are serializer internals (_task_module, _is_mapped, ...) +// - keys in the definition's "required" list are always written by the +// serializer itself (task_type, task_id, ui_color, ...) +// - "has_on_*" keys are flags derived from Python callbacks, not settable +// - non-scalar keys (arrays, objects, refs other than timedelta/datetime, +// anyOf) cannot be expressed as a scalar spec field +// - excludedFields entries are Python-only concerns, documented per key +// +// A new scalar key added to the schema therefore shows up in the regenerated +// spec — visible in the spec.gen.go diff — or must gain an excludedFields +// entry, instead of going silently missing. // // Field defaults are read from the schema, never hard-coded here: a field is // serialized only when it is set (non-zero) and differs from its schema @@ -45,41 +57,29 @@ import ( "strings" ) -// fieldDef selects one schema property for the generated spec struct. goType -// overrides the mechanical schema→Go type mapping for the few keys whose -// schema type is too loose (e.g. "number" that is semantically a float). -type fieldDef struct { - key string - goType string +// excludedFields lists the scalar "operator" keys deliberately not exposed on +// TaskSpec, each with the reason; every other eligible scalar key becomes a +// field. An entry that no longer matches an eligible schema key fails +// generation, so the list cannot go stale. +var excludedFields = map[string]string{ + "doc": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_json": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_yaml": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_rst": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "allow_nested_operators": "Python runtime concern: warns when an operator executes inside another operator", + "multiple_outputs": "TaskFlow (@task) dict-unpacking concern, meaningless outside Python", + "start_from_trigger": "deferrable-operator machinery; its start_trigger_args counterpart is an object the spec cannot express", + "is_setup": "setup/teardown flags carry trigger-rule invariants the SDK does not model yet", + "is_teardown": "setup/teardown flags carry trigger-rule invariants the SDK does not model yet", + "on_failure_fail_dagrun": "only valid on teardown tasks, which the SDK does not model yet", } -// taskSpecFields lists the "operator" keys exposed on TaskSpec, in struct -// order. Adding a task attribute means adding its key here and regenerating. -var taskSpecFields = []fieldDef{ - {key: "queue"}, - {key: "pool"}, - {key: "pool_slots"}, - {key: "retries"}, - {key: "retry_delay"}, - {key: "max_retry_delay"}, - {key: "retry_exponential_backoff", goType: "float64"}, - {key: "priority_weight"}, - {key: "weight_rule"}, - {key: "trigger_rule"}, - {key: "owner"}, - {key: "execution_timeout"}, - {key: "executor"}, - {key: "start_date"}, - {key: "end_date"}, - {key: "depends_on_past"}, - {key: "wait_for_downstream"}, - {key: "do_xcom_push"}, - {key: "email_on_failure"}, - {key: "email_on_retry"}, - {key: "doc_md"}, - {key: "map_index_template"}, - {key: "max_active_tis_per_dag"}, - {key: "max_active_tis_per_dagrun"}, +// goTypeOverrides forces the Go type for keys whose schema type is too loose +// for the mechanical mapping. retry_exponential_backoff is "number" with an +// integral default, but Python declares it float (a backoff multiplier), so +// the mechanical mapping would pick int. +var goTypeOverrides = map[string]string{ + "retry_exponential_backoff": "float64", } // initialisms maps snake_case segments that must keep non-Title capitalization @@ -98,10 +98,42 @@ type propSchema struct { Default any `json:"default"` } +// orderedProps keeps schema property order, which encoding/json's map +// decoding discards; the generated struct follows the schema's field order. +type orderedProps struct { + keys []string + props map[string]propSchema +} + +func (o *orderedProps) UnmarshalJSON(data []byte) error { + if err := json.Unmarshal(data, &o.props); err != nil { + return err + } + dec := json.NewDecoder(bytes.NewReader(data)) + if _, err := dec.Token(); err != nil { + return err + } + for dec.More() { + tok, err := dec.Token() + if err != nil { + return err + } + o.keys = append(o.keys, tok.(string)) + var skipped json.RawMessage + if err := dec.Decode(&skipped); err != nil { + return err + } + } + return nil +} + +type definition struct { + Properties orderedProps `json:"properties"` + Required []string `json:"required"` +} + type schemaDoc struct { - Definitions map[string]struct { - Properties map[string]propSchema `json:"properties"` - } `json:"definitions"` + Definitions map[string]definition `json:"definitions"` } // field is one resolved spec field: schema key + default joined with the Go @@ -137,17 +169,9 @@ func main() { log.Fatal(`gen: schema has no "operator" definition`) } - fields := make([]field, 0, len(taskSpecFields)) - for _, fd := range taskSpecFields { - prop, ok := operator.Properties[fd.key] - if !ok { - log.Fatalf("gen: schema operator definition has no property %q", fd.key) - } - f, err := resolveField(fd, prop) - if err != nil { - log.Fatalf("gen: property %q: %v", fd.key, err) - } - fields = append(fields, f) + fields, err := selectFields(operator) + if err != nil { + log.Fatalf("gen: %v", err) } src, err := render(*pkg, fields) @@ -159,10 +183,66 @@ func main() { } } -func resolveField(fd fieldDef, prop propSchema) (field, error) { +// selectFields derives the TaskSpec fields from the operator definition, in +// schema order, applying the selection rules in the package comment. +func selectFields(op definition) ([]field, error) { + required := make(map[string]bool, len(op.Required)) + for _, key := range op.Required { + required[key] = true + } + + excludedSeen := make(map[string]bool, len(excludedFields)) + fields := make([]field, 0, len(op.Properties.keys)) + for _, key := range op.Properties.keys { + if strings.HasPrefix(key, "_") || required[key] || strings.HasPrefix(key, "has_on_") { + continue + } + if _, ok := excludedFields[key]; ok { + excludedSeen[key] = true + continue + } + f, ok := resolveField(key, op.Properties.props[key]) + if !ok { + continue + } + fields = append(fields, f) + } + + for key := range excludedFields { + if !excludedSeen[key] { + return nil, fmt.Errorf( + "excludedFields entry %q matches no eligible schema property; remove or fix it", + key, + ) + } + } + for key := range goTypeOverrides { + if !containsKey(fields, key) { + return nil, fmt.Errorf( + "goTypeOverrides entry %q matches no generated field; remove or fix it", + key, + ) + } + } + return fields, nil +} + +func containsKey(fields []field, key string) bool { + for _, f := range fields { + if f.key == key { + return true + } + } + return false +} + +// resolveField maps one schema property to a spec field; ok is false when the +// property is not a scalar the spec can express (array, object, anyOf, or a +// ref other than timedelta/datetime). +func resolveField(key string, prop propSchema) (field, bool) { f := field{ - key: fd.key, - goName: goName(fd.key), + key: key, + goName: goName(key), def: prop.Default, hasDef: prop.Default != nil, } @@ -171,8 +251,8 @@ func resolveField(fd fieldDef, prop propSchema) (field, error) { } switch { - case fd.goType != "": - f.goType = fd.goType + case goTypeOverrides[key] != "": + f.goType = goTypeOverrides[key] case f.refName == "timedelta": f.goType = "time.Duration" case f.refName == "datetime": @@ -197,10 +277,10 @@ func resolveField(fd fieldDef, prop propSchema) (field, error) { f.goType = "*bool" } default: - return field{}, fmt.Errorf("unsupported schema type %v (ref %q)", prop.Type, prop.Ref) + return field{}, false } } - return f, nil + return f, true } func goName(key string) string { diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go index a3557ae841570..9dcf4ee832584 100644 --- a/go-sdk/bundle/bundlev1/spec.gen.go +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -26,40 +26,44 @@ import "time" // field maps to the same-named key of the "operator" definition in // airflow-core/src/airflow/serialization/schema.json. type TaskSpec struct { + // Owner maps to the schema key "owner" (schema default "airflow"). + Owner string + // StartDate maps to the schema key "start_date". + StartDate time.Time + // EndDate maps to the schema key "end_date". + EndDate time.Time + // TriggerRule maps to the schema key "trigger_rule" (schema default "all_success"). + TriggerRule string + // DependsOnPast maps to the schema key "depends_on_past" (schema default false). + DependsOnPast bool + // IgnoreFirstDependsOnPast maps to the schema key "ignore_first_depends_on_past" (schema default false). + IgnoreFirstDependsOnPast bool + // WaitForPastDependsBeforeSkipping maps to the schema key "wait_for_past_depends_before_skipping" (schema default false). + WaitForPastDependsBeforeSkipping bool + // WaitForDownstream maps to the schema key "wait_for_downstream" (schema default false). + WaitForDownstream bool + // Retries maps to the schema key "retries" (schema default 0). + Retries int // Queue maps to the schema key "queue" (schema default "default"). Queue string // Pool maps to the schema key "pool" (schema default "default_pool"). Pool string // PoolSlots maps to the schema key "pool_slots" (schema default 1). PoolSlots int - // Retries maps to the schema key "retries" (schema default 0). - Retries int + // ExecutionTimeout maps to the schema key "execution_timeout". + ExecutionTimeout time.Duration // RetryDelay maps to the schema key "retry_delay" (schema default 300 seconds). RetryDelay time.Duration - // MaxRetryDelay maps to the schema key "max_retry_delay". - MaxRetryDelay time.Duration // RetryExponentialBackoff maps to the schema key "retry_exponential_backoff" (schema default 0). RetryExponentialBackoff float64 + // MaxRetryDelay maps to the schema key "max_retry_delay". + MaxRetryDelay time.Duration // PriorityWeight maps to the schema key "priority_weight" (schema default 1). PriorityWeight int // WeightRule maps to the schema key "weight_rule" (schema default "downstream"). WeightRule string - // TriggerRule maps to the schema key "trigger_rule" (schema default "all_success"). - TriggerRule string - // Owner maps to the schema key "owner" (schema default "airflow"). - Owner string - // ExecutionTimeout maps to the schema key "execution_timeout". - ExecutionTimeout time.Duration // Executor maps to the schema key "executor". Executor string - // StartDate maps to the schema key "start_date". - StartDate time.Time - // EndDate maps to the schema key "end_date". - EndDate time.Time - // DependsOnPast maps to the schema key "depends_on_past" (schema default false). - DependsOnPast bool - // WaitForDownstream maps to the schema key "wait_for_downstream" (schema default false). - WaitForDownstream bool // DoXComPush maps to the schema key "do_xcom_push" (schema default true). DoXComPush *bool // EmailOnFailure maps to the schema key "email_on_failure" (schema default true). @@ -82,6 +86,33 @@ type TaskSpec struct { // values are returned as-is; the serializer owns the wire encoding. func (s TaskSpec) SchemaFields() map[string]any { m := map[string]any{} + if s.Owner != "" && s.Owner != "airflow" { + m["owner"] = s.Owner + } + if !s.StartDate.IsZero() { + m["start_date"] = s.StartDate + } + if !s.EndDate.IsZero() { + m["end_date"] = s.EndDate + } + if s.TriggerRule != "" && s.TriggerRule != "all_success" { + m["trigger_rule"] = s.TriggerRule + } + if s.DependsOnPast { + m["depends_on_past"] = s.DependsOnPast + } + if s.IgnoreFirstDependsOnPast { + m["ignore_first_depends_on_past"] = s.IgnoreFirstDependsOnPast + } + if s.WaitForPastDependsBeforeSkipping { + m["wait_for_past_depends_before_skipping"] = s.WaitForPastDependsBeforeSkipping + } + if s.WaitForDownstream { + m["wait_for_downstream"] = s.WaitForDownstream + } + if s.Retries != 0 { + m["retries"] = s.Retries + } if s.Queue != "" && s.Queue != "default" { m["queue"] = s.Queue } @@ -91,48 +122,27 @@ func (s TaskSpec) SchemaFields() map[string]any { if s.PoolSlots != 0 && s.PoolSlots != 1 { m["pool_slots"] = s.PoolSlots } - if s.Retries != 0 { - m["retries"] = s.Retries + if s.ExecutionTimeout != 0 { + m["execution_timeout"] = s.ExecutionTimeout } if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { m["retry_delay"] = s.RetryDelay } - if s.MaxRetryDelay != 0 { - m["max_retry_delay"] = s.MaxRetryDelay - } if s.RetryExponentialBackoff != 0 { m["retry_exponential_backoff"] = s.RetryExponentialBackoff } + if s.MaxRetryDelay != 0 { + m["max_retry_delay"] = s.MaxRetryDelay + } if s.PriorityWeight != 0 && s.PriorityWeight != 1 { m["priority_weight"] = s.PriorityWeight } if s.WeightRule != "" && s.WeightRule != "downstream" { m["weight_rule"] = s.WeightRule } - if s.TriggerRule != "" && s.TriggerRule != "all_success" { - m["trigger_rule"] = s.TriggerRule - } - if s.Owner != "" && s.Owner != "airflow" { - m["owner"] = s.Owner - } - if s.ExecutionTimeout != 0 { - m["execution_timeout"] = s.ExecutionTimeout - } if s.Executor != "" { m["executor"] = s.Executor } - if !s.StartDate.IsZero() { - m["start_date"] = s.StartDate - } - if !s.EndDate.IsZero() { - m["end_date"] = s.EndDate - } - if s.DependsOnPast { - m["depends_on_past"] = s.DependsOnPast - } - if s.WaitForDownstream { - m["wait_for_downstream"] = s.WaitForDownstream - } if s.DoXComPush != nil && !*s.DoXComPush { m["do_xcom_push"] = *s.DoXComPush } diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index a883626b3cac5..83e7df6a67751 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -180,28 +180,30 @@ func applySpecFields(data map[string]any, s bundlev1.TaskSpec) { func TestTaskSpecSchemaFields_EmitsAndOmits(t *testing.T) { spec := bundlev1.TaskSpec{ - Queue: "gpu", - Pool: "gpu_pool", - PoolSlots: 4, - Retries: 3, - RetryDelay: 60 * time.Second, - MaxRetryDelay: 10 * time.Minute, - RetryExponentialBackoff: 2.0, - PriorityWeight: 5, - WeightRule: "upstream", - TriggerRule: "all_done", - Owner: "data-eng", - ExecutionTimeout: 45 * time.Second, - Executor: "KubernetesExecutor", - DependsOnPast: true, - WaitForDownstream: true, - DoXComPush: bundlev1.Bool(false), - EmailOnFailure: bundlev1.Bool(false), - EmailOnRetry: bundlev1.Bool(false), - DocMD: "## task", - MapIndexTemplate: "{{ task.task_id }}", - MaxActiveTisPerDag: 2, - MaxActiveTisPerDagrun: 1, + Queue: "gpu", + Pool: "gpu_pool", + PoolSlots: 4, + Retries: 3, + RetryDelay: 60 * time.Second, + MaxRetryDelay: 10 * time.Minute, + RetryExponentialBackoff: 2.0, + PriorityWeight: 5, + WeightRule: "upstream", + TriggerRule: "all_done", + Owner: "data-eng", + ExecutionTimeout: 45 * time.Second, + Executor: "KubernetesExecutor", + DependsOnPast: true, + IgnoreFirstDependsOnPast: true, + WaitForPastDependsBeforeSkipping: true, + WaitForDownstream: true, + DoXComPush: bundlev1.Bool(false), + EmailOnFailure: bundlev1.Bool(false), + EmailOnRetry: bundlev1.Bool(false), + DocMD: "## task", + MapIndexTemplate: "{{ task.task_id }}", + MaxActiveTisPerDag: 2, + MaxActiveTisPerDagrun: 1, } data := map[string]any{} applySpecFields(data, spec) @@ -220,6 +222,8 @@ func TestTaskSpecSchemaFields_EmitsAndOmits(t *testing.T) { assert.Equal(t, 45.0, data["execution_timeout"]) assert.Equal(t, "KubernetesExecutor", data["executor"]) assert.Equal(t, true, data["depends_on_past"]) + assert.Equal(t, true, data["ignore_first_depends_on_past"]) + assert.Equal(t, true, data["wait_for_past_depends_before_skipping"]) assert.Equal(t, true, data["wait_for_downstream"]) assert.Equal(t, false, data["do_xcom_push"]) assert.Equal(t, false, data["email_on_failure"]) From 1fee90e94e3b9e103b4cd7850a1f15ea3dcce27a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 07:11:53 +0000 Subject: [PATCH 8/9] Go-SDK: Wrap over-length example call to satisfy the golines hook --- go-sdk/example/bundle/main.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 59dde6b28ec0a..06535b0bf8be6 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -57,7 +57,12 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently, v1.TaskSpec{}, nil) + concurrentDag.AddTaskWithName( + "pull_xcoms_concurrently", + concurrentxcom.PullXComsConcurrently, + v1.TaskSpec{}, + nil, + ) return nil } From 0ea0905727a9031493fdeda7176a03da0a91e25a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 03:23:13 +0000 Subject: [PATCH 9/9] Go-SDK: TaskFlow-style dag authoring with XCom-injected task inputs Dag.Task(fn, opts...) replaces AddTask/AddTaskWithName and returns a *TaskRef handle. Passing handles via v1.Inputs(...) wires the dependency edge and feeds each upstream's return value into the matching data parameter of the task function at run time: the runtime pulls the upstream's return-value XCom and strictly decodes it into the declared parameter type, so a Go dag reads like plain function composition: d := reg.AddDag(v1.DagSpec{DagId: "etl", Schedule: "@daily"}) extracted := d.Task(extract) d.Task(transform, v1.Inputs(extracted)) v1.After(...) declares ordering-only edges. Task ids move into the specs: TaskSpec.TaskId (generated from schema.json as an identity field) and DagSpec.DagId replace positional ids and AddTaskWithName. Mismatched input counts or types, refs from another dag, and non-decodable parameters all panic at registration, i.e. dag-parse time. Binding is shared by the coordinator and Edge worker paths. --- go-sdk/README.md | 30 +- go-sdk/bundle/bundlev1/gen/main.go | 107 ++++--- go-sdk/bundle/bundlev1/registry.go | 338 ++++++++++++++++------- go-sdk/bundle/bundlev1/registry_test.go | 254 ++++++++++++----- go-sdk/bundle/bundlev1/schemas.go | 6 +- go-sdk/bundle/bundlev1/spec.gen.go | 20 +- go-sdk/bundle/bundlev1/spec.go | 43 +-- go-sdk/bundle/bundlev1/task.go | 210 +++++++++++++- go-sdk/bundle/bundlev1/task_test.go | 250 ++++++++++++++++- go-sdk/example/bundle/main.go | 83 ++++-- go-sdk/example/bundle/main_test.go | 9 +- go-sdk/pkg/execution/integration_test.go | 61 ++-- go-sdk/pkg/worker/runner_test.go | 24 +- 13 files changed, 1106 insertions(+), 329 deletions(-) diff --git a/go-sdk/README.md b/go-sdk/README.md index 7bd6b3d13c81a..0f1f538acb498 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -87,9 +87,9 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { } func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) + simpleDag := dagbag.AddDag(v1.DagSpec{DagId: "simple_dag"}) + extracted := simpleDag.Task(extract) + simpleDag.Task(transform, v1.Inputs(extracted)) return nil } @@ -100,24 +100,32 @@ func main() { } ``` -A task is an ordinary Go function. The runtime inspects its signature and injects arguments by type: -`sdk.TIRunContext`, `*slog.Logger`, and an `sdk.Client` (or a narrower interface such as -`sdk.VariableClient`). An optional `(any, error)` return becomes the task's XCom; an `error` return marks -the task failed. +A task is an ordinary Go function. The runtime inspects its signature and fills parameters by kind: +runtime values (`sdk.TIRunContext`, `*slog.Logger`, and an `sdk.Client` or a narrower interface such as +`sdk.VariableClient`) are injected by type, and every other parameter is a **data parameter**, bound +positionally to the `v1.Inputs(...)` refs — at run time it receives the upstream task's return value, +pulled from that task's return-value XCom and decoded into the parameter's type. A `(result, error)` +return becomes the task's XCom; an `error` return marks the task failed. ```go -func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { +type ExtractResult struct { + GoVersion string `json:"go_version"` +} + +func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (ExtractResult, error) { conn, err := client.GetConnection(ctx, "test_http") // ... do work, honour ctx cancellation ... - return map[string]any{"go_version": runtime.Version()}, nil + return ExtractResult{GoVersion: runtime.Version()}, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { +// `in` is a data parameter: it receives extract's return value because +// RegisterDags wired v1.Inputs(extracted). +func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger, in ExtractResult) error { val, err := client.GetVariable(ctx, "my_variable") if err != nil { return err } - log.Info("Obtained variable", "my_variable", val) + log.Info("Obtained variable", "my_variable", val, "upstream_go_version", in.GoVersion) return nil } ``` diff --git a/go-sdk/bundle/bundlev1/gen/main.go b/go-sdk/bundle/bundlev1/gen/main.go index c1714dc0c3348..28e16cc83593b 100644 --- a/go-sdk/bundle/bundlev1/gen/main.go +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -28,7 +28,8 @@ // // - "_"-prefixed keys are serializer internals (_task_module, _is_mapped, ...) // - keys in the definition's "required" list are always written by the -// serializer itself (task_type, task_id, ui_color, ...) +// serializer itself (task_type, ui_color, ...), except those re-exposed +// as identity fields by a fieldOverrides entry (task_id) // - "has_on_*" keys are flags derived from Python callbacks, not settable // - non-scalar keys (arrays, objects, refs other than timedelta/datetime, // anyOf) cannot be expressed as a scalar spec field @@ -74,12 +75,31 @@ var excludedFields = map[string]string{ "on_failure_fail_dagrun": "only valid on teardown tasks, which the SDK does not model yet", } -// goTypeOverrides forces the Go type for keys whose schema type is too loose -// for the mechanical mapping. retry_exponential_backoff is "number" with an -// integral default, but Python declares it float (a backoff multiplier), so -// the mechanical mapping would pick int. -var goTypeOverrides = map[string]string{ - "retry_exponential_backoff": "float64", +// fieldOverride carries the per-key tweaks the schema cannot express. goType +// forces the Go type when the schema type is too loose. goName overrides the +// mechanical snake_case→CamelCase name; doc overrides the generated field +// comment. identity marks a field that names the task rather than configuring +// it: it is exposed on the spec even though the schema lists it as required +// (serializer-owned), and it is kept out of SchemaFields because the +// serializer emits it from the registry's resolved TaskInfo.ID instead. +type fieldOverride struct { + goType string + goName string + doc string + identity bool +} + +// fieldOverrides is keyed by schema key; an entry that no longer matches a +// generated field fails generation. retry_exponential_backoff is "number" +// with an integral default, but Python declares it float (a backoff +// multiplier), so the mechanical mapping would pick int. +var fieldOverrides = map[string]fieldOverride{ + "retry_exponential_backoff": {goType: "float64"}, + "task_id": { + goName: "TaskId", + identity: true, + doc: "TaskId sets the task id explicitly; empty derives it from the task function's name", + }, } // initialisms maps snake_case segments that must keep non-Title capitalization @@ -139,12 +159,14 @@ type schemaDoc struct { // field is one resolved spec field: schema key + default joined with the Go // name and type the generated struct uses. type field struct { - key string - goName string - goType string - def any // schema default (nil when absent) - hasDef bool // schema declares a default - refName string + key string + goName string + goType string + doc string // doc-comment override (empty = generated wording) + identity bool // named-identity field, excluded from SchemaFields + def any // schema default (nil when absent) + hasDef bool // schema declares a default + refName string } func main() { @@ -194,7 +216,9 @@ func selectFields(op definition) ([]field, error) { excludedSeen := make(map[string]bool, len(excludedFields)) fields := make([]field, 0, len(op.Properties.keys)) for _, key := range op.Properties.keys { - if strings.HasPrefix(key, "_") || required[key] || strings.HasPrefix(key, "has_on_") { + serializerOwned := strings.HasPrefix(key, "_") || required[key] || + strings.HasPrefix(key, "has_on_") + if serializerOwned && !fieldOverrides[key].identity { continue } if _, ok := excludedFields[key]; ok { @@ -216,10 +240,10 @@ func selectFields(op definition) ([]field, error) { ) } } - for key := range goTypeOverrides { + for key := range fieldOverrides { if !containsKey(fields, key) { return nil, fmt.Errorf( - "goTypeOverrides entry %q matches no generated field; remove or fix it", + "fieldOverrides entry %q matches no generated field; remove or fix it", key, ) } @@ -240,19 +264,25 @@ func containsKey(fields []field, key string) bool { // property is not a scalar the spec can express (array, object, anyOf, or a // ref other than timedelta/datetime). func resolveField(key string, prop propSchema) (field, bool) { + ov := fieldOverrides[key] f := field{ - key: key, - goName: goName(key), - def: prop.Default, - hasDef: prop.Default != nil, + key: key, + goName: ov.goName, + doc: ov.doc, + identity: ov.identity, + def: prop.Default, + hasDef: prop.Default != nil, + } + if f.goName == "" { + f.goName = goName(key) } if ref := strings.TrimPrefix(prop.Ref, "#/definitions/"); ref != prop.Ref { f.refName = ref } switch { - case goTypeOverrides[key] != "": - f.goType = goTypeOverrides[key] + case ov.goType != "": + f.goType = ov.goType case f.refName == "timedelta": f.goType = "time.Duration" case f.refName == "datetime": @@ -322,31 +352,40 @@ func render(pkg string, fields []field) ([]byte, error) { fmt.Fprintf(&b, "\npackage %s\n\n", pkg) b.WriteString("import \"time\"\n\n") - b.WriteString(`// TaskSpec is the optional configuration applied to a task at registration -// time. Every field is optional: the zero value (nil for the *bool fields) -// means "unset" and the scheduler falls back to the schema default. Each -// field maps to the same-named key of the "operator" definition in + b.WriteString(`// TaskSpec describes a task at registration time. Every field is optional: +// the zero value (nil for the *bool fields) means "unset" and the scheduler +// falls back to the schema default. Each field maps to the same-named key of +// the "operator" definition in // airflow-core/src/airflow/serialization/schema.json. type TaskSpec struct { `) for _, f := range fields { - fmt.Fprintf(&b, "\t// %s maps to the schema key %q", f.goName, f.key) - if f.hasDef { - fmt.Fprintf(&b, " (schema default %s)", defaultDoc(f)) + if f.doc != "" { + fmt.Fprintf(&b, "\t// %s. Maps to the schema key %q.\n", f.doc, f.key) + } else { + fmt.Fprintf(&b, "\t// %s maps to the schema key %q", f.goName, f.key) + if f.hasDef { + fmt.Fprintf(&b, " (schema default %s)", defaultDoc(f)) + } + b.WriteString(".\n") } - b.WriteString(".\n") fmt.Fprintf(&b, "\t%s %s\n", f.goName, f.goType) } b.WriteString("}\n\n") - b.WriteString(`// SchemaFields returns the schema-keyed value of every field that is set and -// differs from its schema default, mirroring Python BaseSerialization's -// omission of values the scheduler re-derives. time.Duration and time.Time -// values are returned as-is; the serializer owns the wire encoding. + b.WriteString(`// SchemaFields returns the schema-keyed value of every configuration field +// that is set and differs from its schema default, mirroring Python +// BaseSerialization's omission of values the scheduler re-derives. Identity +// fields (TaskId) are excluded: the serializer emits the id from the +// registry's resolved TaskInfo. time.Duration and time.Time values are +// returned as-is; the serializer owns the wire encoding. func (s TaskSpec) SchemaFields() map[string]any { m := map[string]any{} `) for _, f := range fields { + if f.identity { + continue + } emitCondition(&b, f) } b.WriteString("\treturn m\n}\n") diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 4407bb01f1516..b87c252ffc09a 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -20,6 +20,7 @@ package bundlev1 import ( "fmt" "reflect" + "regexp" "runtime" "strings" "sync" @@ -29,7 +30,7 @@ import ( type ( // Task is one registered task: something the runtime can Execute. Bundle - // authors do not implement this directly; Dag.AddTask wraps a plain Go + // authors do not implement this directly; Dag.Task wraps a plain Go // function into a Task for you. Task = worker.Task @@ -41,31 +42,28 @@ type ( // Dag is the handle returned by Registry.AddDag. Use it to attach the Go // functions that implement the dag's tasks. Dag interface { - // AddTask registers fn as a task in this Dag using fn's Go name as - // the task id (so it must match the @task.stub name in the Python - // dag). spec carries optional per-task configuration (pass TaskSpec{} - // for defaults). depends lists task ids in the same Dag that must run - // before this one; each must already be registered. Pass nil for no - // dependencies. + // Task registers fn as a task and returns its TaskRef, the handle + // downstream registrations pass to Inputs or After. The task id is + // TaskSpec.TaskId when a spec sets it, otherwise fn's own Go name. // - // fn is an ordinary Go function whose parameters are injected by type - // and may appear in any order. Recognised parameters are: - // - context.Context: cancelled when the task is asked to stop - // - *slog.Logger: writes to the task's Airflow log - // - sdk.Client (or a narrower sdk.VariableClient / sdk.ConnectionClient / - // sdk.XComClient): access to Variables, Connections, and XCom + // fn is an ordinary Go function whose parameters are filled by kind: // - // fn must return either error or (result, error): a non-nil error fails - // the task, and a non-nil first result is pushed as the task's - // return-value XCom. Passing a non-function, or a function whose return - // signature does not match, panics at registration time. - AddTask(fn any, spec TaskSpec, depends []string) - - // AddTaskWithName is like AddTask but sets task_id explicitly instead - // of deriving it from the function name. Use it when the Go function - // name cannot match the Python @task.stub id, for example for an - // anonymous function or a differing name. - AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) + // - Injected runtime values, matched by type in any order: + // context.Context (or sdk.TIRunContext), *slog.Logger, and + // sdk.Client (or a narrower sdk.VariableClient / + // sdk.ConnectionClient / sdk.XComClient). + // - Data parameters: every other parameter, bound positionally to + // the Inputs refs. At run time each receives the return value of + // its upstream task, pulled from that task's return-value XCom + // and decoded into the parameter's type. + // + // fn must return either error or (result, error): a non-nil error + // fails the task, and a non-nil first result is pushed as the task's + // return-value XCom (feeding any downstream Inputs). Passing a + // non-function, a function whose signature does not match, a spec / + // Inputs mismatch, or a ref from another dag panics at registration + // (i.e. dag-parse) time. + Task(fn any, opts ...TaskOption) *TaskRef } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -73,11 +71,10 @@ type ( // object serves task lookups at execution time. Registry interface { Bundle - // AddDag registers a dag by its dag_id (matching the Python stub dag) - // and returns a Dag handle for attaching tasks. An optional DagSpec - // configures dag-level attributes such as schedule and tags; passing - // more than one spec panics. Registering the same dag_id twice panics. - AddDag(dagId string, spec ...DagSpec) Dag + // AddDag registers a dag described by spec (spec.DagId is required) + // and returns a Dag handle for attaching tasks. A missing DagId or a + // duplicate registration panics. + AddDag(spec DagSpec) Dag } // EnumerableBundle exposes the dag/task identity recorded by @@ -87,6 +84,12 @@ type ( OrderedDags() []DagInfo } + // TaskOption configures one Dag.Task registration. A TaskSpec value, + // Inputs, and After all implement it. + TaskOption interface { + applyTask(*taskConfig) + } + registry struct { sync.RWMutex taskFuncMap map[string]map[string]Task @@ -97,30 +100,46 @@ type ( } ) -type dagShim struct { - dagId string - registry *registry +// TaskRef is the handle Dag.Task returns for a registered task. Pass it to +// Inputs (data dependency: the upstream's return value feeds a data +// parameter) or After (ordering-only dependency) when registering downstream +// tasks in the same dag. +type TaskRef struct { + id string + out reflect.Type // fn's first return type; nil when fn returns only error + reg *registry + dagID string } -func (d dagShim) AddTask(fn any, spec TaskSpec, depends []string) { - d.registry.registerTask(d.dagId, fn, spec, depends) -} +// ID returns the registered task id. +func (r *TaskRef) ID() string { return r.id } -func (d dagShim) AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) { - d.registry.registerTaskWithName(d.dagId, taskId, fn, spec, depends) +// taskConfig accumulates the options of one Dag.Task call. +type taskConfig struct { + spec TaskSpec + specSet bool + inputs []*TaskRef + after []*TaskRef } -func optionalSpec[T any](specs []T, caller string) T { - switch len(specs) { - case 0: - var zero T - return zero - case 1: - return specs[0] - default: - panic(fmt.Errorf("%s accepts at most one spec, got %d", caller, len(specs))) - } -} +type inputsOption []*TaskRef + +func (o inputsOption) applyTask(c *taskConfig) { c.inputs = append(c.inputs, o...) } + +// Inputs declares data dependencies: each ref's task must run first, and its +// return value is pulled from XCom and passed to the matching data parameter +// of the task function, in order. The number of refs must equal the number of +// data parameters, and each upstream's return type must fit the parameter it +// feeds (checked at registration). +func Inputs(refs ...*TaskRef) TaskOption { return inputsOption(refs) } + +type afterOption []*TaskRef + +func (o afterOption) applyTask(c *taskConfig) { c.after = append(c.after, o...) } + +// After declares ordering-only dependencies: each ref's task must complete +// before this one runs, without feeding it any value. +func After(refs ...*TaskRef) TaskOption { return afterOption(refs) } // New returns an empty Registry on which dags and tasks can be registered. The // runtime creates one and hands it to BundleProvider.RegisterDags, so bundle @@ -135,6 +154,28 @@ func New() Registry { } } +type dagShim struct { + dagId string + registry *registry +} + +func (d dagShim) Task(fn any, opts ...TaskOption) *TaskRef { + return d.registry.registerTask(d.dagId, fn, opts) +} + +// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / +// DagSpec where nil means "leave at schema default": +// +// v1.TaskSpec{DoXComPush: v1.Bool(false)} +func Bool(b bool) *bool { + return &b +} + +// anonFnName matches the names the Go runtime gives anonymous functions once +// splitFullName has taken the segment after the last dot: "funcN" for a +// top-level closure, or a bare number for a nested one ("pkg.fn.func1.2"). +var anonFnName = regexp.MustCompile(`^(func)?\d+$`) + func splitFullName(fullName string) (typeName, pkgPath string) { // fullName looks like "main.extract" or "github.com/x/y.MyTask"; method // values get a "-fm" suffix. @@ -145,96 +186,185 @@ func splitFullName(fullName string) (typeName, pkgPath string) { return strings.TrimSuffix(fullName[lastDot+1:], "-fm"), fullName[:lastDot] } -func getFnName(fn reflect.Value) string { - fullName := runtime.FuncForPC(fn.Pointer()).Name() - name, _ := splitFullName(fullName) - return name -} - -func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { - dagSpec := optionalSpec(spec, "AddDag") +func (r *registry) AddDag(spec DagSpec) Dag { + if spec.DagId == "" { + panic(fmt.Errorf("AddDag requires DagSpec.DagId to be set")) + } r.RWMutex.Lock() defer r.RWMutex.Unlock() - if _, exists := r.taskFuncMap[dagId]; exists { - panic(fmt.Errorf("Dag %q already exists in bundle", dagId)) + if _, exists := r.taskFuncMap[spec.DagId]; exists { + panic(fmt.Errorf("Dag %q already exists in bundle", spec.DagId)) } - r.taskFuncMap[dagId] = make(map[string]Task) - r.taskInfo[dagId] = make(map[string]TaskInfo) - r.dagSpec[dagId] = dagSpec - r.dagOrder = append(r.dagOrder, dagId) - return dagShim{dagId, r} + r.taskFuncMap[spec.DagId] = make(map[string]Task) + r.taskInfo[spec.DagId] = make(map[string]TaskInfo) + r.dagSpec[spec.DagId] = spec + r.dagOrder = append(r.dagOrder, spec.DagId) + return dagShim{spec.DagId, r} } -func (r *registry) registerTask(dagId string, fn any, spec TaskSpec, depends []string) { - val := reflect.ValueOf(fn) +func (r *registry) registerTask(dagId string, fn any, opts []TaskOption) *TaskRef { + var cfg taskConfig + for _, o := range opts { + o.applyTask(&cfg) + } + val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { panic(fmt.Errorf("task fn was a %s, not a func", val.Kind())) } + fullName := runtime.FuncForPC(val.Pointer()).Name() + typeName, pkgPath := splitFullName(fullName) - fnName := getFnName(val) + taskId := cfg.spec.TaskId + if taskId == "" { + if anonFnName.MatchString(typeName) { + panic(fmt.Errorf( + "task function in DAG %q is anonymous (its derived name would be %q); set TaskSpec.TaskId explicitly", + dagId, + typeName, + )) + } + taskId = typeName + } - r.registerTaskWithName(dagId, fnName, fn, spec, depends) -} + // Pair the fn's data parameters with the Inputs refs before wrapping, so + // every mismatch is a registration (dag-parse) error, not a runtime one. + dataParams := dataParamTypes(val.Type()) + if len(cfg.inputs) != len(dataParams) { + panic(fmt.Errorf( + "task %q in DAG %q declares %d data parameter(s) but Inputs supplies %d", + taskId, dagId, len(dataParams), len(cfg.inputs), + )) + } + bindings := make([]InputBinding, len(cfg.inputs)) + for i, ref := range cfg.inputs { + r.validateRef(ref, dagId, taskId, "Inputs") + if ref.out == nil { + panic(fmt.Errorf( + "task %q in DAG %q: input task %q returns no value; use After for an ordering-only dependency", + taskId, + dagId, + ref.id, + )) + } + if !inputCompatible(ref.out, dataParams[i]) { + panic(fmt.Errorf( + "task %q in DAG %q: input %d: task %q returns %s, which cannot fill parameter type %s", + taskId, + dagId, + i, + ref.id, + ref.out, + dataParams[i], + )) + } + bindings[i] = InputBinding{TaskID: ref.id} + } + for _, ref := range cfg.after { + r.validateRef(ref, dagId, taskId, "After") + } -func (r *registry) registerTaskWithName( - dagId, taskId string, - fn any, - spec TaskSpec, - depends []string, -) { - task, err := NewTaskFunction(fn) + task, err := NewTaskFunction(fn, bindings...) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) } - val := reflect.ValueOf(fn) - fullName := runtime.FuncForPC(val.Pointer()).Name() - typeName, pkgPath := splitFullName(fullName) + var outType reflect.Type + if fnType := val.Type(); fnType.NumOut() > 1 { + outType = fnType.Out(0) + } + + var inputIDs []string + if len(cfg.inputs) > 0 { + inputIDs = make([]string, len(cfg.inputs)) + for i, ref := range cfg.inputs { + inputIDs[i] = ref.id + } + } - info := TaskInfo{ID: taskId, TypeName: typeName, PkgPath: pkgPath, Spec: spec} + info := TaskInfo{ + ID: taskId, + TypeName: typeName, + PkgPath: pkgPath, + Spec: cfg.spec, + Inputs: inputIDs, + } r.RWMutex.Lock() defer r.RWMutex.Unlock() dagTasks, exists := r.taskFuncMap[dagId] if !exists { - dagTasks = make(map[string]Task) - r.taskFuncMap[dagId] = dagTasks - r.taskInfo[dagId] = make(map[string]TaskInfo) - r.dagOrder = append(r.dagOrder, dagId) + panic(fmt.Errorf("DAG %q is not registered", dagId)) } - if _, exists := dagTasks[taskId]; exists { panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId)) } - // Resolve depends to upstream TaskInfo entries, validating each exists. - // We dedupe so a repeated id in `depends` only records one downstream - // edge on the parent. - seen := make(map[string]bool, len(depends)) - for _, dep := range depends { - if dep == taskId { - panic(fmt.Errorf("task %q cannot depend on itself in DAG %q", taskId, dagId)) - } - if seen[dep] { + // Record one downstream edge on each distinct upstream, whether the + // dependency carries data (Inputs) or only ordering (After). + seen := make(map[string]bool, len(cfg.inputs)+len(cfg.after)) + for _, ref := range append(append([]*TaskRef{}, cfg.inputs...), cfg.after...) { + if seen[ref.id] { continue } - seen[dep] = true - parent, ok := r.taskInfo[dagId][dep] - if !ok { - panic(fmt.Errorf( - "task %q depends on unknown task %q in DAG %q; register upstream tasks first", - taskId, dep, dagId, - )) - } + seen[ref.id] = true + parent := r.taskInfo[dagId][ref.id] parent.Downstream = append(parent.Downstream, taskId) - r.taskInfo[dagId][dep] = parent + r.taskInfo[dagId][ref.id] = parent } dagTasks[taskId] = task r.taskInfo[dagId][taskId] = info r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) + + return &TaskRef{id: taskId, out: outType, reg: r, dagID: dagId} +} + +// validateRef panics unless ref was returned by this registry for the same +// dag. A nil ref, a ref from another registry, or a ref from a different dag +// is a registration error. +func (r *registry) validateRef(ref *TaskRef, dagId, taskId, opt string) { + switch { + case ref == nil: + panic(fmt.Errorf("task %q in DAG %q: %s got a nil TaskRef", taskId, dagId, opt)) + case ref.reg != r: + panic(fmt.Errorf( + "task %q in DAG %q: %s ref %q belongs to a different registry", + taskId, dagId, opt, ref.id, + )) + case ref.dagID != dagId: + panic(fmt.Errorf( + "task %q in DAG %q: %s ref %q belongs to DAG %q; dependencies cannot cross dags", + taskId, dagId, opt, ref.id, ref.dagID, + )) + } +} + +// inputCompatible reports whether an upstream return type can fill a data +// parameter. Identical or assignable types always fit; a map or empty +// interface parameter opts into loose decoding of any shape; an upstream that +// returns any is accepted and checked structurally when the value is decoded +// at run time. Pointers are stripped from both sides: the value crosses XCom +// as its JSON shape, so *T and T are interchangeable in either position. +func inputCompatible(out, param reflect.Type) bool { + if out.Kind() == reflect.Pointer { + out = out.Elem() + } + if param.Kind() == reflect.Pointer { + param = param.Elem() + } + switch { + case out.AssignableTo(param): + return true + case param.Kind() == reflect.Interface && param.NumMethod() == 0: + return true + case param.Kind() == reflect.Map: + return true + case out.Kind() == reflect.Interface && out.NumMethod() == 0: + return true + } + return false } func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { @@ -250,8 +380,8 @@ func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { } // OrderedDags returns the registered dags in the order AddDag was called, -// each with its tasks in the order AddTask / AddTaskWithName was called. The -// returned slice is freshly allocated; callers may mutate it freely. +// each with its tasks in the order Dag.Task was called. The returned slice is +// freshly allocated; callers may mutate it freely. func (r *registry) OrderedDags() []DagInfo { r.RLock() defer r.RUnlock() diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index 364b7b8a8dc37..dc310e7d5b4b8 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -41,6 +41,18 @@ func NotErrorRet() int { return 0 } +// producer / consumer fixtures for Inputs wiring. +type payload struct { + Value int `json:"value"` +} + +func produce() (payload, error) { return payload{Value: 1}, nil } +func consume(in payload) error { _ = in; return nil } +func relay(in payload) (payload, error) { return in, nil } +func consumeTwo(a payload, b payload) error { _, _ = a, b; return nil } +func consumeLoose(in map[string]any) error { _ = in; return nil } +func consumeString(in string) error { _ = in; return nil } + type RegistrySuite struct { suite.Suite reg Registry @@ -53,72 +65,85 @@ func TestRegistrySuite(t *testing.T) { func (s *RegistrySuite) SetupTest() { s.reg = New() - s.dag = s.reg.AddDag("dag1") + s.dag = s.reg.AddDag(DagSpec{DagId: "dag1"}) } func (s *RegistrySuite) TestAddDag_NewDagRegisters() { s.NotNil(s.dag) } +func (s *RegistrySuite) TestAddDag_MissingIdPanics() { + s.PanicsWithError("AddDag requires DagSpec.DagId to be set", func() { + s.reg.AddDag(DagSpec{}) + }) +} + func (s *RegistrySuite) TestAddDag_DuplicatePanics() { s.PanicsWithError(`Dag "dag1" already exists in bundle`, func() { - s.reg.AddDag("dag1") + s.reg.AddDag(DagSpec{DagId: "dag1"}) }) } -func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask, TaskSpec{}, nil) +func (s *RegistrySuite) TestTask_RegistersAndFindsTask() { + ref := s.dag.Task(myTask) + s.Equal("myTask", ref.ID()) task, exists := s.reg.LookupTask("dag1", "myTask") s.True(exists) s.NotNil(task) } -func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) - task, exists := s.reg.LookupTask("dag1", "special") +func (s *RegistrySuite) TestTask_ExplicitIdViaSpec() { + ref := s.dag.Task(myTask, TaskSpec{TaskId: "special"}) + s.Equal("special", ref.ID()) + _, exists := s.reg.LookupTask("dag1", "special") s.True(exists) - s.NotNil(task) // Lets just make sure it didn't exist under the fn name _, exists = s.reg.LookupTask("dag1", "myTask") s.False(exists) } -func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) +func (s *RegistrySuite) TestTask_DuplicateIdPanics() { + s.dag.Task(myTask, TaskSpec{TaskId: "special"}) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) + s.dag.Task(myTask, TaskSpec{TaskId: "special"}) }) } -func (s *RegistrySuite) TestAddTask_NonFuncPanics() { +func (s *RegistrySuite) TestTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func", TaskSpec{}, nil) + s.dag.Task("not a func") }) } -func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs, TaskSpec{}, nil) +func (s *RegistrySuite) TestTask_InjectedArgsBind() { + s.dag.Task(myTaskWithArgs) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) } -func (s *RegistrySuite) TestAddTask_InvalidReturnType() { +func (s *RegistrySuite) TestTask_InvalidReturnType() { s.PanicsWithError( "error registering task \"NotErrorRet\" for DAG \"dag1\": expected task function github.com/apache/airflow/go-sdk/bundle/bundlev1.NotErrorRet last return value to return error but found int", func() { - s.dag.AddTask(NotErrorRet, TaskSpec{}, nil) + s.dag.Task(NotErrorRet) }, ) } -func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask, TaskSpec{}, nil) +func (s *RegistrySuite) TestTask_ErrorReturnType() { + s.dag.Task(errorTask) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } +func (s *RegistrySuite) TestTask_TooManySpecsPanics() { + s.PanicsWithError("Task accepts at most one TaskSpec", func() { + s.dag.Task(myTask, TaskSpec{}, TaskSpec{}) + }) +} + func (s *RegistrySuite) TestOrderedDags_Empty() { enum, ok := New().(EnumerableBundle) s.Require().True(ok) @@ -129,14 +154,14 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { reg := New() // Register dags out of alphabetical order so the test fails if OrderedDags // ever sorts instead of preserving AddDag order. - zeta := reg.AddDag("zeta") - zeta.AddTaskWithName("z1", myTask, TaskSpec{}, nil) - zeta.AddTaskWithName("z2", myTask, TaskSpec{}, nil) + zeta := reg.AddDag(DagSpec{DagId: "zeta"}) + zeta.Task(myTask, TaskSpec{TaskId: "z1"}) + zeta.Task(myTask, TaskSpec{TaskId: "z2"}) - alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask, TaskSpec{}, nil) + alpha := reg.AddDag(DagSpec{DagId: "alpha"}) + alpha.Task(myTask, TaskSpec{TaskId: "a1"}) - reg.AddDag("mid") // dag with no tasks + reg.AddDag(DagSpec{DagId: "mid"}) // dag with no tasks enum, ok := reg.(EnumerableBundle) s.Require().True(ok) @@ -162,8 +187,8 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { s.Empty(got[2].Tasks) } -func (s *RegistrySuite) TestAddTask_WithSpec() { - s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}, nil) +func (s *RegistrySuite) TestTask_WithSpec() { + s.dag.Task(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -177,84 +202,134 @@ func (s *RegistrySuite) TestAddTask_WithSpec() { s.False(*got.Spec.DoXComPush) } -func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}, nil) - enum, ok := s.reg.(EnumerableBundle) - s.Require().True(ok) - dags := enum.OrderedDags() - s.Require().Len(dags, 1) - s.Require().Len(dags[0].Tasks, 1) - got := dags[0].Tasks[0] - s.Equal("special", got.ID) - s.Equal("gpu", got.Spec.Queue) - s.Equal("gpu_pool", got.Spec.Pool) -} - -func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { - s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) - s.dag.AddTaskWithName("transform", myTask, TaskSpec{}, []string{"extract"}) - s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform"}) - +func (s *RegistrySuite) infoByID() map[string]TaskInfo { enum := s.reg.(EnumerableBundle) tasks := enum.OrderedDags()[0].Tasks byID := make(map[string]TaskInfo, len(tasks)) for _, t := range tasks { byID[t.ID] = t } + return byID +} + +func (s *RegistrySuite) TestTask_InputsRecordEdgesAndBindings() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + transform := s.dag.Task(relay, TaskSpec{TaskId: "transform"}, Inputs(extract)) + s.dag.Task(consume, TaskSpec{TaskId: "load"}, Inputs(transform)) + + byID := s.infoByID() s.Equal([]string{"transform"}, byID["extract"].Downstream) s.Equal([]string{"load"}, byID["transform"].Downstream) s.Nil(byID["load"].Downstream) + + s.Nil(byID["extract"].Inputs) + s.Equal([]string{"extract"}, byID["transform"].Inputs) + s.Equal([]string{"transform"}, byID["load"].Inputs) } -func (s *RegistrySuite) TestAddTask_FanOutFanIn() { - s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) - s.dag.AddTaskWithName("transform_a", myTask, TaskSpec{}, []string{"extract"}) - s.dag.AddTaskWithName("transform_b", myTask, TaskSpec{}, []string{"extract"}) - s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform_a", "transform_b"}) +func (s *RegistrySuite) TestTask_FanOutFanIn() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + a := s.dag.Task(relay, TaskSpec{TaskId: "transform_a"}, Inputs(extract)) + b := s.dag.Task(relay, TaskSpec{TaskId: "transform_b"}, Inputs(extract)) + s.dag.Task(consumeTwo, TaskSpec{TaskId: "load"}, Inputs(a, b)) - enum := s.reg.(EnumerableBundle) - tasks := enum.OrderedDags()[0].Tasks - byID := make(map[string]TaskInfo, len(tasks)) - for _, t := range tasks { - byID[t.ID] = t - } + byID := s.infoByID() s.ElementsMatch([]string{"transform_a", "transform_b"}, byID["extract"].Downstream) s.Equal([]string{"load"}, byID["transform_a"].Downstream) s.Equal([]string{"load"}, byID["transform_b"].Downstream) + s.Equal([]string{"transform_a", "transform_b"}, byID["load"].Inputs) } -func (s *RegistrySuite) TestAddTask_DependsDuplicatesIgnored() { - s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) - s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract", "extract"}) +func (s *RegistrySuite) TestTask_AfterRecordsOrderingEdge() { + first := s.dag.Task(myTask, TaskSpec{TaskId: "first"}) + s.dag.Task(myTask, TaskSpec{TaskId: "second"}, After(first)) - enum := s.reg.(EnumerableBundle) - tasks := enum.OrderedDags()[0].Tasks - byID := make(map[string]TaskInfo, len(tasks)) - for _, t := range tasks { - byID[t.ID] = t - } + byID := s.infoByID() + s.Equal([]string{"second"}, byID["first"].Downstream) + s.Nil(byID["second"].Inputs) +} + +func (s *RegistrySuite) TestTask_InputsAndAfterDeduplicateEdges() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + // The same upstream named as both a data input and an After ref records + // only one downstream edge. + s.dag.Task(consume, TaskSpec{TaskId: "load"}, Inputs(extract), After(extract)) + + byID := s.infoByID() s.Equal([]string{"load"}, byID["extract"].Downstream) } -func (s *RegistrySuite) TestAddTask_DependsUnknownPanics() { +func (s *RegistrySuite) TestTask_InputsCountMismatchPanics() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) s.PanicsWithError( - `task "load" depends on unknown task "extract" in DAG "dag1"; register upstream tasks first`, + `task "consumeTwo" in DAG "dag1" declares 2 data parameter(s) but Inputs supplies 1`, func() { - s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract"}) + s.dag.Task(consumeTwo, Inputs(extract)) }, ) } -func (s *RegistrySuite) TestAddTask_DependsOnSelfPanics() { - s.PanicsWithError(`task "self" cannot depend on itself in DAG "dag1"`, func() { - s.dag.AddTaskWithName("self", myTask, TaskSpec{}, []string{"self"}) +func (s *RegistrySuite) TestTask_DataParamWithoutInputsPanics() { + s.PanicsWithError( + `task "consume" in DAG "dag1" declares 1 data parameter(s) but Inputs supplies 0`, + func() { + s.dag.Task(consume) + }, + ) +} + +func (s *RegistrySuite) TestTask_InputTypeMismatchPanics() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consumeString" in DAG "dag1": input 0: task "extract" returns bundlev1.payload, which cannot fill parameter type string`, + func() { + s.dag.Task(consumeString, Inputs(extract)) + }, + ) +} + +func (s *RegistrySuite) TestTask_InputLooseMapAccepted() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + s.NotPanics(func() { + s.dag.Task(consumeLoose, Inputs(extract)) }) } +func (s *RegistrySuite) TestTask_InputFromValuelessTaskPanics() { + first := s.dag.Task(myTask, TaskSpec{TaskId: "first"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": input task "first" returns no value; use After for an ordering-only dependency`, + func() { + s.dag.Task(consume, Inputs(first)) + }, + ) +} + +func (s *RegistrySuite) TestTask_CrossDagRefPanics() { + other := s.reg.AddDag(DagSpec{DagId: "dag2"}) + upstream := other.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": Inputs ref "extract" belongs to DAG "dag2"; dependencies cannot cross dags`, + func() { + s.dag.Task(consume, Inputs(upstream)) + }, + ) +} + +func (s *RegistrySuite) TestTask_ForeignRegistryRefPanics() { + foreign := New().AddDag(DagSpec{DagId: "dag1"}) + upstream := foreign.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": Inputs ref "extract" belongs to a different registry`, + func() { + s.dag.Task(consume, Inputs(upstream)) + }, + ) +} + func (s *RegistrySuite) TestAddDag_WithSpec() { dag2 := s.reg.AddDag( - "dag2", - DagSpec{Schedule: "@daily", Tags: []string{"etl"}, MaxActiveRuns: 4}, + DagSpec{DagId: "dag2", Schedule: "@daily", Tags: []string{"etl"}, MaxActiveRuns: 4}, ) s.NotNil(dag2) enum, ok := s.reg.(EnumerableBundle) @@ -274,8 +349,35 @@ func (s *RegistrySuite) TestAddDag_WithSpec() { s.Equal(4, got.Spec.MaxActiveRuns) } -func (s *RegistrySuite) TestAddDag_TooManySpecsPanics() { - s.PanicsWithError("AddDag accepts at most one spec, got 2", func() { - s.reg.AddDag("dag3", DagSpec{}, DagSpec{}) +func producePtr() (*payload, error) { return &payload{Value: 1}, nil } +func consumePtr(in *payload) error { _ = in; return nil } + +// Pointer and value shapes must be interchangeable on both sides of a data +// edge: the value crosses XCom as its JSON shape, so *T and T are equivalent. +func (s *RegistrySuite) TestTask_InputPointerValueShapesCompatible() { + val := s.dag.Task(produce, TaskSpec{TaskId: "val"}) + ptr := s.dag.Task(producePtr, TaskSpec{TaskId: "ptr"}) + + s.NotPanics(func() { + s.dag.Task(consume, TaskSpec{TaskId: "val_to_val"}, Inputs(val)) + s.dag.Task(consumePtr, TaskSpec{TaskId: "val_to_ptr"}, Inputs(val)) + s.dag.Task(consume, TaskSpec{TaskId: "ptr_to_val"}, Inputs(ptr)) + s.dag.Task(consumePtr, TaskSpec{TaskId: "ptr_to_ptr"}, Inputs(ptr)) }) } + +func (s *RegistrySuite) TestTask_AnonymousFnWithoutIdPanics() { + s.PanicsWithError( + `task function in DAG "dag1" is anonymous (its derived name would be "1"); set TaskSpec.TaskId explicitly`, + func() { + s.dag.Task(func() error { return nil }) + }, + ) +} + +func (s *RegistrySuite) TestTask_AnonymousFnWithExplicitIdRegisters() { + ref := s.dag.Task(func() error { return nil }, TaskSpec{TaskId: "inline"}) + s.Equal("inline", ref.ID()) + _, exists := s.reg.LookupTask("dag1", "inline") + s.True(exists) +} diff --git a/go-sdk/bundle/bundlev1/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index 46bf5cc28793f..b42021a2390d6 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -34,9 +34,9 @@ type BundleProvider interface { // Registry, for example: // // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { - // dag := dagbag.AddDag("simple_dag") - // dag.AddTask(extract) - // dag.AddTask(transform) + // dag := dagbag.AddDag(bundlev1.DagSpec{DagId: "simple_dag"}) + // extracted := dag.Task(extract) + // dag.Task(transform, bundlev1.Inputs(extracted)) // return nil // } // diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go index 9dcf4ee832584..2a8ee38c02323 100644 --- a/go-sdk/bundle/bundlev1/spec.gen.go +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -20,12 +20,14 @@ package bundlev1 import "time" -// TaskSpec is the optional configuration applied to a task at registration -// time. Every field is optional: the zero value (nil for the *bool fields) -// means "unset" and the scheduler falls back to the schema default. Each -// field maps to the same-named key of the "operator" definition in +// TaskSpec describes a task at registration time. Every field is optional: +// the zero value (nil for the *bool fields) means "unset" and the scheduler +// falls back to the schema default. Each field maps to the same-named key of +// the "operator" definition in // airflow-core/src/airflow/serialization/schema.json. type TaskSpec struct { + // TaskId sets the task id explicitly; empty derives it from the task function's name. Maps to the schema key "task_id". + TaskId string // Owner maps to the schema key "owner" (schema default "airflow"). Owner string // StartDate maps to the schema key "start_date". @@ -80,10 +82,12 @@ type TaskSpec struct { MaxActiveTisPerDagrun int } -// SchemaFields returns the schema-keyed value of every field that is set and -// differs from its schema default, mirroring Python BaseSerialization's -// omission of values the scheduler re-derives. time.Duration and time.Time -// values are returned as-is; the serializer owns the wire encoding. +// SchemaFields returns the schema-keyed value of every configuration field +// that is set and differs from its schema default, mirroring Python +// BaseSerialization's omission of values the scheduler re-derives. Identity +// fields (TaskId) are excluded: the serializer emits the id from the +// registry's resolved TaskInfo. time.Duration and time.Time values are +// returned as-is; the serializer owns the wire encoding. func (s TaskSpec) SchemaFields() map[string]any { m := map[string]any{} if s.Owner != "" && s.Owner != "airflow" { diff --git a/go-sdk/bundle/bundlev1/spec.go b/go-sdk/bundle/bundlev1/spec.go index d9af9e1316f72..cd23f1393fc5e 100644 --- a/go-sdk/bundle/bundlev1/spec.go +++ b/go-sdk/bundle/bundlev1/spec.go @@ -17,20 +17,25 @@ package bundlev1 -import "time" +import ( + "fmt" + "time" +) // This file holds the hand-written companions of the generated spec.gen.go: -// DagSpec (whose Schedule field and always-emit/config-fallback keys the -// schema cannot express mechanically) and the registration Info structs that -// carry Go-side identity the schema knows nothing about. +// DagSpec (whose DagId/Schedule fields and always-emit/config-fallback keys +// the schema cannot express mechanically) and the registration Info structs +// that carry Go-side identity the schema knows nothing about. type ( - // DagSpec is the optional configuration applied to a DAG at registration - // time. Every field is optional: a zero value means "unset" and the + // DagSpec describes a dag at registration time. DagId is required; every + // other field is optional, with a zero value meaning "unset" so the // scheduler falls back to its serialization-schema default. The field // names mirror the keys defined under "dag" in // airflow-core/src/airflow/serialization/schema.json. DagSpec struct { + // DagId is the dag id. Required: AddDag panics when it is empty. + DagId string // Schedule is "@once", "@continuous", a cron expression, or "" for // NullTimetable (no schedule). Schedule string @@ -56,8 +61,8 @@ type ( // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses // it to render the per-task block of a DagFileParsingResult. TaskInfo struct { - // ID is the user-visible task id (the function name unless overridden - // via AddTaskWithName). + // ID is the user-visible task id (TaskSpec.TaskId when set, otherwise + // the task function's name). ID string // TypeName is the unqualified Go function name (e.g. "extract"). TypeName string @@ -66,8 +71,11 @@ type ( // Spec carries the optional per-task configuration supplied at // registration. The zero value means "no overrides". Spec TaskSpec + // Inputs lists the upstream task ids whose return values feed this + // task's data parameters, in parameter order (the Inputs option). + Inputs []string // Downstream lists task ids that depend on this task, populated as - // later tasks declare this id in their AddTask `depends` argument. + // later registrations reference this task via Inputs or After. // Order is registration order; the serializer sorts before emit. Downstream []string } @@ -76,17 +84,18 @@ type ( // registration order. DagInfo struct { DagID string - // Spec carries the optional per-dag configuration supplied at - // registration. The zero value means "no overrides". + // Spec carries the per-dag configuration supplied at registration. Spec DagSpec Tasks []TaskInfo } ) -// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / -// DagSpec where nil means "leave at schema default": -// -// v1.TaskSpec{DoXComPush: v1.Bool(false)} -func Bool(b bool) *bool { - return &b +// applyTask lets a TaskSpec value be passed directly as a TaskOption to +// Dag.Task. Passing more than one spec panics at registration time. +func (s TaskSpec) applyTask(c *taskConfig) { + if c.specSet { + panic(fmt.Errorf("Task accepts at most one TaskSpec")) + } + c.spec = s + c.specSet = true } diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index d31fea84b73f3..97830d8e53ddc 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -18,7 +18,9 @@ package bundlev1 import ( + "bytes" "context" + "encoding/json" "fmt" "log/slog" "reflect" @@ -29,22 +31,106 @@ import ( "github.com/apache/airflow/go-sdk/sdk" ) +// InputBinding names the upstream task whose return-value XCom fills one data +// parameter of a task function. Bindings are positional: the i-th binding +// fills the i-th data parameter. +type InputBinding struct { + TaskID string +} + +// boundInput is one resolved data parameter: which argument slot to fill, +// from which upstream task, decoding into which type. +type boundInput struct { + argIndex int + taskID string + typ reflect.Type +} + type taskFunction struct { fn reflect.Value fullName string + inputs []boundInput } var _ Task = (*taskFunction)(nil) -// NewTaskFunction wraps a plain Go function as a Task, validating its signature -// (injectable parameters, and a return of error or (result, error)). Bundle -// authors normally use Dag.AddTask, which calls this for them; use it directly -// only when building a Task outside the registry. -func NewTaskFunction(fn any) (Task, error) { +// NewTaskFunction wraps a plain Go function as a Task, validating its +// signature (injectable parameters, data parameters matching inputs, and a +// return of error or (result, error)). Bundle authors normally use Dag.Task, +// which calls this for them; use it directly only when building a Task +// outside the registry. +// +// inputs bind the function's data parameters (every parameter that is not an +// injectable runtime type) positionally to upstream tasks: at execution each +// data parameter receives the named task's return-value XCom, decoded into +// the parameter's type. The number of inputs must equal the number of data +// parameters. +func NewTaskFunction(fn any, inputs ...InputBinding) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{v, fullName} - return f, f.validateFn(v.Type()) + f := &taskFunction{fn: v, fullName: fullName} + if err := f.validateFn(v.Type()); err != nil { + return nil, err + } + if err := f.bindInputs(v.Type(), inputs); err != nil { + return nil, err + } + return f, nil +} + +// dataParamTypes returns the types of fnType's data parameters — every +// parameter Execute does not inject — in declaration order. +func dataParamTypes(fnType reflect.Type) []reflect.Type { + if fnType.Kind() != reflect.Func { + return nil + } + var out []reflect.Type + for i := range fnType.NumIn() { + if in := fnType.In(i); !isInjectable(in) { + out = append(out, in) + } + } + return out +} + +// isInjectable reports whether Execute fills a parameter of this type from +// the runtime (rather than from an upstream task's XCom). +func isInjectable(in reflect.Type) bool { + return isTIRunContext(in) || isContext(in) || isLogger(in) || isClient(in) +} + +// bindInputs pairs the fn's data parameters with the supplied input bindings +// and records the pull-and-decode plan Execute runs before calling fn. +func (f *taskFunction) bindInputs(fnType reflect.Type, inputs []InputBinding) error { + var dataIdx []int + for i := range fnType.NumIn() { + if !isInjectable(fnType.In(i)) { + dataIdx = append(dataIdx, i) + } + } + if len(dataIdx) != len(inputs) { + return fmt.Errorf( + "task function %s declares %d data parameter(s) but %d input binding(s) were supplied", + f.fullName, len(dataIdx), len(inputs), + ) + } + f.inputs = make([]boundInput, len(inputs)) + for i, in := range inputs { + typ := fnType.In(dataIdx[i]) + if in.TaskID == "" { + return fmt.Errorf( + "task function %s: input binding %d has an empty task id", f.fullName, i, + ) + } + if !isDecodableType(typ) { + return fmt.Errorf( + "task function %s: data parameter %d has type %s, which cannot hold an XCom value", + f.fullName, dataIdx[i], typ, + ) + } + f.inputs[i] = boundInput{argIndex: dataIdx[i], taskID: in.TaskID, typ: typ} + } + return nil } func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { @@ -83,10 +169,16 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { case isClient(in): reflectArgs[i] = reflect.ValueOf(sdkClient) default: - // TODO: deal with other value types. For now they will all be Zero values unless it's a context + // A data parameter: placeholder until resolveInputs overwrites it + // below (validation guarantees each has an input binding). reflectArgs[i] = reflect.Zero(in) } } + + if err := f.resolveInputs(ctx, sdkClient, logger, reflectArgs); err != nil { + return err + } + slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs) retValues := f.fn.Call(reflectArgs) @@ -108,6 +200,98 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { return err } +// resolveInputs fills the data-parameter argument slots by pulling each bound +// upstream task's return-value XCom for the current dag run and decoding it +// into the parameter's type. Any failure fails the task before its body runs. +func (f *taskFunction) resolveInputs( + ctx context.Context, + c sdk.XComClient, + logger *slog.Logger, + args []reflect.Value, +) error { + if len(f.inputs) == 0 { + return nil + } + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return fmt.Errorf( + "task function %s: no workload in context; cannot resolve task inputs", f.fullName, + ) + } + ti := workload.TI + for _, in := range f.inputs { + // Pull from the upstream's unmapped instance; mapped upstream fan-in + // is out of scope for now. + raw, err := c.GetXCom(ctx, ti.DagId, ti.RunId, in.taskID, nil, api.XComReturnValueKey, nil) + if err != nil { + return fmt.Errorf( + "task function %s: pulling input from task %q: %w", f.fullName, in.taskID, err, + ) + } + decoded, err := decodeXCom(raw, in.typ) + if err != nil { + return fmt.Errorf( + "task function %s: decoding input from task %q into %s: %w", + f.fullName, in.taskID, in.typ, err, + ) + } + logger.DebugContext( + ctx, + "Resolved task input", + "from_task", + in.taskID, + "type", + in.typ.String(), + ) + args[in.argIndex] = decoded + } + return nil +} + +// decodeXCom decodes a raw (generically deserialised) XCom value into target. +// Decoding into a struct is strict: unknown/renamed keys fail rather than +// silently leaving fields zero. Decoding into a map / interface accepts any +// shape, so authors opt into loose decoding by typing the parameter +// map[string]any or any. A null value is allowed only for a nilable target. +func decodeXCom(raw any, target reflect.Type) (reflect.Value, error) { + out := reflect.New(target) + + if raw == nil { + switch target.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface: + return out.Elem(), nil + default: + return reflect.Value{}, fmt.Errorf( + "xcom value is null but parameter type %s is not nilable", target, + ) + } + } + + blob, err := json.Marshal(raw) + if err != nil { + return reflect.Value{}, err + } + dec := json.NewDecoder(bytes.NewReader(blob)) + dec.DisallowUnknownFields() + if err := dec.Decode(out.Interface()); err != nil { + return reflect.Value{}, err + } + return out.Elem(), nil +} + +// isDecodableType reports whether a value can be JSON-decoded into inType. It +// rejects kinds json cannot target (func, chan, unsafe pointer) and non-empty +// interfaces (only the empty interface `any` is a valid decode target). +func isDecodableType(inType reflect.Type) bool { + switch inType.Kind() { + case reflect.Func, reflect.Chan, reflect.UnsafePointer: + return false + case reflect.Interface: + return inType.NumMethod() == 0 + } + return true +} + func (f *taskFunction) sendXcom( ctx context.Context, value any, @@ -200,8 +384,10 @@ func isClient(inType reflect.Type) bool { inType.NumMethod() > 0 && clientType.Implements(inType) } -// validateParam rejects interface parameters Execute cannot inject; they -// would be bound to nil and panic on first use. +// validateParam rejects interface parameters Execute cannot inject and that +// cannot serve as data parameters either; they would be bound to nil and +// panic on first use. Non-interface parameters are data parameters, filled +// from upstream XComs via the Inputs bindings. func validateParam(in reflect.Type) error { if in.Kind() != reflect.Interface || isTIRunContext(in) || isClient(in) { return nil @@ -216,6 +402,10 @@ func validateParam(in reflect.Type) error { in, ) } + if in.NumMethod() == 0 { + // The empty interface is a valid (loose) data-parameter target. + return nil + } return fmt.Errorf( "interface %s is not injectable (want context.Context, sdk.TIRunContext, or a subset of sdk.Client): %s", in, diff --git a/go-sdk/bundle/bundlev1/task_test.go b/go-sdk/bundle/bundlev1/task_test.go index 3f58e930b8721..9f902bfea8d71 100644 --- a/go-sdk/bundle/bundlev1/task_test.go +++ b/go-sdk/bundle/bundlev1/task_test.go @@ -19,12 +19,14 @@ package bundlev1 import ( "context" + "errors" "log/slog" "reflect" "testing" "github.com/stretchr/testify/suite" + "github.com/apache/airflow/go-sdk/pkg/api" "github.com/apache/airflow/go-sdk/pkg/logging" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" @@ -174,10 +176,6 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { }, "method GetVariable is func(context.Context, string) (string, error) on sdk.Client", }, - "empty-interface": { - func(x any) error { return nil }, - "empty interfaces cannot be injected", - }, "context-with-extra-methods": { func(x interface { context.Context @@ -276,3 +274,247 @@ func (s *TaskSuite) TestTIRunContextInjectionWithoutRuntimeContext() { "the injected context must be backed by the one passed to Execute", ) } + +// --- Data-parameter (TaskFlow Inputs) binding --- + +// fakeClient implements sdk.Client with canned per-task XCom values, recording +// the pulls Execute performs to resolve data parameters. +type fakeClient struct { + xcoms map[string]any // upstream task id -> stored return_value + getErr error + pulls []string + pushed map[string]any + pushErr error +} + +func (f *fakeClient) GetVariable(context.Context, string) (string, error) { return "", nil } +func (f *fakeClient) UnmarshalJSONVariable(context.Context, string, any) error { + return nil +} + +func (f *fakeClient) GetConnection(context.Context, string) (sdk.Connection, error) { + return sdk.Connection{}, nil +} + +func (f *fakeClient) GetXCom( + _ context.Context, + dagId, runId, taskId string, + mapIndex *int, + key string, + _ any, +) (any, error) { + f.pulls = append(f.pulls, taskId) + if f.getErr != nil { + return nil, f.getErr + } + if dagId != "dag1" || runId != "run1" || key != api.XComReturnValueKey || mapIndex != nil { + return nil, errors.New("unexpected xcom pull identifiers") + } + return f.xcoms[taskId], nil +} + +func (f *fakeClient) PushXCom(_ context.Context, _ api.TaskInstance, key string, value any) error { + if f.pushed == nil { + f.pushed = map[string]any{} + } + f.pushed[key] = value + return f.pushErr +} + +// bindingResult is the struct shape the fake upstream produces. +type bindingResult struct { + Value int `json:"value"` + Name string `json:"name"` +} + +func (s *TaskSuite) executionContext(client sdk.Client) context.Context { + ctx := context.WithValue(context.Background(), sdkcontext.SdkClientContextKey, client) + return context.WithValue(ctx, sdkcontext.WorkloadContextKey, api.ExecuteTaskWorkload{ + TI: api.TaskInstance{DagId: "dag1", RunId: "run1", TaskId: "consumer"}, + }) +} + +func (s *TaskSuite) TestDataParamInjection() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"value": 42, "name": "answer"}, + }} + + var got bindingResult + task, err := NewTaskFunction(func(in bindingResult) error { + got = in + return nil + }, InputBinding{TaskID: "extract"}) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + s.Require().NoError(err) + s.Equal(bindingResult{Value: 42, Name: "answer"}, got) + s.Equal([]string{"extract"}, client.pulls) +} + +func (s *TaskSuite) TestDataParamsBindPositionallyAmongInjectables() { + client := &fakeClient{xcoms: map[string]any{ + "first": map[string]any{"value": 1, "name": "a"}, + "second": map[string]any{"value": 2, "name": "b"}, + }} + + var gotA, gotB bindingResult + task, err := NewTaskFunction( + func(ctx context.Context, a bindingResult, logger *slog.Logger, b bindingResult, c sdk.XComClient) error { + s.NotNil(ctx) + s.NotNil(logger) + s.NotNil(c) + gotA, gotB = a, b + return nil + }, + InputBinding{TaskID: "first"}, + InputBinding{TaskID: "second"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + s.Require().NoError(err) + s.Equal(bindingResult{Value: 1, Name: "a"}, gotA) + s.Equal(bindingResult{Value: 2, Name: "b"}, gotB) + s.Equal([]string{"first", "second"}, client.pulls) +} + +func (s *TaskSuite) TestDataParamStrictDecodeRejectsUnknownKeys() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"value": 1, "renamed_field": true}, + }} + + task, err := NewTaskFunction( + func(in bindingResult) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), `decoding input from task "extract"`) + s.Contains(err.Error(), "renamed_field") + } +} + +func (s *TaskSuite) TestDataParamLooseMapDecodesAnyShape() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"anything": "goes", "extra": 1}, + }} + + var got map[string]any + task, err := NewTaskFunction( + func(in map[string]any) error { + got = in + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + s.Require().NoError(task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger()))) + s.Equal("goes", got["anything"]) +} + +func (s *TaskSuite) TestDataParamNullValue() { + client := &fakeClient{xcoms: map[string]any{"extract": nil}} + + s.Run("nilable-target-gets-zero", func() { + var got *bindingResult + task, err := NewTaskFunction( + func(in *bindingResult) error { + got = in + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + s.Require(). + NoError(task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger()))) + s.Nil(got) + }) + + s.Run("non-nilable-target-fails", func() { + task, err := NewTaskFunction( + func(in bindingResult) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), "not nilable") + } + }) +} + +func (s *TaskSuite) TestDataParamPullFailureFailsTask() { + client := &fakeClient{getErr: errors.New("supervisor unreachable")} + + ran := false + task, err := NewTaskFunction( + func(in bindingResult) error { + ran = true + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), `pulling input from task "extract"`) + } + s.False(ran, "the task body must not run when an input pull fails") +} + +func (s *TaskSuite) TestDataParamWithoutWorkloadFails() { + client := &fakeClient{xcoms: map[string]any{"extract": map[string]any{"value": 1}}} + task, err := NewTaskFunction( + func(in map[string]any) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + // Client injected but no workload on the context. + ctx := context.WithValue( + context.Background(), + sdkcontext.SdkClientContextKey, + sdk.Client(client), + ) + err = task.Execute(ctx, slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), "no workload in context") + } +} + +func (s *TaskSuite) TestInputBindingValidation() { + s.Run("count-mismatch", func() { + _, err := NewTaskFunction(func(x any) error { return nil }) + if s.Error(err) { + s.Contains( + err.Error(), + "declares 1 data parameter(s) but 0 input binding(s) were supplied", + ) + } + }) + + s.Run("empty-task-id", func() { + _, err := NewTaskFunction( + func(x any) error { return nil }, + InputBinding{}, + ) + if s.Error(err) { + s.Contains(err.Error(), "input binding 0 has an empty task id") + } + }) + + s.Run("undecodable-param", func() { + _, err := NewTaskFunction( + func(x chan int) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + if s.Error(err) { + s.Contains(err.Error(), "cannot hold an XCom value") + } + }) +} diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 06535b0bf8be6..1970a32e4aece 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -45,23 +45,39 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { return v1.BundleInfo{Name: bundleName, Version: &bundleVersion} } +// ExtractResult is extract's return value. Returning it pushes it as the +// task's return_value XCom; declaring it via Inputs feeds it to a downstream +// task's matching data parameter. +type ExtractResult struct { + GoVersion string `json:"go_version"` + Timestamp int64 `json:"timestamp"` +} + +// TransformResult is transform's return value, consumed by load. +type TransformResult struct { + Variable string `json:"variable"` + Extracted ExtractResult `json:"extracted"` +} + func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag", v1.DagSpec{ + simpleDag := dagbag.AddDag(v1.DagSpec{ + DagId: "simple_dag", Schedule: "@daily", Description: "Example Go-authored Dag", Tags: []string{"example", "go-sdk"}, }) - simpleDag.AddTask(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}, nil) - simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}, []string{"extract"}) - simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}, []string{"transform"}) + // TaskFlow style: each Task call returns a ref; passing refs via + // v1.Inputs both orders the tasks and feeds the upstream's return value + // into the downstream function's data parameter. + extracted := simpleDag.Task(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}) + transformed := simpleDag.Task(transform, v1.TaskSpec{Queue: "go-task"}, v1.Inputs(extracted)) + simpleDag.Task(load, v1.TaskSpec{Queue: "go-task"}, v1.Inputs(transformed)) // Tasks defined in other packages register through the same dagbag. - concurrentDag := dagbag.AddDag("concurrent_xcom_dag") - concurrentDag.AddTaskWithName( - "pull_xcoms_concurrently", + concurrentDag := dagbag.AddDag(v1.DagSpec{DagId: "concurrent_xcom_dag"}) + concurrentDag.Task( concurrentxcom.PullXComsConcurrently, - v1.TaskSpec{}, - nil, + v1.TaskSpec{TaskId: "pull_xcoms_concurrently"}, ) return nil @@ -73,7 +89,7 @@ func main() { } } -func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { +func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (ExtractResult, error) { log.Info("Hello from task") // ctx behaves as a context.Context and also carries the task instance @@ -120,7 +136,7 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er // Once per loop,.check if we've been asked to cancel! select { case <-ctx.Done(): - return nil, ctx.Err() + return ExtractResult{}, ctx.Err() default: } log.Info("After the beep the time will be", "time", time.Now()) @@ -128,40 +144,51 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er } log.Info("Goodbye from task") - ret := map[string]any{ - "go_version": runtime.Version(), - "timestamp": time.Now().UnixNano(), - } - - return ret, nil + return ExtractResult{ + GoVersion: runtime.Version(), + Timestamp: time.Now().UnixNano(), + }, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { - // This function takes a VariableClient and not a Client to make unit testing it easier. See - // `./main_test.go` for an example unit of this task fn. Functionally taking a `sdk.Client` is the same (as - // Client includes VariableClient) but by using the dedicated type it can be easier to write unit tests. - // - // It also gives a better indication of what features the tasks use +// transform receives extract's return value as its `in` data parameter: the +// runtime pulls extract's return-value XCom and decodes it into ExtractResult +// before the body runs, because RegisterDags wired v1.Inputs(extracted). +// +// It takes a VariableClient and not a full Client to make unit testing easier +// (see ./main_test.go); functionally `sdk.Client` works the same, but the +// dedicated type documents what the task actually uses. +func transform( + ctx sdk.TIRunContext, + client sdk.VariableClient, + log *slog.Logger, + in ExtractResult, +) (TransformResult, error) { + log.Info("Received upstream result", "go_version", in.GoVersion, "timestamp", in.Timestamp) + key := "my_variable" val, err := client.GetVariable(ctx, key) if err != nil { - return err + return TransformResult{}, err } log.Info("Obtained variable", key, val) - return nil + return TransformResult{Variable: val, Extracted: in}, nil } // load fails on its first attempt and succeeds on the retry. With retries -// configured on the stub task, the first failure makes the supervisor mark the +// configured on the task, the first failure makes the supervisor mark the // task UP_FOR_RETRY -- which only works because the Go SDK now emits a // RetryTask frame (instead of a terminal FAILED) when ti_context.should_retry // is set. The retry then runs this task again and it returns nil. -func load(ctx sdk.TIRunContext, log *slog.Logger) error { +func load(ctx sdk.TIRunContext, log *slog.Logger, in TransformResult) error { tryNumber := ctx.TaskInstance().TryNumber if tryNumber == 1 { log.InfoContext(ctx, "Please fail", "try_number", tryNumber) return fmt.Errorf("Please fail") } - log.InfoContext(ctx, "Recovered on retry", "try_number", tryNumber) + log.InfoContext(ctx, "Recovered on retry", + "try_number", tryNumber, + "variable", in.Variable, + "extracted_go_version", in.Extracted.GoVersion, + ) return nil } diff --git a/go-sdk/example/bundle/main_test.go b/go-sdk/example/bundle/main_test.go index 474a8406d6224..75773784bb502 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/go-sdk/example/bundle/main_test.go @@ -51,9 +51,14 @@ var _ sdk.VariableClient = (*mockVars)(nil) func Test_transform(t *testing.T) { log := slog.Default() - // This is not the best test, but it is a good proof of concept -- you can just call the function. + // This is not the best test, but it is a good proof of concept -- you can + // just call the function, passing the upstream result directly instead of + // letting the runtime pull it from XCom. // sdk.NewTIRunContext wraps any context to build a TIRunContext in a test. ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - err := transform(ctx, &mockVars{}, log) + in := ExtractResult{GoVersion: "go1.24", Timestamp: 42} + got, err := transform(ctx, &mockVars{}, log, in) assert.NoError(t, err) + assert.Equal(t, "value1", got.Variable) + assert.Equal(t, in, got.Extracted) } diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 6abb963f391dd..e6e5e4be5c7b5 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -88,8 +88,8 @@ func buildBundle(t *testing.T, register func(bundlev1.Registry)) bundlev1.Bundle func TestDagParsing(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - d := r.AddDag("test_dag") - d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + d := r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}) + d.Task(simpleTask) }) req := &genmodels.DagFileParseRequest{ @@ -127,8 +127,8 @@ func TestDagParsing(t *testing.T) { func TestDagParsingMultipleDagsPreservesOrder(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) - r.AddDag("dag2").AddTask(failingTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "dag1"}).Task(simpleTask) + r.AddDag(bundlev1.DagSpec{DagId: "dag2"}).Task(failingTask) }) req := &genmodels.DagFileParseRequest{File: "/bundle/main.go", BundlePath: "/bundle"} @@ -146,7 +146,7 @@ func TestDagParsingMultipleDagsPreservesOrder(t *testing.T) { func TestTaskRunnerSuccess(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(simpleTask) }) details := &genmodels.StartupDetails{ @@ -169,7 +169,7 @@ func TestTaskRunnerSuccess(t *testing.T) { func TestTaskRunnerFailure(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(failingTask) }) details := &genmodels.StartupDetails{ @@ -192,7 +192,7 @@ func TestTaskRunnerFailure(t *testing.T) { func TestTaskRunnerRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(failingTask) }) details := &genmodels.StartupDetails{ @@ -219,7 +219,7 @@ func TestTaskRunnerRetry(t *testing.T) { func TestTaskRunnerTaskNotFound(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(simpleTask) }) details := &genmodels.StartupDetails{ @@ -241,7 +241,7 @@ func TestTaskRunnerTaskNotFound(t *testing.T) { func TestTaskRunnerPanic(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(panicTask) }) details := &genmodels.StartupDetails{ @@ -264,7 +264,7 @@ func TestTaskRunnerPanic(t *testing.T) { func TestTaskRunnerPanicRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task(panicTask) }) details := &genmodels.StartupDetails{ @@ -291,8 +291,9 @@ func TestTaskRunnerPanicRetry(t *testing.T) { func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxcheck", - func(ctx context.Context) error { return ctx.Err() }, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( + func(ctx context.Context) error { return ctx.Err() }, + bundlev1.TaskSpec{TaskId: "ctxcheck"}) }) details := &genmodels.StartupDetails{ @@ -325,11 +326,11 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { var got sdk.TIRunContext bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxgrab", + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( func(ctx sdk.TIRunContext) error { got = ctx return nil - }, bundlev1.TaskSpec{}, nil) + }, bundlev1.TaskSpec{TaskId: "ctxgrab"}) }) details := &genmodels.StartupDetails{ @@ -385,11 +386,11 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { var got sdk.TIRunContext bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxgrab", + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( func(ctx sdk.TIRunContext) error { got = ctx return nil - }, bundlev1.TaskSpec{}, nil) + }, bundlev1.TaskSpec{TaskId: "ctxgrab"}) }) details := &genmodels.StartupDetails{ @@ -471,8 +472,16 @@ func TestServeDagFileParseEndToEnd(t *testing.T) { provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - d := r.AddDag("simple_dag") - d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + d := r.AddDag(bundlev1.DagSpec{DagId: "simple_dag"}) + extract := d.Task( + func() (string, error) { return "data", nil }, + bundlev1.TaskSpec{TaskId: "extract"}, + ) + d.Task( + func(in string) error { return nil }, + bundlev1.TaskSpec{TaskId: "transform"}, + bundlev1.Inputs(extract), + ) return nil }, } @@ -510,6 +519,16 @@ func TestServeDagFileParseEndToEnd(t *testing.T) { dag := dags[0].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) assert.Equal(t, "simple_dag", dag["dag_id"]) + // The Inputs wiring must surface as downstream_task_ids on the upstream. + tasks := dag["tasks"].([]any) + require.Len(t, tasks, 2) + extractData := tasks[0].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "extract", extractData["task_id"]) + assert.Equal(t, []any{"transform"}, extractData["downstream_task_ids"]) + transformData := tasks[1].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "transform", transformData["task_id"]) + assert.Nil(t, transformData["downstream_task_ids"]) + select { case err := <-done: require.NoError(t, err) @@ -524,7 +543,7 @@ func TestServeStartupDetailsEndToEnd(t *testing.T) { provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag(bundlev1.DagSpec{DagId: "dag1"}).Task(simpleTask) return nil }, } @@ -582,7 +601,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { var gotVar string provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTaskWithName("getvar", + r.AddDag(bundlev1.DagSpec{DagId: "dag1"}).Task( func(ctx context.Context, c sdk.Client) (string, error) { v, err := c.GetVariable(ctx, varKey) if err != nil { @@ -590,7 +609,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }, bundlev1.TaskSpec{}, nil) + }, bundlev1.TaskSpec{TaskId: "getvar"}) return nil }, } diff --git a/go-sdk/pkg/worker/runner_test.go b/go-sdk/pkg/worker/runner_test.go index 90cdc2e9ffc77..43260c1db0156 100644 --- a/go-sdk/pkg/worker/runner_test.go +++ b/go-sdk/pkg/worker/runner_test.go @@ -160,10 +160,11 @@ func (s *WorkerSuite) TestStartContextErrorTaskDoesntStart() { wasCalled := false // Register a task that should NOT be called if everything works - s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { - wasCalled = true - return nil - }, bundlev1.TaskSpec{}, nil) + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func() error { + wasCalled = true + return nil + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) // Setup the mock s.ti.EXPECT(). @@ -193,10 +194,11 @@ func (s *WorkerSuite) TestTaskHeartbeatsWhileRunning() { id := uuid.New().String() var callCount atomic.Int32 testWorkload := newTestWorkLoad(id, id[:8]) - s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { - time.Sleep(time.Second) - return nil - }, bundlev1.TaskSpec{}, nil) + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func() error { + time.Sleep(time.Second) + return nil + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateSuccess) @@ -224,15 +226,15 @@ func (s *WorkerSuite) TestTaskHeartbeatConflictStopsTask() { id := uuid.New().String() testWorkload := newTestWorkLoad(id, id[:8]) - s.registry.AddDag(testWorkload.TI.DagId). - AddTaskWithName(testWorkload.TI.TaskId, func(ctx context.Context) error { + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func(ctx context.Context) error { select { case <-ctx.Done(): return nil case <-time.After(2 * time.Second): return fmt.Errorf("task context was not cancelled") } - }, bundlev1.TaskSpec{}, nil) + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateFailed)