From d2a25d85eeccffa766c885db4d6454d9bab40f53 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 5 May 2026 11:16:55 +0800 Subject: [PATCH 1/7] Go-SDK: Add coordinator-mode Dag parsing --- go-sdk/bundle/bundlev1/registry.go | 231 +++++++++-- go-sdk/bundle/bundlev1/registry_test.go | 152 +++++++- go-sdk/example/bundle/main.go | 14 +- go-sdk/pkg/execution/dag_parser.go | 54 +++ go-sdk/pkg/execution/integration_test.go | 134 ++++++- go-sdk/pkg/execution/serde.go | 463 +++++++++++++++++++++++ go-sdk/pkg/execution/serde_test.go | 447 ++++++++++++++++++++++ go-sdk/pkg/execution/server.go | 17 +- 8 files changed, 1443 insertions(+), 69 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 563cbca356639..409cd4b913ebe 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,14 +23,19 @@ import ( "runtime" "strings" "sync" + "time" ) 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). + // 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: @@ -43,13 +48,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, 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) + // 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 @@ -58,26 +63,107 @@ 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 } - // TaskInfo describes a registered task by its user-visible id. + // 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; 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 } @@ -85,6 +171,8 @@ type ( registry struct { sync.RWMutex taskFuncMap map[string]map[string]Task + taskInfo map[string]map[string]TaskInfo + dagSpec map[string]DagSpec dagOrder []string taskOrder map[string][]string } @@ -95,12 +183,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, depends []string) { + d.registry.registerTask(d.dagId, fn, spec, depends) } -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, 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 / +// 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 @@ -110,31 +218,43 @@ 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), + dagSpec: make(map[string]DagSpec), 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() - +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 { 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.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, depends []string) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -143,31 +263,66 @@ func (r *registry) registerTask(dagId string, fn any) { fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn) + r.registerTaskWithName(dagId, fnName, fn, spec, depends) } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { +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)) } + val := reflect.ValueOf(fn) + 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() 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)) } + + // 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) } @@ -183,9 +338,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() @@ -195,9 +350,9 @@ 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}) + 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..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 @@ -144,12 +144,138 @@ 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)}, 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("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"}, 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"}) + + 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"}) + }) +} + +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 6d6569889008c..c439164ecaac6 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -37,14 +37,18 @@ type myBundle struct{} var _ v1.BundleProvider = (*myBundle)(nil) 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}, 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") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently, v1.TaskSpec{}, nil) bindingDag := dagbag.AddDag("taskflow_binding_dag") bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go new file mode 100644 index 0000000000000..4b6836ab4d4b4 --- /dev/null +++ b/go-sdk/pkg/execution/dag_parser.go @@ -0,0 +1,54 @@ +// 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" + "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 *genmodels.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/integration_test.go b/go-sdk/pkg/execution/integration_test.go index da6b10784fe0a..ba2c9b0ee2dab 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" @@ -104,9 +105,67 @@ func newStartupDetails( return details } +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) + r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails("simpleTask") @@ -120,7 +179,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 := newStartupDetails("failingTask") @@ -134,7 +193,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 := newStartupDetails("failingTask") @@ -150,7 +209,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 := newStartupDetails("nonexistent") @@ -164,7 +223,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 := newStartupDetails("panicTask") @@ -178,7 +237,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 := newStartupDetails("panicTask") @@ -463,7 +522,7 @@ func TestTaskRunnerMalformedSpecHonorsShouldRetry(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 := newStartupDetails("ctxcheck") @@ -491,7 +550,7 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -551,7 +610,7 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails("ctxgrab") @@ -614,13 +673,66 @@ 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() provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTask(simpleTask) + r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) return nil }, } @@ -686,7 +798,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }) + }, bundlev1.TaskSpec{}, nil) return nil }, } diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go new file mode 100644 index 0000000000000..0aca23993d99e --- /dev/null +++ b/go-sdk/pkg/execution/serde.go @@ -0,0 +1,463 @@ +// 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" +) + +// 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 +// - 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. +// +// 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{ + "__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. 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 + } + pkgPath := info.PkgPath + if pkgPath == "" { + pkgPath = "main" + } + data := map[string]any{ + "task_id": info.ID, + "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 { + sorted := make([]string, len(info.Downstream)) + copy(sorted, info.Downstream) + sort.Strings(sorted) + data["downstream_task_ids"] = sorted + } + return map[string]any{ + "__type": "operator", + "__var": data, + } +} + +// 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, +// 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 + } + 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 { + // 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 + } + if s.DagDisplayName != "" { + data["dag_display_name"] = s.DagDisplayName + } + if s.DocMD != "" { + data["doc_md"] = s.DocMD + } + // 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)) + } + // 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.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)) + 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. 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)) + for i, t := range info.Tasks { + taskIDs[i] = t.ID + tasks[i] = serializeTask(t) + } + + 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(schedule), + "tasks": tasks, + "dag_dependencies": []any{}, + "task_group": serializeTaskGroup(taskIDs), + "edge_info": map[string]any{}, + "params": serializeParams(nil), + "deadline": nil, + "allowed_run_types": nil, + } + applyDagSpec(result, info.Spec) + return result +} + +// 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..e7aa19c70ae66 --- /dev/null +++ b/go-sdk/pkg/execution/serde_test.go @@ -0,0 +1,447 @@ +// 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(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"]) + 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"]) + // 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") +} + +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"}) + 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"}, + }) + 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"}, + }) + 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"]) + 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"]) + + // Optional spec fields stay omitted when unset. + _, hasDesc := result["description"] + assert.False(t, hasDesc) + _, 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) { + info := bundlev1.DagInfo{ + DagID: "etl", + Tasks: []bundlev1.TaskInfo{ + { + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"load"}, + }, + { + ID: "load", TypeName: "load", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }, + }, + } + 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"]) + _, 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) + assert.Contains(t, children, "extract") + 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"]) + // 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"]) + 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_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, 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) { + 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) + } +} diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index d4a4d43f1570d..2bcc2d2a21cb2 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,11 @@ // 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, From ddee3952cd8a194e91c89b6815294638b679bc48 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 16:36:57 +0000 Subject: [PATCH 2/7] Go-SDK: Generate TaskSpec from the Dag serialization schema --- go-sdk/Justfile | 4 + go-sdk/bundle/bundlev1/gen.go | 30 ++ go-sdk/bundle/bundlev1/gen/main.go | 429 +++++++++++++++++++++++++++++ go-sdk/bundle/bundlev1/registry.go | 101 ------- go-sdk/bundle/bundlev1/spec.gen.go | 168 +++++++++++ go-sdk/bundle/bundlev1/spec.go | 92 +++++++ go-sdk/example/bundle/main.go | 7 +- go-sdk/pkg/execution/serde.go | 92 +------ go-sdk/pkg/execution/serde_test.go | 66 +++-- 9 files changed, 776 insertions(+), 213 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 7094a394835db..9417e24b377cb 100644 --- a/go-sdk/Justfile +++ b/go-sdk/Justfile @@ -42,3 +42,7 @@ docs port="6060": # 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..59846166bd759 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen.go @@ -0,0 +1,30 @@ +// 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 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. +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..c1714dc0c3348 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -0,0 +1,429 @@ +// 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 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 +// 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" +) + +// 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", +} + +// 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 +// 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"` +} + +// 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]definition `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, err := selectFields(operator) + if err != nil { + log.Fatalf("gen: %v", err) + } + + 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) + } +} + +// 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: key, + goName: goName(key), + def: prop.Default, + hasDef: prop.Default != nil, + } + if ref := strings.TrimPrefix(prop.Ref, "#/definitions/"); ref != prop.Ref { + f.refName = ref + } + + switch { + case goTypeOverrides[key] != "": + f.goType = goTypeOverrides[key] + 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{}, false + } + } + return f, true +} + +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 409cd4b913ebe..db8fed5ad1a37 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,7 +23,6 @@ import ( "runtime" "strings" "sync" - "time" ) type ( @@ -69,98 +68,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. @@ -191,14 +98,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..9dcf4ee832584 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -0,0 +1,168 @@ +// 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 { + // 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 + // 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 + // 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 + // Executor maps to the schema key "executor". + Executor string + // 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.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 + } + 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.ExecutionTimeout != 0 { + m["execution_timeout"] = s.ExecutionTimeout + } + if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { + m["retry_delay"] = s.RetryDelay + } + 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.Executor != "" { + m["executor"] = s.Executor + } + 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/example/bundle/main.go b/go-sdk/example/bundle/main.go index c439164ecaac6..02b82ba67740d 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -48,7 +48,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, + ) bindingDag := dagbag.AddDag("taskflow_binding_dag") bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) 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..83e7df6a67751 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -172,33 +172,41 @@ 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", - 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{} - applyTaskSpec(data, spec) + applySpecFields(data, spec) assert.Equal(t, "gpu", data["queue"]) assert.Equal(t, "gpu_pool", data["pool"]) @@ -214,6 +222,8 @@ func TestApplyTaskSpec_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"]) @@ -224,7 +234,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 +251,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 6b978864f2aaea4f1b8c4ff2e1c36cccaa3fa874 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 10:09:53 +0000 Subject: [PATCH 3/7] Go-SDK: Generate DagSpec from the Dag serialization schema --- go-sdk/Justfile | 2 +- go-sdk/bundle/bundlev1/gen.go | 12 +- go-sdk/bundle/bundlev1/gen/main.go | 437 +++++++++++++++++------ go-sdk/bundle/bundlev1/registry.go | 36 +- go-sdk/bundle/bundlev1/schemas.go | 4 +- go-sdk/bundle/bundlev1/spec.gen.go | 143 ++++++-- go-sdk/bundle/bundlev1/spec.go | 63 +--- go-sdk/example/bundle/main.go | 35 +- go-sdk/pkg/execution/dag_parser.go | 5 +- go-sdk/pkg/execution/integration_test.go | 19 +- go-sdk/pkg/execution/serde.go | 123 +------ go-sdk/pkg/execution/serde_test.go | 102 +++++- go-sdk/pkg/execution/server.go | 7 +- 13 files changed, 597 insertions(+), 391 deletions(-) diff --git a/go-sdk/Justfile b/go-sdk/Justfile index 9417e24b377cb..9923daaea477e 100644 --- a/go-sdk/Justfile +++ b/go-sdk/Justfile @@ -43,6 +43,6 @@ docs port="6060": generate-models: go generate ./pkg/execution/genmodels/... -# Regenerate the bundle spec structs from the Airflow dag-serialization schema +# Regenerate bundle specs generate-specs: go generate ./bundle/bundlev1/... diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go index 59846166bd759..2a918c1d5b0e7 100644 --- a/go-sdk/bundle/bundlev1/gen.go +++ b/go-sdk/bundle/bundlev1/gen.go @@ -15,16 +15,8 @@ // 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 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. +// spec.gen.go is generated from the Dag serialization schema. +// Edit ./gen, then run `just generate-specs`; do not edit the output. 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 index c1714dc0c3348..ca156a417b2c2 100644 --- a/go-sdk/bundle/bundlev1/gen/main.go +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -15,34 +15,7 @@ // 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 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 -// 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. +// Command gen generates bundle specs from the Dag serialization schema. package main import ( @@ -57,33 +30,97 @@ import ( "strings" ) -// 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", +type behaviorKind int + +const ( + behaviorAlwaysEmit behaviorKind = iota + 1 + behaviorEmitWhenSet +) + +type fieldBehavior struct { + kind behaviorKind + zeroFallback string + fallbackDoc string +} + +type extraField struct { + goName string + goType string + doc []string +} + +type specConfig struct { + defName string + typeName string + structDoc string + excluded map[string]string + typeOverrides map[string]string + behaviors map[string]fieldBehavior + extraFields []extraField } -// 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", +var specConfigs = []specConfig{ + { + defName: "operator", + typeName: "TaskSpec", + structDoc: `TaskSpec configures a registered task. Zero values use schema defaults.`, + excluded: 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", + }, + // Python treats this integral schema default as a float multiplier. + typeOverrides: map[string]string{ + "retry_exponential_backoff": "float64", + }, + }, + { + defName: "dag", + typeName: "DagSpec", + structDoc: `DagSpec configures a registered Dag. Zero values use schema defaults.`, + excluded: map[string]string{ + "relative_fileloc": "computed by the serializer from fileloc and the bundle path", + }, + // Python stores tags as a set. + typeOverrides: map[string]string{ + "tags": "[]string", + }, + behaviors: map[string]fieldBehavior{ + "catchup": {kind: behaviorAlwaysEmit}, + "disable_bundle_versioning": {kind: behaviorAlwaysEmit}, + "max_consecutive_failed_dag_runs": {kind: behaviorAlwaysEmit}, + "max_active_tasks": { + kind: behaviorAlwaysEmit, + zeroFallback: "16", + fallbackDoc: "[core] max_active_tasks_per_dag", + }, + "max_active_runs": { + kind: behaviorAlwaysEmit, + zeroFallback: "16", + fallbackDoc: "[core] max_active_runs_per_dag", + }, + "is_paused_upon_creation": {kind: behaviorEmitWhenSet}, + }, + extraFields: []extraField{ + { + goName: "Schedule", + goType: "string", + doc: []string{ + `Schedule accepts "@once", "@continuous", a cron expression, or "".`, + }, + }, + }, + }, } -// initialisms maps snake_case segments that must keep non-Title capitalization -// in Go names (do_xcom_push -> DoXComPush, doc_md -> DocMD). +// initialisms preserves Go capitalization in schema names. var initialisms = map[string]string{ "id": "ID", "md": "MD", @@ -98,8 +135,7 @@ 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. +// orderedProps preserves schema field order. type orderedProps struct { keys []string props map[string]propSchema @@ -136,17 +172,20 @@ type schemaDoc struct { Definitions map[string]definition `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 + def any + hasDef bool refName string } +type resolvedSpec struct { + cfg specConfig + fields []field +} + 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") @@ -164,17 +203,21 @@ func main() { 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, err := selectFields(operator) - if err != nil { - log.Fatalf("gen: %v", err) + specs := make([]resolvedSpec, 0, len(specConfigs)) + for _, cfg := range specConfigs { + def, ok := doc.Definitions[cfg.defName] + if !ok { + log.Fatalf("gen: schema has no %q definition", cfg.defName) + } + fields, err := selectFields(cfg, def) + if err != nil { + log.Fatalf("gen: %s: %v", cfg.typeName, err) + } + specs = append(specs, resolvedSpec{cfg: cfg, fields: fields}) } - src, err := render(*pkg, fields) + src, err := render(*pkg, specs) if err != nil { log.Fatalf("gen: rendering: %v", err) } @@ -183,50 +226,99 @@ func main() { } } -// 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 { +func selectFields(cfg specConfig, def definition) ([]field, error) { + required := make(map[string]bool, len(def.Required)) + for _, key := range def.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 { + excludedSeen := make(map[string]bool, len(cfg.excluded)) + fields := make([]field, 0, len(def.Properties.keys)) + for _, key := range def.Properties.keys { if strings.HasPrefix(key, "_") || required[key] || strings.HasPrefix(key, "has_on_") { continue } - if _, ok := excludedFields[key]; ok { + if _, ok := cfg.excluded[key]; ok { excludedSeen[key] = true continue } - f, ok := resolveField(key, op.Properties.props[key]) + f, ok := resolveField(cfg, key, def.Properties.props[key]) if !ok { continue } fields = append(fields, f) } - for key := range excludedFields { + for key := range cfg.excluded { if !excludedSeen[key] { return nil, fmt.Errorf( - "excludedFields entry %q matches no eligible schema property; remove or fix it", + "excluded entry %q matches no eligible schema property; remove or fix it", key, ) } } - for key := range goTypeOverrides { + for key := range cfg.typeOverrides { if !containsKey(fields, key) { return nil, fmt.Errorf( - "goTypeOverrides entry %q matches no generated field; remove or fix it", + "typeOverrides entry %q matches no generated field; remove or fix it", key, ) } } + if err := validateBehaviors(cfg, fields); err != nil { + return nil, err + } return fields, nil } +func validateBehaviors(cfg specConfig, fields []field) error { + byKey := make(map[string]field, len(fields)) + for _, f := range fields { + byKey[f.key] = f + } + for key, beh := range cfg.behaviors { + f, ok := byKey[key] + if !ok { + return fmt.Errorf( + "behaviors entry %q matches no generated field; remove or fix it", + key, + ) + } + switch beh.kind { + case behaviorAlwaysEmit: + if strings.HasPrefix(f.goType, "*") || strings.HasPrefix(f.goType, "[]") { + return fmt.Errorf( + "behaviors entry %q: always-emit requires a non-pointer scalar field, got %s", + key, + f.goType, + ) + } + if beh.zeroFallback != "" && f.goType != "int" && f.goType != "float64" { + return fmt.Errorf( + "behaviors entry %q: zeroFallback requires a numeric field, got %s", + key, + f.goType, + ) + } + case behaviorEmitWhenSet: + if beh.zeroFallback != "" { + return fmt.Errorf( + "behaviors entry %q: zeroFallback is only valid with always-emit", + key, + ) + } + if f.goType != "*bool" { + return fmt.Errorf( + "behaviors entry %q: emit-when-set requires a boolean key, got %s", + key, + f.goType, + ) + } + } + } + return nil +} + func containsKey(fields []field, key string) bool { for _, f := range fields { if f.key == key { @@ -236,10 +328,8 @@ func containsKey(fields []field, key string) bool { 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) { +// resolveField rejects unsupported schema properties. +func resolveField(cfg specConfig, key string, prop propSchema) (field, bool) { f := field{ key: key, goName: goName(key), @@ -251,8 +341,8 @@ func resolveField(key string, prop propSchema) (field, bool) { } switch { - case goTypeOverrides[key] != "": - f.goType = goTypeOverrides[key] + case cfg.typeOverrides[key] != "": + f.goType = cfg.typeOverrides[key] case f.refName == "timedelta": f.goType = "time.Duration" case f.refName == "datetime": @@ -272,14 +362,17 @@ func resolveField(key string, prop propSchema) (field, bool) { 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. + // Preserve false against a true default. f.goType = "*bool" } default: return field{}, false } } + if beh, ok := cfg.behaviors[key]; ok && beh.kind == behaviorEmitWhenSet && f.goType == "bool" { + // Preserve explicit false. + f.goType = "*bool" + } return f, true } @@ -313,48 +406,129 @@ const licenseHeader = `// Licensed to the Apache Software Foundation (ASF) under // under the License. ` -func render(pkg string, fields []field) ([]byte, error) { +func render(pkg string, specs []resolvedSpec) ([]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)) + writeImports(&b, specs) + + for i, sp := range specs { + if i > 0 { + b.WriteString("\n") + } + renderSpec(&b, sp) + } + return format.Source(b.Bytes()) +} + +func writeImports(b *bytes.Buffer, specs []resolvedSpec) { + usesTime, usesSlices := false, false + for _, sp := range specs { + for _, f := range sp.fields { + switch f.goType { + case "time.Duration", "time.Time": + usesTime = true + case "[]string": + usesSlices = true + } } - b.WriteString(".\n") - fmt.Fprintf(&b, "\t%s %s\n", f.goName, f.goType) + } + switch { + case usesTime && usesSlices: + b.WriteString("import (\n\t\"slices\"\n\t\"time\"\n)\n\n") + case usesTime: + b.WriteString("import \"time\"\n\n") + case usesSlices: + b.WriteString("import \"slices\"\n\n") + } +} + +func renderSpec(b *bytes.Buffer, sp resolvedSpec) { + cfg := sp.cfg + for _, line := range strings.Split(cfg.structDoc, "\n") { + fmt.Fprintf(b, "// %s\n", line) + } + fmt.Fprintf(b, "type %s struct {\n", cfg.typeName) + for _, ef := range cfg.extraFields { + for _, line := range ef.doc { + fmt.Fprintf(b, "\t// %s\n", line) + } + fmt.Fprintf(b, "\t%s %s\n", ef.goName, ef.goType) + } + for _, f := range sp.fields { + writeFieldDoc(b, cfg, f) + 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("// SchemaFields returns set fields that differ from schema defaults.\n") + fmt.Fprintf( + b, + "func (s %s) SchemaFields() map[string]any {\n\tm := map[string]any{}\n", + cfg.typeName, + ) + for _, f := range sp.fields { + emitCondition(b, cfg, f) } b.WriteString("\treturn m\n}\n") +} - return format.Source(b.Bytes()) +func writeFieldDoc(b *bytes.Buffer, cfg specConfig, f field) { + text := fmt.Sprintf("%s maps to the schema key %q", f.goName, f.key) + if f.hasDef { + text += fmt.Sprintf(" (schema default %s)", defaultDoc(f)) + } + for _, line := range wrapComment(text+fieldDocExtra(cfg, f)+".", 72) { + fmt.Fprintf(b, "\t// %s\n", line) + } +} + +func fieldDocExtra(cfg specConfig, f field) string { + if beh, ok := cfg.behaviors[f.key]; ok { + switch beh.kind { + case behaviorAlwaysEmit: + extra := "; always serialized" + if beh.zeroFallback != "" { + extra += fmt.Sprintf( + ", defaulting to %s from %s", + beh.zeroFallback, + beh.fallbackDoc, + ) + } + return extra + case behaviorEmitWhenSet: + return `; nil omits it` + } + } + if f.goType == "[]string" { + return "; serialized sorted" + } + return "" +} + +func wrapComment(text string, width int) []string { + var lines []string + cur := "" + for _, word := range strings.Fields(text) { + switch { + case cur == "": + cur = word + case len(cur)+1+len(word) > width: + lines = append(lines, cur) + cur = word + default: + cur += " " + word + } + } + if cur != "" { + lines = append(lines, cur) + } + return lines } -// 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) @@ -365,11 +539,40 @@ func defaultDoc(f field) string { 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) { +func emitCondition(b *bytes.Buffer, cfg specConfig, f field) { name := "s." + f.goName key := f.key + if beh, ok := cfg.behaviors[key]; ok { + switch beh.kind { + case behaviorAlwaysEmit: + if beh.zeroFallback == "" { + fmt.Fprintf(b, "\tm[%q] = %s\n", key, name) + return + } + fmt.Fprintf( + b, + "\tif %s != 0 {\n\t\tm[%q] = %s\n\t} else {\n\t\tm[%q] = %s // %s\n\t}\n", + name, + key, + name, + key, + beh.zeroFallback, + beh.fallbackDoc, + ) + case behaviorEmitWhenSet: + fmt.Fprintf(b, "\tif %s != nil {\n\t\tm[%q] = *%s\n\t}\n", name, key, name) + } + return + } switch { + case f.goType == "[]string": + fmt.Fprintf( + b, + "\tif len(%s) > 0 {\n\t\tm[%q] = slices.Sorted(slices.Values(%s))\n\t}\n", + name, + key, + name, + ) 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": @@ -412,8 +615,6 @@ func emitCondition(b *bytes.Buffer, f field) { } } -// 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)) diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index db8fed5ad1a37..e1d481f09a6fa 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -29,12 +29,8 @@ 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. + // AddTask registers fn under its Go name. spec configures the task; + // depends names already-registered upstream tasks. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -49,10 +45,10 @@ type ( // 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 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) } @@ -61,16 +57,11 @@ 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 registers dagId with one optional spec. Duplicates and extra specs panic. AddDag(dagId string, spec ...DagSpec) Dag } - // 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 exposes registered Dag and task metadata. EnumerableBundle interface { OrderedDags() []DagInfo } @@ -124,8 +115,7 @@ func New() Registry { } func splitFullName(fullName string) (typeName, pkgPath string) { - // fullName looks like "main.extract" or "github.com/x/y.MyTask"; method - // values get a "-fm" suffix. + // Method values end in "-fm". lastDot := strings.LastIndex(fullName, ".") if lastDot < 0 { return strings.TrimSuffix(fullName, "-fm"), "" @@ -197,9 +187,7 @@ func (r *registry) registerTaskWithName( 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. + // Record one downstream edge per dependency. seen := make(map[string]bool, len(depends)) for _, dep := range depends { if dep == taskId { @@ -237,9 +225,7 @@ func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { return task, exists } -// 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. +// OrderedDags returns a mutable copy of Dags and tasks in registration order. func (r *registry) OrderedDags() []DagInfo { r.RLock() defer r.RUnlock() diff --git a/go-sdk/bundle/bundlev1/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index a9e2fc7350dd9..89be2768bc740 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -31,8 +31,8 @@ type BundleProvider interface { // // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { // dag := dagbag.AddDag("simple_dag") - // dag.AddTask(extract) - // dag.AddTask(transform) + // dag.AddTask(extract, bundlev1.TaskSpec{}, nil) + // dag.AddTask(transform, bundlev1.TaskSpec{}, nil) // return nil // } // diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go index 9dcf4ee832584..752486fe80cd3 100644 --- a/go-sdk/bundle/bundlev1/spec.gen.go +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -18,13 +18,12 @@ package bundlev1 -import "time" +import ( + "slices" + "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. +// TaskSpec configures a registered task. Zero values use schema defaults. type TaskSpec struct { // Owner maps to the schema key "owner" (schema default "airflow"). Owner string @@ -32,15 +31,20 @@ type TaskSpec struct { 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 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 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 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 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 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 @@ -52,23 +56,29 @@ type TaskSpec struct { PoolSlots 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 maps to the schema key "retry_delay" (schema default 300 + // seconds). RetryDelay time.Duration - // RetryExponentialBackoff maps to the schema key "retry_exponential_backoff" (schema default 0). + // 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 maps to the schema key "priority_weight" (schema default + // 1). PriorityWeight int - // WeightRule maps to the schema key "weight_rule" (schema default "downstream"). + // WeightRule maps to the schema key "weight_rule" (schema default + // "downstream"). WeightRule string // Executor maps to the schema key "executor". Executor string // 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 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 maps to the schema key "email_on_retry" (schema default + // true). EmailOnRetry *bool // DocMD maps to the schema key "doc_md". DocMD string @@ -76,14 +86,12 @@ type TaskSpec struct { 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 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. +// SchemaFields returns set fields that differ from schema defaults. func (s TaskSpec) SchemaFields() map[string]any { m := map[string]any{} if s.Owner != "" && s.Owner != "airflow" { @@ -166,3 +174,94 @@ func (s TaskSpec) SchemaFields() map[string]any { } return m } + +// DagSpec configures a registered Dag. Zero values use schema defaults. +type DagSpec struct { + // Schedule accepts "@once", "@continuous", a cron expression, or "". + Schedule string + // Catchup maps to the schema key "catchup"; always serialized. + Catchup bool + // FailFast maps to the schema key "fail_fast" (schema default false). + FailFast bool + // DagDisplayName maps to the schema key "dag_display_name". + DagDisplayName string + // Description maps to the schema key "description". + Description string + // MaxActiveTasks maps to the schema key "max_active_tasks"; always + // serialized, defaulting to 16 from [core] max_active_tasks_per_dag. + MaxActiveTasks int + // MaxActiveRuns maps to the schema key "max_active_runs"; always + // serialized, defaulting to 16 from [core] max_active_runs_per_dag. + MaxActiveRuns int + // MaxConsecutiveFailedDagRuns maps to the schema key + // "max_consecutive_failed_dag_runs"; always serialized. + MaxConsecutiveFailedDagRuns int + // StartDate maps to the schema key "start_date". + StartDate time.Time + // EndDate maps to the schema key "end_date". + EndDate time.Time + // DagrunTimeout maps to the schema key "dagrun_timeout". + DagrunTimeout time.Duration + // DocMD maps to the schema key "doc_md". + DocMD string + // IsPausedUponCreation maps to the schema key "is_paused_upon_creation"; + // nil omits it. + IsPausedUponCreation *bool + // RenderTemplateAsNativeObj maps to the schema key + // "render_template_as_native_obj" (schema default false). + RenderTemplateAsNativeObj bool + // Tags maps to the schema key "tags"; serialized sorted. + Tags []string + // DisableBundleVersioning maps to the schema key + // "disable_bundle_versioning"; always serialized. + DisableBundleVersioning bool +} + +// SchemaFields returns set fields that differ from schema defaults. +func (s DagSpec) SchemaFields() map[string]any { + m := map[string]any{} + m["catchup"] = s.Catchup + if s.FailFast { + m["fail_fast"] = s.FailFast + } + if s.DagDisplayName != "" { + m["dag_display_name"] = s.DagDisplayName + } + if s.Description != "" { + m["description"] = s.Description + } + if s.MaxActiveTasks != 0 { + m["max_active_tasks"] = s.MaxActiveTasks + } else { + m["max_active_tasks"] = 16 // [core] max_active_tasks_per_dag + } + if s.MaxActiveRuns != 0 { + m["max_active_runs"] = s.MaxActiveRuns + } else { + m["max_active_runs"] = 16 // [core] max_active_runs_per_dag + } + m["max_consecutive_failed_dag_runs"] = s.MaxConsecutiveFailedDagRuns + if !s.StartDate.IsZero() { + m["start_date"] = s.StartDate + } + if !s.EndDate.IsZero() { + m["end_date"] = s.EndDate + } + if s.DagrunTimeout != 0 { + m["dagrun_timeout"] = s.DagrunTimeout + } + if s.DocMD != "" { + m["doc_md"] = s.DocMD + } + if s.IsPausedUponCreation != nil { + m["is_paused_upon_creation"] = *s.IsPausedUponCreation + } + if s.RenderTemplateAsNativeObj { + m["render_template_as_native_obj"] = s.RenderTemplateAsNativeObj + } + if len(s.Tags) > 0 { + m["tags"] = slices.Sorted(slices.Values(s.Tags)) + } + m["disable_bundle_versioning"] = s.DisableBundleVersioning + return m +} diff --git a/go-sdk/bundle/bundlev1/spec.go b/go-sdk/bundle/bundlev1/spec.go index d9af9e1316f72..17a61dd7487b6 100644 --- a/go-sdk/bundle/bundlev1/spec.go +++ b/go-sdk/bundle/bundlev1/spec.go @@ -17,76 +17,31 @@ 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 describes a registered task. TaskInfo struct { - // ID is the user-visible task id (the function name unless overridden - // via AddTaskWithName). + // ID is the task ID. ID string - // TypeName is the unqualified Go function name (e.g. "extract"). + // TypeName is the unqualified Go function name. TypeName string - // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). + // PkgPath is the Go package path. PkgPath string - // Spec carries the optional per-task configuration supplied at - // registration. The zero value means "no overrides". + // Spec configures the task. 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 lists dependent task IDs in registration order. Downstream []string } - // DagInfo describes a registered dag together with its tasks in - // registration order. + // DagInfo describes a registered Dag and its tasks. DagInfo struct { DagID string - // Spec carries the optional per-dag configuration supplied at - // registration. The zero value means "no overrides". + // Spec configures the Dag. 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)} +// Bool returns a pointer for nullable spec fields. func Bool(b bool) *bool { return &b } diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 02b82ba67740d..59596dead4340 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -56,16 +56,31 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { ) bindingDag := dagbag.AddDag("taskflow_binding_dag") - bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) - bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) - bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion) - bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) - bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) - bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) - bindingDag.AddTaskWithName("via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg) - bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap) - bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap) - bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap) + bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName( + "via_struct_no_tags", + taskflowbinding.ViaStructNoTags, + v1.TaskSpec{}, + nil, + ) + bindingDag.AddTaskWithName( + "via_struct_arg_tag", + taskflowbinding.ViaStructArgTag, + v1.TaskSpec{}, + nil, + ) + bindingDag.AddTaskWithName( + "via_struct_unmatched_arg", + taskflowbinding.ViaStructUnmatchedArg, + v1.TaskSpec{}, + nil, + ) + bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap, v1.TaskSpec{}, nil) return nil } diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go index 4b6836ab4d4b4..177f2788119ec 100644 --- a/go-sdk/pkg/execution/dag_parser.go +++ b/go-sdk/pkg/execution/dag_parser.go @@ -22,10 +22,7 @@ import ( "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. +// ParseDags serializes registered Dags into a DagFileParsingResult. func ParseDags(bundle bundlev1.Bundle, req *genmodels.DagFileParseRequest) map[string]any { fileloc := req.File bundlePath := req.BundlePath diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index ba2c9b0ee2dab..8c87781a163da 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -260,7 +260,7 @@ func TestTaskRunnerBindsArgs(t *testing.T) { gotCountry = country gotMeta = meta return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -295,7 +295,7 @@ func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { func(country string, meta map[string]any) error { ran = true return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -327,7 +327,7 @@ func TestTaskRunnerBindsStructArgs(t *testing.T) { func(input regionInput) error { got = input return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -355,7 +355,7 @@ func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { func(input regionInput) error { got = input return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -386,7 +386,7 @@ func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", - func(count int) error { return nil }) + func(count int) error { return nil }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -413,7 +413,7 @@ func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { func(country string) error { ran = true return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -436,7 +436,7 @@ func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { func(country string) error { ran = true return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -485,7 +485,7 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { func(country string) error { ran = true return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails("transform", tc.spec) @@ -503,7 +503,7 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { return nil }) + func(country string) error { return nil }, bundlev1.TaskSpec{}, nil) }) details := newStartupDetails( @@ -695,7 +695,6 @@ func TestServeDagFileParseEndToEnd(t *testing.T) { 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", diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go index b323688a68ab2..624c5a9b66ed2 100644 --- a/go-sdk/pkg/execution/serde.go +++ b/go-sdk/pkg/execution/serde.go @@ -27,21 +27,7 @@ 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 -// - 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 +// serializeValue matches Python's type/var encoding. func serializeValue(value any) any { if value == nil { return nil @@ -96,8 +82,6 @@ func serializeValue(value any) any { } 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: @@ -121,8 +105,7 @@ func serializeValue(value any) any { } } -// unwrapTypeEncoding extracts the "__var" part from a type-encoded value. -// In Python's serialize_to_json, non-decorated fields are serialized then unwrapped. +// unwrapTypeEncoding strips type/var metadata. func unwrapTypeEncoding(value any) any { m, ok := value.(map[string]any) if !ok { @@ -138,16 +121,7 @@ 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. +// TODO: honor scheduler interval flags once the supervisor sends them. // Tracked at https://github.com/apache/airflow/issues/67938 func serializeTimetable(schedule *string) map[string]any { if schedule == nil { @@ -180,9 +154,7 @@ func serializeTimetable(schedule *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. +// serializeTask builds a serialized task. func serializeTask(info bundlev1.TaskInfo) map[string]any { typeName := info.TypeName if typeName == "" { @@ -197,17 +169,9 @@ 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. + // Python always emits this; Go tasks have no template fields. "template_fields": []any{}, } - // 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)) } @@ -223,72 +187,10 @@ func serializeTask(info bundlev1.TaskInfo) map[string]any { } } -// 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. +// applyDagSpec writes serialized Dag fields. 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 { - // 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 - } - if s.DagDisplayName != "" { - data["dag_display_name"] = s.DagDisplayName - } - if s.DocMD != "" { - data["doc_md"] = s.DocMD - } - // 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)) - } - // 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.IsPausedUponCreation != nil { - data["is_paused_upon_creation"] = *s.IsPausedUponCreation + for key, value := range s.SchemaFields() { + data[key] = unwrapTypeEncoding(serializeValue(value)) } } @@ -313,7 +215,7 @@ func serializeTaskGroup(taskIDs []string) map[string]any { } } -// serializeParams converts DAG params to Airflow's serialization format. +// serializeParams builds serialized Dag params. func serializeParams(params map[string]any) []any { if len(params) == 0 { return []any{} @@ -334,9 +236,7 @@ func serializeParams(params map[string]any) []any { return result } -// SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 -// format. Required fields are always present; spec-driven fields are emitted -// per the rules in applyDagSpec (some always, some only when set). +// SerializeDag converts DagInfo to DagSerialization v3. func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { taskIDs := make([]string, len(info.Tasks)) tasks := make([]any, len(info.Tasks)) @@ -352,7 +252,6 @@ func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[st } result := map[string]any{ - // Required fields (always present) "dag_id": info.DagID, "fileloc": fileloc, "relative_fileloc": relativeFileloc, @@ -370,7 +269,7 @@ func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[st return result } -// computeRelativeFileloc computes the relative file location from the bundle path. +// computeRelativeFileloc returns a path relative to the bundle. func computeRelativeFileloc(fileloc, bundlePath string) string { if fileloc == "" { return "" diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go index 83e7df6a67751..beb675236eb7e 100644 --- a/go-sdk/pkg/execution/serde_test.go +++ b/go-sdk/pkg/execution/serde_test.go @@ -131,7 +131,6 @@ 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") @@ -235,7 +234,6 @@ func TestTaskSpecSchemaFields_EmitsAndOmits(t *testing.T) { } func TestTaskSpecSchemaFields_OmitsSchemaDefaults(t *testing.T) { - // Values equal to schema defaults must be dropped. spec := bundlev1.TaskSpec{ Queue: "default", Pool: "default_pool", @@ -301,18 +299,15 @@ 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) _, 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, 16, result["max_active_tasks"]) + assert.Equal(t, 16, result["max_active_runs"]) assert.Equal(t, 0, result["max_consecutive_failed_dag_runs"]) } @@ -385,7 +380,6 @@ func TestSerializeDagWithSpec(t *testing.T) { assert.Equal(t, "@daily", v["expression"]) assert.Equal(t, "Extract, transform, load", result["description"]) - // 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"]) @@ -399,25 +393,20 @@ func TestSerializeDagWithSpec(t *testing.T) { 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_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, 16, data["max_active_tasks"]) + assert.Equal(t, 16, 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", @@ -429,8 +418,6 @@ func TestApplyDagSpec_AlwaysEmitsNonSchemaDefaultFields(t *testing.T) { } 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, bundlev1.DagSpec{FailFast: false, RenderTemplateAsNativeObj: false}) _, hasFailFast := data["fail_fast"] @@ -439,6 +426,87 @@ func TestApplyDagSpec_OmitsFalseSchemaDefaultBooleans(t *testing.T) { assert.False(t, hasNative) } +func TestDagSpecSchemaFields_EmitsAllSetFields(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + spec := bundlev1.DagSpec{ + Schedule: "@daily", + Catchup: true, + FailFast: true, + DagDisplayName: "ETL Pipeline", + Description: "Extract, transform, load", + MaxActiveTasks: 32, + MaxActiveRuns: 4, + MaxConsecutiveFailedDagRuns: 3, + StartDate: start, + EndDate: end, + DagrunTimeout: 2 * time.Hour, + DocMD: "## ETL", + IsPausedUponCreation: bundlev1.Bool(true), + RenderTemplateAsNativeObj: true, + Tags: []string{"prod", "etl"}, + DisableBundleVersioning: true, + } + m := spec.SchemaFields() + + assert.Equal(t, true, m["catchup"]) + assert.Equal(t, true, m["fail_fast"]) + assert.Equal(t, "ETL Pipeline", m["dag_display_name"]) + assert.Equal(t, "Extract, transform, load", m["description"]) + assert.Equal(t, 32, m["max_active_tasks"]) + assert.Equal(t, 4, m["max_active_runs"]) + assert.Equal(t, 3, m["max_consecutive_failed_dag_runs"]) + assert.Equal(t, start, m["start_date"]) + assert.Equal(t, end, m["end_date"]) + assert.Equal(t, 2*time.Hour, m["dagrun_timeout"]) + assert.Equal(t, "## ETL", m["doc_md"]) + assert.Equal(t, true, m["is_paused_upon_creation"]) + assert.Equal(t, true, m["render_template_as_native_obj"]) + assert.Equal(t, []string{"etl", "prod"}, m["tags"]) + assert.Equal(t, true, m["disable_bundle_versioning"]) + _, hasSchedule := m["schedule"] + assert.False(t, hasSchedule) +} + +func TestDagSpecSchemaFields_EmptySpecEmitsOnlyAlwaysKeys(t *testing.T) { + assert.Equal(t, map[string]any{ + "catchup": false, + "disable_bundle_versioning": false, + "max_active_tasks": 16, + "max_active_runs": 16, + "max_consecutive_failed_dag_runs": 0, + }, bundlev1.DagSpec{}.SchemaFields()) +} + +func TestDagSpecSchemaFields_NullableIsPausedUponCreation(t *testing.T) { + tests := []struct { + name string + value *bool + want any // nil means the key must be absent + }{ + {"nil is omitted", nil, nil}, + {"explicit false serializes", bundlev1.Bool(false), false}, + {"explicit true serializes", bundlev1.Bool(true), true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := bundlev1.DagSpec{IsPausedUponCreation: tt.value}.SchemaFields()["is_paused_upon_creation"] + if tt.want == nil { + assert.False(t, ok) + return + } + assert.Equal(t, tt.want, got) + }) + } +} + +func TestDagSpecSchemaFields_TagsSortedCopy(t *testing.T) { + tags := []string{"prod", "etl", "hourly"} + m := bundlev1.DagSpec{Tags: tags}.SchemaFields() + assert.Equal(t, []string{"etl", "hourly", "prod"}, m["tags"]) + assert.Equal(t, []string{"prod", "etl", "hourly"}, tags, "input slice must not be mutated") +} + func TestComputeRelativeFileloc(t *testing.T) { tests := []struct { fileloc string diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 2bcc2d2a21cb2..4bb8992583b87 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,11 +20,7 @@ // the Airflow supervisor (Python ExecutableCoordinator), bundlev1server.Serve // dispatches here. // -// 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. +// The first frame is either DagFileParseRequest or StartupDetails. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -169,7 +165,6 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { 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) From 0245dcbd7fdb75ff821fc5a6fa7bc17981831c7c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 20 Aug 2026 08:44:51 +0000 Subject: [PATCH 4/7] Go-SDK: Enforce native Dag execution settings Native bundle registration bypasses Python-side validation, so invalid scheduling, identifiers, and task execution flags must be enforced in the Go runtime. --- go-sdk/bundle/bundlev1/registry.go | 57 ++++++++++++- go-sdk/bundle/bundlev1/registry_test.go | 90 +++++++++++++++++++- go-sdk/bundle/bundlev1/task.go | 27 ++++-- kubernetes-tests/lang_sdk/go_example/main.go | 4 +- 4 files changed, 166 insertions(+), 12 deletions(-) diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index e1d481f09a6fa..ff87be09a7b79 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,6 +23,8 @@ import ( "runtime" "strings" "sync" + "unicode" + "unicode/utf8" ) type ( @@ -129,8 +131,57 @@ func getFnName(fn reflect.Value) string { return name } +const maxIDLength = 250 + +func validateID(kind, id string) { + if length := utf8.RuneCountInString(id); length > maxIDLength { + panic( + fmt.Errorf( + "%s %q must be at most %d characters, got %d", + kind, + id, + maxIDLength, + length, + ), + ) + } + for _, r := range id { + if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' || r == '-' || r == '.' { + continue + } + panic( + fmt.Errorf( + "%s %q may contain only letters, numbers, dashes, dots, and underscores", + kind, + id, + ), + ) + } + if id == "" { + panic(fmt.Errorf("%s must not be empty", kind)) + } +} + +func normalizeDagSpec(dagID string, spec DagSpec) DagSpec { + if spec.Schedule != "@continuous" { + return spec + } + if spec.MaxActiveRuns > 1 { + panic(fmt.Errorf( + "Dag %q uses @continuous and requires MaxActiveRuns <= 1, got %d", + dagID, + spec.MaxActiveRuns, + )) + } + if spec.MaxActiveRuns == 0 { + spec.MaxActiveRuns = 1 + } + return spec +} + func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { - dagSpec := optionalSpec(spec, "AddDag") + validateID("Dag ID", dagId) + dagSpec := normalizeDagSpec(dagId, optionalSpec(spec, "AddDag")) r.RWMutex.Lock() defer r.RWMutex.Unlock() if _, exists := r.taskFuncMap[dagId]; exists { @@ -161,7 +212,9 @@ func (r *registry) registerTaskWithName( spec TaskSpec, depends []string, ) { - task, err := NewTaskFunction(fn) + validateID("task ID", taskId) + doXComPush := spec.DoXComPush == nil || *spec.DoXComPush + task, err := newTaskFunction(fn, doXComPush) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) } diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index 364b7b8a8dc37..0f27ae90d5926 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -21,10 +21,12 @@ import ( "context" "errors" "log/slog" + "strings" "testing" "github.com/stretchr/testify/suite" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" ) @@ -35,7 +37,34 @@ func myTaskWithArgs(ctx context.Context, logger *slog.Logger, client sdk.Client) } return nil } -func errorTask() error { return errors.New("fail") } +func errorTask() error { return errors.New("fail") } +func resultTask() (string, error) { return "result", nil } +func nilResultTask() (*string, error) { return nil, nil } + +type recordingClient struct { + sdk.Client + values []any +} + +func (c *recordingClient) PushXCom( + _ context.Context, + _ sdk.TaskInstance, + _ string, + value any, +) error { + c.values = append(c.values, value) + return nil +} + +func taskContext(client sdk.Client) context.Context { + ctx := context.WithValue(context.Background(), sdkcontext.SdkClientContextKey, client) + ti := sdk.TaskInstance{DagID: "dag1", RunID: "run1", TaskID: "task1"} + runtimeContext := sdk.NewTIRunContext(context.Background(), ti, sdk.DagRun{ + DagID: ti.DagID, + RunID: ti.RunID, + }) + return context.WithValue(ctx, sdkcontext.RuntimeContextKey, runtimeContext) +} func NotErrorRet() int { return 0 @@ -66,6 +95,18 @@ func (s *RegistrySuite) TestAddDag_DuplicatePanics() { }) } +func (s *RegistrySuite) TestAddDag_ValidatesID() { + for name, id := range map[string]string{ + "empty": "", + "invalid-char": "bad/id", + "over-max-size": strings.Repeat("a", maxIDLength+1), + } { + s.Run(name, func() { + s.Panics(func() { New().AddDag(id) }) + }) + } +} + func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { s.dag.AddTask(myTask, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "myTask") @@ -84,6 +125,18 @@ func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { s.False(exists) } +func (s *RegistrySuite) TestAddTaskWithName_ValidatesID() { + for name, id := range map[string]string{ + "empty": "", + "invalid-char": "bad/id", + "over-max-size": strings.Repeat("a", maxIDLength+1), + } { + s.Run(name, func() { + s.Panics(func() { s.dag.AddTaskWithName(id, myTask, TaskSpec{}, nil) }) + }) + } +} + func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { @@ -177,6 +230,25 @@ func (s *RegistrySuite) TestAddTask_WithSpec() { s.False(*got.Spec.DoXComPush) } +func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { + s.dag.AddTask(resultTask, TaskSpec{DoXComPush: Bool(false)}, nil) + task, exists := s.reg.LookupTask("dag1", "resultTask") + s.Require().True(exists) + client := &recordingClient{} + s.Require().NoError(task.Execute(taskContext(client), slog.Default(), nil)) + s.Empty(client.values) +} + +func (s *RegistrySuite) TestAddTask_PushesNilResult() { + s.dag.AddTask(nilResultTask, TaskSpec{}, nil) + task, exists := s.reg.LookupTask("dag1", "nilResultTask") + s.Require().True(exists) + client := &recordingClient{} + s.Require().NoError(task.Execute(taskContext(client), slog.Default(), nil)) + s.Require().Len(client.values, 1) + s.Nil(client.values[0]) +} + func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}, nil) enum, ok := s.reg.(EnumerableBundle) @@ -274,6 +346,22 @@ func (s *RegistrySuite) TestAddDag_WithSpec() { s.Equal(4, got.Spec.MaxActiveRuns) } +func (s *RegistrySuite) TestAddDag_ContinuousDefaultsToOneActiveRun() { + s.reg.AddDag("continuous", DagSpec{Schedule: "@continuous"}) + dags := s.reg.(EnumerableBundle).OrderedDags() + s.Equal(1, dags[1].Spec.MaxActiveRuns) + s.Equal(1, dags[1].Spec.SchemaFields()["max_active_runs"]) +} + +func (s *RegistrySuite) TestAddDag_ContinuousRejectsConcurrentRuns() { + s.PanicsWithError( + `Dag "continuous" uses @continuous and requires MaxActiveRuns <= 1, got 2`, + func() { + s.reg.AddDag("continuous", DagSpec{Schedule: "@continuous", MaxActiveRuns: 2}) + }, + ) +} + 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/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index 6a14a11e4f0cb..d0c58c64360b6 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -44,18 +44,23 @@ type Bundle interface { } type taskFunction struct { - fn reflect.Value - fullName string - plan *binding.Plan + fn reflect.Value + fullName string + plan *binding.Plan + doXComPush bool } var _ Task = (*taskFunction)(nil) // NewTaskFunction validates and wraps a Go function as a Task. func NewTaskFunction(fn any) (Task, error) { + return newTaskFunction(fn, true) +} + +func newTaskFunction(fn any, doXComPush bool) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{fn: v, fullName: fullName} + f := &taskFunction{fn: v, fullName: fullName, doXComPush: doXComPush} if err := f.validateFn(v.Type()); err != nil { return nil, err } @@ -106,9 +111,17 @@ func (f *taskFunction) call( ) } } - // If there are two results, convert the first only if it's not a nil pointer - if len(retValues) > 1 && (retValues[0].Kind() != reflect.Ptr || !retValues[0].IsNil()) { - res := retValues[0].Interface() + if len(retValues) > 1 && f.doXComPush { + var res any + value := retValues[0] + switch value.Kind() { + case reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + if !value.IsNil() { + res = value.Interface() + } + default: + res = value.Interface() + } f.sendXcom(ctx, res, sdkClient, logger) } return err diff --git a/kubernetes-tests/lang_sdk/go_example/main.go b/kubernetes-tests/lang_sdk/go_example/main.go index d1de46bf50fe6..8875735b92a16 100644 --- a/kubernetes-tests/lang_sdk/go_example/main.go +++ b/kubernetes-tests/lang_sdk/go_example/main.go @@ -44,8 +44,8 @@ func (m *combinedBundle) RegisterDags(dagbag v1.Registry) error { dag := dagbag.AddDag(combinedDagID) // Explicit task ids so the Go tasks are namespaced apart from the Java // tasks that share this dag_id in the Python stub. - dag.AddTaskWithName("go_extract", goExtract) - dag.AddTaskWithName("go_transform", goTransform) + dag.AddTaskWithName("go_extract", goExtract, v1.TaskSpec{}, nil) + dag.AddTaskWithName("go_transform", goTransform, v1.TaskSpec{}, nil) return nil } From 507c36beb4113b495982e83e5875c4ce356f7c48 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 21 Aug 2026 07:15:35 +0000 Subject: [PATCH 5/7] Go-SDK: Make task specs optional in Dag authoring Dependency-only task declarations should not need an empty TaskSpec value. --- go-sdk/README.md | 4 +- go-sdk/adr/0001-bundle-packing-options.md | 2 +- go-sdk/bundle/bundlev1/registry.go | 29 +++++---- go-sdk/bundle/bundlev1/registry_test.go | 66 ++++++++++++-------- go-sdk/bundle/bundlev1/schemas.go | 4 +- go-sdk/example/bundle/main.go | 24 +++---- go-sdk/pkg/execution/integration_test.go | 48 +++++++------- kubernetes-tests/lang_sdk/go_example/main.go | 4 +- 8 files changed, 98 insertions(+), 83 deletions(-) diff --git a/go-sdk/README.md b/go-sdk/README.md index 360e185a9a546..8bad8464cfa3f 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -84,8 +84,8 @@ var _ v1.BundleProvider = (*myBundle)(nil) func (m *myBundle) RegisterDags(dagbag v1.Registry) error { simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) + simpleDag.AddTask(extract, nil) + simpleDag.AddTask(transform, []string{"extract"}) return nil } diff --git a/go-sdk/adr/0001-bundle-packing-options.md b/go-sdk/adr/0001-bundle-packing-options.md index f504b3c9abe14..36289ab2e3fb0 100644 --- a/go-sdk/adr/0001-bundle-packing-options.md +++ b/go-sdk/adr/0001-bundle-packing-options.md @@ -194,7 +194,7 @@ The SDK settles that sub-decision in favour of a *single* ### Option E: Static AST scan, no introspection Parser-only packer: walk the DAG source AST, find `dagbag.AddDag("X")` -calls and the `.AddTask(fn)` calls chained off them, and synthesise the +calls and the `.AddTask(fn, nil)` calls chained off them, and synthesise the manifest without running the binary. - **Pros:** no runtime dependency on the binary (works even if it diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index ff87be09a7b79..814a941b18f16 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -31,8 +31,9 @@ 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 under its Go name. spec configures the task; - // depends names already-registered upstream tasks. + // AddTask registers fn under its Go name. depends names + // already-registered upstream tasks; the optional spec configures the + // task. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -45,13 +46,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, depends []string) + AddTask(fn any, depends []string, 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) + AddTaskWithName(taskId string, fn any, depends []string, spec ...TaskSpec) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -83,12 +84,18 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any, spec TaskSpec, depends []string) { - d.registry.registerTask(d.dagId, fn, spec, depends) +func (d dagShim) AddTask(fn any, depends []string, spec ...TaskSpec) { + d.registry.registerTask(d.dagId, fn, depends, optionalSpec(spec, "AddTask")) } -func (d dagShim) AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) { - d.registry.registerTaskWithName(d.dagId, taskId, fn, spec, depends) +func (d dagShim) AddTaskWithName(taskId string, fn any, depends []string, spec ...TaskSpec) { + d.registry.registerTaskWithName( + d.dagId, + taskId, + fn, + depends, + optionalSpec(spec, "AddTaskWithName"), + ) } func optionalSpec[T any](specs []T, caller string) T { @@ -194,7 +201,7 @@ func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { return dagShim{dagId, r} } -func (r *registry) registerTask(dagId string, fn any, spec TaskSpec, depends []string) { +func (r *registry) registerTask(dagId string, fn any, depends []string, spec TaskSpec) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -203,14 +210,14 @@ func (r *registry) registerTask(dagId string, fn any, spec TaskSpec, depends []s fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn, spec, depends) + r.registerTaskWithName(dagId, fnName, fn, depends, spec) } func (r *registry) registerTaskWithName( dagId, taskId string, fn any, - spec TaskSpec, depends []string, + spec TaskSpec, ) { validateID("task ID", taskId) doXComPush := spec.DoXComPush == nil || *spec.DoXComPush diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index 0f27ae90d5926..47292ead7506f 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -108,14 +108,14 @@ func (s *RegistrySuite) TestAddDag_ValidatesID() { } func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask, TaskSpec{}, nil) + s.dag.AddTask(myTask, nil) 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) + s.dag.AddTaskWithName("special", myTask, nil) task, exists := s.reg.LookupTask("dag1", "special") s.True(exists) s.NotNil(task) @@ -132,26 +132,26 @@ func (s *RegistrySuite) TestAddTaskWithName_ValidatesID() { "over-max-size": strings.Repeat("a", maxIDLength+1), } { s.Run(name, func() { - s.Panics(func() { s.dag.AddTaskWithName(id, myTask, TaskSpec{}, nil) }) + s.Panics(func() { s.dag.AddTaskWithName(id, myTask, nil) }) }) } } func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("special", myTask, nil) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("special", myTask, nil) }) } func (s *RegistrySuite) TestAddTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func", TaskSpec{}, nil) + s.dag.AddTask("not a func", nil) }) } func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs, TaskSpec{}, nil) + s.dag.AddTask(myTaskWithArgs, nil) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) @@ -161,13 +161,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, TaskSpec{}, nil) + s.dag.AddTask(NotErrorRet, nil) }, ) } func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask, TaskSpec{}, nil) + s.dag.AddTask(errorTask, nil) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } @@ -183,11 +183,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, TaskSpec{}, nil) - zeta.AddTaskWithName("z2", myTask, TaskSpec{}, nil) + zeta.AddTaskWithName("z1", myTask, nil) + zeta.AddTaskWithName("z2", myTask, nil) alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask, TaskSpec{}, nil) + alpha.AddTaskWithName("a1", myTask, nil) reg.AddDag("mid") // dag with no tasks @@ -216,7 +216,7 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { } func (s *RegistrySuite) TestAddTask_WithSpec() { - s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}, nil) + s.dag.AddTask(myTask, nil, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -231,7 +231,7 @@ func (s *RegistrySuite) TestAddTask_WithSpec() { } func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { - s.dag.AddTask(resultTask, TaskSpec{DoXComPush: Bool(false)}, nil) + s.dag.AddTask(resultTask, nil, TaskSpec{DoXComPush: Bool(false)}) task, exists := s.reg.LookupTask("dag1", "resultTask") s.Require().True(exists) client := &recordingClient{} @@ -240,7 +240,7 @@ func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { } func (s *RegistrySuite) TestAddTask_PushesNilResult() { - s.dag.AddTask(nilResultTask, TaskSpec{}, nil) + s.dag.AddTask(nilResultTask, nil) task, exists := s.reg.LookupTask("dag1", "nilResultTask") s.Require().True(exists) client := &recordingClient{} @@ -250,7 +250,7 @@ func (s *RegistrySuite) TestAddTask_PushesNilResult() { } func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { - s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}, nil) + s.dag.AddTaskWithName("special", myTask, nil, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -263,9 +263,9 @@ func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { } 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"}) + s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("transform", myTask, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, []string{"transform"}) enum := s.reg.(EnumerableBundle) tasks := enum.OrderedDags()[0].Tasks @@ -279,10 +279,10 @@ func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { } 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"}) + s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("transform_a", myTask, []string{"extract"}) + s.dag.AddTaskWithName("transform_b", myTask, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, []string{"transform_a", "transform_b"}) enum := s.reg.(EnumerableBundle) tasks := enum.OrderedDags()[0].Tasks @@ -296,8 +296,8 @@ func (s *RegistrySuite) TestAddTask_FanOutFanIn() { } func (s *RegistrySuite) TestAddTask_DependsDuplicatesIgnored() { - s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) - s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract", "extract"}) + s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("load", myTask, []string{"extract", "extract"}) enum := s.reg.(EnumerableBundle) tasks := enum.OrderedDags()[0].Tasks @@ -312,14 +312,26 @@ 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"}) + s.dag.AddTaskWithName("load", myTask, []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"}) + s.dag.AddTaskWithName("self", myTask, []string{"self"}) + }) +} + +func (s *RegistrySuite) TestAddTask_TooManySpecsPanics() { + s.PanicsWithError("AddTask accepts at most one spec, got 2", func() { + s.dag.AddTask(myTask, nil, TaskSpec{}, TaskSpec{}) + }) +} + +func (s *RegistrySuite) TestAddTaskWithName_TooManySpecsPanics() { + s.PanicsWithError("AddTaskWithName accepts at most one spec, got 2", func() { + s.dag.AddTaskWithName("special", myTask, nil, TaskSpec{}, TaskSpec{}) }) } diff --git a/go-sdk/bundle/bundlev1/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index 89be2768bc740..e61b3547224e2 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -31,8 +31,8 @@ type BundleProvider interface { // // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { // dag := dagbag.AddDag("simple_dag") - // dag.AddTask(extract, bundlev1.TaskSpec{}, nil) - // dag.AddTask(transform, bundlev1.TaskSpec{}, nil) + // dag.AddTask(extract, nil) + // dag.AddTask(transform, []string{"extract"}) // return nil // } // diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 59596dead4340..d317a02e1d986 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -42,45 +42,41 @@ 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}, nil) - simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}, []string{"extract"}) - simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}, []string{"transform"}) + simpleDag.AddTask(extract, nil, v1.TaskSpec{Queue: "go-task", Retries: 2}) + simpleDag.AddTask(transform, []string{"extract"}, v1.TaskSpec{Queue: "go-task"}) + simpleDag.AddTask(load, []string{"transform"}, v1.TaskSpec{Queue: "go-task"}) // 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, ) bindingDag := dagbag.AddDag("taskflow_binding_dag") - bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig, v1.TaskSpec{}, nil) - bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers, v1.TaskSpec{}, nil) - bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion, v1.TaskSpec{}, nil) - bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig, nil) + bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers, nil) + bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion, nil) + bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs, nil) bindingDag.AddTaskWithName( "via_struct_no_tags", taskflowbinding.ViaStructNoTags, - v1.TaskSpec{}, nil, ) bindingDag.AddTaskWithName( "via_struct_arg_tag", taskflowbinding.ViaStructArgTag, - v1.TaskSpec{}, nil, ) bindingDag.AddTaskWithName( "via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg, - v1.TaskSpec{}, nil, ) - bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap, v1.TaskSpec{}, nil) - bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap, v1.TaskSpec{}, nil) - bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap, v1.TaskSpec{}, nil) + bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap, nil) + bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap, nil) + bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap, nil) return nil } diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 8c87781a163da..93fa5ca825aaa 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -108,7 +108,7 @@ func newStartupDetails( func TestDagParsing(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { d := r.AddDag("test_dag") - d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + d.AddTask(simpleTask, nil) }) req := &genmodels.DagFileParseRequest{ @@ -146,8 +146,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("dag1").AddTask(simpleTask, nil) + r.AddDag("dag2").AddTask(failingTask, nil) }) req := &genmodels.DagFileParseRequest{File: "/bundle/main.go", BundlePath: "/bundle"} @@ -165,7 +165,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("test_dag").AddTask(simpleTask, nil) }) details := newStartupDetails("simpleTask") @@ -179,7 +179,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("test_dag").AddTask(failingTask, nil) }) details := newStartupDetails("failingTask") @@ -193,7 +193,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("test_dag").AddTask(failingTask, nil) }) details := newStartupDetails("failingTask") @@ -209,7 +209,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("test_dag").AddTask(simpleTask, nil) }) details := newStartupDetails("nonexistent") @@ -223,7 +223,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("test_dag").AddTask(panicTask, nil) }) details := newStartupDetails("panicTask") @@ -237,7 +237,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("test_dag").AddTask(panicTask, nil) }) details := newStartupDetails("panicTask") @@ -260,7 +260,7 @@ func TestTaskRunnerBindsArgs(t *testing.T) { gotCountry = country gotMeta = meta return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -295,7 +295,7 @@ func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { func(country string, meta map[string]any) error { ran = true return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -327,7 +327,7 @@ func TestTaskRunnerBindsStructArgs(t *testing.T) { func(input regionInput) error { got = input return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -355,7 +355,7 @@ func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { func(input regionInput) error { got = input return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -386,7 +386,7 @@ func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", - func(count int) error { return nil }, bundlev1.TaskSpec{}, nil) + func(count int) error { return nil }, nil) }) details := newStartupDetails( @@ -413,7 +413,7 @@ func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { func(country string) error { ran = true return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -436,7 +436,7 @@ func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { func(country string) error { ran = true return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails( @@ -485,7 +485,7 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { func(country string) error { ran = true return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails("transform", tc.spec) @@ -503,7 +503,7 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { return nil }, bundlev1.TaskSpec{}, nil) + func(country string) error { return nil }, nil) }) details := newStartupDetails( @@ -522,7 +522,7 @@ func TestTaskRunnerMalformedSpecHonorsShouldRetry(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) + func(ctx context.Context) error { return ctx.Err() }, nil) }) details := newStartupDetails("ctxcheck") @@ -550,7 +550,7 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := &genmodels.StartupDetails{ @@ -610,7 +610,7 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }, bundlev1.TaskSpec{}, nil) + }, nil) }) details := newStartupDetails("ctxgrab") @@ -680,7 +680,7 @@ 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.AddTask(simpleTask, nil) return nil }, } @@ -731,7 +731,7 @@ func TestServeStartupDetailsEndToEnd(t *testing.T) { provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag("dag1").AddTask(simpleTask, nil) return nil }, } @@ -797,7 +797,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }, bundlev1.TaskSpec{}, nil) + }, nil) return nil }, } diff --git a/kubernetes-tests/lang_sdk/go_example/main.go b/kubernetes-tests/lang_sdk/go_example/main.go index 8875735b92a16..f4d6a0973d34c 100644 --- a/kubernetes-tests/lang_sdk/go_example/main.go +++ b/kubernetes-tests/lang_sdk/go_example/main.go @@ -44,8 +44,8 @@ func (m *combinedBundle) RegisterDags(dagbag v1.Registry) error { dag := dagbag.AddDag(combinedDagID) // Explicit task ids so the Go tasks are namespaced apart from the Java // tasks that share this dag_id in the Python stub. - dag.AddTaskWithName("go_extract", goExtract, v1.TaskSpec{}, nil) - dag.AddTaskWithName("go_transform", goTransform, v1.TaskSpec{}, nil) + dag.AddTaskWithName("go_extract", goExtract, nil) + dag.AddTaskWithName("go_transform", goTransform, nil) return nil } From a0bd425dc03668f8429e6d13974247a890be9191 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 21 Aug 2026 12:22:05 +0000 Subject: [PATCH 6/7] Go-SDK: Add a separate native Dag example --- .../airflow-go-pack/pack_integration_test.go | 5 +++ go-sdk/example/bundle/main.go | 36 ++++++++++++++++--- go-sdk/example/bundle/main_test.go | 30 ++++++++++++++++ 3 files changed, 67 insertions(+), 4 deletions(-) diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index ace867461e2cc..951c3d7bc9bb7 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -149,6 +149,11 @@ dags: concurrent_xcom_dag: tasks: - "pull_xcoms_concurrently" + native_dag: + tasks: + - "nativeExtract" + - "nativeTransform" + - "nativeLoad" simple_dag: tasks: - "extract" diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index d317a02e1d986..706cf13a4c9c3 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -37,14 +37,27 @@ type myBundle struct{} var _ v1.BundleProvider = (*myBundle)(nil) func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag", v1.DagSpec{ + simpleDag := dagbag.AddDag("simple_dag") + simpleDag.AddTask(extract, nil) + simpleDag.AddTask(transform, nil) + simpleDag.AddTask(load, nil) + + nativeDag := dagbag.AddDag("native_dag", v1.DagSpec{ Schedule: "@daily", Description: "Example Go-authored Dag", Tags: []string{"example", "go-sdk"}, }) - simpleDag.AddTask(extract, nil, v1.TaskSpec{Queue: "go-task", Retries: 2}) - simpleDag.AddTask(transform, []string{"extract"}, v1.TaskSpec{Queue: "go-task"}) - simpleDag.AddTask(load, []string{"transform"}, v1.TaskSpec{Queue: "go-task"}) + nativeDag.AddTask(nativeExtract, nil, v1.TaskSpec{Queue: "go-task"}) + nativeDag.AddTask( + nativeTransform, + []string{"nativeExtract"}, + v1.TaskSpec{Queue: "go-task"}, + ) + nativeDag.AddTask( + nativeLoad, + []string{"nativeTransform"}, + v1.TaskSpec{Queue: "go-task"}, + ) // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") @@ -81,6 +94,21 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { return nil } +func nativeExtract(log *slog.Logger) error { + log.Info("Extracting native Dag data") + return nil +} + +func nativeTransform(log *slog.Logger) error { + log.Info("Transforming native Dag data") + return nil +} + +func nativeLoad(log *slog.Logger) error { + log.Info("Loading native Dag data") + return nil +} + func main() { if err := bundlev1server.Serve(&myBundle{}); err != nil { log.Fatal(err) diff --git a/go-sdk/example/bundle/main_test.go b/go-sdk/example/bundle/main_test.go index 16be1e22e92e0..02cb452cedcef 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/go-sdk/example/bundle/main_test.go @@ -23,7 +23,9 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + v1 "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/sdk" ) @@ -49,6 +51,34 @@ func (m *mockVars) UnmarshalJSONVariable(ctx context.Context, key string, pointe var _ sdk.VariableClient = (*mockVars)(nil) +func TestRegisterDags(t *testing.T) { + registry := v1.New() + require.NoError(t, (&myBundle{}).RegisterDags(registry)) + + dags := registry.(v1.EnumerableBundle).OrderedDags() + require.Len(t, dags, 4) + + simpleDag := dags[0] + assert.Equal(t, "simple_dag", simpleDag.DagID) + assert.Equal(t, v1.DagSpec{}, simpleDag.Spec) + require.Len(t, simpleDag.Tasks, 3) + assert.Equal(t, "extract", simpleDag.Tasks[0].ID) + assert.Empty(t, simpleDag.Tasks[0].Downstream) + assert.Equal(t, "transform", simpleDag.Tasks[1].ID) + assert.Empty(t, simpleDag.Tasks[1].Downstream) + assert.Equal(t, "load", simpleDag.Tasks[2].ID) + assert.Empty(t, simpleDag.Tasks[2].Downstream) + + nativeDag := dags[1] + assert.Equal(t, "native_dag", nativeDag.DagID) + assert.Equal(t, "@daily", nativeDag.Spec.Schedule) + assert.Equal(t, []string{"example", "go-sdk"}, nativeDag.Spec.Tags) + require.Len(t, nativeDag.Tasks, 3) + assert.Equal(t, []string{"nativeTransform"}, nativeDag.Tasks[0].Downstream) + assert.Equal(t, []string{"nativeLoad"}, nativeDag.Tasks[1].Downstream) + assert.Empty(t, nativeDag.Tasks[2].Downstream) +} + 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. From 7a7f3368dd55c3fd03fbe3d2d25b573139dc5e14 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 21 Aug 2026 12:27:21 +0000 Subject: [PATCH 7/7] Go-SDK: Make task dependencies optional --- go-sdk/README.md | 2 +- go-sdk/adr/0001-bundle-packing-options.md | 2 +- go-sdk/bundle/bundlev1/registry.go | 73 ++++++++++++++++--- go-sdk/bundle/bundlev1/registry_test.go | 74 +++++++++++++------- go-sdk/bundle/bundlev1/schemas.go | 2 +- go-sdk/example/bundle/main.go | 26 +++---- kubernetes-tests/lang_sdk/go_example/main.go | 4 +- 7 files changed, 129 insertions(+), 54 deletions(-) diff --git a/go-sdk/README.md b/go-sdk/README.md index 8bad8464cfa3f..035470cca0ae5 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -84,7 +84,7 @@ var _ v1.BundleProvider = (*myBundle)(nil) func (m *myBundle) RegisterDags(dagbag v1.Registry) error { simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract, nil) + simpleDag.AddTask(extract) simpleDag.AddTask(transform, []string{"extract"}) return nil } diff --git a/go-sdk/adr/0001-bundle-packing-options.md b/go-sdk/adr/0001-bundle-packing-options.md index 36289ab2e3fb0..f504b3c9abe14 100644 --- a/go-sdk/adr/0001-bundle-packing-options.md +++ b/go-sdk/adr/0001-bundle-packing-options.md @@ -194,7 +194,7 @@ The SDK settles that sub-decision in favour of a *single* ### Option E: Static AST scan, no introspection Parser-only packer: walk the DAG source AST, find `dagbag.AddDag("X")` -calls and the `.AddTask(fn, nil)` calls chained off them, and synthesise the +calls and the `.AddTask(fn)` calls chained off them, and synthesise the manifest without running the binary. - **Pros:** no runtime dependency on the binary (works even if it diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 814a941b18f16..8830bcb4ec085 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -31,9 +31,9 @@ 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 under its Go name. depends names - // already-registered upstream tasks; the optional spec configures the - // task. + // AddTask registers fn under its Go name. Its optional arguments are an + // upstream task-id slice followed by a TaskSpec. Either or both may be + // omitted. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -46,13 +46,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, depends []string, spec ...TaskSpec) + AddTask(fn any, args ...any) // 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, depends []string, spec ...TaskSpec) + AddTaskWithName(taskId string, fn any, args ...any) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -84,20 +84,75 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any, depends []string, spec ...TaskSpec) { - d.registry.registerTask(d.dagId, fn, depends, optionalSpec(spec, "AddTask")) +func (d dagShim) AddTask(fn any, args ...any) { + depends, spec := taskArgs(args, "AddTask") + d.registry.registerTask(d.dagId, fn, depends, spec) } -func (d dagShim) AddTaskWithName(taskId string, fn any, depends []string, spec ...TaskSpec) { +func (d dagShim) AddTaskWithName(taskId string, fn any, args ...any) { + depends, spec := taskArgs(args, "AddTaskWithName") d.registry.registerTaskWithName( d.dagId, taskId, fn, depends, - optionalSpec(spec, "AddTaskWithName"), + spec, ) } +func taskArgs(args []any, caller string) (depends []string, spec TaskSpec) { + if len(args) > 2 { + panic(fmt.Errorf( + "%s accepts at most two optional arguments: depends []string, then TaskSpec; got %d", + caller, + len(args), + )) + } + + if len(args) == 0 { + return nil, TaskSpec{} + } + + if len(args) == 1 { + switch arg := args[0].(type) { + case nil: + return nil, TaskSpec{} + case []string: + return arg, TaskSpec{} + case TaskSpec: + return nil, arg + default: + panic(fmt.Errorf( + "%s optional argument must be depends []string or TaskSpec, got %T", + caller, + arg, + )) + } + } + + if args[0] != nil { + var ok bool + depends, ok = args[0].([]string) + if !ok { + panic(fmt.Errorf( + "%s first optional argument must be depends []string, got %T", + caller, + args[0], + )) + } + } + var ok bool + spec, ok = args[1].(TaskSpec) + if !ok { + panic(fmt.Errorf( + "%s second optional argument must be TaskSpec, got %T", + caller, + args[1], + )) + } + return depends, spec +} + func optionalSpec[T any](specs []T, caller string) T { switch len(specs) { case 0: diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index 47292ead7506f..5651bc10a92fc 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -108,14 +108,14 @@ func (s *RegistrySuite) TestAddDag_ValidatesID() { } func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask, nil) + s.dag.AddTask(myTask) task, exists := s.reg.LookupTask("dag1", "myTask") s.True(exists) s.NotNil(task) } func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { - s.dag.AddTaskWithName("special", myTask, nil) + s.dag.AddTaskWithName("special", myTask) task, exists := s.reg.LookupTask("dag1", "special") s.True(exists) s.NotNil(task) @@ -132,26 +132,26 @@ func (s *RegistrySuite) TestAddTaskWithName_ValidatesID() { "over-max-size": strings.Repeat("a", maxIDLength+1), } { s.Run(name, func() { - s.Panics(func() { s.dag.AddTaskWithName(id, myTask, nil) }) + s.Panics(func() { s.dag.AddTaskWithName(id, myTask) }) }) } } func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask, nil) + s.dag.AddTaskWithName("special", myTask) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask, nil) + s.dag.AddTaskWithName("special", myTask) }) } func (s *RegistrySuite) TestAddTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func", nil) + s.dag.AddTask("not a func") }) } func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs, nil) + s.dag.AddTask(myTaskWithArgs) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) @@ -161,13 +161,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, nil) + s.dag.AddTask(NotErrorRet) }, ) } func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask, nil) + s.dag.AddTask(errorTask) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } @@ -183,11 +183,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, nil) - zeta.AddTaskWithName("z2", myTask, nil) + zeta.AddTaskWithName("z1", myTask) + zeta.AddTaskWithName("z2", myTask) alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask, nil) + alpha.AddTaskWithName("a1", myTask) reg.AddDag("mid") // dag with no tasks @@ -216,7 +216,7 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { } func (s *RegistrySuite) TestAddTask_WithSpec() { - s.dag.AddTask(myTask, nil, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) + 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() @@ -231,7 +231,7 @@ func (s *RegistrySuite) TestAddTask_WithSpec() { } func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { - s.dag.AddTask(resultTask, nil, TaskSpec{DoXComPush: Bool(false)}) + s.dag.AddTask(resultTask, TaskSpec{DoXComPush: Bool(false)}) task, exists := s.reg.LookupTask("dag1", "resultTask") s.Require().True(exists) client := &recordingClient{} @@ -240,7 +240,7 @@ func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { } func (s *RegistrySuite) TestAddTask_PushesNilResult() { - s.dag.AddTask(nilResultTask, nil) + s.dag.AddTask(nilResultTask) task, exists := s.reg.LookupTask("dag1", "nilResultTask") s.Require().True(exists) client := &recordingClient{} @@ -250,7 +250,7 @@ func (s *RegistrySuite) TestAddTask_PushesNilResult() { } func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { - s.dag.AddTaskWithName("special", myTask, nil, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}) + s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}) enum, ok := s.reg.(EnumerableBundle) s.Require().True(ok) dags := enum.OrderedDags() @@ -263,7 +263,7 @@ func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { } func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { - s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("extract", myTask) s.dag.AddTaskWithName("transform", myTask, []string{"extract"}) s.dag.AddTaskWithName("load", myTask, []string{"transform"}) @@ -279,7 +279,7 @@ func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { } func (s *RegistrySuite) TestAddTask_FanOutFanIn() { - s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("extract", myTask) s.dag.AddTaskWithName("transform_a", myTask, []string{"extract"}) s.dag.AddTaskWithName("transform_b", myTask, []string{"extract"}) s.dag.AddTaskWithName("load", myTask, []string{"transform_a", "transform_b"}) @@ -296,7 +296,7 @@ func (s *RegistrySuite) TestAddTask_FanOutFanIn() { } func (s *RegistrySuite) TestAddTask_DependsDuplicatesIgnored() { - s.dag.AddTaskWithName("extract", myTask, nil) + s.dag.AddTaskWithName("extract", myTask) s.dag.AddTaskWithName("load", myTask, []string{"extract", "extract"}) enum := s.reg.(EnumerableBundle) @@ -324,15 +324,39 @@ func (s *RegistrySuite) TestAddTask_DependsOnSelfPanics() { } func (s *RegistrySuite) TestAddTask_TooManySpecsPanics() { - s.PanicsWithError("AddTask accepts at most one spec, got 2", func() { - s.dag.AddTask(myTask, nil, TaskSpec{}, TaskSpec{}) - }) + s.PanicsWithError( + "AddTask accepts at most two optional arguments: depends []string, then TaskSpec; got 3", + func() { + s.dag.AddTask(myTask, nil, TaskSpec{}, TaskSpec{}) + }, + ) } func (s *RegistrySuite) TestAddTaskWithName_TooManySpecsPanics() { - s.PanicsWithError("AddTaskWithName accepts at most one spec, got 2", func() { - s.dag.AddTaskWithName("special", myTask, nil, TaskSpec{}, TaskSpec{}) - }) + s.PanicsWithError( + "AddTaskWithName accepts at most two optional arguments: depends []string, then TaskSpec; got 3", + func() { + s.dag.AddTaskWithName("special", myTask, nil, TaskSpec{}, TaskSpec{}) + }, + ) +} + +func (s *RegistrySuite) TestAddTask_RejectsSpecBeforeDepends() { + s.PanicsWithError( + "AddTask first optional argument must be depends []string, got bundlev1.TaskSpec", + func() { + s.dag.AddTask(myTask, TaskSpec{}, []string{"upstream"}) + }, + ) +} + +func (s *RegistrySuite) TestAddTask_RejectsInvalidOptionalArgument() { + s.PanicsWithError( + "AddTask optional argument must be depends []string or TaskSpec, got string", + func() { + s.dag.AddTask(myTask, "upstream") + }, + ) } func (s *RegistrySuite) TestAddDag_WithSpec() { diff --git a/go-sdk/bundle/bundlev1/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index e61b3547224e2..d170a7f9f8e8c 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -31,7 +31,7 @@ type BundleProvider interface { // // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { // dag := dagbag.AddDag("simple_dag") - // dag.AddTask(extract, nil) + // dag.AddTask(extract) // dag.AddTask(transform, []string{"extract"}) // return nil // } diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 706cf13a4c9c3..01bb2a0087706 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -38,16 +38,16 @@ var _ v1.BundleProvider = (*myBundle)(nil) func (m *myBundle) RegisterDags(dagbag v1.Registry) error { simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract, nil) - simpleDag.AddTask(transform, nil) - simpleDag.AddTask(load, nil) + simpleDag.AddTask(extract) + simpleDag.AddTask(transform) + simpleDag.AddTask(load) nativeDag := dagbag.AddDag("native_dag", v1.DagSpec{ Schedule: "@daily", Description: "Example Go-authored Dag", Tags: []string{"example", "go-sdk"}, }) - nativeDag.AddTask(nativeExtract, nil, v1.TaskSpec{Queue: "go-task"}) + nativeDag.AddTask(nativeExtract, v1.TaskSpec{Queue: "go-task"}) nativeDag.AddTask( nativeTransform, []string{"nativeExtract"}, @@ -64,32 +64,28 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { concurrentDag.AddTaskWithName( "pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently, - nil, ) bindingDag := dagbag.AddDag("taskflow_binding_dag") - bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig, nil) - bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers, nil) - bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion, nil) - bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs, nil) + bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) + bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) + bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion) + bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) bindingDag.AddTaskWithName( "via_struct_no_tags", taskflowbinding.ViaStructNoTags, - nil, ) bindingDag.AddTaskWithName( "via_struct_arg_tag", taskflowbinding.ViaStructArgTag, - nil, ) bindingDag.AddTaskWithName( "via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg, - nil, ) - bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap, nil) - bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap, nil) - bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap, nil) + bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap) + bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap) + bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap) return nil } diff --git a/kubernetes-tests/lang_sdk/go_example/main.go b/kubernetes-tests/lang_sdk/go_example/main.go index f4d6a0973d34c..d1de46bf50fe6 100644 --- a/kubernetes-tests/lang_sdk/go_example/main.go +++ b/kubernetes-tests/lang_sdk/go_example/main.go @@ -44,8 +44,8 @@ func (m *combinedBundle) RegisterDags(dagbag v1.Registry) error { dag := dagbag.AddDag(combinedDagID) // Explicit task ids so the Go tasks are namespaced apart from the Java // tasks that share this dag_id in the Python stub. - dag.AddTaskWithName("go_extract", goExtract, nil) - dag.AddTaskWithName("go_transform", goTransform, nil) + dag.AddTaskWithName("go_extract", goExtract) + dag.AddTaskWithName("go_transform", goTransform) return nil }