diff --git a/go-sdk/Justfile b/go-sdk/Justfile index a88599969634b..8312d483335ee 100644 --- a/go-sdk/Justfile +++ b/go-sdk/Justfile @@ -50,3 +50,7 @@ protoc: # Regenerate the coordinator-protocol data models from the task-sdk supervisor schema generate-models: go generate ./pkg/execution/genmodels/... + +# Regenerate the bundle spec structs from the Airflow dag-serialization schema +generate-specs: + go generate ./bundle/bundlev1/... diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go new file mode 100644 index 0000000000000..f0fb4115605f2 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen.go @@ -0,0 +1,32 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// spec.gen.go is generated from the Airflow dag-serialization schema by the +// local gen tool: the TaskSpec and DagSpec field sets, their Go types, and +// their SchemaFields emit rules derive from the "operator" and "dag" +// definitions, so they cannot drift from what the scheduler deserializes. +// Don't edit spec.gen.go by hand; change the schema (or the per-definition +// config in ./gen) and re-run `just generate-specs` (go generate). What the +// schema cannot express is declared in that config and validated against the +// schema on every run: the always-emit dag keys with their [core]-config +// fallbacks, the nullable is_paused_upon_creation *bool, sorted tags, and +// the SDK-only Schedule field the serializer turns into the timetable +// object. Only the registration Info structs stay hand-written in spec.go: +// they carry Go-side identity the schema knows nothing about. +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..6f5b1ae7a83aa --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -0,0 +1,811 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Command gen generates spec.gen.go from the Airflow dag-serialization schema +// (airflow-core/src/airflow/serialization/schema.json). +// +// It reads the "operator" and "dag" definitions and emits the TaskSpec and +// DagSpec structs plus their SchemaFields methods, so the exposed fields, +// their Go types, and the emit rules the serializer relies on cannot drift +// from the schema. Each field set is derived from its definition rather than +// hand-listed: every scalar property (string/integer/number/boolean, or a +// timedelta/datetime ref) becomes a spec field, in schema order, unless one +// of these rules skips it: +// +// - "_"-prefixed keys are serializer internals (_task_module, _concurrency, ...) +// - keys in the definition's "required" list are always written by the +// serializer itself (task_type, task_id, dag_id, fileloc, ...) +// - "has_on_*" keys are flags derived from Python callbacks, not settable +// - non-scalar keys (arrays, objects, refs other than timedelta/datetime, +// anyOf) cannot be expressed as a scalar spec field, unless the config's +// typeOverrides maps them explicitly (the dag "tags" key -> []string) +// - the config's excluded entries are Python-only or serializer-owned +// concerns, documented per key +// +// A new scalar key added to the schema therefore shows up in the regenerated +// spec — visible in the spec.gen.go diff — or must gain an excluded entry, +// instead of going silently missing. +// +// Field defaults are read from the schema, never hard-coded here: a field is +// serialized only when it is set (non-zero) and differs from its schema +// default, mirroring Python BaseSerialization's behaviour of omitting values +// the scheduler will re-derive. Booleans whose schema default is true become +// *bool so nil can mean "unset" while an explicit false still serializes. +// +// Emit rules the schema cannot express live in each config's behaviors table +// and are validated against the schema on every run (a stale entry fails +// generation): always-emit keys have no schema default and Python's +// serializer always writes their resolved value, some falling back to a +// [core] config default when unset; emit-when-set booleans also have no +// schema default but are nullable, becoming *bool where nil omits the key +// while an explicit false still serializes. Extra fields (DagSpec.Schedule) +// are SDK-only concepts with no schema key; they are declared in the config +// and never emitted by SchemaFields. +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "go/format" + "log" + "math" + "os" + "strings" +) + +type behaviorKind int + +const ( + // behaviorAlwaysEmit marks a key with no schema default that Python's + // serializer always writes: the field is emitted unconditionally, with + // zeroFallback (when set) substituted for the zero value. + behaviorAlwaysEmit behaviorKind = iota + 1 + // behaviorEmitWhenSet marks a nullable boolean with no schema default: + // the field becomes *bool, nil omits the key, and an explicit false + // still serializes. + behaviorEmitWhenSet +) + +// fieldBehavior is one behaviors-table entry: an emit rule the schema cannot +// express, keyed by schema property name and validated for staleness. +type fieldBehavior struct { + kind behaviorKind + // zeroFallback is the literal emitted when an always-emit field is zero, + // e.g. "16"; fallbackDoc names its source, e.g. the [core] config option. + zeroFallback string + fallbackDoc string +} + +// extraField is a spec field with no schema key at all (an SDK-only concept +// like DagSpec.Schedule); it is rendered first and never emitted by +// SchemaFields. +type extraField struct { + goName string + goType string + doc []string +} + +// specConfig describes one schema definition rendered into spec.gen.go. +type specConfig struct { + defName string + typeName string + // structDoc is the doc comment above the generated struct, pre-wrapped. + structDoc string + // excluded lists the eligible scalar keys deliberately not exposed, each + // with the reason; every other eligible scalar key becomes a field. An + // entry that no longer matches an eligible schema key fails generation, + // so the list cannot go stale. + excluded map[string]string + // typeOverrides forces the Go type for keys the mechanical mapping gets + // wrong or cannot express; entries matching no field fail generation. + typeOverrides map[string]string + // behaviors holds the emit rules the schema cannot express; entries + // matching no field fail generation. + behaviors map[string]fieldBehavior + // extraFields are SDK-only fields with no schema key. + extraFields []extraField +} + +// specConfigs lists the definitions rendered into spec.gen.go, in output +// order. +var specConfigs = []specConfig{ + { + defName: "operator", + typeName: "TaskSpec", + structDoc: `TaskSpec is the optional configuration applied to a task at registration +time. Every field is optional: the zero value (nil for the *bool fields) +means "unset" and the scheduler falls back to the schema default. Each +field maps to the same-named key of the "operator" definition in +airflow-core/src/airflow/serialization/schema.json.`, + 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", + }, + // retry_exponential_backoff is "number" with an integral default, but + // Python declares it float (a backoff multiplier), so the mechanical + // mapping would pick int. + typeOverrides: map[string]string{ + "retry_exponential_backoff": "float64", + }, + }, + { + defName: "dag", + typeName: "DagSpec", + structDoc: `DagSpec is the optional configuration applied to a Dag at registration +time. Every field is optional: the zero value (nil for the *bool fields) +means "unset", which either omits the key or, for the always-serialized +keys, emits its resolved default. Except for Schedule, each field maps to +the same-named key of the "dag" definition in +airflow-core/src/airflow/serialization/schema.json.`, + excluded: map[string]string{ + "relative_fileloc": "computed by the serializer from fileloc and the bundle path", + }, + // tags is a bare "array" in the schema; Python stores a set of + // strings, expressed here as []string and serialized sorted. + 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 is "@once", "@continuous", a cron expression, or "" for`, + `NullTimetable (no schedule). It has no schema key: the serializer`, + `derives the timetable object from it, so SchemaFields never emits it.`, + }, + }, + }, + }, +} + +// initialisms maps snake_case segments that must keep non-Title capitalization +// in Go names (do_xcom_push -> DoXComPush, doc_md -> DocMD). +var initialisms = map[string]string{ + "id": "ID", + "md": "MD", + "ui": "UI", + "uri": "URI", + "xcom": "XCom", +} + +type propSchema struct { + Type any `json:"type"` + Ref string `json:"$ref"` + Default any `json:"default"` +} + +// orderedProps keeps schema property order, which encoding/json's map +// decoding discards; the generated struct follows the schema's field order. +type orderedProps struct { + keys []string + props map[string]propSchema +} + +func (o *orderedProps) UnmarshalJSON(data []byte) error { + if err := json.Unmarshal(data, &o.props); err != nil { + return err + } + dec := json.NewDecoder(bytes.NewReader(data)) + if _, err := dec.Token(); err != nil { + return err + } + for dec.More() { + tok, err := dec.Token() + if err != nil { + return err + } + o.keys = append(o.keys, tok.(string)) + var skipped json.RawMessage + if err := dec.Decode(&skipped); err != nil { + return err + } + } + return nil +} + +type definition struct { + Properties orderedProps `json:"properties"` + Required []string `json:"required"` +} + +type schemaDoc struct { + Definitions map[string]definition `json:"definitions"` +} + +// field is one resolved spec field: schema key + default joined with the Go +// name and type the generated struct uses. +type field struct { + key string + goName string + goType string + def any // schema default (nil when absent) + hasDef bool // schema declares a default + refName string +} + +// resolvedSpec pairs a config with the fields selected from its definition. +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) + } +} + +// selectFields derives the spec fields from the definition, in schema order, +// applying the selection rules in the package comment. +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 +} + +// validateBehaviors rejects behaviors-table entries that match no generated +// field or whose kind cannot apply to the field's resolved type. +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 maps one schema property to a spec field; ok is false when the +// property is not a scalar the spec can express (array, object, anyOf, or a +// ref other than timedelta/datetime) and no typeOverrides entry rescues it. +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 { + // A plain bool cannot distinguish "unset" from an explicit + // false when the schema default is true. + f.goType = "*bool" + } + default: + return field{}, false + } + } + if beh, ok := cfg.behaviors[key]; ok && beh.kind == behaviorEmitWhenSet && f.goType == "bool" { + // Nullable in Python: nil must mean "unset" while an explicit false + // still serializes. + 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()) +} + +// writeImports emits the import block the resolved field types need. +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") + } +} + +// renderSpec writes one config's struct and SchemaFields method. +func renderSpec(b *bytes.Buffer, sp resolvedSpec) { + cfg := sp.cfg + for _, line := range strings.Split(cfg.structDoc, "\n") { + fmt.Fprintf(b, "// %s\n", line) + } + fmt.Fprintf(b, "type %s struct {\n", cfg.typeName) + for _, ef := range cfg.extraFields { + for _, line := range ef.doc { + fmt.Fprintf(b, "\t// %s\n", line) + } + fmt.Fprintf(b, "\t%s %s\n", ef.goName, ef.goType) + } + for _, f := range sp.fields { + writeFieldDoc(b, cfg, f) + fmt.Fprintf(b, "\t%s %s\n", f.goName, f.goType) + } + b.WriteString("}\n\n") + + b.WriteString(`// SchemaFields returns the schema-keyed value of every field that is set and +// differs from its schema default, mirroring Python BaseSerialization's +// omission of values the scheduler re-derives. time.Duration and time.Time +// values are returned as-is; the serializer owns the wire encoding. +`) + if extra := schemaFieldsExtraDoc(cfg, sp.fields); extra != "" { + b.WriteString("//\n") + for _, line := range wrapComment(extra, 76) { + fmt.Fprintf(b, "// %s\n", line) + } + } + 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") +} + +// writeFieldDoc renders a field's doc comment: the plain schema-key line, or +// a wrapped multi-line comment when a behavior or override adds policy text. +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)) + } + extra := fieldDocExtra(cfg, f) + if extra == "" { + fmt.Fprintf(b, "\t// %s.\n", text) + return + } + for _, line := range wrapComment(text+extra+".", 72) { + fmt.Fprintf(b, "\t// %s\n", line) + } +} + +// fieldDocExtra renders the policy sentence a behavior or override entry adds +// to a field's doc comment. +func fieldDocExtra(cfg specConfig, f field) string { + if beh, ok := cfg.behaviors[f.key]; ok { + switch beh.kind { + case behaviorAlwaysEmit: + extra := "; it has no schema default and is always serialized" + if beh.zeroFallback != "" { + extra += fmt.Sprintf( + ", falling back to %s (the %s default) when unset", + beh.zeroFallback, + beh.fallbackDoc, + ) + } + return extra + case behaviorEmitWhenSet: + return `; it has no schema default: nil means "unset", Bool(true) / Bool(false) set it explicitly` + } + } + if f.goType == "[]string" { + return "; serialized as a sorted copy (Python stores this value in a set)" + } + return "" +} + +// schemaFieldsExtraDoc renders the SchemaFields doc sentences for the emit +// policies that deviate from plain omit-if-default, derived from the same +// tables that emit the code so the two cannot drift. +func schemaFieldsExtraDoc(cfg specConfig, fields []field) string { + var always, fallbacks, whenSet, sorted []string + for _, f := range fields { + if beh, ok := cfg.behaviors[f.key]; ok { + switch beh.kind { + case behaviorAlwaysEmit: + always = append(always, f.key) + if beh.zeroFallback != "" { + fallbacks = append(fallbacks, f.key) + } + case behaviorEmitWhenSet: + whenSet = append(whenSet, f.key) + } + continue + } + if f.goType == "[]string" { + sorted = append(sorted, f.key) + } + } + var parts []string + if len(always) > 0 { + s := fmt.Sprintf("Keys with no schema default (%s) are always emitted", joinAnd(always)) + if len(fallbacks) > 0 { + s += fmt.Sprintf( + "; %s fall back to their [core] config defaults when unset", + joinAnd(fallbacks), + ) + } + parts = append(parts, s+".") + } + if len(whenSet) > 0 { + parts = append(parts, fmt.Sprintf( + "%s is emitted whenever non-nil, so an explicit false still serializes.", + joinAnd(whenSet), + )) + } + if len(sorted) > 0 { + parts = append(parts, fmt.Sprintf( + "%s is emitted as a sorted copy, matching Python's set-backed serialization.", + joinAnd(sorted), + )) + } + for _, ef := range cfg.extraFields { + parts = append( + parts, + fmt.Sprintf("%s has no schema key and is never emitted here.", ef.goName), + ) + } + return strings.Join(parts, " ") +} + +func joinAnd(items []string) string { + if len(items) <= 1 { + return strings.Join(items, "") + } + return strings.Join(items[:len(items)-1], ", ") + " and " + items[len(items)-1] +} + +// wrapComment breaks text on spaces into lines of at most width characters +// for comment rendering. +func wrapComment(text string, width int) []string { + var lines []string + cur := "" + for _, word := range strings.Fields(text) { + switch { + case cur == "": + cur = word + case len(cur)+1+len(word) > width: + lines = append(lines, cur) + cur = word + default: + cur += " " + word + } + } + if cur != "" { + lines = append(lines, cur) + } + return lines +} + +// defaultDoc renders a field's schema default for its doc comment. +func defaultDoc(f field) string { + if f.refName == "timedelta" { + return fmt.Sprintf("%v seconds", f.def) + } + if s, ok := f.def.(string); ok { + return fmt.Sprintf("%q", s) + } + return fmt.Sprintf("%v", f.def) +} + +// emitCondition writes one field's SchemaFields statement: its behavior-table +// rule when one exists, otherwise the "set and not schema default" guard. +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) + } +} + +// durationLit renders a schema timedelta default (seconds) as a Go duration +// expression, e.g. 300 -> "300 * time.Second". +func durationLit(secs float64) string { + if secs == math.Trunc(secs) { + return fmt.Sprintf("%d * time.Second", int64(secs)) + } + return fmt.Sprintf("time.Duration(%v * float64(time.Second))", secs) +} + +func numLit(v float64, goType string) string { + if goType == "int" { + return fmt.Sprintf("%d", int64(v)) + } + return fmt.Sprintf("%v", v) +} diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 80253fb36653b..4407bb01f1516 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -41,8 +41,12 @@ type ( // Dag is the handle returned by Registry.AddDag. Use it to attach the Go // functions that implement the dag's tasks. Dag interface { - // AddTask registers fn as a task, deriving the task_id from fn's own - // name (so it must match the @task.stub name in the Python dag). + // AddTask registers fn as a task in this Dag using fn's Go name as + // the task id (so it must match the @task.stub name in the Python + // dag). spec carries optional per-task configuration (pass TaskSpec{} + // for defaults). depends lists task ids in the same Dag that must run + // before this one; each must already be registered. Pass nil for no + // dependencies. // // fn is an ordinary Go function whose parameters are injected by type // and may appear in any order. Recognised parameters are: @@ -55,13 +59,13 @@ type ( // the task, and a non-nil first result is pushed as the task's // return-value XCom. Passing a non-function, or a function whose return // signature does not match, panics at registration time. - AddTask(fn any) + AddTask(fn any, spec TaskSpec, depends []string) - // AddTaskWithName is like AddTask but sets task_id explicitly instead of - // deriving it from the function name. Use it when the Go function name - // cannot match the Python @task.stub id, for example for an anonymous - // function or a differing name. - AddTaskWithName(taskId string, fn any) + // AddTaskWithName is like AddTask but sets task_id explicitly instead + // of deriving it from the function name. Use it when the Go function + // name cannot match the Python @task.stub id, for example for an + // anonymous function or a differing name. + AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -70,26 +74,15 @@ type ( Registry interface { Bundle // AddDag registers a dag by its dag_id (matching the Python stub dag) - // and returns a Dag handle for attaching tasks. Registering the same - // dag_id twice panics. - AddDag(dagId string) Dag + // and returns a Dag handle for attaching tasks. An optional DagSpec + // configures dag-level attributes such as schedule and tags; passing + // more than one spec panics. Registering the same dag_id twice panics. + AddDag(dagId string, spec ...DagSpec) Dag } - // TaskInfo describes a registered task by its user-visible id. - 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 the dag/task identity recorded by + // RegisterDags. The default registry implements it; the coordinator-mode + // runtime relies on it for the DAG-parse one-shot. EnumerableBundle interface { OrderedDags() []DagInfo } @@ -97,6 +90,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 } @@ -107,12 +102,24 @@ type dagShim struct { registry *registry } -func (d dagShim) AddTask(fn any) { - d.registry.registerTask(d.dagId, fn) +func (d dagShim) AddTask(fn any, spec TaskSpec, depends []string) { + d.registry.registerTask(d.dagId, fn, spec, depends) +} + +func (d dagShim) AddTaskWithName(taskId string, fn any, spec TaskSpec, depends []string) { + d.registry.registerTaskWithName(d.dagId, taskId, fn, spec, depends) } -func (d dagShim) AddTaskWithName(taskId string, fn any) { - d.registry.registerTaskWithName(d.dagId, taskId, fn) +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 @@ -122,31 +129,43 @@ func (d dagShim) AddTaskWithName(taskId string, fn any) { func New() Registry { return ®istry{ taskFuncMap: make(map[string]map[string]Task), + taskInfo: make(map[string]map[string]TaskInfo), + dagSpec: make(map[string]DagSpec), taskOrder: make(map[string][]string), } } +func splitFullName(fullName string) (typeName, pkgPath string) { + // fullName looks like "main.extract" or "github.com/x/y.MyTask"; method + // values get a "-fm" suffix. + lastDot := strings.LastIndex(fullName, ".") + if lastDot < 0 { + return strings.TrimSuffix(fullName, "-fm"), "" + } + return strings.TrimSuffix(fullName[lastDot+1:], "-fm"), fullName[:lastDot] +} + func getFnName(fn reflect.Value) string { fullName := runtime.FuncForPC(fn.Pointer()).Name() - parts := strings.Split(fullName, ".") - fnName := parts[len(parts)-1] - // Go adds `-fm` suffix to a method names - return strings.TrimSuffix(fnName, "-fm") + name, _ := splitFullName(fullName) + return name } -func (r *registry) AddDag(dagId string) Dag { - r.Lock() - defer r.Unlock() - +func (r *registry) AddDag(dagId string, spec ...DagSpec) Dag { + dagSpec := optionalSpec(spec, "AddDag") + r.RWMutex.Lock() + defer r.RWMutex.Unlock() if _, exists := r.taskFuncMap[dagId]; exists { panic(fmt.Errorf("Dag %q already exists in bundle", dagId)) } r.taskFuncMap[dagId] = make(map[string]Task) + r.taskInfo[dagId] = make(map[string]TaskInfo) + r.dagSpec[dagId] = dagSpec r.dagOrder = append(r.dagOrder, dagId) return dagShim{dagId, r} } -func (r *registry) registerTask(dagId string, fn any) { +func (r *registry) registerTask(dagId string, fn any, spec TaskSpec, depends []string) { val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { @@ -155,31 +174,66 @@ func (r *registry) registerTask(dagId string, fn any) { fnName := getFnName(val) - r.registerTaskWithName(dagId, fnName, fn) + r.registerTaskWithName(dagId, fnName, fn, spec, depends) } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { +func (r *registry) registerTaskWithName( + dagId, taskId string, + fn any, + spec TaskSpec, + depends []string, +) { task, err := NewTaskFunction(fn) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) } + val := reflect.ValueOf(fn) + fullName := runtime.FuncForPC(val.Pointer()).Name() + typeName, pkgPath := splitFullName(fullName) + + info := TaskInfo{ID: taskId, TypeName: typeName, PkgPath: pkgPath, Spec: spec} + r.RWMutex.Lock() defer r.RWMutex.Unlock() dagTasks, exists := r.taskFuncMap[dagId] - if !exists { dagTasks = make(map[string]Task) r.taskFuncMap[dagId] = dagTasks + r.taskInfo[dagId] = make(map[string]TaskInfo) r.dagOrder = append(r.dagOrder, dagId) } - _, exists = dagTasks[taskId] - if exists { + if _, exists := dagTasks[taskId]; exists { panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId)) } + + // Resolve depends to upstream TaskInfo entries, validating each exists. + // We dedupe so a repeated id in `depends` only records one downstream + // edge on the parent. + seen := make(map[string]bool, len(depends)) + for _, dep := range depends { + if dep == taskId { + panic(fmt.Errorf("task %q cannot depend on itself in DAG %q", taskId, dagId)) + } + if seen[dep] { + continue + } + seen[dep] = true + parent, ok := r.taskInfo[dagId][dep] + if !ok { + panic(fmt.Errorf( + "task %q depends on unknown task %q in DAG %q; register upstream tasks first", + taskId, dep, dagId, + )) + } + parent.Downstream = append(parent.Downstream, taskId) + r.taskInfo[dagId][dep] = parent + } + dagTasks[taskId] = task + r.taskInfo[dagId][taskId] = info r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) } @@ -195,9 +249,9 @@ func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { return task, exists } -// OrderedDags returns the registered dags in AddDag order, each with its tasks -// in registration order. The returned slice is freshly allocated; callers may -// mutate it freely. +// OrderedDags returns the registered dags in the order AddDag was called, +// each with its tasks in the order AddTask / AddTaskWithName was called. The +// returned slice is freshly allocated; callers may mutate it freely. func (r *registry) OrderedDags() []DagInfo { r.RLock() defer r.RUnlock() @@ -207,9 +261,9 @@ func (r *registry) OrderedDags() []DagInfo { taskIDs := r.taskOrder[dagID] tasks := make([]TaskInfo, 0, len(taskIDs)) for _, tid := range taskIDs { - tasks = append(tasks, TaskInfo{ID: tid}) + tasks = append(tasks, r.taskInfo[dagID][tid]) } - out = append(out, DagInfo{DagID: dagID, Tasks: tasks}) + out = append(out, DagInfo{DagID: dagID, Spec: r.dagSpec[dagID], Tasks: tasks}) } return out } diff --git a/go-sdk/bundle/bundlev1/registry_test.go b/go-sdk/bundle/bundlev1/registry_test.go index f405c8f531578..364b7b8a8dc37 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -67,14 +67,14 @@ func (s *RegistrySuite) TestAddDag_DuplicatePanics() { } func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask) + s.dag.AddTask(myTask, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "myTask") s.True(exists) s.NotNil(task) } func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "special") s.True(exists) s.NotNil(task) @@ -85,20 +85,20 @@ func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { } func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask) + s.dag.AddTaskWithName("special", myTask, TaskSpec{}, nil) }) } func (s *RegistrySuite) TestAddTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func") + s.dag.AddTask("not a func", TaskSpec{}, nil) }) } func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs) + s.dag.AddTask(myTaskWithArgs, TaskSpec{}, nil) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) @@ -108,13 +108,13 @@ func (s *RegistrySuite) TestAddTask_InvalidReturnType() { s.PanicsWithError( "error registering task \"NotErrorRet\" for DAG \"dag1\": expected task function github.com/apache/airflow/go-sdk/bundle/bundlev1.NotErrorRet last return value to return error but found int", func() { - s.dag.AddTask(NotErrorRet) + s.dag.AddTask(NotErrorRet, TaskSpec{}, nil) }, ) } func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask) + s.dag.AddTask(errorTask, TaskSpec{}, nil) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } @@ -130,11 +130,11 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { // Register dags out of alphabetical order so the test fails if OrderedDags // ever sorts instead of preserving AddDag order. zeta := reg.AddDag("zeta") - zeta.AddTaskWithName("z1", myTask) - zeta.AddTaskWithName("z2", myTask) + zeta.AddTaskWithName("z1", myTask, TaskSpec{}, nil) + zeta.AddTaskWithName("z2", myTask, TaskSpec{}, nil) alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask) + alpha.AddTaskWithName("a1", myTask, TaskSpec{}, nil) reg.AddDag("mid") // dag with no tasks @@ -144,12 +144,138 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { got := enum.OrderedDags() s.Require().Len(got, 3) + taskIDs := func(tasks []TaskInfo) []string { + ids := make([]string, 0, len(tasks)) + for _, t := range tasks { + ids = append(ids, t.ID) + } + return ids + } + s.Equal("zeta", got[0].DagID) - s.Equal([]TaskInfo{{ID: "z1"}, {ID: "z2"}}, got[0].Tasks) + s.Equal([]string{"z1", "z2"}, taskIDs(got[0].Tasks)) s.Equal("alpha", got[1].DagID) - s.Equal([]TaskInfo{{ID: "a1"}}, got[1].Tasks) + s.Equal([]string{"a1"}, taskIDs(got[1].Tasks)) s.Equal("mid", got[2].DagID) s.Empty(got[2].Tasks) } + +func (s *RegistrySuite) TestAddTask_WithSpec() { + s.dag.AddTask(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}, nil) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 1) + s.Require().Len(dags[0].Tasks, 1) + got := dags[0].Tasks[0] + s.Equal("myTask", got.ID) + s.Equal("high_mem", got.Spec.Queue) + s.Equal(3, got.Spec.Retries) + s.Require().NotNil(got.Spec.DoXComPush) + s.False(*got.Spec.DoXComPush) +} + +func (s *RegistrySuite) TestAddTaskWithName_WithSpec() { + s.dag.AddTaskWithName("special", myTask, TaskSpec{Queue: "gpu", Pool: "gpu_pool"}, nil) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 1) + s.Require().Len(dags[0].Tasks, 1) + got := dags[0].Tasks[0] + s.Equal("special", got.ID) + s.Equal("gpu", got.Spec.Queue) + s.Equal("gpu_pool", got.Spec.Pool) +} + +func (s *RegistrySuite) TestAddTask_DependsRecordsDownstream() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("transform", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.Equal([]string{"transform"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform"].Downstream) + s.Nil(byID["load"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_FanOutFanIn() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("transform_a", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("transform_b", myTask, TaskSpec{}, []string{"extract"}) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"transform_a", "transform_b"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.ElementsMatch([]string{"transform_a", "transform_b"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform_a"].Downstream) + s.Equal([]string{"load"}, byID["transform_b"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_DependsDuplicatesIgnored() { + s.dag.AddTaskWithName("extract", myTask, TaskSpec{}, nil) + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract", "extract"}) + + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + s.Equal([]string{"load"}, byID["extract"].Downstream) +} + +func (s *RegistrySuite) TestAddTask_DependsUnknownPanics() { + s.PanicsWithError( + `task "load" depends on unknown task "extract" in DAG "dag1"; register upstream tasks first`, + func() { + s.dag.AddTaskWithName("load", myTask, TaskSpec{}, []string{"extract"}) + }, + ) +} + +func (s *RegistrySuite) TestAddTask_DependsOnSelfPanics() { + s.PanicsWithError(`task "self" cannot depend on itself in DAG "dag1"`, func() { + s.dag.AddTaskWithName("self", myTask, TaskSpec{}, []string{"self"}) + }) +} + +func (s *RegistrySuite) TestAddDag_WithSpec() { + dag2 := s.reg.AddDag( + "dag2", + DagSpec{Schedule: "@daily", Tags: []string{"etl"}, MaxActiveRuns: 4}, + ) + s.NotNil(dag2) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + s.Require().Len(dags, 2) + var got DagInfo + for _, d := range dags { + if d.DagID == "dag2" { + got = d + break + } + } + s.Equal("dag2", got.DagID) + s.Equal("@daily", got.Spec.Schedule) + s.Equal([]string{"etl"}, got.Spec.Tags) + s.Equal(4, got.Spec.MaxActiveRuns) +} + +func (s *RegistrySuite) TestAddDag_TooManySpecsPanics() { + s.PanicsWithError("AddDag accepts at most one spec, got 2", func() { + s.reg.AddDag("dag3", DagSpec{}, DagSpec{}) + }) +} diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go new file mode 100644 index 0000000000000..7bf5c66c3219b --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -0,0 +1,286 @@ +// 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 is the optional configuration applied to a task at registration +// time. Every field is optional: the zero value (nil for the *bool fields) +// means "unset" and the scheduler falls back to the schema default. Each +// field maps to the same-named key of the "operator" definition in +// airflow-core/src/airflow/serialization/schema.json. +type TaskSpec struct { + // Owner maps to the schema key "owner" (schema default "airflow"). + Owner string + // StartDate maps to the schema key "start_date". + StartDate time.Time + // EndDate maps to the schema key "end_date". + EndDate time.Time + // TriggerRule maps to the schema key "trigger_rule" (schema default "all_success"). + TriggerRule string + // DependsOnPast maps to the schema key "depends_on_past" (schema default false). + DependsOnPast bool + // IgnoreFirstDependsOnPast maps to the schema key "ignore_first_depends_on_past" (schema default false). + IgnoreFirstDependsOnPast bool + // WaitForPastDependsBeforeSkipping maps to the schema key "wait_for_past_depends_before_skipping" (schema default false). + WaitForPastDependsBeforeSkipping bool + // WaitForDownstream maps to the schema key "wait_for_downstream" (schema default false). + WaitForDownstream bool + // Retries maps to the schema key "retries" (schema default 0). + Retries int + // Queue maps to the schema key "queue" (schema default "default"). + Queue string + // Pool maps to the schema key "pool" (schema default "default_pool"). + Pool string + // PoolSlots maps to the schema key "pool_slots" (schema default 1). + PoolSlots int + // ExecutionTimeout maps to the schema key "execution_timeout". + ExecutionTimeout time.Duration + // RetryDelay maps to the schema key "retry_delay" (schema default 300 seconds). + RetryDelay time.Duration + // RetryExponentialBackoff maps to the schema key "retry_exponential_backoff" (schema default 0). + RetryExponentialBackoff float64 + // MaxRetryDelay maps to the schema key "max_retry_delay". + MaxRetryDelay time.Duration + // PriorityWeight maps to the schema key "priority_weight" (schema default 1). + PriorityWeight int + // WeightRule maps to the schema key "weight_rule" (schema default "downstream"). + WeightRule string + // Executor maps to the schema key "executor". + Executor string + // DoXComPush maps to the schema key "do_xcom_push" (schema default true). + DoXComPush *bool + // EmailOnFailure maps to the schema key "email_on_failure" (schema default true). + EmailOnFailure *bool + // EmailOnRetry maps to the schema key "email_on_retry" (schema default true). + EmailOnRetry *bool + // DocMD maps to the schema key "doc_md". + DocMD string + // MapIndexTemplate maps to the schema key "map_index_template". + MapIndexTemplate string + // MaxActiveTisPerDag maps to the schema key "max_active_tis_per_dag". + MaxActiveTisPerDag int + // MaxActiveTisPerDagrun maps to the schema key "max_active_tis_per_dagrun". + MaxActiveTisPerDagrun int +} + +// SchemaFields returns the schema-keyed value of every field that is set and +// differs from its schema default, mirroring Python BaseSerialization's +// omission of values the scheduler re-derives. time.Duration and time.Time +// values are returned as-is; the serializer owns the wire encoding. +func (s TaskSpec) SchemaFields() map[string]any { + m := map[string]any{} + if s.Owner != "" && s.Owner != "airflow" { + m["owner"] = s.Owner + } + if !s.StartDate.IsZero() { + m["start_date"] = s.StartDate + } + if !s.EndDate.IsZero() { + m["end_date"] = s.EndDate + } + if s.TriggerRule != "" && s.TriggerRule != "all_success" { + m["trigger_rule"] = s.TriggerRule + } + if s.DependsOnPast { + m["depends_on_past"] = s.DependsOnPast + } + if s.IgnoreFirstDependsOnPast { + m["ignore_first_depends_on_past"] = s.IgnoreFirstDependsOnPast + } + if s.WaitForPastDependsBeforeSkipping { + m["wait_for_past_depends_before_skipping"] = s.WaitForPastDependsBeforeSkipping + } + if s.WaitForDownstream { + m["wait_for_downstream"] = s.WaitForDownstream + } + if s.Retries != 0 { + m["retries"] = s.Retries + } + if s.Queue != "" && s.Queue != "default" { + m["queue"] = s.Queue + } + if s.Pool != "" && s.Pool != "default_pool" { + m["pool"] = s.Pool + } + if s.PoolSlots != 0 && s.PoolSlots != 1 { + m["pool_slots"] = s.PoolSlots + } + if s.ExecutionTimeout != 0 { + m["execution_timeout"] = s.ExecutionTimeout + } + if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { + m["retry_delay"] = s.RetryDelay + } + if s.RetryExponentialBackoff != 0 { + m["retry_exponential_backoff"] = s.RetryExponentialBackoff + } + if s.MaxRetryDelay != 0 { + m["max_retry_delay"] = s.MaxRetryDelay + } + if s.PriorityWeight != 0 && s.PriorityWeight != 1 { + m["priority_weight"] = s.PriorityWeight + } + if s.WeightRule != "" && s.WeightRule != "downstream" { + m["weight_rule"] = s.WeightRule + } + if s.Executor != "" { + m["executor"] = s.Executor + } + if s.DoXComPush != nil && !*s.DoXComPush { + m["do_xcom_push"] = *s.DoXComPush + } + if s.EmailOnFailure != nil && !*s.EmailOnFailure { + m["email_on_failure"] = *s.EmailOnFailure + } + if s.EmailOnRetry != nil && !*s.EmailOnRetry { + m["email_on_retry"] = *s.EmailOnRetry + } + if s.DocMD != "" { + m["doc_md"] = s.DocMD + } + if s.MapIndexTemplate != "" { + m["map_index_template"] = s.MapIndexTemplate + } + if s.MaxActiveTisPerDag != 0 { + m["max_active_tis_per_dag"] = s.MaxActiveTisPerDag + } + if s.MaxActiveTisPerDagrun != 0 { + m["max_active_tis_per_dagrun"] = s.MaxActiveTisPerDagrun + } + return m +} + +// DagSpec is the optional configuration applied to a Dag at registration +// time. Every field is optional: the zero value (nil for the *bool fields) +// means "unset", which either omits the key or, for the always-serialized +// keys, emits its resolved default. Except for Schedule, each field maps to +// the same-named key of the "dag" definition in +// airflow-core/src/airflow/serialization/schema.json. +type DagSpec struct { + // Schedule is "@once", "@continuous", a cron expression, or "" for + // NullTimetable (no schedule). It has no schema key: the serializer + // derives the timetable object from it, so SchemaFields never emits it. + Schedule string + // Catchup maps to the schema key "catchup"; it has no schema default and + // is 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"; it has no + // schema default and is always serialized, falling back to 16 (the [core] + // max_active_tasks_per_dag default) when unset. + MaxActiveTasks int + // MaxActiveRuns maps to the schema key "max_active_runs"; it has no schema + // default and is always serialized, falling back to 16 (the [core] + // max_active_runs_per_dag default) when unset. + MaxActiveRuns int + // MaxConsecutiveFailedDagRuns maps to the schema key + // "max_consecutive_failed_dag_runs"; it has no schema default and is + // 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"; + // it has no schema default: nil means "unset", Bool(true) / Bool(false) + // set it explicitly. + 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 as a sorted copy (Python + // stores this value in a set). + Tags []string + // DisableBundleVersioning maps to the schema key + // "disable_bundle_versioning"; it has no schema default and is always + // serialized. + DisableBundleVersioning bool +} + +// SchemaFields returns the schema-keyed value of every field that is set and +// differs from its schema default, mirroring Python BaseSerialization's +// omission of values the scheduler re-derives. time.Duration and time.Time +// values are returned as-is; the serializer owns the wire encoding. +// +// Keys with no schema default (catchup, max_active_tasks, max_active_runs, +// max_consecutive_failed_dag_runs and disable_bundle_versioning) are always +// emitted; max_active_tasks and max_active_runs fall back to their [core] +// config defaults when unset. is_paused_upon_creation is emitted whenever +// non-nil, so an explicit false still serializes. tags is emitted as a sorted +// copy, matching Python's set-backed serialization. Schedule has no schema key +// and is never emitted here. +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..940296504401c --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.go @@ -0,0 +1,61 @@ +// 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 + +// This file holds the hand-written companions of the generated spec.gen.go: +// the registration Info structs that carry Go-side identity the schema knows +// nothing about. + +type ( + // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses + // it to render the per-task block of a DagFileParsingResult. + TaskInfo struct { + // ID is the user-visible task id (the function name unless overridden + // via AddTaskWithName). + ID string + // TypeName is the unqualified Go function name (e.g. "extract"). + TypeName string + // PkgPath is the Go package path (e.g. "main", "github.com/x/y"). + PkgPath string + // Spec carries the optional per-task configuration supplied at + // registration. The zero value means "no overrides". + Spec TaskSpec + // Downstream lists task ids that depend on this task, populated as + // later tasks declare this id in their AddTask `depends` argument. + // Order is registration order; the serializer sorts before emit. + Downstream []string + } + + // DagInfo describes a registered dag together with its tasks in + // registration order. + DagInfo struct { + DagID string + // Spec carries the optional per-dag configuration supplied at + // registration. The zero value means "no overrides". + Spec DagSpec + Tasks []TaskInfo + } +) + +// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / +// DagSpec where nil means "leave at schema default": +// +// v1.TaskSpec{DoXComPush: v1.Bool(false)} +func Bool(b bool) *bool { + return &b +} diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..06535b0bf8be6 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -46,14 +46,23 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { } func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) - simpleDag.AddTask(load) + simpleDag := dagbag.AddDag("simple_dag", v1.DagSpec{ + Schedule: "@daily", + Description: "Example Go-authored Dag", + Tags: []string{"example", "go-sdk"}, + }) + simpleDag.AddTask(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}, nil) + simpleDag.AddTask(transform, v1.TaskSpec{Queue: "go-task"}, []string{"extract"}) + simpleDag.AddTask(load, v1.TaskSpec{Queue: "go-task"}, []string{"transform"}) // Tasks defined in other packages register through the same dagbag. concurrentDag := dagbag.AddDag("concurrent_xcom_dag") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + concurrentDag.AddTaskWithName( + "pull_xcoms_concurrently", + concurrentxcom.PullXComsConcurrently, + v1.TaskSpec{}, + nil, + ) return nil } diff --git a/go-sdk/pkg/execution/dag_parser.go b/go-sdk/pkg/execution/dag_parser.go new file mode 100644 index 0000000000000..4b6836ab4d4b4 --- /dev/null +++ b/go-sdk/pkg/execution/dag_parser.go @@ -0,0 +1,54 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package execution + +import ( + "github.com/apache/airflow/go-sdk/bundle/bundlev1" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// ParseDags processes a DagFileParseRequest by serialising every dag +// registered on bundle to DagSerialization v3 and returning the result as a +// DagFileParsingResult body. bundle is the materialised registry produced by +// running BundleProvider.RegisterDags. +func ParseDags(bundle bundlev1.Bundle, req *genmodels.DagFileParseRequest) map[string]any { + fileloc := req.File + bundlePath := req.BundlePath + relativeFileloc := computeRelativeFileloc(fileloc, bundlePath) + + var dags []bundlev1.DagInfo + if enum, ok := bundle.(bundlev1.EnumerableBundle); ok { + dags = enum.OrderedDags() + } + + serializedDags := make([]any, 0, len(dags)) + for _, d := range dags { + serializedDags = append(serializedDags, map[string]any{ + "data": map[string]any{ + "__version": 3, + "dag": SerializeDag(d, fileloc, relativeFileloc), + }, + }) + } + + return map[string]any{ + "type": "DagFileParsingResult", + "fileloc": fileloc, + "serialized_dags": serializedDags, + } +} diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2559359453866..6abb963f391dd 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -29,6 +29,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/vmihailenco/msgpack/v5" "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" @@ -85,9 +86,67 @@ func buildBundle(t *testing.T, register func(bundlev1.Registry)) bundlev1.Bundle // --- Tests --- +func TestDagParsing(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + d := r.AddDag("test_dag") + d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + }) + + req := &genmodels.DagFileParseRequest{ + File: "/bundles/test/main.go", + BundlePath: "/bundles/test", + } + + result := ParseDags(bundle, req) + + assert.Equal(t, "DagFileParsingResult", result["type"]) + assert.Equal(t, "/bundles/test/main.go", result["fileloc"]) + + serializedDags, ok := result["serialized_dags"].([]any) + require.True(t, ok) + require.Len(t, serializedDags, 1) + + dagEntry := serializedDags[0].(map[string]any) + data := dagEntry["data"].(map[string]any) + assert.Equal(t, 3, data["__version"]) + + dagMap := data["dag"].(map[string]any) + assert.Equal(t, "test_dag", dagMap["dag_id"]) + + tt := dagMap["timetable"].(map[string]any) + assert.Equal(t, "airflow.timetables.simple.NullTimetable", tt["__type"]) + + tasks := dagMap["tasks"].([]any) + require.Len(t, tasks, 1) + taskMap := tasks[0].(map[string]any) + assert.Equal(t, "operator", taskMap["__type"]) + taskData := taskMap["__var"].(map[string]any) + assert.Equal(t, "simpleTask", taskData["task_id"]) + assert.Equal(t, "go", taskData["language"]) +} + +func TestDagParsingMultipleDagsPreservesOrder(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + r.AddDag("dag2").AddTask(failingTask, bundlev1.TaskSpec{}, nil) + }) + + req := &genmodels.DagFileParseRequest{File: "/bundle/main.go", BundlePath: "/bundle"} + result := ParseDags(bundle, req) + + serializedDags := result["serialized_dags"].([]any) + require.Len(t, serializedDags, 2) + + dag1Data := serializedDags[0].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "dag1", dag1Data["dag_id"]) + + dag2Data := serializedDags[1].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "dag2", dag2Data["dag_id"]) +} + func TestTaskRunnerSuccess(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask) + r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -110,7 +169,7 @@ func TestTaskRunnerSuccess(t *testing.T) { func TestTaskRunnerFailure(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask) + r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -133,7 +192,7 @@ func TestTaskRunnerFailure(t *testing.T) { func TestTaskRunnerRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(failingTask) + r.AddDag("test_dag").AddTask(failingTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -160,7 +219,7 @@ func TestTaskRunnerRetry(t *testing.T) { func TestTaskRunnerTaskNotFound(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(simpleTask) + r.AddDag("test_dag").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -182,7 +241,7 @@ func TestTaskRunnerTaskNotFound(t *testing.T) { func TestTaskRunnerPanic(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask) + r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -205,7 +264,7 @@ func TestTaskRunnerPanic(t *testing.T) { func TestTaskRunnerPanicRetry(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTask(panicTask) + r.AddDag("test_dag").AddTask(panicTask, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -233,7 +292,7 @@ func TestTaskRunnerPanicRetry(t *testing.T) { func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("ctxcheck", - func(ctx context.Context) error { return ctx.Err() }) + func(ctx context.Context) error { return ctx.Err() }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -270,7 +329,7 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -330,7 +389,7 @@ func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{}, nil) }) details := &genmodels.StartupDetails{ @@ -406,13 +465,66 @@ func startSupervisor( return commLn.Addr().String(), logsLn.Addr().String(), commCh, logsCh, cleanup } +func TestServeDagFileParseEndToEnd(t *testing.T) { + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + defer cleanup() + + provider := &fakeProvider{ + register: func(r bundlev1.Registry) error { + d := r.AddDag("simple_dag") + d.AddTask(simpleTask, bundlev1.TaskSpec{}, nil) + return nil + }, + } + + done := make(chan error, 1) + go func() { done <- Serve(provider, commAddr, logsAddr) }() + + commConn := <-commCh + require.NotNil(t, commConn) + defer commConn.Close() + logsConn := <-logsCh + require.NotNil(t, logsConn) + defer logsConn.Close() + + // Send DagFileParseRequest as a request frame. + payload, err := encodeRequest(0, map[string]any{ + "type": "DagFileParseRequest", + "file": "/bundle/main.go", + "bundle_path": "/bundle", + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, payload)) + + frame, err := readFrame(commConn) + require.NoError(t, err) + assert.Equal(t, int64(0), frame.ID) + require.True(t, isNilRaw(frame.Err)) + assert.Equal(t, "DagFileParsingResult", peekBodyType(frame.Body)) + + var body map[string]any + require.NoError(t, msgpack.Unmarshal(frame.Body, &body)) + + dags := body["serialized_dags"].([]any) + require.Len(t, dags, 1) + dag := dags[0].(map[string]any)["data"].(map[string]any)["dag"].(map[string]any) + assert.Equal(t, "simple_dag", dag["dag_id"]) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after parse result") + } +} + func TestServeStartupDetailsEndToEnd(t *testing.T) { commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) defer cleanup() provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTask(simpleTask) + r.AddDag("dag1").AddTask(simpleTask, bundlev1.TaskSpec{}, nil) return nil }, } @@ -478,7 +590,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }) + }, bundlev1.TaskSpec{}, nil) return nil }, } diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go new file mode 100644 index 0000000000000..eeca5fdc3f200 --- /dev/null +++ b/go-sdk/pkg/execution/serde.go @@ -0,0 +1,324 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package execution + +import ( + "fmt" + "path/filepath" + "reflect" + "sort" + "time" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// serializeValue recursively serializes a value with Airflow's type/var encoding. +// This matches Python's BaseSerialization.serialize() output: +// - primitives (string, bool, int, float) pass through unchanged +// - time.Time -> {"__type": "datetime", "__var": epoch_seconds_float} +// - time.Duration -> {"__type": "timedelta", "__var": total_seconds_float} +// - map[string]any -> {"__type": "dict", "__var": {k: serialize(v), ...}} +// - []any -> direct array with each element serialized +func serializeValue(value any) any { + if value == nil { + return nil + } + switch v := value.(type) { + case string, bool: + return v + case int: + return v + case int8: + return int(v) + case int16: + return int(v) + case int32: + return int(v) + case int64: + return v + case float32: + return float64(v) + case float64: + return v + case time.Time: + epochSec := float64(v.Unix()) + float64(v.Nanosecond())/1e9 + return map[string]any{ + "__type": "datetime", + "__var": epochSec, + } + case time.Duration: + return map[string]any{ + "__type": "timedelta", + "__var": v.Seconds(), + } + case map[string]any: + serialized := make(map[string]any, len(v)) + for k, val := range v { + serialized[k] = serializeValue(val) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case []string: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + case []any: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + default: + // Use reflection to handle typed maps and slices that don't match + // the concrete types above (e.g., map[string]map[string][]string). + rv := reflect.ValueOf(v) + switch rv.Kind() { + case reflect.Map: + serialized := make(map[string]any, rv.Len()) + for _, key := range rv.MapKeys() { + serialized[fmt.Sprint(key.Interface())] = serializeValue(rv.MapIndex(key).Interface()) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case reflect.Slice, reflect.Array: + result := make([]any, rv.Len()) + for i := range result { + result[i] = serializeValue(rv.Index(i).Interface()) + } + return result + default: + return v + } + } +} + +// unwrapTypeEncoding extracts the "__var" part from a type-encoded value. +// In Python's serialize_to_json, non-decorated fields are serialized then unwrapped. +func unwrapTypeEncoding(value any) any { + m, ok := value.(map[string]any) + if !ok { + return value + } + if _, hasType := m["__type"]; !hasType { + return value + } + if v, hasVar := m["__var"]; hasVar { + return v + } + return value +} + +// serializeTimetable converts a schedule string to the Airflow timetable format. +// +// TODO: respect [scheduler] create_cron_data_intervals (and, once timedelta +// schedules are supported, create_delta_data_intervals). Python's +// _create_timetable selects CronDataIntervalTimetable when +// create_cron_data_intervals is True and CronTriggerTimetable when False; we +// hardcode CronTriggerTimetable, which matches only the default (False). The +// Go bundle binary cannot read airflow.cfg, so the supervisor (Python +// ExecutableCoordinator) must send these scheduler flags to the lang-SDK over +// the coordinator protocol (e.g. on DagFileParseRequest) before we can honor +// non-default deployments. Until that channel exists this stays default-only. +// Tracked at https://github.com/apache/airflow/issues/67938 +func serializeTimetable(schedule *string) map[string]any { + if schedule == nil { + return map[string]any{ + "__type": "airflow.timetables.simple.NullTimetable", + "__var": map[string]any{}, + } + } + switch *schedule { + case "@once": + return map[string]any{ + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": map[string]any{}, + } + case "@continuous": + return map[string]any{ + "__type": "airflow.timetables.simple.ContinuousTimetable", + "__var": map[string]any{}, + } + default: + return map[string]any{ + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": map[string]any{ + "expression": *schedule, + "timezone": "UTC", + "interval": 0.0, + "run_immediately": false, + }, + } + } +} + +// serializeTask converts a task to the Airflow serialization format. The +// downstream_task_ids slice is read from info.Downstream (populated by the +// registry from each task's `depends` argument) and sorted for stable JSON. +func serializeTask(info bundlev1.TaskInfo) map[string]any { + typeName := info.TypeName + if typeName == "" { + typeName = info.ID + } + pkgPath := info.PkgPath + if pkgPath == "" { + pkgPath = "main" + } + data := map[string]any{ + "task_id": info.ID, + "task_type": typeName, + "_task_module": pkgPath, + "language": "go", + // Python's operator serializer always emits template_fields (its + // list value never matches the tuple default it is compared against), + // so it is unconditional here too. Go tasks have no template fields. + "template_fields": []any{}, + } + // TaskSpec.SchemaFields (generated from schema.json) returns only the + // fields that are set and differ from their schema default, so this + // mirrors Python BaseSerialization's "omit hard-coded default" behavior. + // serializeValue converts time.Duration / time.Time to their wire floats; + // unwrapTypeEncoding strips the __type wrapper because operator fields are + // stored unwrapped. + for key, value := range info.Spec.SchemaFields() { + data[key] = unwrapTypeEncoding(serializeValue(value)) + } + if len(info.Downstream) > 0 { + sorted := make([]string, len(info.Downstream)) + copy(sorted, info.Downstream) + sort.Strings(sorted) + data["downstream_task_ids"] = sorted + } + return map[string]any{ + "__type": "operator", + "__var": data, + } +} + +// applyDagSpec writes DAG-level fields onto data. The emit policy — which +// keys are omitted when unset, always emitted, or fall back to a [core] +// config default — lives in the generated DagSpec.SchemaFields; this only +// converts the returned values to the wire encoding, as serializeTask does +// for TaskSpec. +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 converts DAG params to Airflow's serialization format. +func serializeParams(params map[string]any) []any { + if len(params) == 0 { + return []any{} + } + result := make([]any, 0, len(params)) + for k, v := range params { + result = append(result, []any{ + k, + map[string]any{ + "__class": "airflow.sdk.definitions.param.Param", + "default": serializeValue(v), + "description": nil, + "schema": serializeValue(map[string]any{}), + "source": nil, + }, + }) + } + return result +} + +// SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 +// format. Required fields are always present; spec-driven fields are emitted +// per the generated DagSpec.SchemaFields policy (some always, some only when +// set). +func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { + taskIDs := make([]string, len(info.Tasks)) + tasks := make([]any, len(info.Tasks)) + for i, t := range info.Tasks { + taskIDs[i] = t.ID + tasks[i] = serializeTask(t) + } + + var schedule *string + if info.Spec.Schedule != "" { + s := info.Spec.Schedule + schedule = &s + } + + result := map[string]any{ + // Required fields (always present) + "dag_id": info.DagID, + "fileloc": fileloc, + "relative_fileloc": relativeFileloc, + "timezone": "UTC", + "timetable": serializeTimetable(schedule), + "tasks": tasks, + "dag_dependencies": []any{}, + "task_group": serializeTaskGroup(taskIDs), + "edge_info": map[string]any{}, + "params": serializeParams(nil), + "deadline": nil, + "allowed_run_types": nil, + } + applyDagSpec(result, info.Spec) + return result +} + +// computeRelativeFileloc computes the relative file location from the bundle path. +func computeRelativeFileloc(fileloc, bundlePath string) string { + if fileloc == "" { + return "" + } + if bundlePath == "" { + return "." + } + rel, err := filepath.Rel(bundlePath, fileloc) + if err != nil { + return "." + } + if rel == "" { + return "." + } + return rel +} diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go new file mode 100644 index 0000000000000..18823b0045ec1 --- /dev/null +++ b/go-sdk/pkg/execution/serde_test.go @@ -0,0 +1,545 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package execution + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +func TestSerializeValuePrimitives(t *testing.T) { + assert.Nil(t, serializeValue(nil)) + assert.Equal(t, "hello", serializeValue("hello")) + assert.Equal(t, true, serializeValue(true)) + assert.Equal(t, 42, serializeValue(42)) + assert.Equal(t, float64(3.14), serializeValue(3.14)) +} + +func TestSerializeValueDatetime(t *testing.T) { + ts := time.Date(2024, 1, 15, 10, 30, 0, 500000000, time.UTC) + result := serializeValue(ts) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "datetime", m["__type"]) + epochSec := m["__var"].(float64) + expected := float64(ts.Unix()) + 0.5 + assert.InDelta(t, expected, epochSec, 0.001) +} + +func TestSerializeValueTimedelta(t *testing.T) { + dur := 90 * time.Second + result := serializeValue(dur) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "timedelta", m["__type"]) + assert.Equal(t, 90.0, m["__var"]) +} + +func TestSerializeValueMap(t *testing.T) { + input := map[string]any{ + "key1": "val1", + "key2": 42, + } + result := serializeValue(input) + m, ok := result.(map[string]any) + require.True(t, ok) + assert.Equal(t, "dict", m["__type"]) + inner := m["__var"].(map[string]any) + assert.Equal(t, "val1", inner["key1"]) + assert.Equal(t, 42, inner["key2"]) +} + +func TestSerializeValueSlice(t *testing.T) { + input := []any{"a", 1, true} + result := serializeValue(input) + arr, ok := result.([]any) + require.True(t, ok) + assert.Len(t, arr, 3) + assert.Equal(t, "a", arr[0]) +} + +func TestUnwrapTypeEncoding(t *testing.T) { + wrapped := map[string]any{ + "__type": "datetime", + "__var": 1705313400.5, + } + assert.Equal(t, 1705313400.5, unwrapTypeEncoding(wrapped)) + + assert.Equal(t, "hello", unwrapTypeEncoding("hello")) + assert.Equal(t, 42, unwrapTypeEncoding(42)) +} + +func TestSerializeTimetable(t *testing.T) { + t.Run("nil schedule", func(t *testing.T) { + result := serializeTimetable(nil) + assert.Equal(t, "airflow.timetables.simple.NullTimetable", result["__type"]) + }) + + t.Run("@once", func(t *testing.T) { + s := "@once" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.simple.OnceTimetable", result["__type"]) + }) + + t.Run("@continuous", func(t *testing.T) { + s := "@continuous" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.simple.ContinuousTimetable", result["__type"]) + }) + + t.Run("cron expression", func(t *testing.T) { + s := "0 12 * * *" + result := serializeTimetable(&s) + assert.Equal(t, "airflow.timetables.trigger.CronTriggerTimetable", result["__type"]) + v := result["__var"].(map[string]any) + assert.Equal(t, "0 12 * * *", v["expression"]) + assert.Equal(t, "UTC", v["timezone"]) + assert.Equal(t, 0.0, v["interval"]) + assert.Equal(t, false, v["run_immediately"]) + }) +} + +func TestSerializeTask(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"transform"}, + }) + assert.Equal(t, "operator", result["__type"]) + data := result["__var"].(map[string]any) + assert.Equal(t, "extract", data["task_id"]) + assert.Equal(t, "extract", data["task_type"]) + assert.Equal(t, "main", data["_task_module"]) + assert.Equal(t, "go", data["language"]) + assert.Equal(t, []string{"transform"}, data["downstream_task_ids"]) + // template_fields is always present (matches Python), empty for Go tasks. + assert.Equal(t, []any{}, data["template_fields"]) + _, hasQueue := data["queue"] + assert.False(t, hasQueue, "queue should be omitted when unset") +} + +func TestSerializeTaskDownstreamSorted(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"transform", "audit", "load"}, + }) + data := result["__var"].(map[string]any) + assert.Equal(t, []string{"audit", "load", "transform"}, data["downstream_task_ids"]) +} + +func TestSerializeTaskNoDownstream(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ID: "load", TypeName: "load", PkgPath: "main"}) + data := result["__var"].(map[string]any) + _, hasDownstream := data["downstream_task_ids"] + assert.False(t, hasDownstream) +} + +func TestSerializeTaskCustomQueue(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }) + data := result["__var"].(map[string]any) + assert.Equal(t, "high_mem", data["queue"]) +} + +func TestSerializeTaskDefaultQueueOmitted(t *testing.T) { + result := serializeTask(bundlev1.TaskInfo{ + ID: "extract", TypeName: "extract", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "default"}, + }) + data := result["__var"].(map[string]any) + _, hasQueue := data["queue"] + assert.False(t, hasQueue, "queue=\"default\" matches the schema default and should be omitted") +} + +func 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) { + // Values equal to schema defaults must be dropped. + spec := bundlev1.TaskSpec{ + Queue: "default", + Pool: "default_pool", + PoolSlots: 1, + Retries: 0, + RetryDelay: 300 * time.Second, + PriorityWeight: 1, + WeightRule: "downstream", + TriggerRule: "all_success", + Owner: "airflow", + DoXComPush: bundlev1.Bool(true), + EmailOnFailure: bundlev1.Bool(true), + EmailOnRetry: bundlev1.Bool(true), + } + data := map[string]any{} + 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"]) + + // Optional spec fields stay omitted when unset. + _, hasDesc := result["description"] + assert.False(t, hasDesc) + _, hasTags := result["tags"] + assert.False(t, hasTags) + + // Python always serializes these (no schema default to omit against), so + // a minimal Dag still carries them at their resolved [core] defaults — + // 16 is [core] max_active_tasks_per_dag / max_active_runs_per_dag, pinned + // in the generated DagSpec.SchemaFields. + 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"]) + // Tags are emitted sorted, matching Python's set-backed serialization. + assert.Equal(t, []any{"etl", "prod"}, result["tags"]) + assert.Equal(t, "ETL Pipeline", result["dag_display_name"]) + assert.Equal(t, "## ETL", result["doc_md"]) + assert.Equal(t, 32, result["max_active_tasks"]) + assert.Equal(t, 4, result["max_active_runs"]) + assert.Equal(t, 3, result["max_consecutive_failed_dag_runs"]) + assert.Equal(t, (2 * time.Hour).Seconds(), result["dagrun_timeout"]) + assert.Equal(t, true, result["catchup"]) + assert.Equal(t, true, result["fail_fast"]) + assert.Equal(t, true, result["render_template_as_native_obj"]) + assert.Equal(t, true, result["disable_bundle_versioning"]) + assert.Equal(t, true, result["is_paused_upon_creation"]) + + // start_date is a raw epoch number, not the type-wrapped form. + startDate := result["start_date"].(float64) + assert.InDelta(t, float64(start.Unix()), startDate, 0.001) +} + +func TestApplyDagSpec_AlwaysEmitsNonSchemaDefaultFields(t *testing.T) { + // catchup, disable_bundle_versioning and the three max_* fields have no + // JSON-schema default, so Python always serializes them. An empty spec + // must still emit them at their resolved [core] defaults. + data := map[string]any{} + applyDagSpec(data, bundlev1.DagSpec{}) + + assert.Equal(t, false, data["catchup"]) + assert.Equal(t, false, data["disable_bundle_versioning"]) + // 16 is [core] max_active_tasks_per_dag / max_active_runs_per_dag, pinned + // in the generated DagSpec.SchemaFields. + assert.Equal(t, 16, data["max_active_tasks"]) + assert.Equal(t, 16, data["max_active_runs"]) + assert.Equal(t, 0, data["max_consecutive_failed_dag_runs"]) + + // Fields with a false/empty schema default stay omitted when unset. + for _, k := range []string{ + "description", "start_date", "end_date", "tags", "dag_display_name", + "doc_md", "dagrun_timeout", "fail_fast", "render_template_as_native_obj", + "is_paused_upon_creation", + } { + _, ok := data[k] + assert.Falsef(t, ok, "%q should be omitted when unset", k) + } +} + +func TestApplyDagSpec_OmitsFalseSchemaDefaultBooleans(t *testing.T) { + // fail_fast and render_template_as_native_obj default to false in the + // schema, so Python omits them unless explicitly true. + data := map[string]any{} + applyDagSpec(data, bundlev1.DagSpec{FailFast: false, RenderTemplateAsNativeObj: false}) + _, hasFailFast := data["fail_fast"] + assert.False(t, hasFailFast) + _, hasNative := data["render_template_as_native_obj"] + assert.False(t, hasNative) +} + +func 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"]) + // time values come back raw; the serializer owns the wire encoding. + 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"]) + // Schedule has no schema key; the serializer derives the timetable from it. + _, hasSchedule := m["schedule"] + assert.False(t, hasSchedule) +} + +func TestDagSpecSchemaFields_EmptySpecEmitsOnlyAlwaysKeys(t *testing.T) { + // 16 is [core] max_active_tasks_per_dag / max_active_runs_per_dag. + 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 7e6e05ef94dc9..71ec3fc82616c 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,11 @@ // the bundle binary is launched with --comm/--logs by the Airflow supervisor // (Python ExecutableCoordinator), bundlev1server.Serve dispatches here. // -// The first inbound frame on the comm socket is a StartupDetails message -// that drives multi-round task execution. +// The first inbound frame on the comm socket selects between two +// sub-protocols: +// +// - DagFileParseRequest: one-shot, returns DagFileParsingResult and exits. +// - StartupDetails: multi-round task execution. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -163,6 +166,16 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { } switch msg := body.(type) { + case *genmodels.DagFileParseRequest: + logger.Debug("DAG parsing mode", "file", msg.File) + result := ParseDags(bundle, msg) + // Bound the terminal write so a wedged socket cannot hang shutdown. + _ = commConn.SetWriteDeadline(time.Now().Add(terminalSendTimeout)) + if err := comm.SendRequest(frame.ID, result); err != nil { + return fmt.Errorf("sending parse result: %w", err) + } + logger.Debug("DAG parsing complete") + case *genmodels.StartupDetails: logger.Debug("Task execution mode", "dag_id", msg.TI.DagID, diff --git a/go-sdk/pkg/worker/runner_test.go b/go-sdk/pkg/worker/runner_test.go index 01ab4ba62b29a..90cdc2e9ffc77 100644 --- a/go-sdk/pkg/worker/runner_test.go +++ b/go-sdk/pkg/worker/runner_test.go @@ -163,7 +163,7 @@ func (s *WorkerSuite) TestStartContextErrorTaskDoesntStart() { s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { wasCalled = true return nil - }) + }, bundlev1.TaskSpec{}, nil) // Setup the mock s.ti.EXPECT(). @@ -196,7 +196,7 @@ func (s *WorkerSuite) TestTaskHeartbeatsWhileRunning() { s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { time.Sleep(time.Second) return nil - }) + }, bundlev1.TaskSpec{}, nil) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateSuccess) @@ -232,7 +232,7 @@ func (s *WorkerSuite) TestTaskHeartbeatConflictStopsTask() { case <-time.After(2 * time.Second): return fmt.Errorf("task context was not cancelled") } - }) + }, bundlev1.TaskSpec{}, nil) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateFailed)