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/README.md b/go-sdk/README.md index 7bd6b3d13c81a..0f1f538acb498 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -87,9 +87,9 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { } func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) + simpleDag := dagbag.AddDag(v1.DagSpec{DagId: "simple_dag"}) + extracted := simpleDag.Task(extract) + simpleDag.Task(transform, v1.Inputs(extracted)) return nil } @@ -100,24 +100,32 @@ func main() { } ``` -A task is an ordinary Go function. The runtime inspects its signature and injects arguments by type: -`sdk.TIRunContext`, `*slog.Logger`, and an `sdk.Client` (or a narrower interface such as -`sdk.VariableClient`). An optional `(any, error)` return becomes the task's XCom; an `error` return marks -the task failed. +A task is an ordinary Go function. The runtime inspects its signature and fills parameters by kind: +runtime values (`sdk.TIRunContext`, `*slog.Logger`, and an `sdk.Client` or a narrower interface such as +`sdk.VariableClient`) are injected by type, and every other parameter is a **data parameter**, bound +positionally to the `v1.Inputs(...)` refs — at run time it receives the upstream task's return value, +pulled from that task's return-value XCom and decoded into the parameter's type. A `(result, error)` +return becomes the task's XCom; an `error` return marks the task failed. ```go -func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { +type ExtractResult struct { + GoVersion string `json:"go_version"` +} + +func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (ExtractResult, error) { conn, err := client.GetConnection(ctx, "test_http") // ... do work, honour ctx cancellation ... - return map[string]any{"go_version": runtime.Version()}, nil + return ExtractResult{GoVersion: runtime.Version()}, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { +// `in` is a data parameter: it receives extract's return value because +// RegisterDags wired v1.Inputs(extracted). +func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger, in ExtractResult) error { val, err := client.GetVariable(ctx, "my_variable") if err != nil { return err } - log.Info("Obtained variable", "my_variable", val) + log.Info("Obtained variable", "my_variable", val, "upstream_go_version", in.GoVersion) return nil } ``` diff --git a/go-sdk/bundle/bundlev1/gen.go b/go-sdk/bundle/bundlev1/gen.go new file mode 100644 index 0000000000000..59846166bd759 --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen.go @@ -0,0 +1,30 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// spec.gen.go is generated from the Airflow dag-serialization schema by the +// local gen tool: the TaskSpec field set, its Go types, and its SchemaFields +// omit-if-default rules all derive from the "operator" definition, so they +// cannot drift from what the scheduler deserializes. Don't edit spec.gen.go +// by hand; change the schema (or the exclusion rules in ./gen) and re-run +// `just generate-specs` (go generate). DagSpec and the registration Info +// structs stay hand-written +// in spec.go: Schedule is an SDK-level concept the schema has no scalar for, +// and several dag keys are always-emitted with [core]-config fallbacks the +// schema cannot express. +package bundlev1 + +//go:generate go run ./gen -schema ../../../airflow-core/src/airflow/serialization/schema.json -out spec.gen.go diff --git a/go-sdk/bundle/bundlev1/gen/main.go b/go-sdk/bundle/bundlev1/gen/main.go new file mode 100644 index 0000000000000..28e16cc83593b --- /dev/null +++ b/go-sdk/bundle/bundlev1/gen/main.go @@ -0,0 +1,468 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Command gen generates spec.gen.go from the Airflow dag-serialization schema +// (airflow-core/src/airflow/serialization/schema.json). +// +// It reads the "operator" definition and emits the TaskSpec struct plus its +// SchemaFields method, so the exposed fields, their Go types, and the +// omit-if-default rules the serializer relies on cannot drift from the schema. +// The field set is derived from the schema rather than hand-listed: every +// scalar property (string/integer/number/boolean, or a timedelta/datetime ref) +// becomes a TaskSpec field, in schema order, unless one of these rules skips +// it: +// +// - "_"-prefixed keys are serializer internals (_task_module, _is_mapped, ...) +// - keys in the definition's "required" list are always written by the +// serializer itself (task_type, ui_color, ...), except those re-exposed +// as identity fields by a fieldOverrides entry (task_id) +// - "has_on_*" keys are flags derived from Python callbacks, not settable +// - non-scalar keys (arrays, objects, refs other than timedelta/datetime, +// anyOf) cannot be expressed as a scalar spec field +// - excludedFields entries are Python-only concerns, documented per key +// +// A new scalar key added to the schema therefore shows up in the regenerated +// spec — visible in the spec.gen.go diff — or must gain an excludedFields +// entry, instead of going silently missing. +// +// Field defaults are read from the schema, never hard-coded here: a field is +// serialized only when it is set (non-zero) and differs from its schema +// default, mirroring Python BaseSerialization's behaviour of omitting values +// the scheduler will re-derive. Booleans whose schema default is true become +// *bool so nil can mean "unset" while an explicit false still serializes. +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "go/format" + "log" + "math" + "os" + "strings" +) + +// excludedFields lists the scalar "operator" keys deliberately not exposed on +// TaskSpec, each with the reason; every other eligible scalar key becomes a +// field. An entry that no longer matches an eligible schema key fails +// generation, so the list cannot go stale. +var excludedFields = map[string]string{ + "doc": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_json": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_yaml": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "doc_rst": "legacy doc attribute; the UI renders Markdown, so only doc_md is exposed", + "allow_nested_operators": "Python runtime concern: warns when an operator executes inside another operator", + "multiple_outputs": "TaskFlow (@task) dict-unpacking concern, meaningless outside Python", + "start_from_trigger": "deferrable-operator machinery; its start_trigger_args counterpart is an object the spec cannot express", + "is_setup": "setup/teardown flags carry trigger-rule invariants the SDK does not model yet", + "is_teardown": "setup/teardown flags carry trigger-rule invariants the SDK does not model yet", + "on_failure_fail_dagrun": "only valid on teardown tasks, which the SDK does not model yet", +} + +// fieldOverride carries the per-key tweaks the schema cannot express. goType +// forces the Go type when the schema type is too loose. goName overrides the +// mechanical snake_case→CamelCase name; doc overrides the generated field +// comment. identity marks a field that names the task rather than configuring +// it: it is exposed on the spec even though the schema lists it as required +// (serializer-owned), and it is kept out of SchemaFields because the +// serializer emits it from the registry's resolved TaskInfo.ID instead. +type fieldOverride struct { + goType string + goName string + doc string + identity bool +} + +// fieldOverrides is keyed by schema key; an entry that no longer matches a +// generated field fails generation. retry_exponential_backoff is "number" +// with an integral default, but Python declares it float (a backoff +// multiplier), so the mechanical mapping would pick int. +var fieldOverrides = map[string]fieldOverride{ + "retry_exponential_backoff": {goType: "float64"}, + "task_id": { + goName: "TaskId", + identity: true, + doc: "TaskId sets the task id explicitly; empty derives it from the task function's name", + }, +} + +// initialisms maps snake_case segments that must keep non-Title capitalization +// 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 + doc string // doc-comment override (empty = generated wording) + identity bool // named-identity field, excluded from SchemaFields + def any // schema default (nil when absent) + hasDef bool // schema declares a default + refName string +} + +func main() { + schemaPath := flag.String("schema", "", "path to the dag-serialization schema.json") + outPath := flag.String("out", "spec.gen.go", "path to write the generated spec file") + pkg := flag.String("package", "bundlev1", "package name for the generated file") + flag.Parse() + + if *schemaPath == "" { + log.Fatal("gen: -schema is required") + } + raw, err := os.ReadFile(*schemaPath) + if err != nil { + log.Fatalf("gen: reading schema: %v", err) + } + var doc schemaDoc + if err := json.Unmarshal(raw, &doc); err != nil { + log.Fatalf("gen: parsing schema: %v", err) + } + operator, ok := doc.Definitions["operator"] + if !ok { + log.Fatal(`gen: schema has no "operator" definition`) + } + + fields, err := selectFields(operator) + if err != nil { + log.Fatalf("gen: %v", err) + } + + src, err := render(*pkg, fields) + if err != nil { + log.Fatalf("gen: rendering: %v", err) + } + if err := os.WriteFile(*outPath, src, 0o644); err != nil { + log.Fatalf("gen: writing %s: %v", *outPath, err) + } +} + +// selectFields derives the TaskSpec fields from the operator definition, in +// schema order, applying the selection rules in the package comment. +func selectFields(op definition) ([]field, error) { + required := make(map[string]bool, len(op.Required)) + for _, key := range op.Required { + required[key] = true + } + + excludedSeen := make(map[string]bool, len(excludedFields)) + fields := make([]field, 0, len(op.Properties.keys)) + for _, key := range op.Properties.keys { + serializerOwned := strings.HasPrefix(key, "_") || required[key] || + strings.HasPrefix(key, "has_on_") + if serializerOwned && !fieldOverrides[key].identity { + continue + } + if _, ok := excludedFields[key]; ok { + excludedSeen[key] = true + continue + } + f, ok := resolveField(key, op.Properties.props[key]) + if !ok { + continue + } + fields = append(fields, f) + } + + for key := range excludedFields { + if !excludedSeen[key] { + return nil, fmt.Errorf( + "excludedFields entry %q matches no eligible schema property; remove or fix it", + key, + ) + } + } + for key := range fieldOverrides { + if !containsKey(fields, key) { + return nil, fmt.Errorf( + "fieldOverrides entry %q matches no generated field; remove or fix it", + key, + ) + } + } + return fields, nil +} + +func containsKey(fields []field, key string) bool { + for _, f := range fields { + if f.key == key { + return true + } + } + return false +} + +// resolveField maps one schema property to a spec field; ok is false when the +// property is not a scalar the spec can express (array, object, anyOf, or a +// ref other than timedelta/datetime). +func resolveField(key string, prop propSchema) (field, bool) { + ov := fieldOverrides[key] + f := field{ + key: key, + goName: ov.goName, + doc: ov.doc, + identity: ov.identity, + def: prop.Default, + hasDef: prop.Default != nil, + } + if f.goName == "" { + f.goName = goName(key) + } + if ref := strings.TrimPrefix(prop.Ref, "#/definitions/"); ref != prop.Ref { + f.refName = ref + } + + switch { + case ov.goType != "": + f.goType = ov.goType + case f.refName == "timedelta": + f.goType = "time.Duration" + case f.refName == "datetime": + f.goType = "time.Time" + default: + typ, _ := prop.Type.(string) + switch typ { + case "string": + f.goType = "string" + case "integer": + f.goType = "int" + case "number": + f.goType = "int" + if n, ok := prop.Default.(float64); ok && n != math.Trunc(n) { + f.goType = "float64" + } + case "boolean": + f.goType = "bool" + if d, ok := prop.Default.(bool); ok && d { + // A plain bool cannot distinguish "unset" from an explicit + // false when the schema default is true. + f.goType = "*bool" + } + default: + return field{}, false + } + } + return f, true +} + +func goName(key string) string { + parts := strings.Split(key, "_") + for i, p := range parts { + if init, ok := initialisms[p]; ok { + parts[i] = init + continue + } + parts[i] = strings.ToUpper(p[:1]) + p[1:] + } + return strings.Join(parts, "") +} + +const licenseHeader = `// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +` + +func render(pkg string, fields []field) ([]byte, error) { + var b bytes.Buffer + b.WriteString( + "// Code generated by bundlev1/gen from airflow-core/src/airflow/serialization/schema.json, DO NOT EDIT.\n", + ) + b.WriteString(licenseHeader) + fmt.Fprintf(&b, "\npackage %s\n\n", pkg) + b.WriteString("import \"time\"\n\n") + + b.WriteString(`// TaskSpec describes a task at registration time. Every field is optional: +// the zero value (nil for the *bool fields) means "unset" and the scheduler +// falls back to the schema default. Each field maps to the same-named key of +// the "operator" definition in +// airflow-core/src/airflow/serialization/schema.json. +type TaskSpec struct { +`) + for _, f := range fields { + if f.doc != "" { + fmt.Fprintf(&b, "\t// %s. Maps to the schema key %q.\n", f.doc, f.key) + } else { + fmt.Fprintf(&b, "\t// %s maps to the schema key %q", f.goName, f.key) + if f.hasDef { + fmt.Fprintf(&b, " (schema default %s)", defaultDoc(f)) + } + b.WriteString(".\n") + } + 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 configuration field +// that is set and differs from its schema default, mirroring Python +// BaseSerialization's omission of values the scheduler re-derives. Identity +// fields (TaskId) are excluded: the serializer emits the id from the +// registry's resolved TaskInfo. time.Duration and time.Time values are +// returned as-is; the serializer owns the wire encoding. +func (s TaskSpec) SchemaFields() map[string]any { + m := map[string]any{} +`) + for _, f := range fields { + if f.identity { + continue + } + emitCondition(&b, f) + } + b.WriteString("\treturn m\n}\n") + + return format.Source(b.Bytes()) +} + +// defaultDoc renders a field's schema default for its doc comment. +func defaultDoc(f field) string { + if f.refName == "timedelta" { + return fmt.Sprintf("%v seconds", f.def) + } + if s, ok := f.def.(string); ok { + return fmt.Sprintf("%q", s) + } + return fmt.Sprintf("%v", f.def) +} + +// emitCondition writes the "set and not schema default" guard for one field. +func emitCondition(b *bytes.Buffer, f field) { + name := "s." + f.goName + key := f.key + switch { + case f.goType == "time.Time": + fmt.Fprintf(b, "\tif !%s.IsZero() {\n\t\tm[%q] = %s\n\t}\n", name, key, name) + case f.goType == "time.Duration": + cond := fmt.Sprintf("%s != 0", name) + if secs, ok := f.def.(float64); ok && secs != 0 { + cond += fmt.Sprintf(" && %s != %s", name, durationLit(secs)) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + case f.goType == "*bool": + notDefault := "*" + name + if def, _ := f.def.(bool); def { + notDefault = "!*" + name + } + fmt.Fprintf( + b, + "\tif %s != nil && %s {\n\t\tm[%q] = *%s\n\t}\n", + name, + notDefault, + key, + name, + ) + case f.goType == "bool": + notDefault := name + if def, _ := f.def.(bool); def { + notDefault = "!" + name + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", notDefault, key, name) + case f.goType == "string": + cond := fmt.Sprintf("%s != \"\"", name) + if def, ok := f.def.(string); ok && def != "" { + cond += fmt.Sprintf(" && %s != %q", name, def) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + default: // int / float64 + cond := fmt.Sprintf("%s != 0", name) + if def, ok := f.def.(float64); ok && def != 0 { + cond += fmt.Sprintf(" && %s != %s", name, numLit(def, f.goType)) + } + fmt.Fprintf(b, "\tif %s {\n\t\tm[%q] = %s\n\t}\n", cond, key, name) + } +} + +// durationLit renders a schema timedelta default (seconds) as a Go duration +// expression, e.g. 300 -> "300 * time.Second". +func durationLit(secs float64) string { + if secs == math.Trunc(secs) { + return fmt.Sprintf("%d * time.Second", int64(secs)) + } + return fmt.Sprintf("time.Duration(%v * float64(time.Second))", secs) +} + +func numLit(v float64, goType string) string { + if goType == "int" { + return fmt.Sprintf("%d", int64(v)) + } + return fmt.Sprintf("%v", v) +} diff --git a/go-sdk/bundle/bundlev1/registry.go b/go-sdk/bundle/bundlev1/registry.go index 80253fb36653b..b87c252ffc09a 100644 --- a/go-sdk/bundle/bundlev1/registry.go +++ b/go-sdk/bundle/bundlev1/registry.go @@ -20,6 +20,7 @@ package bundlev1 import ( "fmt" "reflect" + "regexp" "runtime" "strings" "sync" @@ -29,7 +30,7 @@ import ( type ( // Task is one registered task: something the runtime can Execute. Bundle - // authors do not implement this directly; Dag.AddTask wraps a plain Go + // authors do not implement this directly; Dag.Task wraps a plain Go // function into a Task for you. Task = worker.Task @@ -41,27 +42,28 @@ type ( // Dag is the handle returned by Registry.AddDag. Use it to attach the Go // functions that implement the dag's tasks. Dag interface { - // AddTask registers fn as a task, deriving the task_id from fn's own - // name (so it must match the @task.stub name in the Python dag). + // Task registers fn as a task and returns its TaskRef, the handle + // downstream registrations pass to Inputs or After. The task id is + // TaskSpec.TaskId when a spec sets it, otherwise fn's own Go name. // - // fn is an ordinary Go function whose parameters are injected by type - // and may appear in any order. Recognised parameters are: - // - context.Context: cancelled when the task is asked to stop - // - *slog.Logger: writes to the task's Airflow log - // - sdk.Client (or a narrower sdk.VariableClient / sdk.ConnectionClient / - // sdk.XComClient): access to Variables, Connections, and XCom + // fn is an ordinary Go function whose parameters are filled by kind: // - // fn must return either error or (result, error): a non-nil error fails - // the task, and a non-nil first result is pushed as the task's - // return-value XCom. Passing a non-function, or a function whose return - // signature does not match, panics at registration time. - AddTask(fn any) - - // 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) + // - Injected runtime values, matched by type in any order: + // context.Context (or sdk.TIRunContext), *slog.Logger, and + // sdk.Client (or a narrower sdk.VariableClient / + // sdk.ConnectionClient / sdk.XComClient). + // - Data parameters: every other parameter, bound positionally to + // the Inputs refs. At run time each receives the return value of + // its upstream task, pulled from that task's return-value XCom + // and decoded into the parameter's type. + // + // fn must return either error or (result, error): a non-nil error + // fails the task, and a non-nil first result is pushed as the task's + // return-value XCom (feeding any downstream Inputs). Passing a + // non-function, a function whose signature does not match, a spec / + // Inputs mismatch, or a ref from another dag panics at registration + // (i.e. dag-parse) time. + Task(fn any, opts ...TaskOption) *TaskRef } // Registry is the recorder passed to BundleProvider.RegisterDags. Use it to @@ -69,52 +71,76 @@ type ( // object serves task lookups at execution time. Registry interface { Bundle - // AddDag registers a dag by its dag_id (matching the Python stub dag) - // and returns a Dag handle for attaching tasks. Registering the same - // dag_id twice panics. - AddDag(dagId string) Dag + // AddDag registers a dag described by spec (spec.DagId is required) + // and returns a Dag handle for attaching tasks. A missing DagId or a + // duplicate registration panics. + AddDag(spec DagSpec) Dag } - // 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 } + // TaskOption configures one Dag.Task registration. A TaskSpec value, + // Inputs, and After all implement it. + TaskOption interface { + applyTask(*taskConfig) + } + registry struct { sync.RWMutex taskFuncMap map[string]map[string]Task + taskInfo map[string]map[string]TaskInfo + dagSpec map[string]DagSpec dagOrder []string taskOrder map[string][]string } ) -type dagShim struct { - dagId string - registry *registry +// TaskRef is the handle Dag.Task returns for a registered task. Pass it to +// Inputs (data dependency: the upstream's return value feeds a data +// parameter) or After (ordering-only dependency) when registering downstream +// tasks in the same dag. +type TaskRef struct { + id string + out reflect.Type // fn's first return type; nil when fn returns only error + reg *registry + dagID string } -func (d dagShim) AddTask(fn any) { - d.registry.registerTask(d.dagId, fn) -} +// ID returns the registered task id. +func (r *TaskRef) ID() string { return r.id } -func (d dagShim) AddTaskWithName(taskId string, fn any) { - d.registry.registerTaskWithName(d.dagId, taskId, fn) +// taskConfig accumulates the options of one Dag.Task call. +type taskConfig struct { + spec TaskSpec + specSet bool + inputs []*TaskRef + after []*TaskRef } +type inputsOption []*TaskRef + +func (o inputsOption) applyTask(c *taskConfig) { c.inputs = append(c.inputs, o...) } + +// Inputs declares data dependencies: each ref's task must run first, and its +// return value is pulled from XCom and passed to the matching data parameter +// of the task function, in order. The number of refs must equal the number of +// data parameters, and each upstream's return type must fit the parameter it +// feeds (checked at registration). +func Inputs(refs ...*TaskRef) TaskOption { return inputsOption(refs) } + +type afterOption []*TaskRef + +func (o afterOption) applyTask(c *taskConfig) { c.after = append(c.after, o...) } + +// After declares ordering-only dependencies: each ref's task must complete +// before this one runs, without feeding it any value. +func After(refs ...*TaskRef) TaskOption { return afterOption(refs) } + // New returns an empty Registry on which dags and tasks can be registered. The // runtime creates one and hands it to BundleProvider.RegisterDags, so bundle // authors rarely call this directly; it is handy for unit-testing a @@ -122,65 +148,223 @@ 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 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") +type dagShim struct { + dagId string + registry *registry +} + +func (d dagShim) Task(fn any, opts ...TaskOption) *TaskRef { + return d.registry.registerTask(d.dagId, fn, opts) +} + +// Bool returns a pointer to b. Use it for the *bool fields on TaskSpec / +// DagSpec where nil means "leave at schema default": +// +// v1.TaskSpec{DoXComPush: v1.Bool(false)} +func Bool(b bool) *bool { + return &b } -func (r *registry) AddDag(dagId string) Dag { - r.Lock() - defer r.Unlock() +// anonFnName matches the names the Go runtime gives anonymous functions once +// splitFullName has taken the segment after the last dot: "funcN" for a +// top-level closure, or a bare number for a nested one ("pkg.fn.func1.2"). +var anonFnName = regexp.MustCompile(`^(func)?\d+$`) + +func splitFullName(fullName string) (typeName, pkgPath string) { + // fullName looks like "main.extract" or "github.com/x/y.MyTask"; method + // values get a "-fm" suffix. + lastDot := strings.LastIndex(fullName, ".") + if lastDot < 0 { + return strings.TrimSuffix(fullName, "-fm"), "" + } + return strings.TrimSuffix(fullName[lastDot+1:], "-fm"), fullName[:lastDot] +} - if _, exists := r.taskFuncMap[dagId]; exists { - panic(fmt.Errorf("Dag %q already exists in bundle", dagId)) +func (r *registry) AddDag(spec DagSpec) Dag { + if spec.DagId == "" { + panic(fmt.Errorf("AddDag requires DagSpec.DagId to be set")) + } + r.RWMutex.Lock() + defer r.RWMutex.Unlock() + if _, exists := r.taskFuncMap[spec.DagId]; exists { + panic(fmt.Errorf("Dag %q already exists in bundle", spec.DagId)) } - r.taskFuncMap[dagId] = make(map[string]Task) - r.dagOrder = append(r.dagOrder, dagId) - return dagShim{dagId, r} + r.taskFuncMap[spec.DagId] = make(map[string]Task) + r.taskInfo[spec.DagId] = make(map[string]TaskInfo) + r.dagSpec[spec.DagId] = spec + r.dagOrder = append(r.dagOrder, spec.DagId) + return dagShim{spec.DagId, r} } -func (r *registry) registerTask(dagId string, fn any) { - val := reflect.ValueOf(fn) +func (r *registry) registerTask(dagId string, fn any, opts []TaskOption) *TaskRef { + var cfg taskConfig + for _, o := range opts { + o.applyTask(&cfg) + } + val := reflect.ValueOf(fn) if val.Kind() != reflect.Func { panic(fmt.Errorf("task fn was a %s, not a func", val.Kind())) } + fullName := runtime.FuncForPC(val.Pointer()).Name() + typeName, pkgPath := splitFullName(fullName) + + taskId := cfg.spec.TaskId + if taskId == "" { + if anonFnName.MatchString(typeName) { + panic(fmt.Errorf( + "task function in DAG %q is anonymous (its derived name would be %q); set TaskSpec.TaskId explicitly", + dagId, + typeName, + )) + } + taskId = typeName + } - fnName := getFnName(val) - - r.registerTaskWithName(dagId, fnName, fn) -} + // Pair the fn's data parameters with the Inputs refs before wrapping, so + // every mismatch is a registration (dag-parse) error, not a runtime one. + dataParams := dataParamTypes(val.Type()) + if len(cfg.inputs) != len(dataParams) { + panic(fmt.Errorf( + "task %q in DAG %q declares %d data parameter(s) but Inputs supplies %d", + taskId, dagId, len(dataParams), len(cfg.inputs), + )) + } + bindings := make([]InputBinding, len(cfg.inputs)) + for i, ref := range cfg.inputs { + r.validateRef(ref, dagId, taskId, "Inputs") + if ref.out == nil { + panic(fmt.Errorf( + "task %q in DAG %q: input task %q returns no value; use After for an ordering-only dependency", + taskId, + dagId, + ref.id, + )) + } + if !inputCompatible(ref.out, dataParams[i]) { + panic(fmt.Errorf( + "task %q in DAG %q: input %d: task %q returns %s, which cannot fill parameter type %s", + taskId, + dagId, + i, + ref.id, + ref.out, + dataParams[i], + )) + } + bindings[i] = InputBinding{TaskID: ref.id} + } + for _, ref := range cfg.after { + r.validateRef(ref, dagId, taskId, "After") + } -func (r *registry) registerTaskWithName(dagId, taskId string, fn any) { - task, err := NewTaskFunction(fn) + task, err := NewTaskFunction(fn, bindings...) if err != nil { panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err)) } + var outType reflect.Type + if fnType := val.Type(); fnType.NumOut() > 1 { + outType = fnType.Out(0) + } + + var inputIDs []string + if len(cfg.inputs) > 0 { + inputIDs = make([]string, len(cfg.inputs)) + for i, ref := range cfg.inputs { + inputIDs[i] = ref.id + } + } + + info := TaskInfo{ + ID: taskId, + TypeName: typeName, + PkgPath: pkgPath, + Spec: cfg.spec, + Inputs: inputIDs, + } + r.RWMutex.Lock() defer r.RWMutex.Unlock() dagTasks, exists := r.taskFuncMap[dagId] - if !exists { - dagTasks = make(map[string]Task) - r.taskFuncMap[dagId] = dagTasks - r.dagOrder = append(r.dagOrder, dagId) + panic(fmt.Errorf("DAG %q is not registered", dagId)) } - - _, exists = dagTasks[taskId] - if exists { + if _, exists := dagTasks[taskId]; exists { panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId)) } + + // Record one downstream edge on each distinct upstream, whether the + // dependency carries data (Inputs) or only ordering (After). + seen := make(map[string]bool, len(cfg.inputs)+len(cfg.after)) + for _, ref := range append(append([]*TaskRef{}, cfg.inputs...), cfg.after...) { + if seen[ref.id] { + continue + } + seen[ref.id] = true + parent := r.taskInfo[dagId][ref.id] + parent.Downstream = append(parent.Downstream, taskId) + r.taskInfo[dagId][ref.id] = parent + } + dagTasks[taskId] = task + r.taskInfo[dagId][taskId] = info r.taskOrder[dagId] = append(r.taskOrder[dagId], taskId) + + return &TaskRef{id: taskId, out: outType, reg: r, dagID: dagId} +} + +// validateRef panics unless ref was returned by this registry for the same +// dag. A nil ref, a ref from another registry, or a ref from a different dag +// is a registration error. +func (r *registry) validateRef(ref *TaskRef, dagId, taskId, opt string) { + switch { + case ref == nil: + panic(fmt.Errorf("task %q in DAG %q: %s got a nil TaskRef", taskId, dagId, opt)) + case ref.reg != r: + panic(fmt.Errorf( + "task %q in DAG %q: %s ref %q belongs to a different registry", + taskId, dagId, opt, ref.id, + )) + case ref.dagID != dagId: + panic(fmt.Errorf( + "task %q in DAG %q: %s ref %q belongs to DAG %q; dependencies cannot cross dags", + taskId, dagId, opt, ref.id, ref.dagID, + )) + } +} + +// inputCompatible reports whether an upstream return type can fill a data +// parameter. Identical or assignable types always fit; a map or empty +// interface parameter opts into loose decoding of any shape; an upstream that +// returns any is accepted and checked structurally when the value is decoded +// at run time. Pointers are stripped from both sides: the value crosses XCom +// as its JSON shape, so *T and T are interchangeable in either position. +func inputCompatible(out, param reflect.Type) bool { + if out.Kind() == reflect.Pointer { + out = out.Elem() + } + if param.Kind() == reflect.Pointer { + param = param.Elem() + } + switch { + case out.AssignableTo(param): + return true + case param.Kind() == reflect.Interface && param.NumMethod() == 0: + return true + case param.Kind() == reflect.Map: + return true + case out.Kind() == reflect.Interface && out.NumMethod() == 0: + return true + } + return false } func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) { @@ -195,9 +379,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 Dag.Task was called. The returned slice is +// freshly allocated; callers may mutate it freely. func (r *registry) OrderedDags() []DagInfo { r.RLock() defer r.RUnlock() @@ -207,9 +391,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..dc310e7d5b4b8 100644 --- a/go-sdk/bundle/bundlev1/registry_test.go +++ b/go-sdk/bundle/bundlev1/registry_test.go @@ -41,6 +41,18 @@ func NotErrorRet() int { return 0 } +// producer / consumer fixtures for Inputs wiring. +type payload struct { + Value int `json:"value"` +} + +func produce() (payload, error) { return payload{Value: 1}, nil } +func consume(in payload) error { _ = in; return nil } +func relay(in payload) (payload, error) { return in, nil } +func consumeTwo(a payload, b payload) error { _, _ = a, b; return nil } +func consumeLoose(in map[string]any) error { _ = in; return nil } +func consumeString(in string) error { _ = in; return nil } + type RegistrySuite struct { suite.Suite reg Registry @@ -53,72 +65,85 @@ func TestRegistrySuite(t *testing.T) { func (s *RegistrySuite) SetupTest() { s.reg = New() - s.dag = s.reg.AddDag("dag1") + s.dag = s.reg.AddDag(DagSpec{DagId: "dag1"}) } func (s *RegistrySuite) TestAddDag_NewDagRegisters() { s.NotNil(s.dag) } +func (s *RegistrySuite) TestAddDag_MissingIdPanics() { + s.PanicsWithError("AddDag requires DagSpec.DagId to be set", func() { + s.reg.AddDag(DagSpec{}) + }) +} + func (s *RegistrySuite) TestAddDag_DuplicatePanics() { s.PanicsWithError(`Dag "dag1" already exists in bundle`, func() { - s.reg.AddDag("dag1") + s.reg.AddDag(DagSpec{DagId: "dag1"}) }) } -func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() { - s.dag.AddTask(myTask) +func (s *RegistrySuite) TestTask_RegistersAndFindsTask() { + ref := s.dag.Task(myTask) + s.Equal("myTask", ref.ID()) task, exists := s.reg.LookupTask("dag1", "myTask") s.True(exists) s.NotNil(task) } -func (s *RegistrySuite) TestAddTaskWithName_RegistersAndFindsTask() { - s.dag.AddTaskWithName("special", myTask) - task, exists := s.reg.LookupTask("dag1", "special") +func (s *RegistrySuite) TestTask_ExplicitIdViaSpec() { + ref := s.dag.Task(myTask, TaskSpec{TaskId: "special"}) + s.Equal("special", ref.ID()) + _, exists := s.reg.LookupTask("dag1", "special") s.True(exists) - s.NotNil(task) // Lets just make sure it didn't exist under the fn name _, exists = s.reg.LookupTask("dag1", "myTask") s.False(exists) } -func (s *RegistrySuite) TestRegisterTaskWithName_DuplicatePanics() { - s.dag.AddTaskWithName("special", myTask) +func (s *RegistrySuite) TestTask_DuplicateIdPanics() { + s.dag.Task(myTask, TaskSpec{TaskId: "special"}) s.PanicsWithError("taskId \"special\" is already registered for DAG \"dag1\"", func() { - s.dag.AddTaskWithName("special", myTask) + s.dag.Task(myTask, TaskSpec{TaskId: "special"}) }) } -func (s *RegistrySuite) TestAddTask_NonFuncPanics() { +func (s *RegistrySuite) TestTask_NonFuncPanics() { s.PanicsWithError("task fn was a string, not a func", func() { - s.dag.AddTask("not a func") + s.dag.Task("not a func") }) } -func (s *RegistrySuite) TestAddTaskWithArgs_BindsCorrectArgs() { - s.dag.AddTask(myTaskWithArgs) +func (s *RegistrySuite) TestTask_InjectedArgsBind() { + s.dag.Task(myTaskWithArgs) task, exists := s.reg.LookupTask("dag1", "myTaskWithArgs") s.True(exists) s.NotNil(task) } -func (s *RegistrySuite) TestAddTask_InvalidReturnType() { +func (s *RegistrySuite) TestTask_InvalidReturnType() { s.PanicsWithError( "error registering task \"NotErrorRet\" for DAG \"dag1\": expected task function github.com/apache/airflow/go-sdk/bundle/bundlev1.NotErrorRet last return value to return error but found int", func() { - s.dag.AddTask(NotErrorRet) + s.dag.Task(NotErrorRet) }, ) } -func (s *RegistrySuite) TestAddTask_ErrorReturnType() { - s.dag.AddTask(errorTask) +func (s *RegistrySuite) TestTask_ErrorReturnType() { + s.dag.Task(errorTask) _, exists := s.reg.LookupTask("dag1", "errorTask") s.True(exists) } +func (s *RegistrySuite) TestTask_TooManySpecsPanics() { + s.PanicsWithError("Task accepts at most one TaskSpec", func() { + s.dag.Task(myTask, TaskSpec{}, TaskSpec{}) + }) +} + func (s *RegistrySuite) TestOrderedDags_Empty() { enum, ok := New().(EnumerableBundle) s.Require().True(ok) @@ -129,14 +154,14 @@ func (s *RegistrySuite) TestOrderedDags_PreservesRegistrationOrder() { reg := New() // Register dags out of alphabetical order so the test fails if OrderedDags // ever sorts instead of preserving AddDag order. - zeta := reg.AddDag("zeta") - zeta.AddTaskWithName("z1", myTask) - zeta.AddTaskWithName("z2", myTask) + zeta := reg.AddDag(DagSpec{DagId: "zeta"}) + zeta.Task(myTask, TaskSpec{TaskId: "z1"}) + zeta.Task(myTask, TaskSpec{TaskId: "z2"}) - alpha := reg.AddDag("alpha") - alpha.AddTaskWithName("a1", myTask) + alpha := reg.AddDag(DagSpec{DagId: "alpha"}) + alpha.Task(myTask, TaskSpec{TaskId: "a1"}) - reg.AddDag("mid") // dag with no tasks + reg.AddDag(DagSpec{DagId: "mid"}) // dag with no tasks enum, ok := reg.(EnumerableBundle) s.Require().True(ok) @@ -144,12 +169,215 @@ 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) TestTask_WithSpec() { + s.dag.Task(myTask, TaskSpec{Queue: "high_mem", Retries: 3, DoXComPush: Bool(false)}) + enum, ok := s.reg.(EnumerableBundle) + s.Require().True(ok) + dags := enum.OrderedDags() + 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) infoByID() map[string]TaskInfo { + enum := s.reg.(EnumerableBundle) + tasks := enum.OrderedDags()[0].Tasks + byID := make(map[string]TaskInfo, len(tasks)) + for _, t := range tasks { + byID[t.ID] = t + } + return byID +} + +func (s *RegistrySuite) TestTask_InputsRecordEdgesAndBindings() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + transform := s.dag.Task(relay, TaskSpec{TaskId: "transform"}, Inputs(extract)) + s.dag.Task(consume, TaskSpec{TaskId: "load"}, Inputs(transform)) + + byID := s.infoByID() + s.Equal([]string{"transform"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform"].Downstream) + s.Nil(byID["load"].Downstream) + + s.Nil(byID["extract"].Inputs) + s.Equal([]string{"extract"}, byID["transform"].Inputs) + s.Equal([]string{"transform"}, byID["load"].Inputs) +} + +func (s *RegistrySuite) TestTask_FanOutFanIn() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + a := s.dag.Task(relay, TaskSpec{TaskId: "transform_a"}, Inputs(extract)) + b := s.dag.Task(relay, TaskSpec{TaskId: "transform_b"}, Inputs(extract)) + s.dag.Task(consumeTwo, TaskSpec{TaskId: "load"}, Inputs(a, b)) + + byID := s.infoByID() + s.ElementsMatch([]string{"transform_a", "transform_b"}, byID["extract"].Downstream) + s.Equal([]string{"load"}, byID["transform_a"].Downstream) + s.Equal([]string{"load"}, byID["transform_b"].Downstream) + s.Equal([]string{"transform_a", "transform_b"}, byID["load"].Inputs) +} + +func (s *RegistrySuite) TestTask_AfterRecordsOrderingEdge() { + first := s.dag.Task(myTask, TaskSpec{TaskId: "first"}) + s.dag.Task(myTask, TaskSpec{TaskId: "second"}, After(first)) + + byID := s.infoByID() + s.Equal([]string{"second"}, byID["first"].Downstream) + s.Nil(byID["second"].Inputs) +} + +func (s *RegistrySuite) TestTask_InputsAndAfterDeduplicateEdges() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + // The same upstream named as both a data input and an After ref records + // only one downstream edge. + s.dag.Task(consume, TaskSpec{TaskId: "load"}, Inputs(extract), After(extract)) + + byID := s.infoByID() + s.Equal([]string{"load"}, byID["extract"].Downstream) +} + +func (s *RegistrySuite) TestTask_InputsCountMismatchPanics() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consumeTwo" in DAG "dag1" declares 2 data parameter(s) but Inputs supplies 1`, + func() { + s.dag.Task(consumeTwo, Inputs(extract)) + }, + ) +} + +func (s *RegistrySuite) TestTask_DataParamWithoutInputsPanics() { + s.PanicsWithError( + `task "consume" in DAG "dag1" declares 1 data parameter(s) but Inputs supplies 0`, + func() { + s.dag.Task(consume) + }, + ) +} + +func (s *RegistrySuite) TestTask_InputTypeMismatchPanics() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consumeString" in DAG "dag1": input 0: task "extract" returns bundlev1.payload, which cannot fill parameter type string`, + func() { + s.dag.Task(consumeString, Inputs(extract)) + }, + ) +} + +func (s *RegistrySuite) TestTask_InputLooseMapAccepted() { + extract := s.dag.Task(produce, TaskSpec{TaskId: "extract"}) + s.NotPanics(func() { + s.dag.Task(consumeLoose, Inputs(extract)) + }) +} + +func (s *RegistrySuite) TestTask_InputFromValuelessTaskPanics() { + first := s.dag.Task(myTask, TaskSpec{TaskId: "first"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": input task "first" returns no value; use After for an ordering-only dependency`, + func() { + s.dag.Task(consume, Inputs(first)) + }, + ) +} + +func (s *RegistrySuite) TestTask_CrossDagRefPanics() { + other := s.reg.AddDag(DagSpec{DagId: "dag2"}) + upstream := other.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": Inputs ref "extract" belongs to DAG "dag2"; dependencies cannot cross dags`, + func() { + s.dag.Task(consume, Inputs(upstream)) + }, + ) +} + +func (s *RegistrySuite) TestTask_ForeignRegistryRefPanics() { + foreign := New().AddDag(DagSpec{DagId: "dag1"}) + upstream := foreign.Task(produce, TaskSpec{TaskId: "extract"}) + s.PanicsWithError( + `task "consume" in DAG "dag1": Inputs ref "extract" belongs to a different registry`, + func() { + s.dag.Task(consume, Inputs(upstream)) + }, + ) +} + +func (s *RegistrySuite) TestAddDag_WithSpec() { + dag2 := s.reg.AddDag( + DagSpec{DagId: "dag2", 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 producePtr() (*payload, error) { return &payload{Value: 1}, nil } +func consumePtr(in *payload) error { _ = in; return nil } + +// Pointer and value shapes must be interchangeable on both sides of a data +// edge: the value crosses XCom as its JSON shape, so *T and T are equivalent. +func (s *RegistrySuite) TestTask_InputPointerValueShapesCompatible() { + val := s.dag.Task(produce, TaskSpec{TaskId: "val"}) + ptr := s.dag.Task(producePtr, TaskSpec{TaskId: "ptr"}) + + s.NotPanics(func() { + s.dag.Task(consume, TaskSpec{TaskId: "val_to_val"}, Inputs(val)) + s.dag.Task(consumePtr, TaskSpec{TaskId: "val_to_ptr"}, Inputs(val)) + s.dag.Task(consume, TaskSpec{TaskId: "ptr_to_val"}, Inputs(ptr)) + s.dag.Task(consumePtr, TaskSpec{TaskId: "ptr_to_ptr"}, Inputs(ptr)) + }) +} + +func (s *RegistrySuite) TestTask_AnonymousFnWithoutIdPanics() { + s.PanicsWithError( + `task function in DAG "dag1" is anonymous (its derived name would be "1"); set TaskSpec.TaskId explicitly`, + func() { + s.dag.Task(func() error { return nil }) + }, + ) +} + +func (s *RegistrySuite) TestTask_AnonymousFnWithExplicitIdRegisters() { + ref := s.dag.Task(func() error { return nil }, TaskSpec{TaskId: "inline"}) + s.Equal("inline", ref.ID()) + _, exists := s.reg.LookupTask("dag1", "inline") + s.True(exists) +} diff --git a/go-sdk/bundle/bundlev1/schemas.go b/go-sdk/bundle/bundlev1/schemas.go index 46bf5cc28793f..b42021a2390d6 100644 --- a/go-sdk/bundle/bundlev1/schemas.go +++ b/go-sdk/bundle/bundlev1/schemas.go @@ -34,9 +34,9 @@ type BundleProvider interface { // Registry, for example: // // func (m *myBundle) RegisterDags(dagbag bundlev1.Registry) error { - // dag := dagbag.AddDag("simple_dag") - // dag.AddTask(extract) - // dag.AddTask(transform) + // dag := dagbag.AddDag(bundlev1.DagSpec{DagId: "simple_dag"}) + // extracted := dag.Task(extract) + // dag.Task(transform, bundlev1.Inputs(extracted)) // return nil // } // diff --git a/go-sdk/bundle/bundlev1/spec.gen.go b/go-sdk/bundle/bundlev1/spec.gen.go new file mode 100644 index 0000000000000..2a8ee38c02323 --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.gen.go @@ -0,0 +1,172 @@ +// Code generated by bundlev1/gen from airflow-core/src/airflow/serialization/schema.json, DO NOT EDIT. +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package bundlev1 + +import "time" + +// TaskSpec describes a task at registration time. Every field is optional: +// the zero value (nil for the *bool fields) means "unset" and the scheduler +// falls back to the schema default. Each field maps to the same-named key of +// the "operator" definition in +// airflow-core/src/airflow/serialization/schema.json. +type TaskSpec struct { + // TaskId sets the task id explicitly; empty derives it from the task function's name. Maps to the schema key "task_id". + TaskId string + // Owner maps to the schema key "owner" (schema default "airflow"). + Owner string + // StartDate maps to the schema key "start_date". + 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 configuration field +// that is set and differs from its schema default, mirroring Python +// BaseSerialization's omission of values the scheduler re-derives. Identity +// fields (TaskId) are excluded: the serializer emits the id from the +// registry's resolved TaskInfo. time.Duration and time.Time values are +// returned as-is; the serializer owns the wire encoding. +func (s TaskSpec) SchemaFields() map[string]any { + m := map[string]any{} + if s.Owner != "" && s.Owner != "airflow" { + m["owner"] = s.Owner + } + if !s.StartDate.IsZero() { + m["start_date"] = s.StartDate + } + if !s.EndDate.IsZero() { + m["end_date"] = s.EndDate + } + if s.TriggerRule != "" && s.TriggerRule != "all_success" { + m["trigger_rule"] = s.TriggerRule + } + if s.DependsOnPast { + m["depends_on_past"] = s.DependsOnPast + } + if s.IgnoreFirstDependsOnPast { + m["ignore_first_depends_on_past"] = s.IgnoreFirstDependsOnPast + } + if s.WaitForPastDependsBeforeSkipping { + m["wait_for_past_depends_before_skipping"] = s.WaitForPastDependsBeforeSkipping + } + if s.WaitForDownstream { + m["wait_for_downstream"] = s.WaitForDownstream + } + if s.Retries != 0 { + m["retries"] = s.Retries + } + if s.Queue != "" && s.Queue != "default" { + m["queue"] = s.Queue + } + if s.Pool != "" && s.Pool != "default_pool" { + m["pool"] = s.Pool + } + if s.PoolSlots != 0 && s.PoolSlots != 1 { + m["pool_slots"] = s.PoolSlots + } + if s.ExecutionTimeout != 0 { + m["execution_timeout"] = s.ExecutionTimeout + } + if s.RetryDelay != 0 && s.RetryDelay != 300*time.Second { + m["retry_delay"] = s.RetryDelay + } + if s.RetryExponentialBackoff != 0 { + m["retry_exponential_backoff"] = s.RetryExponentialBackoff + } + if s.MaxRetryDelay != 0 { + m["max_retry_delay"] = s.MaxRetryDelay + } + if s.PriorityWeight != 0 && s.PriorityWeight != 1 { + m["priority_weight"] = s.PriorityWeight + } + if s.WeightRule != "" && s.WeightRule != "downstream" { + m["weight_rule"] = s.WeightRule + } + if s.Executor != "" { + m["executor"] = s.Executor + } + if s.DoXComPush != nil && !*s.DoXComPush { + m["do_xcom_push"] = *s.DoXComPush + } + if s.EmailOnFailure != nil && !*s.EmailOnFailure { + m["email_on_failure"] = *s.EmailOnFailure + } + if s.EmailOnRetry != nil && !*s.EmailOnRetry { + m["email_on_retry"] = *s.EmailOnRetry + } + if s.DocMD != "" { + m["doc_md"] = s.DocMD + } + if s.MapIndexTemplate != "" { + m["map_index_template"] = s.MapIndexTemplate + } + if s.MaxActiveTisPerDag != 0 { + m["max_active_tis_per_dag"] = s.MaxActiveTisPerDag + } + if s.MaxActiveTisPerDagrun != 0 { + m["max_active_tis_per_dagrun"] = s.MaxActiveTisPerDagrun + } + return m +} diff --git a/go-sdk/bundle/bundlev1/spec.go b/go-sdk/bundle/bundlev1/spec.go new file mode 100644 index 0000000000000..cd23f1393fc5e --- /dev/null +++ b/go-sdk/bundle/bundlev1/spec.go @@ -0,0 +1,101 @@ +// 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 ( + "fmt" + "time" +) + +// This file holds the hand-written companions of the generated spec.gen.go: +// DagSpec (whose DagId/Schedule fields and always-emit/config-fallback keys +// the schema cannot express mechanically) and the registration Info structs +// that carry Go-side identity the schema knows nothing about. + +type ( + // DagSpec describes a dag at registration time. DagId is required; every + // other field is optional, with a zero value meaning "unset" so the + // scheduler falls back to its serialization-schema default. The field + // names mirror the keys defined under "dag" in + // airflow-core/src/airflow/serialization/schema.json. + DagSpec struct { + // DagId is the dag id. Required: AddDag panics when it is empty. + DagId string + // Schedule is "@once", "@continuous", a cron expression, or "" for + // NullTimetable (no schedule). + Schedule string + Description string + StartDate time.Time + EndDate time.Time + Tags []string + DagDisplayName string + DocMD string + MaxActiveTasks int + MaxActiveRuns int + MaxConsecutiveFailedDagRuns int + DagrunTimeout time.Duration + Catchup bool + FailFast bool + RenderTemplateAsNativeObj bool + DisableBundleVersioning bool + // IsPausedUponCreation has no schema default. nil means "unset"; pass + // Bool(true) or Bool(false) to set it explicitly. + IsPausedUponCreation *bool + } + + // TaskInfo describes a registered task. Coordinator-mode DAG parsing uses + // it to render the per-task block of a DagFileParsingResult. + TaskInfo struct { + // ID is the user-visible task id (TaskSpec.TaskId when set, otherwise + // the task function's name). + 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 + // Inputs lists the upstream task ids whose return values feed this + // task's data parameters, in parameter order (the Inputs option). + Inputs []string + // Downstream lists task ids that depend on this task, populated as + // later registrations reference this task via Inputs or After. + // 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 per-dag configuration supplied at registration. + Spec DagSpec + Tasks []TaskInfo + } +) + +// applyTask lets a TaskSpec value be passed directly as a TaskOption to +// Dag.Task. Passing more than one spec panics at registration time. +func (s TaskSpec) applyTask(c *taskConfig) { + if c.specSet { + panic(fmt.Errorf("Task accepts at most one TaskSpec")) + } + c.spec = s + c.specSet = true +} diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index d31fea84b73f3..97830d8e53ddc 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -18,7 +18,9 @@ package bundlev1 import ( + "bytes" "context" + "encoding/json" "fmt" "log/slog" "reflect" @@ -29,22 +31,106 @@ import ( "github.com/apache/airflow/go-sdk/sdk" ) +// InputBinding names the upstream task whose return-value XCom fills one data +// parameter of a task function. Bindings are positional: the i-th binding +// fills the i-th data parameter. +type InputBinding struct { + TaskID string +} + +// boundInput is one resolved data parameter: which argument slot to fill, +// from which upstream task, decoding into which type. +type boundInput struct { + argIndex int + taskID string + typ reflect.Type +} + type taskFunction struct { fn reflect.Value fullName string + inputs []boundInput } var _ Task = (*taskFunction)(nil) -// NewTaskFunction wraps a plain Go function as a Task, validating its signature -// (injectable parameters, and a return of error or (result, error)). Bundle -// authors normally use Dag.AddTask, which calls this for them; use it directly -// only when building a Task outside the registry. -func NewTaskFunction(fn any) (Task, error) { +// NewTaskFunction wraps a plain Go function as a Task, validating its +// signature (injectable parameters, data parameters matching inputs, and a +// return of error or (result, error)). Bundle authors normally use Dag.Task, +// which calls this for them; use it directly only when building a Task +// outside the registry. +// +// inputs bind the function's data parameters (every parameter that is not an +// injectable runtime type) positionally to upstream tasks: at execution each +// data parameter receives the named task's return-value XCom, decoded into +// the parameter's type. The number of inputs must equal the number of data +// parameters. +func NewTaskFunction(fn any, inputs ...InputBinding) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{v, fullName} - return f, f.validateFn(v.Type()) + f := &taskFunction{fn: v, fullName: fullName} + if err := f.validateFn(v.Type()); err != nil { + return nil, err + } + if err := f.bindInputs(v.Type(), inputs); err != nil { + return nil, err + } + return f, nil +} + +// dataParamTypes returns the types of fnType's data parameters — every +// parameter Execute does not inject — in declaration order. +func dataParamTypes(fnType reflect.Type) []reflect.Type { + if fnType.Kind() != reflect.Func { + return nil + } + var out []reflect.Type + for i := range fnType.NumIn() { + if in := fnType.In(i); !isInjectable(in) { + out = append(out, in) + } + } + return out +} + +// isInjectable reports whether Execute fills a parameter of this type from +// the runtime (rather than from an upstream task's XCom). +func isInjectable(in reflect.Type) bool { + return isTIRunContext(in) || isContext(in) || isLogger(in) || isClient(in) +} + +// bindInputs pairs the fn's data parameters with the supplied input bindings +// and records the pull-and-decode plan Execute runs before calling fn. +func (f *taskFunction) bindInputs(fnType reflect.Type, inputs []InputBinding) error { + var dataIdx []int + for i := range fnType.NumIn() { + if !isInjectable(fnType.In(i)) { + dataIdx = append(dataIdx, i) + } + } + if len(dataIdx) != len(inputs) { + return fmt.Errorf( + "task function %s declares %d data parameter(s) but %d input binding(s) were supplied", + f.fullName, len(dataIdx), len(inputs), + ) + } + f.inputs = make([]boundInput, len(inputs)) + for i, in := range inputs { + typ := fnType.In(dataIdx[i]) + if in.TaskID == "" { + return fmt.Errorf( + "task function %s: input binding %d has an empty task id", f.fullName, i, + ) + } + if !isDecodableType(typ) { + return fmt.Errorf( + "task function %s: data parameter %d has type %s, which cannot hold an XCom value", + f.fullName, dataIdx[i], typ, + ) + } + f.inputs[i] = boundInput{argIndex: dataIdx[i], taskID: in.TaskID, typ: typ} + } + return nil } func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { @@ -83,10 +169,16 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { case isClient(in): reflectArgs[i] = reflect.ValueOf(sdkClient) default: - // TODO: deal with other value types. For now they will all be Zero values unless it's a context + // A data parameter: placeholder until resolveInputs overwrites it + // below (validation guarantees each has an input binding). reflectArgs[i] = reflect.Zero(in) } } + + if err := f.resolveInputs(ctx, sdkClient, logger, reflectArgs); err != nil { + return err + } + slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs) retValues := f.fn.Call(reflectArgs) @@ -108,6 +200,98 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { return err } +// resolveInputs fills the data-parameter argument slots by pulling each bound +// upstream task's return-value XCom for the current dag run and decoding it +// into the parameter's type. Any failure fails the task before its body runs. +func (f *taskFunction) resolveInputs( + ctx context.Context, + c sdk.XComClient, + logger *slog.Logger, + args []reflect.Value, +) error { + if len(f.inputs) == 0 { + return nil + } + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return fmt.Errorf( + "task function %s: no workload in context; cannot resolve task inputs", f.fullName, + ) + } + ti := workload.TI + for _, in := range f.inputs { + // Pull from the upstream's unmapped instance; mapped upstream fan-in + // is out of scope for now. + raw, err := c.GetXCom(ctx, ti.DagId, ti.RunId, in.taskID, nil, api.XComReturnValueKey, nil) + if err != nil { + return fmt.Errorf( + "task function %s: pulling input from task %q: %w", f.fullName, in.taskID, err, + ) + } + decoded, err := decodeXCom(raw, in.typ) + if err != nil { + return fmt.Errorf( + "task function %s: decoding input from task %q into %s: %w", + f.fullName, in.taskID, in.typ, err, + ) + } + logger.DebugContext( + ctx, + "Resolved task input", + "from_task", + in.taskID, + "type", + in.typ.String(), + ) + args[in.argIndex] = decoded + } + return nil +} + +// decodeXCom decodes a raw (generically deserialised) XCom value into target. +// Decoding into a struct is strict: unknown/renamed keys fail rather than +// silently leaving fields zero. Decoding into a map / interface accepts any +// shape, so authors opt into loose decoding by typing the parameter +// map[string]any or any. A null value is allowed only for a nilable target. +func decodeXCom(raw any, target reflect.Type) (reflect.Value, error) { + out := reflect.New(target) + + if raw == nil { + switch target.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface: + return out.Elem(), nil + default: + return reflect.Value{}, fmt.Errorf( + "xcom value is null but parameter type %s is not nilable", target, + ) + } + } + + blob, err := json.Marshal(raw) + if err != nil { + return reflect.Value{}, err + } + dec := json.NewDecoder(bytes.NewReader(blob)) + dec.DisallowUnknownFields() + if err := dec.Decode(out.Interface()); err != nil { + return reflect.Value{}, err + } + return out.Elem(), nil +} + +// isDecodableType reports whether a value can be JSON-decoded into inType. It +// rejects kinds json cannot target (func, chan, unsafe pointer) and non-empty +// interfaces (only the empty interface `any` is a valid decode target). +func isDecodableType(inType reflect.Type) bool { + switch inType.Kind() { + case reflect.Func, reflect.Chan, reflect.UnsafePointer: + return false + case reflect.Interface: + return inType.NumMethod() == 0 + } + return true +} + func (f *taskFunction) sendXcom( ctx context.Context, value any, @@ -200,8 +384,10 @@ func isClient(inType reflect.Type) bool { inType.NumMethod() > 0 && clientType.Implements(inType) } -// validateParam rejects interface parameters Execute cannot inject; they -// would be bound to nil and panic on first use. +// validateParam rejects interface parameters Execute cannot inject and that +// cannot serve as data parameters either; they would be bound to nil and +// panic on first use. Non-interface parameters are data parameters, filled +// from upstream XComs via the Inputs bindings. func validateParam(in reflect.Type) error { if in.Kind() != reflect.Interface || isTIRunContext(in) || isClient(in) { return nil @@ -216,6 +402,10 @@ func validateParam(in reflect.Type) error { in, ) } + if in.NumMethod() == 0 { + // The empty interface is a valid (loose) data-parameter target. + return nil + } return fmt.Errorf( "interface %s is not injectable (want context.Context, sdk.TIRunContext, or a subset of sdk.Client): %s", in, diff --git a/go-sdk/bundle/bundlev1/task_test.go b/go-sdk/bundle/bundlev1/task_test.go index 3f58e930b8721..9f902bfea8d71 100644 --- a/go-sdk/bundle/bundlev1/task_test.go +++ b/go-sdk/bundle/bundlev1/task_test.go @@ -19,12 +19,14 @@ package bundlev1 import ( "context" + "errors" "log/slog" "reflect" "testing" "github.com/stretchr/testify/suite" + "github.com/apache/airflow/go-sdk/pkg/api" "github.com/apache/airflow/go-sdk/pkg/logging" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" @@ -174,10 +176,6 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { }, "method GetVariable is func(context.Context, string) (string, error) on sdk.Client", }, - "empty-interface": { - func(x any) error { return nil }, - "empty interfaces cannot be injected", - }, "context-with-extra-methods": { func(x interface { context.Context @@ -276,3 +274,247 @@ func (s *TaskSuite) TestTIRunContextInjectionWithoutRuntimeContext() { "the injected context must be backed by the one passed to Execute", ) } + +// --- Data-parameter (TaskFlow Inputs) binding --- + +// fakeClient implements sdk.Client with canned per-task XCom values, recording +// the pulls Execute performs to resolve data parameters. +type fakeClient struct { + xcoms map[string]any // upstream task id -> stored return_value + getErr error + pulls []string + pushed map[string]any + pushErr error +} + +func (f *fakeClient) GetVariable(context.Context, string) (string, error) { return "", nil } +func (f *fakeClient) UnmarshalJSONVariable(context.Context, string, any) error { + return nil +} + +func (f *fakeClient) GetConnection(context.Context, string) (sdk.Connection, error) { + return sdk.Connection{}, nil +} + +func (f *fakeClient) GetXCom( + _ context.Context, + dagId, runId, taskId string, + mapIndex *int, + key string, + _ any, +) (any, error) { + f.pulls = append(f.pulls, taskId) + if f.getErr != nil { + return nil, f.getErr + } + if dagId != "dag1" || runId != "run1" || key != api.XComReturnValueKey || mapIndex != nil { + return nil, errors.New("unexpected xcom pull identifiers") + } + return f.xcoms[taskId], nil +} + +func (f *fakeClient) PushXCom(_ context.Context, _ api.TaskInstance, key string, value any) error { + if f.pushed == nil { + f.pushed = map[string]any{} + } + f.pushed[key] = value + return f.pushErr +} + +// bindingResult is the struct shape the fake upstream produces. +type bindingResult struct { + Value int `json:"value"` + Name string `json:"name"` +} + +func (s *TaskSuite) executionContext(client sdk.Client) context.Context { + ctx := context.WithValue(context.Background(), sdkcontext.SdkClientContextKey, client) + return context.WithValue(ctx, sdkcontext.WorkloadContextKey, api.ExecuteTaskWorkload{ + TI: api.TaskInstance{DagId: "dag1", RunId: "run1", TaskId: "consumer"}, + }) +} + +func (s *TaskSuite) TestDataParamInjection() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"value": 42, "name": "answer"}, + }} + + var got bindingResult + task, err := NewTaskFunction(func(in bindingResult) error { + got = in + return nil + }, InputBinding{TaskID: "extract"}) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + s.Require().NoError(err) + s.Equal(bindingResult{Value: 42, Name: "answer"}, got) + s.Equal([]string{"extract"}, client.pulls) +} + +func (s *TaskSuite) TestDataParamsBindPositionallyAmongInjectables() { + client := &fakeClient{xcoms: map[string]any{ + "first": map[string]any{"value": 1, "name": "a"}, + "second": map[string]any{"value": 2, "name": "b"}, + }} + + var gotA, gotB bindingResult + task, err := NewTaskFunction( + func(ctx context.Context, a bindingResult, logger *slog.Logger, b bindingResult, c sdk.XComClient) error { + s.NotNil(ctx) + s.NotNil(logger) + s.NotNil(c) + gotA, gotB = a, b + return nil + }, + InputBinding{TaskID: "first"}, + InputBinding{TaskID: "second"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + s.Require().NoError(err) + s.Equal(bindingResult{Value: 1, Name: "a"}, gotA) + s.Equal(bindingResult{Value: 2, Name: "b"}, gotB) + s.Equal([]string{"first", "second"}, client.pulls) +} + +func (s *TaskSuite) TestDataParamStrictDecodeRejectsUnknownKeys() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"value": 1, "renamed_field": true}, + }} + + task, err := NewTaskFunction( + func(in bindingResult) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), `decoding input from task "extract"`) + s.Contains(err.Error(), "renamed_field") + } +} + +func (s *TaskSuite) TestDataParamLooseMapDecodesAnyShape() { + client := &fakeClient{xcoms: map[string]any{ + "extract": map[string]any{"anything": "goes", "extra": 1}, + }} + + var got map[string]any + task, err := NewTaskFunction( + func(in map[string]any) error { + got = in + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + s.Require().NoError(task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger()))) + s.Equal("goes", got["anything"]) +} + +func (s *TaskSuite) TestDataParamNullValue() { + client := &fakeClient{xcoms: map[string]any{"extract": nil}} + + s.Run("nilable-target-gets-zero", func() { + var got *bindingResult + task, err := NewTaskFunction( + func(in *bindingResult) error { + got = in + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + s.Require(). + NoError(task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger()))) + s.Nil(got) + }) + + s.Run("non-nilable-target-fails", func() { + task, err := NewTaskFunction( + func(in bindingResult) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), "not nilable") + } + }) +} + +func (s *TaskSuite) TestDataParamPullFailureFailsTask() { + client := &fakeClient{getErr: errors.New("supervisor unreachable")} + + ran := false + task, err := NewTaskFunction( + func(in bindingResult) error { + ran = true + return nil + }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + err = task.Execute(s.executionContext(client), slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), `pulling input from task "extract"`) + } + s.False(ran, "the task body must not run when an input pull fails") +} + +func (s *TaskSuite) TestDataParamWithoutWorkloadFails() { + client := &fakeClient{xcoms: map[string]any{"extract": map[string]any{"value": 1}}} + task, err := NewTaskFunction( + func(in map[string]any) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + s.Require().NoError(err) + + // Client injected but no workload on the context. + ctx := context.WithValue( + context.Background(), + sdkcontext.SdkClientContextKey, + sdk.Client(client), + ) + err = task.Execute(ctx, slog.New(logging.NewTeeLogger())) + if s.Error(err) { + s.Contains(err.Error(), "no workload in context") + } +} + +func (s *TaskSuite) TestInputBindingValidation() { + s.Run("count-mismatch", func() { + _, err := NewTaskFunction(func(x any) error { return nil }) + if s.Error(err) { + s.Contains( + err.Error(), + "declares 1 data parameter(s) but 0 input binding(s) were supplied", + ) + } + }) + + s.Run("empty-task-id", func() { + _, err := NewTaskFunction( + func(x any) error { return nil }, + InputBinding{}, + ) + if s.Error(err) { + s.Contains(err.Error(), "input binding 0 has an empty task id") + } + }) + + s.Run("undecodable-param", func() { + _, err := NewTaskFunction( + func(x chan int) error { return nil }, + InputBinding{TaskID: "extract"}, + ) + if s.Error(err) { + s.Contains(err.Error(), "cannot hold an XCom value") + } + }) +} diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..1970a32e4aece 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -45,15 +45,40 @@ func (m *myBundle) GetBundleVersion() v1.BundleInfo { return v1.BundleInfo{Name: bundleName, Version: &bundleVersion} } +// ExtractResult is extract's return value. Returning it pushes it as the +// task's return_value XCom; declaring it via Inputs feeds it to a downstream +// task's matching data parameter. +type ExtractResult struct { + GoVersion string `json:"go_version"` + Timestamp int64 `json:"timestamp"` +} + +// TransformResult is transform's return value, consumed by load. +type TransformResult struct { + Variable string `json:"variable"` + Extracted ExtractResult `json:"extracted"` +} + func (m *myBundle) RegisterDags(dagbag v1.Registry) error { - simpleDag := dagbag.AddDag("simple_dag") - simpleDag.AddTask(extract) - simpleDag.AddTask(transform) - simpleDag.AddTask(load) + simpleDag := dagbag.AddDag(v1.DagSpec{ + DagId: "simple_dag", + Schedule: "@daily", + Description: "Example Go-authored Dag", + Tags: []string{"example", "go-sdk"}, + }) + // TaskFlow style: each Task call returns a ref; passing refs via + // v1.Inputs both orders the tasks and feeds the upstream's return value + // into the downstream function's data parameter. + extracted := simpleDag.Task(extract, v1.TaskSpec{Queue: "go-task", Retries: 2}) + transformed := simpleDag.Task(transform, v1.TaskSpec{Queue: "go-task"}, v1.Inputs(extracted)) + simpleDag.Task(load, v1.TaskSpec{Queue: "go-task"}, v1.Inputs(transformed)) // Tasks defined in other packages register through the same dagbag. - concurrentDag := dagbag.AddDag("concurrent_xcom_dag") - concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + concurrentDag := dagbag.AddDag(v1.DagSpec{DagId: "concurrent_xcom_dag"}) + concurrentDag.Task( + concurrentxcom.PullXComsConcurrently, + v1.TaskSpec{TaskId: "pull_xcoms_concurrently"}, + ) return nil } @@ -64,7 +89,7 @@ func main() { } } -func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { +func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (ExtractResult, error) { log.Info("Hello from task") // ctx behaves as a context.Context and also carries the task instance @@ -111,7 +136,7 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er // Once per loop,.check if we've been asked to cancel! select { case <-ctx.Done(): - return nil, ctx.Err() + return ExtractResult{}, ctx.Err() default: } log.Info("After the beep the time will be", "time", time.Now()) @@ -119,40 +144,51 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er } log.Info("Goodbye from task") - ret := map[string]any{ - "go_version": runtime.Version(), - "timestamp": time.Now().UnixNano(), - } - - return ret, nil + return ExtractResult{ + GoVersion: runtime.Version(), + Timestamp: time.Now().UnixNano(), + }, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { - // This function takes a VariableClient and not a Client to make unit testing it easier. See - // `./main_test.go` for an example unit of this task fn. Functionally taking a `sdk.Client` is the same (as - // Client includes VariableClient) but by using the dedicated type it can be easier to write unit tests. - // - // It also gives a better indication of what features the tasks use +// transform receives extract's return value as its `in` data parameter: the +// runtime pulls extract's return-value XCom and decodes it into ExtractResult +// before the body runs, because RegisterDags wired v1.Inputs(extracted). +// +// It takes a VariableClient and not a full Client to make unit testing easier +// (see ./main_test.go); functionally `sdk.Client` works the same, but the +// dedicated type documents what the task actually uses. +func transform( + ctx sdk.TIRunContext, + client sdk.VariableClient, + log *slog.Logger, + in ExtractResult, +) (TransformResult, error) { + log.Info("Received upstream result", "go_version", in.GoVersion, "timestamp", in.Timestamp) + key := "my_variable" val, err := client.GetVariable(ctx, key) if err != nil { - return err + return TransformResult{}, err } log.Info("Obtained variable", key, val) - return nil + return TransformResult{Variable: val, Extracted: in}, nil } // load fails on its first attempt and succeeds on the retry. With retries -// configured on the stub task, the first failure makes the supervisor mark the +// configured on the task, the first failure makes the supervisor mark the // task UP_FOR_RETRY -- which only works because the Go SDK now emits a // RetryTask frame (instead of a terminal FAILED) when ti_context.should_retry // is set. The retry then runs this task again and it returns nil. -func load(ctx sdk.TIRunContext, log *slog.Logger) error { +func load(ctx sdk.TIRunContext, log *slog.Logger, in TransformResult) error { tryNumber := ctx.TaskInstance().TryNumber if tryNumber == 1 { log.InfoContext(ctx, "Please fail", "try_number", tryNumber) return fmt.Errorf("Please fail") } - log.InfoContext(ctx, "Recovered on retry", "try_number", tryNumber) + log.InfoContext(ctx, "Recovered on retry", + "try_number", tryNumber, + "variable", in.Variable, + "extracted_go_version", in.Extracted.GoVersion, + ) return nil } diff --git a/go-sdk/example/bundle/main_test.go b/go-sdk/example/bundle/main_test.go index 474a8406d6224..75773784bb502 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/go-sdk/example/bundle/main_test.go @@ -51,9 +51,14 @@ var _ sdk.VariableClient = (*mockVars)(nil) func Test_transform(t *testing.T) { log := slog.Default() - // This is not the best test, but it is a good proof of concept -- you can just call the function. + // This is not the best test, but it is a good proof of concept -- you can + // just call the function, passing the upstream result directly instead of + // letting the runtime pull it from XCom. // sdk.NewTIRunContext wraps any context to build a TIRunContext in a test. ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - err := transform(ctx, &mockVars{}, log) + in := ExtractResult{GoVersion: "go1.24", Timestamp: 42} + got, err := transform(ctx, &mockVars{}, log, in) assert.NoError(t, err) + assert.Equal(t, "value1", got.Variable) + assert.Equal(t, in, got.Extracted) } diff --git a/go-sdk/pkg/execution/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..e6e5e4be5c7b5 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(bundlev1.DagSpec{DagId: "test_dag"}) + d.Task(simpleTask) + }) + + 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(bundlev1.DagSpec{DagId: "dag1"}).Task(simpleTask) + r.AddDag(bundlev1.DagSpec{DagId: "dag2"}).Task(failingTask) + }) + + 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(simpleTask) }) 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(failingTask) }) 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(failingTask) }) 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(simpleTask) }) 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(panicTask) }) 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(bundlev1.DagSpec{DagId: "test_dag"}).Task(panicTask) }) details := &genmodels.StartupDetails{ @@ -232,8 +291,9 @@ func TestTaskRunnerPanicRetry(t *testing.T) { func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxcheck", - func(ctx context.Context) error { return ctx.Err() }) + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( + func(ctx context.Context) error { return ctx.Err() }, + bundlev1.TaskSpec{TaskId: "ctxcheck"}) }) details := &genmodels.StartupDetails{ @@ -266,11 +326,11 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { var got sdk.TIRunContext bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxgrab", + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{TaskId: "ctxgrab"}) }) details := &genmodels.StartupDetails{ @@ -326,11 +386,11 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) { func TestRunTaskRuntimeContextMappedIndex(t *testing.T) { var got sdk.TIRunContext bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("ctxgrab", + r.AddDag(bundlev1.DagSpec{DagId: "test_dag"}).Task( func(ctx sdk.TIRunContext) error { got = ctx return nil - }) + }, bundlev1.TaskSpec{TaskId: "ctxgrab"}) }) details := &genmodels.StartupDetails{ @@ -406,13 +466,84 @@ 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(bundlev1.DagSpec{DagId: "simple_dag"}) + extract := d.Task( + func() (string, error) { return "data", nil }, + bundlev1.TaskSpec{TaskId: "extract"}, + ) + d.Task( + func(in string) error { return nil }, + bundlev1.TaskSpec{TaskId: "transform"}, + bundlev1.Inputs(extract), + ) + return nil + }, + } + + 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"]) + + // The Inputs wiring must surface as downstream_task_ids on the upstream. + tasks := dag["tasks"].([]any) + require.Len(t, tasks, 2) + extractData := tasks[0].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "extract", extractData["task_id"]) + assert.Equal(t, []any{"transform"}, extractData["downstream_task_ids"]) + transformData := tasks[1].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "transform", transformData["task_id"]) + assert.Nil(t, transformData["downstream_task_ids"]) + + select { + case err := <-done: + require.NoError(t, err) + 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(bundlev1.DagSpec{DagId: "dag1"}).Task(simpleTask) return nil }, } @@ -470,7 +601,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { var gotVar string provider := &fakeProvider{ register: func(r bundlev1.Registry) error { - r.AddDag("dag1").AddTaskWithName("getvar", + r.AddDag(bundlev1.DagSpec{DagId: "dag1"}).Task( func(ctx context.Context, c sdk.Client) (string, error) { v, err := c.GetVariable(ctx, varKey) if err != nil { @@ -478,7 +609,7 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { } gotVar = v return "xval", nil - }) + }, bundlev1.TaskSpec{TaskId: "getvar"}) return nil }, } diff --git a/go-sdk/pkg/execution/serde.go b/go-sdk/pkg/execution/serde.go new file mode 100644 index 0000000000000..b323688a68ab2 --- /dev/null +++ b/go-sdk/pkg/execution/serde.go @@ -0,0 +1,389 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package execution + +import ( + "fmt" + "path/filepath" + "reflect" + "sort" + "time" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// Per-DAG defaults that Python resolves from [core] config when the DAG does +// not override them. The serializer always emits these fields (they have no +// JSON-schema default to omit against), so we fall back to the same values. +const ( + defaultMaxActiveTasksPerDag = 16 // [core] max_active_tasks_per_dag + defaultMaxActiveRunsPerDag = 16 // [core] max_active_runs_per_dag +) + +// serializeValue recursively serializes a value with Airflow's type/var encoding. +// This matches Python's BaseSerialization.serialize() output: +// - primitives (string, bool, int, float) pass through unchanged +// - time.Time -> {"__type": "datetime", "__var": epoch_seconds_float} +// - time.Duration -> {"__type": "timedelta", "__var": total_seconds_float} +// - map[string]any -> {"__type": "dict", "__var": {k: serialize(v), ...}} +// - []any -> direct array with each element serialized +func serializeValue(value any) any { + if value == nil { + return nil + } + switch v := value.(type) { + case string, bool: + return v + case int: + return v + case int8: + return int(v) + case int16: + return int(v) + case int32: + return int(v) + case int64: + return v + case float32: + return float64(v) + case float64: + return v + case time.Time: + epochSec := float64(v.Unix()) + float64(v.Nanosecond())/1e9 + return map[string]any{ + "__type": "datetime", + "__var": epochSec, + } + case time.Duration: + return map[string]any{ + "__type": "timedelta", + "__var": v.Seconds(), + } + case map[string]any: + serialized := make(map[string]any, len(v)) + for k, val := range v { + serialized[k] = serializeValue(val) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case []string: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + case []any: + result := make([]any, len(v)) + for i, item := range v { + result[i] = serializeValue(item) + } + return result + default: + // Use reflection to handle typed maps and slices that don't match + // the concrete types above (e.g., map[string]map[string][]string). + rv := reflect.ValueOf(v) + switch rv.Kind() { + case reflect.Map: + serialized := make(map[string]any, rv.Len()) + for _, key := range rv.MapKeys() { + serialized[fmt.Sprint(key.Interface())] = serializeValue(rv.MapIndex(key).Interface()) + } + return map[string]any{ + "__type": "dict", + "__var": serialized, + } + case reflect.Slice, reflect.Array: + result := make([]any, rv.Len()) + for i := range result { + result[i] = serializeValue(rv.Index(i).Interface()) + } + return result + default: + return v + } + } +} + +// unwrapTypeEncoding extracts the "__var" part from a type-encoded value. +// In Python's serialize_to_json, non-decorated fields are serialized then unwrapped. +func unwrapTypeEncoding(value any) any { + m, ok := value.(map[string]any) + if !ok { + return value + } + if _, hasType := m["__type"]; !hasType { + return value + } + if v, hasVar := m["__var"]; hasVar { + return v + } + return value +} + +// serializeTimetable converts a schedule string to the Airflow timetable format. +// +// TODO: respect [scheduler] create_cron_data_intervals (and, once timedelta +// schedules are supported, create_delta_data_intervals). Python's +// _create_timetable selects CronDataIntervalTimetable when +// create_cron_data_intervals is True and CronTriggerTimetable when False; we +// hardcode CronTriggerTimetable, which matches only the default (False). The +// Go bundle binary cannot read airflow.cfg, so the supervisor (Python +// ExecutableCoordinator) must send these scheduler flags to the lang-SDK over +// the coordinator protocol (e.g. on DagFileParseRequest) before we can honor +// non-default deployments. Until that channel exists this stays default-only. +// Tracked at https://github.com/apache/airflow/issues/67938 +func serializeTimetable(schedule *string) map[string]any { + if schedule == nil { + return map[string]any{ + "__type": "airflow.timetables.simple.NullTimetable", + "__var": map[string]any{}, + } + } + switch *schedule { + case "@once": + return map[string]any{ + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": map[string]any{}, + } + case "@continuous": + return map[string]any{ + "__type": "airflow.timetables.simple.ContinuousTimetable", + "__var": map[string]any{}, + } + default: + return map[string]any{ + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": map[string]any{ + "expression": *schedule, + "timezone": "UTC", + "interval": 0.0, + "run_immediately": false, + }, + } + } +} + +// serializeTask converts a task to the Airflow serialization format. The +// downstream_task_ids slice is read from info.Downstream (populated by the +// registry from each task's `depends` argument) and sorted for stable JSON. +func serializeTask(info bundlev1.TaskInfo) map[string]any { + typeName := info.TypeName + if typeName == "" { + typeName = info.ID + } + pkgPath := info.PkgPath + if pkgPath == "" { + pkgPath = "main" + } + data := map[string]any{ + "task_id": info.ID, + "task_type": typeName, + "_task_module": pkgPath, + "language": "go", + // Python's operator serializer always emits template_fields (its + // list value never matches the tuple default it is compared against), + // so it is unconditional here too. Go tasks have no template fields. + "template_fields": []any{}, + } + // 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. Fields with a JSON-schema +// default (description, dates, tags, fail_fast, …) are omitted when unset, as +// in applyTaskSpec. Fields with no schema default (catchup, +// disable_bundle_versioning, max_active_tasks, max_active_runs, +// max_consecutive_failed_dag_runs) are always emitted, because Python's +// serializer never omits them. +func applyDagSpec(data map[string]any, s bundlev1.DagSpec) { + if s.Description != "" { + data["description"] = s.Description + } + if !s.StartDate.IsZero() { + data["start_date"] = unwrapTypeEncoding(serializeValue(s.StartDate)) + } + if !s.EndDate.IsZero() { + data["end_date"] = unwrapTypeEncoding(serializeValue(s.EndDate)) + } + if len(s.Tags) > 0 { + // Python stores tags in a set and serializes them sorted (for stable + // dag_hash); mirror that here regardless of registration order. + sorted := make([]string, len(s.Tags)) + copy(sorted, s.Tags) + sort.Strings(sorted) + tags := make([]any, len(sorted)) + for i, t := range sorted { + tags[i] = t + } + data["tags"] = tags + } + if s.DagDisplayName != "" { + data["dag_display_name"] = s.DagDisplayName + } + if s.DocMD != "" { + data["doc_md"] = s.DocMD + } + // max_active_tasks / max_active_runs / max_consecutive_failed_dag_runs and + // the catchup / disable_bundle_versioning booleans have no schema default, + // so Python's serializer never omits them — it always writes the resolved + // value (the per-DAG default falls back to the matching [core] config). + // Emit them unconditionally to match, using the config defaults for unset + // (zero) fields. + maxActiveTasks := s.MaxActiveTasks + if maxActiveTasks == 0 { + maxActiveTasks = defaultMaxActiveTasksPerDag + } + data["max_active_tasks"] = maxActiveTasks + maxActiveRuns := s.MaxActiveRuns + if maxActiveRuns == 0 { + maxActiveRuns = defaultMaxActiveRunsPerDag + } + data["max_active_runs"] = maxActiveRuns + data["max_consecutive_failed_dag_runs"] = s.MaxConsecutiveFailedDagRuns + data["catchup"] = s.Catchup + data["disable_bundle_versioning"] = s.DisableBundleVersioning + if s.DagrunTimeout != 0 { + data["dagrun_timeout"] = unwrapTypeEncoding(serializeValue(s.DagrunTimeout)) + } + // fail_fast and render_template_as_native_obj both have schema default + // false, so Python omits them when false; keep that behavior. + if s.FailFast { + data["fail_fast"] = true + } + if s.RenderTemplateAsNativeObj { + data["render_template_as_native_obj"] = true + } + if s.IsPausedUponCreation != nil { + data["is_paused_upon_creation"] = *s.IsPausedUponCreation + } +} + +// serializeTaskGroup creates a flat root task group containing all task IDs. +func serializeTaskGroup(taskIDs []string) map[string]any { + children := make(map[string]any, len(taskIDs)) + for _, id := range taskIDs { + children[id] = []any{"operator", id} + } + return map[string]any{ + "_group_id": nil, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": children, + "upstream_group_ids": []any{}, + "downstream_group_ids": []any{}, + "upstream_task_ids": []any{}, + "downstream_task_ids": []any{}, + } +} + +// serializeParams converts DAG params to Airflow's serialization format. +func serializeParams(params map[string]any) []any { + if len(params) == 0 { + return []any{} + } + result := make([]any, 0, len(params)) + for k, v := range params { + result = append(result, []any{ + k, + map[string]any{ + "__class": "airflow.sdk.definitions.param.Param", + "default": serializeValue(v), + "description": nil, + "schema": serializeValue(map[string]any{}), + "source": nil, + }, + }) + } + return result +} + +// SerializeDag converts a bundlev1.DagInfo to Airflow DagSerialization v3 +// format. Required fields are always present; spec-driven fields are emitted +// per the rules in applyDagSpec (some always, some only when set). +func SerializeDag(info bundlev1.DagInfo, fileloc, relativeFileloc string) map[string]any { + taskIDs := make([]string, len(info.Tasks)) + tasks := make([]any, len(info.Tasks)) + for i, t := range info.Tasks { + taskIDs[i] = t.ID + tasks[i] = serializeTask(t) + } + + var schedule *string + if info.Spec.Schedule != "" { + s := info.Spec.Schedule + schedule = &s + } + + result := map[string]any{ + // Required fields (always present) + "dag_id": info.DagID, + "fileloc": fileloc, + "relative_fileloc": relativeFileloc, + "timezone": "UTC", + "timetable": serializeTimetable(schedule), + "tasks": tasks, + "dag_dependencies": []any{}, + "task_group": serializeTaskGroup(taskIDs), + "edge_info": map[string]any{}, + "params": serializeParams(nil), + "deadline": nil, + "allowed_run_types": nil, + } + applyDagSpec(result, info.Spec) + return result +} + +// computeRelativeFileloc computes the relative file location from the bundle path. +func computeRelativeFileloc(fileloc, bundlePath string) string { + if fileloc == "" { + return "" + } + if bundlePath == "" { + return "." + } + rel, err := filepath.Rel(bundlePath, fileloc) + if err != nil { + return "." + } + if rel == "" { + return "." + } + return rel +} diff --git a/go-sdk/pkg/execution/serde_test.go b/go-sdk/pkg/execution/serde_test.go new file mode 100644 index 0000000000000..83e7df6a67751 --- /dev/null +++ b/go-sdk/pkg/execution/serde_test.go @@ -0,0 +1,457 @@ +// 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. + assert.Equal(t, false, result["catchup"]) + assert.Equal(t, false, result["disable_bundle_versioning"]) + assert.Equal(t, defaultMaxActiveTasksPerDag, result["max_active_tasks"]) + assert.Equal(t, defaultMaxActiveRunsPerDag, result["max_active_runs"]) + assert.Equal(t, 0, result["max_consecutive_failed_dag_runs"]) +} + +func TestSerializeDagWithTasks(t *testing.T) { + info := bundlev1.DagInfo{ + DagID: "etl", + Tasks: []bundlev1.TaskInfo{ + { + ID: "extract", TypeName: "extract", PkgPath: "main", + Downstream: []string{"load"}, + }, + { + ID: "load", TypeName: "load", PkgPath: "main", + Spec: bundlev1.TaskSpec{Queue: "high_mem"}, + }, + }, + } + result := SerializeDag(info, "/bundle/main.go", "main.go") + + tasks := result["tasks"].([]any) + require.Len(t, tasks, 2) + first := tasks[0].(map[string]any) + v := first["__var"].(map[string]any) + assert.Equal(t, "extract", v["task_id"]) + assert.Equal(t, "extract", v["task_type"]) + assert.Equal(t, "main", v["_task_module"]) + assert.Equal(t, "go", v["language"]) + _, hasQueue := v["queue"] + assert.False(t, hasQueue, "extract has no queue set; field should be omitted") + assert.Equal(t, []string{"load"}, v["downstream_task_ids"]) + + second := tasks[1].(map[string]any)["__var"].(map[string]any) + assert.Equal(t, "high_mem", second["queue"]) + _, hasDownstream := second["downstream_task_ids"] + assert.False(t, hasDownstream, "leaf task has no downstream") + + tg := result["task_group"].(map[string]any) + children := tg["children"].(map[string]any) + assert.Contains(t, children, "extract") + assert.Contains(t, children, "load") +} + +func TestSerializeDagWithSpec(t *testing.T) { + start := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + info := bundlev1.DagInfo{ + DagID: "etl", + Spec: bundlev1.DagSpec{ + Schedule: "@daily", + Description: "Extract, transform, load", + StartDate: start, + Tags: []string{"prod", "etl"}, + DagDisplayName: "ETL Pipeline", + DocMD: "## ETL", + MaxActiveTasks: 32, + MaxActiveRuns: 4, + MaxConsecutiveFailedDagRuns: 3, + DagrunTimeout: 2 * time.Hour, + Catchup: true, + FailFast: true, + RenderTemplateAsNativeObj: true, + DisableBundleVersioning: true, + IsPausedUponCreation: bundlev1.Bool(true), + }, + } + result := SerializeDag(info, "/bundle/main.go", "main.go") + + tt := result["timetable"].(map[string]any) + assert.Equal(t, "airflow.timetables.trigger.CronTriggerTimetable", tt["__type"]) + v := tt["__var"].(map[string]any) + assert.Equal(t, "@daily", v["expression"]) + + assert.Equal(t, "Extract, transform, load", result["description"]) + // Tags are emitted sorted, matching Python's set-backed serialization. + assert.Equal(t, []any{"etl", "prod"}, result["tags"]) + assert.Equal(t, "ETL Pipeline", result["dag_display_name"]) + assert.Equal(t, "## ETL", result["doc_md"]) + assert.Equal(t, 32, result["max_active_tasks"]) + assert.Equal(t, 4, result["max_active_runs"]) + assert.Equal(t, 3, result["max_consecutive_failed_dag_runs"]) + assert.Equal(t, (2 * time.Hour).Seconds(), result["dagrun_timeout"]) + assert.Equal(t, true, result["catchup"]) + assert.Equal(t, true, result["fail_fast"]) + assert.Equal(t, true, result["render_template_as_native_obj"]) + assert.Equal(t, true, result["disable_bundle_versioning"]) + assert.Equal(t, true, result["is_paused_upon_creation"]) + + // start_date is a raw epoch number, not the type-wrapped form. + startDate := result["start_date"].(float64) + assert.InDelta(t, float64(start.Unix()), startDate, 0.001) +} + +func TestApplyDagSpec_AlwaysEmitsNonSchemaDefaultFields(t *testing.T) { + // catchup, disable_bundle_versioning and the three max_* fields have no + // JSON-schema default, so Python always serializes them. An empty spec + // must still emit them at their resolved [core] defaults. + data := map[string]any{} + applyDagSpec(data, bundlev1.DagSpec{}) + + assert.Equal(t, false, data["catchup"]) + assert.Equal(t, false, data["disable_bundle_versioning"]) + assert.Equal(t, defaultMaxActiveTasksPerDag, data["max_active_tasks"]) + assert.Equal(t, defaultMaxActiveRunsPerDag, data["max_active_runs"]) + assert.Equal(t, 0, data["max_consecutive_failed_dag_runs"]) + + // Fields with a false/empty schema default stay omitted when unset. + for _, k := range []string{ + "description", "start_date", "end_date", "tags", "dag_display_name", + "doc_md", "dagrun_timeout", "fail_fast", "render_template_as_native_obj", + "is_paused_upon_creation", + } { + _, ok := data[k] + assert.Falsef(t, ok, "%q should be omitted when unset", k) + } +} + +func TestApplyDagSpec_OmitsFalseSchemaDefaultBooleans(t *testing.T) { + // fail_fast and render_template_as_native_obj default to false in the + // schema, so Python omits them unless explicitly true. + data := map[string]any{} + applyDagSpec(data, bundlev1.DagSpec{FailFast: false, RenderTemplateAsNativeObj: false}) + _, hasFailFast := data["fail_fast"] + assert.False(t, hasFailFast) + _, hasNative := data["render_template_as_native_obj"] + assert.False(t, hasNative) +} + +func TestComputeRelativeFileloc(t *testing.T) { + tests := []struct { + fileloc string + bundlePath string + want string + }{ + {"", "", ""}, + {"/a/b/c.go", "", "."}, + {"/bundles/my/dags.go", "/bundles/my", "dags.go"}, + {"/bundles/my/sub/dags.go", "/bundles/my", "sub/dags.go"}, + } + for _, tt := range tests { + result := computeRelativeFileloc(tt.fileloc, tt.bundlePath) + assert.Equal(t, tt.want, result, "fileloc=%q bundlePath=%q", tt.fileloc, tt.bundlePath) + } +} diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 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..43260c1db0156 100644 --- a/go-sdk/pkg/worker/runner_test.go +++ b/go-sdk/pkg/worker/runner_test.go @@ -160,10 +160,11 @@ func (s *WorkerSuite) TestStartContextErrorTaskDoesntStart() { wasCalled := false // Register a task that should NOT be called if everything works - s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { - wasCalled = true - return nil - }) + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func() error { + wasCalled = true + return nil + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) // Setup the mock s.ti.EXPECT(). @@ -193,10 +194,11 @@ func (s *WorkerSuite) TestTaskHeartbeatsWhileRunning() { id := uuid.New().String() var callCount atomic.Int32 testWorkload := newTestWorkLoad(id, id[:8]) - s.registry.AddDag(testWorkload.TI.DagId).AddTaskWithName(testWorkload.TI.TaskId, func() error { - time.Sleep(time.Second) - return nil - }) + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func() error { + time.Sleep(time.Second) + return nil + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateSuccess) @@ -224,15 +226,15 @@ func (s *WorkerSuite) TestTaskHeartbeatConflictStopsTask() { id := uuid.New().String() testWorkload := newTestWorkLoad(id, id[:8]) - s.registry.AddDag(testWorkload.TI.DagId). - AddTaskWithName(testWorkload.TI.TaskId, func(ctx context.Context) error { + s.registry.AddDag(bundlev1.DagSpec{DagId: testWorkload.TI.DagId}). + Task(func(ctx context.Context) error { select { case <-ctx.Done(): return nil case <-time.After(2 * time.Second): return fmt.Errorf("task context was not cancelled") } - }) + }, bundlev1.TaskSpec{TaskId: testWorkload.TI.TaskId}) s.ExpectTaskRun(id) s.ExpectTaskState(id, api.TerminalTIStateFailed)