diff --git a/go-sdk/Justfile b/go-sdk/Justfile index 7094a394835db..9923daaea477e 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 bundle specs +generate-specs: + go generate ./bundle/bundlev1/... diff --git a/go-sdk/README.md b/go-sdk/README.md index 360e185a9a546..035470cca0ae5 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -85,7 +85,7 @@ 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(transform, []string{"extract"}) return nil } diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go new file mode 100644 index 0000000000000..2a918c1d5b0e7 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen.go @@ -0,0 +1,22 @@ +// 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 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 new file mode 100644 index 0000000000000..ca156a417b2c2 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -0,0 +1,630 @@ +// 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 bundle specs from the Dag serialization schema. +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "go/format" + "log" + "math" + "os" + "strings" +) + +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 +} + +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 preserves Go capitalization in schema names. +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 preserves schema 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"` +} + +type field struct { + key string + goName string + goType string + 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") + 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) + } + + 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, specs) + if err != nil { + log.Fatalf("gen: rendering: %v", err) + } + if err := os.WriteFile(*outPath, src, 0o644); err != nil { + log.Fatalf("gen: writing %s: %v", *outPath, err) + } +} + +func 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(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 := cfg.excluded[key]; ok { + excludedSeen[key] = true + continue + } + f, ok := resolveField(cfg, key, def.Properties.props[key]) + if !ok { + continue + } + fields = append(fields, f) + } + + for key := range cfg.excluded { + if !excludedSeen[key] { + return nil, fmt.Errorf( + "excluded entry %q matches no eligible schema property; remove or fix it", + key, + ) + } + } + for key := range cfg.typeOverrides { + if !containsKey(fields, key) { + return nil, fmt.Errorf( + "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 { + return true + } + } + return false +} + +// resolveField rejects unsupported schema properties. +func resolveField(cfg specConfig, 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 cfg.typeOverrides[key] != "": + f.goType = cfg.typeOverrides[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 { + // 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 +} + +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, 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) + 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 + } + } + } + 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 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") +} + +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 +} + +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) +} + +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": + 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) + } +} + +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 563cbca356639..8830bcb4ec085 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -23,14 +23,17 @@ import ( "runtime" "strings" "sync" + "unicode" + "unicode/utf8" ) 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 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: @@ -43,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) + 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) + AddTaskWithName(taskId string, fn any, args ...any) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -57,27 +60,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. Registering the same - // dag_id twice panics. - AddDag(dagId string) Dag + // AddDag registers dagId with one optional spec. Duplicates and extra specs panic. + AddDag(dagId string, spec ...DagSpec) Dag } - // TaskInfo describes a registered task by its user-visible id. - TaskInfo struct { - ID string - } - - // DagInfo describes a registered dag together with its tasks in - // registration order. - DagInfo struct { - DagID string - 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 registered Dag and task metadata. EnumerableBundle interface { OrderedDags() []DagInfo } @@ -85,6 +72,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 +84,85 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any) { - d.registry.registerTask(d.dagId, fn) +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) { - d.registry.registerTaskWithName(d.dagId, taskId, fn) +func (d dagShim) AddTaskWithName(taskId string, fn any, args ...any) { + depends, spec := taskArgs(args, "AddTaskWithName") + d.registry.registerTaskWithName( + d.dagId, + taskId, + fn, + depends, + 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: + 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 +172,91 @@ 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) { + // Method values end in "-fm". + 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() +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 { + validateID("Dag ID", dagId) + dagSpec := normalizeDagSpec(dagId, 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, depends []string, spec TaskSpec) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -143,31 +265,66 @@ func (r *registry) registerTask(dagId string, fn any) { fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn) + r.registerTaskWithName(dagId, fnName, fn, depends, spec) } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { - task, err := NewTaskFunction(fn) +func (r *registry) registerTaskWithName( + dagId, taskId string, + fn any, + depends []string, + spec TaskSpec, +) { + 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)) } + 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)) } + + // Record one downstream edge per dependency. + 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 +340,7 @@ 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 a mutable copy of Dags and tasks in registration order. 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..5651bc10a92fc 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) 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) }) + }) + } +} + func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { s.dag.AddTaskWithName("special", myTask) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { @@ -144,12 +197,209 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { got := enum.OrderedDags() s.Require().Len(got, 3) + taskIDs := func(tasks []TaskInfo) []string { + ids := make([]string, 0, len(tasks)) + for _, t := range tasks { + ids = append(ids, t.ID) + } + return ids + } + s.Equal("zeta", got[0].DagID) - s.Equal([]TaskInfo{{ID: "z1"}, {ID: "z2"}}, got[0].Tasks) + s.Equal([]string{"z1", "z2"}, taskIDs(got[0].Tasks)) s.Equal("alpha", got[1].DagID) - s.Equal([]TaskInfo{{ID: "a1"}}, got[1].Tasks) + s.Equal([]string{"a1"}, taskIDs(got[1].Tasks)) s.Equal("mid", got[2].DagID) s.Empty(got[2].Tasks) } + +func (s *RegistrySuite) TestAddTask_WithSpec() { + s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 1) + s.Require().Len(dags[0].Tasks, 1) + got := dags[0].Tasks[0] + s.Equal("myTask", got.ID) + s.Equal("high_mem", got.Spec.Queue) + s.Equal(3, got.Spec.Retries) + s.Require().NotNil(got.Spec.DoXComPush) + s.False(*got.Spec.DoXComPush) +} + +func (s *RegistrySuite) TestAddTask_HonorsDoXComPush() { + s.dag.AddTask(resultTask, TaskSpec{DoXComPush: Bool(false)}) + 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) + 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"}) + 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) + s.dag.AddTaskWithName("transform", myTask, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, []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) + 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 + 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) + s.dag.AddTaskWithName("load", myTask, []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, []string{"extract"}) + }, + ) +} + +func (s *RegistrySuite) TestAddTask_DependsOnSelfPanics() { + s.PanicsWithError(`task "self" cannot depend on itself in DAG "dag1"`, func() { + s.dag.AddTaskWithName("self", myTask, []string{"self"}) + }) +} + +func (s *RegistrySuite) TestAddTask_TooManySpecsPanics() { + 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 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() { + 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_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/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index a9e2fc7350dd9..d170a7f9f8e8c 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -32,7 +32,7 @@ type BundleProvider interface { // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { // dag := dagbag.AddDag("simple_dag") // dag.AddTask(extract) - // dag.AddTask(transform) + // dag.AddTask(transform, []string{"extract"}) // return nil // } // diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go new file mode 100644 index 0000000000000..752486fe80cd3 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -0,0 +1,267 @@ +// 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 ( + "slices" + "time" +) + +// 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 + // 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 set fields that differ from schema defaults. +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 +} + +// 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 new file mode 100644 index 0000000000000..17a61dd7487b6 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.go @@ -0,0 +1,47 @@ +// 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 + +type ( + // TaskInfo describes a registered task. + TaskInfo struct { + // ID is the task ID. + ID string + // TypeName is the unqualified Go function name. + TypeName string + // PkgPath is the Go package path. + PkgPath string + // Spec configures the task. + Spec TaskSpec + // Downstream lists dependent task IDs in registration order. + Downstream []string + } + + // DagInfo describes a registered Dag and its tasks. + DagInfo struct { + DagID string + // Spec configures the Dag. + Spec DagSpec + Tasks []TaskInfo + } +) + +// Bool returns a pointer for nullable spec fields. +func Bool(b bool) *bool { + return &b +} 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/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 6d6569889008c..01bb2a0087706 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -42,18 +42,47 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { 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, 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") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + concurrentDag.AddTaskWithName( + "pull_xcoms_concurrently", + concurrentxcom.PullXComsConcurrently, + ) 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_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) @@ -61,6 +90,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. diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go new file mode 100644 index 0000000000000..177f2788119ec --- /dev/null +++ b/go-sdk/pkg/execution/dag_parser.go @@ -0,0 +1,51 @@ +// 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 serializes registered Dags into a DagFileParsingResult. +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..93fa5ca825aaa 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, 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, nil) + r.AddDag("dag2").AddTask(failingTask, 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, 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, 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, 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, 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, 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, nil) }) details := newStartupDetails("panicTask") @@ -201,7 +260,7 @@ func TestTaskRunnerBindsArgs(t *testing.T) { gotCountry = country gotMeta = meta return nil - }) + }, nil) }) details := newStartupDetails( @@ -236,7 +295,7 @@ func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { func(country string, meta map[string]any) error { ran = true return nil - }) + }, nil) }) details := newStartupDetails( @@ -268,7 +327,7 @@ func TestTaskRunnerBindsStructArgs(t *testing.T) { func(input regionInput) error { got = input return nil - }) + }, nil) }) details := newStartupDetails( @@ -296,7 +355,7 @@ func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { func(input regionInput) error { got = input return nil - }) + }, nil) }) details := newStartupDetails( @@ -327,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 }, nil) }) details := newStartupDetails( @@ -354,7 +413,7 @@ func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { func(country string) error { ran = true return nil - }) + }, nil) }) details := newStartupDetails( @@ -377,7 +436,7 @@ func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { func(country string) error { ran = true return nil - }) + }, nil) }) details := newStartupDetails( @@ -426,7 +485,7 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { func(country string) error { ran = true return nil - }) + }, nil) }) details := newStartupDetails("transform", tc.spec) @@ -444,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 }, nil) }) details := newStartupDetails( @@ -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() }, nil) }) details := newStartupDetails("ctxcheck") @@ -491,7 +550,7 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, nil) }) details := &genmodels.StartupDetails{ @@ -551,7 +610,7 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, nil) }) details := newStartupDetails("ctxgrab") @@ -614,13 +673,65 @@ 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, 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() + + 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, nil) return nil }, } @@ -686,7 +797,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }) + }, 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..624c5a9b66ed2 --- /dev/null +++ b/go-sdk/pkg/execution/serde.go @@ -0,0 +1,288 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package execution + +import ( + "fmt" + "path/filepath" + "reflect" + "sort" + "time" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// serializeValue matches Python's type/var encoding. +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: + 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 strips type/var metadata. +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: 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 { + 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 builds a serialized task. +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 always emits this; Go tasks have no template fields. + "template_fields": []any{}, + } + 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) + sort.Strings(sorted) + data["downstream_task_ids"] = sorted + } + return map[string]any{ + "__type": "operator", + "__var": data, + } +} + +// applyDagSpec writes serialized Dag fields. +func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { + for key, value := range s.SchemaFields() { + data[key] = unwrapTypeEncoding(serializeValue(value)) + } +} + +// 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 builds serialized Dag params. +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 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)) + 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{ + "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 returns a path relative to the bundle. +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..beb675236eb7e --- /dev/null +++ b/go-sdk/pkg/execution/serde_test.go @@ -0,0 +1,525 @@ +// 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"]) + 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 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, + IgnoreFirstDependsOnPast: true, + WaitForPastDependsBeforeSkipping: true, + WaitForDownstream: true, + DoXComPush: bundlev1.Bool(false), + EmailOnFailure: bundlev1.Bool(false), + EmailOnRetry: bundlev1.Bool(false), + DocMD: "## task", + MapIndexTemplate: "{{ task.task_id }}", + MaxActiveTisPerDag: 2, + MaxActiveTisPerDagrun: 1, + } + data := map[string]any{} + applySpecFields(data, spec) + + 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["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"]) + 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 TestTaskSpecSchemaFields_OmitsSchemaDefaults(t *testing.T) { + 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{} + applySpecFields(data, spec) + assert.Empty(t, data, "all fields equal schema defaults; nothing should be emitted") +} + +func TestTaskSpecSchemaFields_EmptySpecNoOp(t *testing.T) { + data := map[string]any{} + applySpecFields(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"]) + + _, hasDesc := result["description"] + assert.False(t, hasDesc) + _, hasTags := result["tags"] + assert.False(t, hasTags) + + assert.Equal(t, false, result["catchup"]) + assert.Equal(t, false, result["disable_bundle_versioning"]) + 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"]) +} + +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"]) + 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"]) + + startDate := result["start_date"].(float64) + assert.InDelta(t, float64(start.Unix()), startDate, 0.001) +} + +func TestApplyDagSpec_AlwaysEmitsNonSchemaDefaultFields(t *testing.T) { + 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, 16, data["max_active_tasks"]) + assert.Equal(t, 16, data["max_active_runs"]) + assert.Equal(t, 0, data["max_consecutive_failed_dag_runs"]) + + 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) { + 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 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 + 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..4bb8992583b87 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,7 @@ // 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 frame is either DagFileParseRequest or StartupDetails. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -163,6 +162,15 @@ 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) + _ = 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,