From 2cb70850d8209aa1768e09e50f4d0d948ce5fb5f Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Tue, 20 Oct 2026 12:24:27 +0800 Subject: [PATCH 1/2] Go SDK: Write the edge label text before the task Signed-off-by: PoAn Yang --- go-sdk/adr/0008-native-dag-interface.md | 17 ++- go-sdk/airflow/conformance_test.go | 2 +- go-sdk/airflow/dag.go | 6 +- go-sdk/airflow/if_test.go | 2 +- go-sdk/airflow/node.go | 121 +++++++++++++---- go-sdk/airflow/node_test.go | 173 +++++++++++++++++++----- go-sdk/airflow/serialize_test.go | 8 +- go-sdk/airflow/task_group_test.go | 30 ++-- go-sdk/example/native/main.go | 2 +- go-sdk/internal/genspec/authoring.go | 2 +- 10 files changed, 271 insertions(+), 92 deletions(-) diff --git a/go-sdk/adr/0008-native-dag-interface.md b/go-sdk/adr/0008-native-dag-interface.md index e505eb4c9fbaa..e8b4f786c81ad 100644 --- a/go-sdk/adr/0008-native-dag-interface.md +++ b/go-sdk/adr/0008-native-dag-interface.md @@ -35,7 +35,7 @@ Proposed. 5. **`airflow.Inputs(refs...)` declares the data and the edge in one call** for defining graph with TaskFlow syntax. `airflow.Literal(value)` passes a JSON-serializable value in the same call without adding an edge, e.g. `airflow.Inputs(extracted, airflow.Literal("uk"))`. 6. **`Before` and `After` are order-only edges on `airflow.Node`**, which both `*airflow.TaskRef` and `*airflow.TaskGroupRef` satisfy. They are the Go pair for `>>` and `<<`, and both return their argument set as one `Node`, so `a.Before(b, c).Before(d)` is Python's `a >> [b, c] >> d`. -7. **An edge label wraps the endpoint**: `loaded.Before(airflow.Label(notify, "when empty"))` is Python's `loaded >> Label("when empty") >> notify`. Labelling the endpoint rather than the call lets one fan-out give each edge its own label. +7. **An edge label wraps the endpoint, and its text comes first**: `loaded.Before(airflow.Label("when empty").On(notify))` is Python's `loaded >> Label("when empty") >> notify`. Labelling the endpoint rather than the call lets one fan-out give each edge its own label. Because the text comes before the endpoint, it sits between the two tasks in the code, the same place the label sits in the rendered graph. The Java SDK uses the same order: `Label.of("when empty").on(notify)`. 8. **Trigger rules belong to the task**, as `airflow.TaskSpec{TriggerRule: ...}`, never to an edge. 9. **A user-facing enum carries its type in the constant name** — e.g. `airflow.TriggerRuleAllDone`. 10. **Everything an author writes comes from one `airflow` package.** @@ -100,7 +100,7 @@ staging.Before(loaded) // staging >> load loaded.Before(notified, cleaned) // loaded >> [notify, cleanup] cleaned.After(extracted) // cleanup << extracted -loaded.Before(airflow.Label(emptyNotice, "when empty")) // loaded >> Label("when empty") >> notify_empty +loaded.Before(airflow.Label("when empty").On(emptyNotice)) // loaded >> Label("when empty") >> notify_empty ``` ## Signature @@ -153,8 +153,10 @@ type Node interface { node() } -// Label carries an edge label into Before or After. The Node it returns stands for node itself. -func Label(node Node, text string) Node +// Label makes an edge label, and On carries the label into Before or After. The Node that On +// returns stands for node and each node in more. +func Label(text string) *LabelRef +func (l *LabelRef) On(node Node, more ...Node) Node // A user-facing enum is a named string type with its type in each constant name. type TriggerRule string @@ -185,7 +187,7 @@ const ( `airflow.DagRunState` is a type alias of `genmodels.DagRunState`, not a new type. `airflow.DagRun` is an alias of `sdk.DagRun`, and package `sdk` cannot import `airflow`. So a state field that `sdk.DagRun` gains later has a type declared outside `airflow`. `actx.DagRun().State == airflow.DagRunStateFailed` compiles only if both sides have the same type. The cost is that `%T` and package reflect report `airflow.DagRunState` as `genmodels.DagRunState`, and package `airflow` cannot declare a method on that type. - **A value outside an enum is rejected when the task is added.** Go converts a string literal to an enum type without a cast, so the enum type alone does not stop `TaskSpec{TriggerRule: "all_sucess"}`. `dag.Task` panics with a message that names the field and lists the valid values. Python rejects an unknown trigger rule when the operator is constructed. It rejects an unknown weight rule only when the Dag is serialized. -- **A data edge is labelled by redeclaring it.** Declaring an edge that already exists is idempotent, so `extracted.Before(airflow.Label(transformed, "rows"))` labels the edge `Inputs` created. +- **A data edge is labelled by redeclaring it.** Declaring an edge that already exists is idempotent, so `extracted.Before(airflow.Label("rows").On(transformed))` labels the edge `Inputs` created. - **Renaming a Go function renames the task.** The id is derived, and history, clears, and the UI all key on task_id, so renaming a function whose task carries no `TaskSpec` id is a Dag change. `airflow.TaskSpec{TaskID: ...}` pins an id that has to outlive the function's name, and it is also how a Dag gets snake_case ids, since nothing transforms a Go name. Renaming a task group renames every task whose task_id the group_id prefixes, for the same reason. @@ -196,8 +198,8 @@ const ( - **A group with no task is stepped over.** An edge to it continues along each edge from it, and an edge from it back along each edge to it, whenever those were declared, as the TypeScript SDK does. So `extracted.Before(empty).Before(loaded)` runs `load` after `extract`, as Python's `extract >> empty >> load` does. What Python does with an empty group depends on how and when its edges are declared: `load << empty << extract` adds no edge, and an empty group with no task before it falls back to the last tasks of the enclosing group or of the whole Dag, which can make a task depend on itself. - **An edge cannot connect a group to what it holds.** Python's `group >> node_in_group` orders the last tasks of the group before the first tasks of a node inside it, which fails as a cycle whenever that node holds a task. The edge verb rejects the edge when it is declared. -- **A label stays on the edge it is declared on.** A label on a group edge labels none of the task edges that the group edge stands for, and a labelled edge keeps its ends. - Python's behavior depends on which end is the receiver and which groups hold the ends. When no group holds `extract`, `extract >> Label("rows") >> group` labels each of those task edges as well. When `a` and `b` are in different groups, `a >> Label("x") >> b` replaces `a` with its group, so every last task of `a`'s group runs before `b`. The Go SDK does neither. +- **A label stays on the edge it is declared on, with one exception.** A label on a group edge labels that group edge. A group edge from a task in no group to a group with tasks also puts its label on the task edges from that task to the first tasks of the group. Python's `extract >> Label("rows") >> group` does the same. If one of those task edges was also declared on its own, or an earlier group edge added it, that task edge keeps the label it has. If it has no label, it stays unlabelled. Python labels those task edges when `>>` runs, so in Python the label written last wins. + A label does not change the ends of its edge. When `a` and `b` are in different groups, Python's `a >> Label("x") >> b` replaces `a` with its group, so all the last tasks of `a`'s group run before `b`. The Go SDK keeps the edge between `a` and `b`. - **`TaskGroup` takes `spec ...TaskGroupSpec`, as `Dag` takes `spec ...DagSpec`.** A group has one kind of option, so it needs no sealed option interface. A mapped task group would come through a method of its own, not through this variadic. ## Alternatives @@ -208,5 +210,6 @@ const ( callable per run, so an edge existing only in execution order cannot be parsed without running the program to completion. `Inputs` keeps the typed outputs that style is reached for. - **Labelling the call**, `loaded.Before(notify).Label("when empty")`. Rejected: one call fans out, so a single label on the call cannot give `a.Before(b, c)` a different label per edge. +- **The endpoint first**, `airflow.Label(notify, "when empty")`. The SDK had this form before it followed the Java SDK, and no release included this form. Rejected: the text comes after the second task in the code, while the rendered graph shows the label between the two tasks. - **A positional task_id**, `dag.Task(taskID, fn, opts...)`. Rejected: It's more straightforward and native for Go user to define a Task without defining the task_id explicitly. They could still set the task_id other than the function name in the TaskSpec. - **Separate `Dag` and `MixedLangDag` types.** Rejected: Python has one Dag class, and the Mixed Lang case is not a Dag at all ([ADR 7](0007-mixed-lang-task-handler-interface.md)). diff --git a/go-sdk/airflow/conformance_test.go b/go-sdk/airflow/conformance_test.go index 149ad6aed4000..a86b9209531c4 100644 --- a/go-sdk/airflow/conformance_test.go +++ b/go-sdk/airflow/conformance_test.go @@ -242,7 +242,7 @@ func buildConformanceDag(t *testing.T, dagCase conformanceDag) *airflow.DagRef { for _, edge := range dagCase.OrderEdges { downstream := node(edge[1]) if len(edge) == 3 { - downstream = airflow.Label(downstream, edge[2]) + downstream = airflow.Label(edge[2]).On(downstream) } node(edge[0]).Before(downstream) } diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go index 31950d51a0c51..63c606ed45ee5 100644 --- a/go-sdk/airflow/dag.go +++ b/go-sdk/airflow/dag.go @@ -50,9 +50,9 @@ type DagRef struct { // in common. groups []*TaskGroupRef groupsByID map[string]*TaskGroupRef - // edgeLabels holds every edge between two tasks of the Dag, whether Inputs, Before or After - // declared it, and the label that Label put on it. An edge with no label maps to the empty - // string, so a lookup reports whether the edge has been declared. + // edgeLabels maps each edge between two tasks of the Dag to its label, whether Inputs, Before or + // After declared the edge. An edge with no label maps to the empty string, so a lookup reports + // whether the edge has been declared. edgeLabels map[edgeKey]string // groupEdges holds the edges that have a task group at one end or both, in the order they // were first declared, and groupEdgeLabels holds their labels as edgeLabels does. Registration diff --git a/go-sdk/airflow/if_test.go b/go-sdk/airflow/if_test.go index 8571002e70845..ebd982ed2506b 100644 --- a/go-sdk/airflow/if_test.go +++ b/go-sdk/airflow/if_test.go @@ -284,7 +284,7 @@ func TestDeclaringTheEdgeOfAConditionAgainChangesNothing(t *testing.T) { gate := dag.If(isReady) gate.Then(loaded) - gate.task.Before(Label(loaded, "when ready")) + gate.task.Before(Label("when ready").On(loaded)) assertTasks(t, gate.task.downstreams, loaded) assertTasks(t, loaded.upstreams, gate.task) diff --git a/go-sdk/airflow/node.go b/go-sdk/airflow/node.go index 643c599d8145f..5f620ad9aaa8e 100644 --- a/go-sdk/airflow/node.go +++ b/go-sdk/airflow/node.go @@ -20,6 +20,7 @@ package airflow import ( "fmt" "slices" + "strings" ) // Node is what an edge connects: a task or a whole task group. A [TaskRef] and a [TaskGroupRef] @@ -53,9 +54,9 @@ type Node interface { node() } -// nodeEndpoint is one task or one task group that a Node stands for, with the label that [Label] -// carries into the edge verb the Node is passed to. Exactly one of task and group is set, except -// for the endpoint of a nil *TaskRef, which has neither. +// nodeEndpoint is one task or one task group that a Node stands for, with the label that +// [LabelRef.On] carries into the edge verb the Node is passed to. Exactly one of task and group is +// set, except for the endpoint of a nil *TaskRef, which has neither. type nodeEndpoint struct { task *TaskRef group *TaskGroupRef @@ -96,8 +97,8 @@ func (e nodeEndpoint) describe() string { } // nodeSet is the argument set that Before and After return as one Node, and the Node that -// [Label] returns. A label belongs to the edge of the verb the Node is passed to, so a nodeSet -// carries its labels no further once it is the receiver of the next verb. +// [LabelRef.On] returns. The labels of a nodeSet go on the edges of the verb that the nodeSet is +// passed to, so a nodeSet with labels cannot be the receiver of a verb. type nodeSet []nodeEndpoint func (nodeSet) node() {} @@ -150,7 +151,7 @@ func endpoints(where string, node Node) []nodeEndpoint { // // Before returns the nodes it was given as one Node, not the receiver, so a chain fans out from // them: loaded.Before(notified, cleaned).Before(done) is loaded >> [notify, cleanup] >> done. -// Wrap a node in [Label] to label the edge that reaches it. +// To label the edge that reaches a node, pass Label(text).On(node) in place of the node. // // Declaring an edge that the Dag already has changes nothing, other than to apply a [Label]. // @@ -178,20 +179,51 @@ func (t *TaskRef) Before(nodes ...Node) Node { return declareEdges(t, nodes, dir // cleanup << extract << start. func (t *TaskRef) After(nodes ...Node) Node { return declareEdges(t, nodes, dirAfter) } -// Label puts text on the edge that an edge verb declares to node, as Python's -// loaded >> Label("when empty") >> notify_empty does: +// LabelRef is an edge label made by [Label]. [LabelRef.On] returns a Node. To label edges, pass +// that Node to Before or After. +type LabelRef struct { + text string +} + +// Label returns a label with the given text. The Python code +// loaded >> Label("when empty") >> notify_empty becomes: +// +// loaded.Before(Label("when empty").On(emptyNotice)) +// +// The text comes before the node passed to On, so it sits between the two tasks in the code. The +// Airflow UI draws the label between the same two tasks in the graph. +// +// Label panics if text is empty or holds only white space. +func Label(text string) *LabelRef { + if strings.TrimSpace(text) == "" { + panic("airflow.Label: got a blank label; pass the text to put on the edge") + } + return &LabelRef{text: text} +} + +// On returns node, and each node in more, as one Node. Pass that Node to an edge verb, and the verb +// puts the label on each edge between its receiver and one of those nodes. With After, those edges +// point at the receiver: +// +// cleaned.After(Label("always").On(loaded)) // loaded >> Label("always") >> cleanup // -// loaded.Before(Label(emptyNotice, "when empty")) +// Each Node passed to Before or After can carry its own label, so the edges that one call declares +// can have different labels: // -// The label wraps the endpoint rather than the call, so each edge of a fan-out can carry a label -// of its own: +// checked.Before(Label("rows found").On(processed), Label("no rows").On(emptyNotice)) // -// checked.Before(Label(processed, "rows found"), Label(emptyNotice, "no rows")) +// When On is given more than one node, the edge to or from each of those nodes gets the same +// label. Python's loaded >> Label("forks") >> [notify_empty, notify_suspicious] becomes: // -// A label is the Node's only in the verb it is passed to. The Node that Label returns stands for -// node itself, and the Node a verb returns carries no label on either side of the next verb. So a -// label on the receiver of a verb has no edge to land on and is dropped, the way Python's -// Label("x") >> b alone sets no label. +// loaded.Before(Label("forks").On(emptyNotice, suspiciousNotice)) +// +// A label is the Node's only in the verb it is passed to. The Node that On returns stands for the +// nodes themselves. The Node that a verb returns carries no label into the next verb, either as +// its receiver or as an argument. Before and After panic when their receiver carries a label, +// because that label has no edge to go on. Pass the Node that On returns only as an argument. +// Python's Label("x") >> b drops the label without an error instead. When labels nest, the outer +// label is used: loaded.Before(Label("outer").On(Label("inner").On(cleaned))) labels the edge +// "outer". // // A label declared on an edge that already carries one replaces it, as Python's DAG.set_edge_info // does. @@ -199,25 +231,38 @@ func (t *TaskRef) After(nodes ...Node) Node { return declareEdges(t, nodes, dirA // An [Inputs] edge is labelled by declaring it again, which is idempotent: // // transformed := dag.Task(transform, Inputs(extracted)) -// extracted.Before(Label(transformed, "rows")) +// extracted.Before(Label("rows").On(transformed)) // // A label on an edge to or from a task group stays on that edge, as [TaskGroupRef.Before] // describes. A label on an edge between two tasks stays on that edge too, even when the tasks are // in different task groups. In that case Python can replace the receiver of the verb with a group // that holds it, which makes the edge connect other tasks. // -// Label panics if node is nil or text is empty. -func Label(node Node, text string) Node { - if node == nil { - panic("airflow.Label: got a nil airflow.Node") - } - if text == "" { - panic("airflow.Label: got an empty label; pass the text to put on the edge") +// On panics if: +// - l is nil or points to a zero LabelRef +// - node or a node in more is nil, is a nil *TaskGroupRef, or is a struct that embeds a Node +// +// On does not panic on a nil *TaskRef. The edge verb that takes the Node panics on it instead. +func (l *LabelRef) On(node Node, more ...Node) Node { + const where = "airflow.LabelRef.On" + if l == nil || l.text == "" { + panic(where + ": airflow.Label did not return the *airflow.LabelRef") } - labelling := endpoints("airflow.Label: node", node) - labelled := make(nodeSet, len(labelling)) - for i, endpoint := range labelling { - labelled[i] = nodeEndpoint{task: endpoint.task, group: endpoint.group, label: text} + nodes := append([]Node{node}, more...) + labelled := make(nodeSet, 0, len(nodes)) + for i, labelling := range nodes { + name := "node" + if i > 0 { + name = fmt.Sprintf("more[%d]", i-1) + } + if labelling == nil { + panic(fmt.Sprintf("%s: %s is a nil airflow.Node", where, name)) + } + for _, endpoint := range endpoints(where+": "+name, labelling) { + labelled = append( + labelled, nodeEndpoint{task: endpoint.task, group: endpoint.group, label: l.text}, + ) + } } return labelled } @@ -281,6 +326,7 @@ func declareEdges(receiver Node, nodes []Node, dir edgeDir) Node { } dag := edgeDag(where, all) + checkReceiverUnlabelled(where, dag, dir, ends) dag.mu.Lock() defer dag.mu.Unlock() @@ -347,6 +393,25 @@ func declareEdges(receiver Node, nodes []Node, dir edgeDir) Node { return args.unlabelled() } +// checkReceiverUnlabelled panics if the receiver of an edge verb carries a label. Only the Node +// that [LabelRef.On] returns carries labels. A label goes on the edges between the receiver and +// the nodes that the verb is given, so a label on the receiver has no edge to go on. +func checkReceiverUnlabelled(where string, dag *DagRef, dir edgeDir, ends []nodeEndpoint) { + var labelled []string + for _, end := range ends { + if end.label != "" { + labelled = append(labelled, fmt.Sprintf("%s labelled %q", end.describe(), end.label)) + } + } + if len(labelled) > 0 { + panic(fmt.Sprintf( + "%s: Dag %q: cannot call %s on %s; "+ + "pass airflow.Label(...).On(...) to Before or After as an argument instead", + where, dag.dagID, dir, strings.Join(labelled, ", "), + )) + } +} + // checkGroupEdgeEnds panics if one end of an edge is a task group that holds the other end. The // edge would order the group against part of itself: in Python, group >> task_in_group makes the // last tasks of the group upstreams of a task that may be one of them. diff --git a/go-sdk/airflow/node_test.go b/go-sdk/airflow/node_test.go index 474011af37cf4..7e6dfcff77aad 100644 --- a/go-sdk/airflow/node_test.go +++ b/go-sdk/airflow/node_test.go @@ -172,7 +172,7 @@ func TestLabelLabelsTheEdgeToTheNodeItWraps(t *testing.T) { loaded := orderedTask(t, dag, "load") emptyNotice := orderedTask(t, dag, "notify_empty") - loaded.Before(Label(emptyNotice, "when empty")) + loaded.Before(Label("when empty").On(emptyNotice)) assertTasks(t, loaded.downstreams, emptyNotice) assertEdgeLabel(t, dag, "load", "notify_empty", "when empty") @@ -184,18 +184,36 @@ func TestLabelLabelsEachEdgeOfAFanOut(t *testing.T) { processed := orderedTask(t, dag, "process") emptyNotice := orderedTask(t, dag, "notify_empty") - checked.Before(Label(processed, "rows found"), Label(emptyNotice, "no rows")) + checked.Before(Label("rows found").On(processed), Label("no rows").On(emptyNotice)) assertEdgeLabel(t, dag, "check", "process", "rows found") assertEdgeLabel(t, dag, "check", "notify_empty", "no rows") } +// TestLabelOnLabelsTheEdgeToEachNodeItIsGiven checks the Go form of Python's +// loaded >> Label("forks") >> [notify_empty, staging, notify_suspicious]. +func TestLabelOnLabelsTheEdgeToEachNodeItIsGiven(t *testing.T) { + dag := Dag("etl") + loaded := orderedTask(t, dag, "load") + emptyNotice := orderedTask(t, dag, "notify_empty") + staging := dag.TaskGroup("staging") + groupTask(t, staging, "stage") + suspiciousNotice := orderedTask(t, dag, "notify_suspicious") + + loaded.Before(Label("forks").On(emptyNotice, staging, suspiciousNotice)) + + assertTasks(t, loaded.downstreams, emptyNotice, suspiciousNotice) + assertEdgeLabel(t, dag, "load", "notify_empty", "forks") + assertGroupEdgeLabel(t, dag, "load", "staging", "forks") + assertEdgeLabel(t, dag, "load", "notify_suspicious", "forks") +} + func TestAfterLabelsTheEdgeToTheNodeItWasGiven(t *testing.T) { dag := Dag("etl") extracted := orderedTask(t, dag, "extract") cleaned := orderedTask(t, dag, "cleanup") - cleaned.After(Label(extracted, "always")) + cleaned.After(Label("always").On(extracted)) assertEdgeLabel(t, dag, "extract", "cleanup", "always") } @@ -207,7 +225,7 @@ func TestRedeclaringAnInputsEdgeLabelsIt(t *testing.T) { read := dag.Task(readRows) counted := dag.Task(countRows, Inputs(read)) - read.Before(Label(counted, "rows")) + read.Before(Label("rows").On(counted)) assertTasks(t, read.downstreams, counted) assertTasks(t, counted.upstreams, read) @@ -219,9 +237,9 @@ func TestRedeclaringALabelledEdgeKeepsTheLabel(t *testing.T) { loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") - loaded.Before(Label(cleaned, "always")) + loaded.Before(Label("always").On(cleaned)) loaded.Before(cleaned) - cleaned.After(Label(loaded, "always")) + cleaned.After(Label("always").On(loaded)) assertEdgeLabel(t, dag, "load", "cleanup", "always") } @@ -233,8 +251,8 @@ func TestASecondLabelOnAnEdgeReplacesTheFirst(t *testing.T) { loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") - loaded.Before(Label(cleaned, "always")) - loaded.Before(Label(cleaned, "when empty")) + loaded.Before(Label("always").On(cleaned)) + loaded.Before(Label("when empty").On(cleaned)) assertEdgeLabel(t, dag, "load", "cleanup", "when empty") assertTasks(t, loaded.downstreams, cleaned) @@ -247,24 +265,66 @@ func TestTheLastLabelOfOneCallWins(t *testing.T) { loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") - loaded.Before(Label(cleaned, "always"), Label(cleaned, "when empty")) + loaded.Before(Label("always").On(cleaned), Label("when empty").On(cleaned)) assertEdgeLabel(t, dag, "load", "cleanup", "when empty") assertTasks(t, loaded.downstreams, cleaned) } -// TestALabelOnTheReceiverOfAVerbIsDropped pins the silence Python has too: Label marks the edge -// that reaches a node, so a label on the node an edge leaves from has no edge to land on. -// Python's Label("x") >> b, with nothing upstream of the label, sets no label either. -func TestALabelOnTheReceiverOfAVerbIsDropped(t *testing.T) { +// TestEdgeVerbsRejectALabelledReceiver covers a verb whose receiver carries a label. That label has +// no edge to go on. The verb panics before it declares any edge. The Java SDK throws in the same +// case. Python's Label("x") >> b drops the label without an error instead. When the label is on a +// nil *TaskRef, the verb's panic reports the nil task. +func TestEdgeVerbsRejectALabelledReceiver(t *testing.T) { dag := Dag("etl") loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") - - Label(loaded, "always").Before(cleaned) - - assertTasks(t, loaded.downstreams, cleaned) - assertEdgeLabel(t, dag, "load", "cleanup", "") + audited := orderedTask(t, dag, "audit") + staging := dag.TaskGroup("staging") + groupTask(t, staging, "stage") + + const fix = "pass airflow.Label(...).On(...) to Before or After as an argument instead" + for _, tc := range []struct { + name string + call func() + want string + }{ + { + name: "Before", + call: func() { Label("always").On(loaded).Before(cleaned) }, + want: `airflow.Node.Before: Dag "etl": cannot call Before on task "load" ` + + `labelled "always"; ` + fix, + }, + { + name: "After", + call: func() { Label("always").On(loaded).After(cleaned) }, + want: `airflow.Node.After: Dag "etl": cannot call After on task "load" ` + + `labelled "always"; ` + fix, + }, + { + name: "a verb with no node", + call: func() { Label("always").On(loaded).Before() }, + want: `airflow.Node.Before: Dag "etl": cannot call Before on task "load" ` + + `labelled "always"; ` + fix, + }, + { + name: "a label on several nodes", + call: func() { Label("always").On(loaded, staging).Before(audited) }, + want: `airflow.Node.Before: Dag "etl": cannot call Before on task "load" ` + + `labelled "always", task group "staging" labelled "always"; ` + fix, + }, + { + name: "a label on a nil task", + call: func() { Label("always").On((*TaskRef)(nil)).Before(cleaned) }, + want: "airflow.Node.Before: got a nil *airflow.TaskRef", + }, + } { + t.Run(tc.name, func(t *testing.T) { + assert.PanicsWithValue(t, tc.want, tc.call) + }) + } + assert.Empty(t, dag.edgeLabels) + assert.Empty(t, dag.groupEdges) } // TestTheOuterLabelOfANestedLabelWins pins what nesting means: the outer Label relabels the @@ -274,7 +334,7 @@ func TestTheOuterLabelOfANestedLabelWins(t *testing.T) { loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") - loaded.Before(Label(Label(cleaned, "inner"), "outer")) + loaded.Before(Label("outer").On(Label("inner").On(cleaned))) assertEdgeLabel(t, dag, "load", "cleanup", "outer") } @@ -289,7 +349,7 @@ func TestLabelBelongsToTheVerbItIsPassedTo(t *testing.T) { done := orderedTask(t, dag, "done") extracted := orderedTask(t, dag, "extract") - labelled := loaded.Before(Label(cleaned, "always")) + labelled := loaded.Before(Label("always").On(cleaned)) labelled.Before(done) extracted.Before(labelled) @@ -306,27 +366,57 @@ func TestLabelPassesThroughNoVerbTwice(t *testing.T) { cleaned := orderedTask(t, dag, "cleanup") notified := orderedTask(t, dag, "notify") - notified.After(cleaned.After(Label(loaded, "always"))) + notified.After(cleaned.After(Label("always").On(loaded))) assertEdgeLabel(t, dag, "load", "cleanup", "always") assertEdgeLabel(t, dag, "load", "notify", "") } -func TestLabelPanicsOnANilNode(t *testing.T) { +func TestLabelOnPanicsOnANilNode(t *testing.T) { + dag := Dag("etl") + loaded := orderedTask(t, dag, "load") + + assert.PanicsWithValue(t, + "airflow.LabelRef.On: node is a nil airflow.Node", + func() { Label("when empty").On(nil) }, + ) assert.PanicsWithValue(t, - "airflow.Label: got a nil airflow.Node", - func() { Label(nil, "when empty") }, + "airflow.LabelRef.On: more[1] is a nil airflow.Node", + func() { Label("when empty").On(loaded, loaded, nil) }, ) } -func TestLabelPanicsOnAnEmptyLabel(t *testing.T) { +func TestLabelPanicsOnABlankLabel(t *testing.T) { + for _, text := range []string{"", " ", "\t\n", " ", "　"} { + assert.PanicsWithValue(t, + "airflow.Label: got a blank label; pass the text to put on the edge", + func() { Label(text) }, + "%q", text, + ) + } +} + +// TestLabelOnPanicsOnALabelThatLabelDidNotReturn covers a nil LabelRef and a zero LabelRef. A zero +// LabelRef has empty text, and an edge verb reads empty text as no label. Without the panic in On, +// the verb would silently ignore that LabelRef. +func TestLabelOnPanicsOnALabelThatLabelDidNotReturn(t *testing.T) { dag := Dag("etl") loaded := orderedTask(t, dag, "load") - assert.PanicsWithValue(t, - "airflow.Label: got an empty label; pass the text to put on the edge", - func() { Label(loaded, "") }, - ) + for _, tc := range []struct { + name string + label *LabelRef + }{ + {name: "a nil label", label: nil}, + {name: "a zero label", label: &LabelRef{}}, + } { + t.Run(tc.name, func(t *testing.T) { + assert.PanicsWithValue(t, + "airflow.LabelRef.On: airflow.Label did not return the *airflow.LabelRef", + func() { tc.label.On(loaded) }, + ) + }) + } } func TestEdgeVerbsRejectANilNode(t *testing.T) { @@ -442,6 +532,22 @@ func TestEdgeVerbsWithNoNodeDeclareNoEdge(t *testing.T) { assert.Empty(t, dag.edgeLabels) } +// TestAVerbThatDeclaresNoEdgeReturnsNoLabel covers a verb whose receiver stands for no task. The +// verb declares no edge. The Node that the verb returns carries none of the labels on the verb's +// arguments, so a chain can go on from that Node. +func TestAVerbThatDeclaresNoEdgeReturnsNoLabel(t *testing.T) { + dag := Dag("etl") + loaded := orderedTask(t, dag, "load") + cleaned := orderedTask(t, dag, "cleanup") + done := orderedTask(t, dag, "done") + + var none []Node + loaded.Before(none...).Before(Label("always").On(cleaned)).Before(done) + + assert.Empty(t, cleaned.upstreams) + assertEdgeLabel(t, dag, "cleanup", "done", "") +} + // TestEdgeVerbsCheckTheTasksEvenWithNoNode pins that a verb with no node to point at still checks // the tasks it was called on. It declares no edge, but a Dag that can no longer take one, or a // task the Dag never returned, is a mistake either way. @@ -506,9 +612,14 @@ func TestEdgeVerbsRejectANodeTheyDoNotDefine(t *testing.T) { func() { loaded.Before(wrappedNode{cleaned}) }, ) assert.PanicsWithValue(t, - "airflow.Label: node has type airflow.wrappedNode, "+ + "airflow.LabelRef.On: node has type airflow.wrappedNode, "+ + "which is not a node that package airflow defines", + func() { Label("when empty").On(wrappedNode{cleaned}) }, + ) + assert.PanicsWithValue(t, + "airflow.LabelRef.On: more[0] has type airflow.wrappedNode, "+ "which is not a node that package airflow defines", - func() { Label(wrappedNode{cleaned}, "when empty") }, + func() { Label("when empty").On(cleaned, wrappedNode{cleaned}) }, ) assert.Empty(t, loaded.downstreams) } diff --git a/go-sdk/airflow/serialize_test.go b/go-sdk/airflow/serialize_test.go index 174791840c84c..4f728ada5d899 100644 --- a/go-sdk/airflow/serialize_test.go +++ b/go-sdk/airflow/serialize_test.go @@ -759,10 +759,10 @@ func TestSerializeWritesTheLabelsOfEdges(t *testing.T) { groupTask(t, transform, "clean") publish := dag.TaskGroup("publish") groupTask(t, publish, "push") - extracted.Before(Label(loaded, "rows"), notified) - extracted.Before(Label(transform, "to transform")) - transform.Before(Label(publish, "to publish")) - publish.Before(Label(notified, "to notify")) + extracted.Before(Label("rows").On(loaded), notified) + extracted.Before(Label("to transform").On(transform)) + transform.Before(Label("to publish").On(publish)) + publish.Before(Label("to notify").On(notified)) got := serializedDag(t, dag) diff --git a/go-sdk/airflow/task_group_test.go b/go-sdk/airflow/task_group_test.go index d136bfc787fcf..0f28b6d5bdb28 100644 --- a/go-sdk/airflow/task_group_test.go +++ b/go-sdk/airflow/task_group_test.go @@ -531,7 +531,7 @@ func TestTaskGroupIsSafeForConcurrentUse(t *testing.T) { task := groupTask(t, group, fmt.Sprintf("task_%d", i)) nested := group.TaskGroup(fmt.Sprintf("group_%d", i)) groupTask(t, nested, "rows") - extracted.Before(Label(nested, fmt.Sprintf("to %d", i))) + extracted.Before(Label(fmt.Sprintf("to %d", i)).On(nested)) nested.After(task) }) } @@ -694,9 +694,9 @@ func TestAGroupEdgeLeavesTheLabelOfATaskEdgeAlone(t *testing.T) { group := dag.TaskGroup("transform") cleaned := groupTask(t, group, "clean") validated := groupTask(t, group, "validate") - extracted.Before(Label(validated, "rows")) + extracted.Before(Label("rows").On(validated)) - extracted.Before(Label(group, "to transform")) + extracted.Before(Label("to transform").On(group)) Bundle().Register(dag) assertGroupEdgeLabel(t, dag, "extract", "transform", "to transform") @@ -718,7 +718,7 @@ func TestALabelOnAGroupEdge(t *testing.T) { { name: "a task before a labelled group", declare: func(extracted *TaskRef, transform, _ *TaskGroupRef) { - extracted.Before(Label(transform, "x")) + extracted.Before(Label("x").On(transform)) }, upstream: "extract", downstream: "transform", taskEdgeLabel: "x", @@ -726,21 +726,21 @@ func TestALabelOnAGroupEdge(t *testing.T) { { name: "a task after a labelled group", declare: func(extracted *TaskRef, transform, _ *TaskGroupRef) { - extracted.After(Label(transform, "x")) + extracted.After(Label("x").On(transform)) }, upstream: "transform", downstream: "extract", }, { name: "a group before a labelled task", declare: func(extracted *TaskRef, transform, _ *TaskGroupRef) { - transform.Before(Label(extracted, "x")) + transform.Before(Label("x").On(extracted)) }, upstream: "transform", downstream: "extract", }, { name: "a group after a labelled task", declare: func(extracted *TaskRef, transform, _ *TaskGroupRef) { - transform.After(Label(extracted, "x")) + transform.After(Label("x").On(extracted)) }, upstream: "extract", downstream: "transform", taskEdgeLabel: "x", @@ -748,7 +748,7 @@ func TestALabelOnAGroupEdge(t *testing.T) { { name: "a group before a labelled group", declare: func(_ *TaskRef, transform, publish *TaskGroupRef) { - transform.Before(Label(publish, "x")) + transform.Before(Label("x").On(publish)) }, upstream: "transform", downstream: "publish", }, @@ -786,7 +786,7 @@ func TestALabelOnAnEdgeBetweenTasksOfDifferentGroupsKeepsTheTaskEdge(t *testing. cleaned.Before(validated) loaded := orderedTask(t, dag, "load") - cleaned.Before(Label(loaded, "rows")) + cleaned.Before(Label("rows").On(loaded)) Bundle().Register(dag) assertEdgeLabel(t, dag, "transform.clean", "load", "rows") @@ -803,7 +803,7 @@ func TestALabelOnAGroupEdgeFromAGroupedTaskStaysOnTheGroupEdge(t *testing.T) { transform := dag.TaskGroup("transform") groupTask(t, transform, "clean") - pulled.Before(Label(transform, "rows")) + pulled.Before(Label("rows").On(transform)) Bundle().Register(dag) assertGroupEdgeLabel(t, dag, "extract.pull", "transform", "rows") @@ -816,7 +816,7 @@ func TestRedeclaringAGroupEdgeIsIdempotent(t *testing.T) { group := dag.TaskGroup("transform") groupTask(t, group, "clean") - extracted.Before(Label(group, "rows")) + extracted.Before(Label("rows").On(group)) group.After(extracted) require.Len(t, dag.groupEdges, 1) @@ -938,8 +938,8 @@ func TestGroupEdgeVerbsRejectAGroupTheDagDidNotReturn(t *testing.T) { }, { name: "a nil group in a label", - call: func() { Label((*TaskGroupRef)(nil), "rows") }, - want: "airflow.Label: node is a nil *airflow.TaskGroupRef", + call: func() { Label("rows").On((*TaskGroupRef)(nil)) }, + want: "airflow.LabelRef.On: node is a nil *airflow.TaskGroupRef", }, } { t.Run(tc.name, func(t *testing.T) { @@ -1329,7 +1329,7 @@ func TestAPanicLeavesTheDagAsItWas(t *testing.T) { { name: "a fan-out that relabels a group edge before a later pair fails", call: func(p dagParts) { - p.extracted.Before(Label(p.transform, "relabelled"), p.extracted) + p.extracted.Before(Label("relabelled").On(p.transform), p.extracted) }, want: `task "extract" cannot depend on itself`, }, @@ -1363,7 +1363,7 @@ func TestAPanicLeavesTheDagAsItWas(t *testing.T) { transform := dag.TaskGroup("transform") cleaned := groupTask(t, transform, "clean") checks := transform.TaskGroup("checks") - extracted.Before(Label(transform, "rows")) + extracted.Before(Label("rows").On(transform)) parts := dagParts{dag, transform, checks, extracted, cleaned} if tc.prepare != nil { tc.prepare(parts) diff --git a/go-sdk/example/native/main.go b/go-sdk/example/native/main.go index fda85d96eb601..ef2638bfacc4e 100644 --- a/go-sdk/example/native/main.go +++ b/go-sdk/example/native/main.go @@ -78,7 +78,7 @@ func newPipelineDag() *airflow.DagRef { airflow.Inputs(airflow.Literal("south"), airflow.Literal(2)), ) summed := extract.Task(summarize, airflow.Inputs(north, south)) - seeded.Before(airflow.Label(extract, "fan out")) + seeded.Before(airflow.Label("fan out").On(extract)) loaded := dag.Task(loadRows, airflow.Inputs(summed, airflow.Literal("s3://bucket/out"))) empty := dag.Task(reportEmpty) diff --git a/go-sdk/internal/genspec/authoring.go b/go-sdk/internal/genspec/authoring.go index 463eef2e44021..d1dedd1717e8c 100644 --- a/go-sdk/internal/genspec/authoring.go +++ b/go-sdk/internal/genspec/authoring.go @@ -79,7 +79,7 @@ var dagShape = authoringShape{ "bundle_name": "the name of the bundle that carries the Dag, not the Dag's", "tasks": "the tasks dag.Task registers", "task_group": "the groups dag.TaskGroup registers", - "edge_info": "the labels airflow.Label carries into an edge verb", + "edge_info": "the edge labels airflow.Label makes", "dag_dependencies": "derived from the edges and the assets a Dag declares", "timezone": "always UTC for a Go Dag, even when StartDate has another location", "timetable": "the serialized form of Schedule, which is injected instead", From 478077bc4c170fc7bf364519410da5ba84ba44a5 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Wed, 21 Oct 2026 22:47:34 +0800 Subject: [PATCH 2/2] Go SDK: Reject labels with no edge and name nil task positions Signed-off-by: PoAn Yang --- go-sdk/adr/0008-native-dag-interface.md | 3 + go-sdk/airflow/node.go | 193 +++++++++++++++--------- go-sdk/airflow/node_test.go | 90 +++++++++-- go-sdk/airflow/task_group_test.go | 3 +- 4 files changed, 200 insertions(+), 89 deletions(-) diff --git a/go-sdk/adr/0008-native-dag-interface.md b/go-sdk/adr/0008-native-dag-interface.md index e8b4f786c81ad..806e430644ac9 100644 --- a/go-sdk/adr/0008-native-dag-interface.md +++ b/go-sdk/adr/0008-native-dag-interface.md @@ -153,6 +153,8 @@ type Node interface { node() } +type LabelRef struct{ /* unexported */ } + // Label makes an edge label, and On carries the label into Before or After. The Node that On // returns stands for node and each node in more. func Label(text string) *LabelRef @@ -200,6 +202,7 @@ const ( - **An edge cannot connect a group to what it holds.** Python's `group >> node_in_group` orders the last tasks of the group before the first tasks of a node inside it, which fails as a cycle whenever that node holds a task. The edge verb rejects the edge when it is declared. - **A label stays on the edge it is declared on, with one exception.** A label on a group edge labels that group edge. A group edge from a task in no group to a group with tasks also puts its label on the task edges from that task to the first tasks of the group. Python's `extract >> Label("rows") >> group` does the same. If one of those task edges was also declared on its own, or an earlier group edge added it, that task edge keeps the label it has. If it has no label, it stays unlabelled. Python labels those task edges when `>>` runs, so in Python the label written last wins. A label does not change the ends of its edge. When `a` and `b` are in different groups, Python's `a >> Label("x") >> b` replaces `a` with its group, so all the last tasks of `a`'s group run before `b`. The Go SDK keeps the edge between `a` and `b`. +- **A label on a task or task group is rejected when it would get no edge.** `Before` and `After` panic when their receiver carries a label. Python's `Label("x") >> b` drops that label without an error. `Before` and `After` also panic when they are given a label but their receiver stands for no task or task group. A chain such as `loaded.Before(shards...).Before(airflow.Label("x").On(merged))` panics when `shards` is empty. The same chain without the label declares nothing and does not panic. When `On` is given only empty sets, the label is on no task or task group, so `Before` and `After` drop it without an error. - **`TaskGroup` takes `spec ...TaskGroupSpec`, as `Dag` takes `spec ...DagSpec`.** A group has one kind of option, so it needs no sealed option interface. A mapped task group would come through a method of its own, not through this variadic. ## Alternatives diff --git a/go-sdk/airflow/node.go b/go-sdk/airflow/node.go index 5f620ad9aaa8e..8e1cfa18aafab 100644 --- a/go-sdk/airflow/node.go +++ b/go-sdk/airflow/node.go @@ -50,13 +50,13 @@ type Node interface { // After makes the receiver a downstream of every node, as Python's << does. After(nodes ...Node) Node // node seals the interface. Only a type that package airflow declares is a Node, so - // endpoints can read what one stands for. + // appendEndpoints can read what one stands for. node() } // nodeEndpoint is one task or one task group that a Node stands for, with the label that // [LabelRef.On] carries into the edge verb the Node is passed to. Exactly one of task and group is -// set, except for the endpoint of a nil *TaskRef, which has neither. +// set. type nodeEndpoint struct { task *TaskRef group *TaskGroupRef @@ -96,47 +96,71 @@ func (e nodeEndpoint) describe() string { return fmt.Sprintf("task %q", e.task.taskID) } -// nodeSet is the argument set that Before and After return as one Node, and the Node that -// [LabelRef.On] returns. The labels of a nodeSet go on the edges of the verb that the nodeSet is -// passed to, so a nodeSet with labels cannot be the receiver of a verb. +// nodeSet is the Node that Before and After return, and it holds the nodes that were passed to +// the verb. It carries no labels, because a label belongs to the edges of the verb call that it +// was passed to. type nodeSet []nodeEndpoint func (nodeSet) node() {} -// unlabelled returns the set with its labels dropped. A label belongs to the one verb it was -// passed to, so the Node a verb returns, which stands for the nodes it pointed at, carries none. -func (s nodeSet) unlabelled() nodeSet { - plain := make(nodeSet, len(s)) - for i, endpoint := range s { - plain[i] = nodeEndpoint{task: endpoint.task, group: endpoint.group} - } - return plain -} - func (s nodeSet) Before(nodes ...Node) Node { return declareEdges(s, nodes, dirBefore) } func (s nodeSet) After(nodes ...Node) Node { return declareEdges(s, nodes, dirAfter) } +// labelledSet is the Node that [LabelRef.On] returns. Every endpoint in a labelledSet carries the +// label of the LabelRef. labelledSet is a separate type so that an edge verb can use a type +// assertion to recognize the Node that On returns, without looking at each endpoint. An empty +// labelledSet carries no label. +type labelledSet []nodeEndpoint + +func (labelledSet) node() {} + +func (s labelledSet) Before(nodes ...Node) Node { return declareEdges(s, nodes, dirBefore) } + +func (s labelledSet) After(nodes ...Node) Node { return declareEdges(s, nodes, dirAfter) } + func (*TaskRef) node() {} -// endpoints returns the tasks and task groups that node stands for. The Node method is -// unexported, so the types of this package are the only Nodes, and a struct that embeds one is all -// that reaches the default case. where names the Node in a panic, such as -// "airflow.Node.Before: nodes[0]". -func endpoints(where string, node Node) []nodeEndpoint { +// position names a Node in a panic message, such as airflow.Node.Before: nodes[1]. String builds +// the name only when a panic needs it, so a call that does not panic builds no string. +type position struct { + where string // the function, such as airflow.Node.Before + param string // the parameter, such as nodes, or the receiver + index int // the index into a variadic parameter, or -1 for a parameter that holds one Node +} + +func (p position) String() string { + if p.index < 0 { + return p.where + ": " + p.param + } + return fmt.Sprintf("%s: %s[%d]", p.where, p.param, p.index) +} + +// appendEndpoints appends to set the tasks and task groups that node stands for, with the labels +// that node carries. It returns the extended set. The Node method is unexported, so the types of +// this package are the only Nodes. Only a struct that embeds one of those types reaches the +// default case. at names node in a panic message. +func appendEndpoints(set []nodeEndpoint, node Node, at position) []nodeEndpoint { switch node := node.(type) { + case nil: + panic(at.String() + " is a nil airflow.Node") case *TaskRef: - return []nodeEndpoint{{task: node}} + if node == nil { + panic(at.String() + " is a nil *airflow.TaskRef") + } + return append(set, nodeEndpoint{task: node}) case *TaskGroupRef: if node == nil { - panic(where + " is a nil *airflow.TaskGroupRef") + panic(at.String() + " is a nil *airflow.TaskGroupRef") } - return []nodeEndpoint{{group: node}} + return append(set, nodeEndpoint{group: node}) case nodeSet: - return node + return append(set, node...) + case labelledSet: + return append(set, node...) default: panic(fmt.Sprintf( - "%s has type %T, which is not a node that package airflow defines", where, node, + "%s has type %T, which is not a node that package airflow defines", at.String(), node, )) } } @@ -219,11 +243,19 @@ func Label(text string) *LabelRef { // // A label is the Node's only in the verb it is passed to. The Node that On returns stands for the // nodes themselves. The Node that a verb returns carries no label into the next verb, either as -// its receiver or as an argument. Before and After panic when their receiver carries a label, -// because that label has no edge to go on. Pass the Node that On returns only as an argument. -// Python's Label("x") >> b drops the label without an error instead. When labels nest, the outer -// label is used: loaded.Before(Label("outer").On(Label("inner").On(cleaned))) labels the edge -// "outer". +// its receiver or as an argument. When labels nest, the outer label is used: +// loaded.Before(Label("outer").On(Label("inner").On(cleaned))) labels the edge "outer". +// +// Before and After panic when a label on a task or task group would get no edge: +// - Their receiver carries a label, so pass the Node that On returns only as an argument. +// Python's Label("x") >> b drops such a label without an error instead. +// - They are given a label, but their receiver, such as the Node that loaded.Before() returns, +// stands for no task or task group. So they declare no edge. A chain that passes an empty +// slice of nodes to one verb reaches this case when the next verb in the chain is given a +// label. +// +// When On is given only empty sets, the Node it returns holds no task or task group, so its +// label is on nothing. Before and After then drop the label without an error. // // A label declared on an edge that already carries one replaces it, as Python's DAG.set_edge_info // does. @@ -233,38 +265,34 @@ func Label(text string) *LabelRef { // transformed := dag.Task(transform, Inputs(extracted)) // extracted.Before(Label("rows").On(transformed)) // -// A label on an edge to or from a task group stays on that edge, as [TaskGroupRef.Before] -// describes. A label on an edge between two tasks stays on that edge too, even when the tasks are -// in different task groups. In that case Python can replace the receiver of the verb with a group +// A label on an edge to or from a task group labels that group edge. On an edge from a task that +// no group holds to a group with tasks, it also labels the edges to the first tasks of the group, +// as [TaskGroupRef.Before] describes. If one of those task edges was also declared on its own, or +// was added by an earlier group edge, that task edge keeps the label it already has, or stays +// unlabelled. A label on an edge between two tasks stays on that edge, even when the tasks are in +// different task groups. In that case Python can replace the receiver of the verb with a group // that holds it, which makes the edge connect other tasks. // // On panics if: // - l is nil or points to a zero LabelRef -// - node or a node in more is nil, is a nil *TaskGroupRef, or is a struct that embeds a Node -// -// On does not panic on a nil *TaskRef. The edge verb that takes the Node panics on it instead. +// - node or a node in more is nil, is a nil *TaskRef or *TaskGroupRef, or is a struct that +// embeds a Node func (l *LabelRef) On(node Node, more ...Node) Node { const where = "airflow.LabelRef.On" if l == nil || l.text == "" { panic(where + ": airflow.Label did not return the *airflow.LabelRef") } - nodes := append([]Node{node}, more...) - labelled := make(nodeSet, 0, len(nodes)) - for i, labelling := range nodes { - name := "node" - if i > 0 { - name = fmt.Sprintf("more[%d]", i-1) - } - if labelling == nil { - panic(fmt.Sprintf("%s: %s is a nil airflow.Node", where, name)) - } - for _, endpoint := range endpoints(where+": "+name, labelling) { - labelled = append( - labelled, nodeEndpoint{task: endpoint.task, group: endpoint.group, label: l.text}, - ) - } + labelled := make([]nodeEndpoint, 0, 1+len(more)) + labelled = appendEndpoints(labelled, node, position{where, "node", -1}) + for i, next := range more { + labelled = appendEndpoints(labelled, next, position{where, "more", i}) + } + // When a Node given to this On came from another On, the label of l replaces the label of that + // Node. + for i := range labelled { + labelled[i].label = l.text } - return labelled + return labelledSet(labelled) } // edgeKey identifies an edge of a Dag. Tasks and task groups share one namespace of IDs, so a @@ -280,6 +308,9 @@ const ( dirAfter ) +// verbNames holds the name that panic messages use for each edge verb. +var verbNames = [...]string{dirBefore: "airflow.Node.Before", dirAfter: "airflow.Node.After"} + func (dir edgeDir) String() string { if dir == dirAfter { return "After" @@ -311,22 +342,19 @@ func (e pendingEdge) betweenTasks() bool { // declareEdges records an edge from every endpoint of receiver to every node it was given, or // the other way round for After, and returns those nodes as one Node. func declareEdges(receiver Node, nodes []Node, dir edgeDir) Node { - where := "airflow.Node." + dir.String() - args := make(nodeSet, 0, len(nodes)) + where := verbNames[dir] + args := make([]nodeEndpoint, 0, len(nodes)) for i, node := range nodes { - if node == nil { - panic(fmt.Sprintf("%s: nodes[%d] is a nil airflow.Node", where, i)) - } - args = append(args, endpoints(fmt.Sprintf("%s: nodes[%d]", where, i), node)...) + args = appendEndpoints(args, node, position{where, "nodes", i}) } - ends := endpoints(where+": the receiver", receiver) + ends := appendEndpoints(nil, receiver, position{where, "the receiver", -1}) all := slices.Concat(ends, args) if len(all) == 0 { return nodeSet(nil) } dag := edgeDag(where, all) - checkReceiverUnlabelled(where, dag, dir, ends) + checkReceiverUnlabelled(where, dag, dir, receiver) dag.mu.Lock() defer dag.mu.Unlock() @@ -349,9 +377,17 @@ func declareEdges(receiver Node, nodes []Node, dir edgeDir) Node { dag.checkGroupLocked(where, endpoint.group) } // A verb with no node to point at, which a spread of an empty slice reaches, declares no - // edge. So does one on the empty set that such a verb returned. + // edge. So does a verb called on the empty set that such a verb returned. A label given to this + // second verb would have no edge to go on. if len(ends) == 0 || len(args) == 0 { - return args.unlabelled() + if labelled := describeLabelled(args); labelled != "" { + panic(fmt.Sprintf( + "%s: Dag %q: the receiver stands for no task or task group, so there is no edge "+ + "for %s", + where, dag.dagID, labelled, + )) + } + return nodeSet(args) } // Check and merge every pair the call declares before it records any of them. @@ -390,28 +426,39 @@ func declareEdges(receiver Node, nodes []Node, dir edgeDir) Node { dag.addGroupEdgeLocked(edge.upstream, edge.downstream, edge.label) } } - return args.unlabelled() + // The Node that a verb returns carries no label into the next verb. args belongs to this call, so + // this loop clears the labels of args in place. + for i := range args { + args[i].label = "" + } + return nodeSet(args) } // checkReceiverUnlabelled panics if the receiver of an edge verb carries a label. Only the Node // that [LabelRef.On] returns carries labels. A label goes on the edges between the receiver and // the nodes that the verb is given, so a label on the receiver has no edge to go on. -func checkReceiverUnlabelled(where string, dag *DagRef, dir edgeDir, ends []nodeEndpoint) { - var labelled []string - for _, end := range ends { - if end.label != "" { - labelled = append(labelled, fmt.Sprintf("%s labelled %q", end.describe(), end.label)) - } - } - if len(labelled) > 0 { +func checkReceiverUnlabelled(where string, dag *DagRef, dir edgeDir, receiver Node) { + if labelled, ok := receiver.(labelledSet); ok && len(labelled) > 0 { panic(fmt.Sprintf( "%s: Dag %q: cannot call %s on %s; "+ "pass airflow.Label(...).On(...) to Before or After as an argument instead", - where, dag.dagID, dir, strings.Join(labelled, ", "), + where, dag.dagID, dir, describeLabelled(labelled), )) } } +// describeLabelled names each endpoint in set that carries a label, as in task "load" labelled +// "always". It returns "" when no endpoint carries one. The result goes into a panic message. +func describeLabelled(set []nodeEndpoint) string { + var labelled []string + for _, end := range set { + if end.label != "" { + labelled = append(labelled, fmt.Sprintf("%s labelled %q", end.describe(), end.label)) + } + } + return strings.Join(labelled, ", ") +} + // checkGroupEdgeEnds panics if one end of an edge is a task group that holds the other end. The // edge would order the group against part of itself: in Python, group >> task_in_group makes the // last tasks of the group upstreams of a task that may be one of them. @@ -444,8 +491,6 @@ func edgeDag(where string, ends []nodeEndpoint) *DagRef { var dag *DagRef for _, end := range ends { switch { - case end.task == nil && end.group == nil: - panic(fmt.Sprintf("%s: got a nil *airflow.TaskRef", where)) case end.group != nil && end.group.dag == nil: panic(fmt.Sprintf( "%s: got a *airflow.TaskGroupRef that DagRef.TaskGroup or "+ diff --git a/go-sdk/airflow/node_test.go b/go-sdk/airflow/node_test.go index 7e6dfcff77aad..a3fa867aeef3c 100644 --- a/go-sdk/airflow/node_test.go +++ b/go-sdk/airflow/node_test.go @@ -206,6 +206,9 @@ func TestLabelOnLabelsTheEdgeToEachNodeItIsGiven(t *testing.T) { assertEdgeLabel(t, dag, "load", "notify_empty", "forks") assertGroupEdgeLabel(t, dag, "load", "staging", "forks") assertEdgeLabel(t, dag, "load", "notify_suspicious", "forks") + + Bundle().Register(dag) + assertEdgeLabel(t, dag, "load", "staging.stage", "forks") } func TestAfterLabelsTheEdgeToTheNodeItWasGiven(t *testing.T) { @@ -274,7 +277,7 @@ func TestTheLastLabelOfOneCallWins(t *testing.T) { // TestEdgeVerbsRejectALabelledReceiver covers a verb whose receiver carries a label. That label has // no edge to go on. The verb panics before it declares any edge. The Java SDK throws in the same // case. Python's Label("x") >> b drops the label without an error instead. When the label is on a -// nil *TaskRef, the verb's panic reports the nil task. +// zero TaskRef, the verb panics about that TaskRef instead of about the label. func TestEdgeVerbsRejectALabelledReceiver(t *testing.T) { dag := Dag("etl") loaded := orderedTask(t, dag, "load") @@ -314,9 +317,10 @@ func TestEdgeVerbsRejectALabelledReceiver(t *testing.T) { `labelled "always", task group "staging" labelled "always"; ` + fix, }, { - name: "a label on a nil task", - call: func() { Label("always").On((*TaskRef)(nil)).Before(cleaned) }, - want: "airflow.Node.Before: got a nil *airflow.TaskRef", + name: "a label on a task that DagRef.Task did not return", + call: func() { Label("always").On(&TaskRef{}).Before(cleaned) }, + want: "airflow.Node.Before: got a *airflow.TaskRef that DagRef.Task or " + + "TaskGroupRef.Task did not return", }, } { t.Run(tc.name, func(t *testing.T) { @@ -358,6 +362,46 @@ func TestLabelBelongsToTheVerbItIsPassedTo(t *testing.T) { assertEdgeLabel(t, dag, "extract", "cleanup", "") } +// TestALabelledNodeCanGoIntoSeveralVerbs checks that a verb does not change the labels of the Node +// that On returned. The same Node can then label the edges of every verb it is passed to. +func TestALabelledNodeCanGoIntoSeveralVerbs(t *testing.T) { + dag := Dag("etl") + loaded := orderedTask(t, dag, "load") + extracted := orderedTask(t, dag, "extract") + cleaned := orderedTask(t, dag, "cleanup") + + always := Label("always").On(cleaned) + loaded.Before(always) + extracted.Before(always) + + assertEdgeLabel(t, dag, "load", "cleanup", "always") + assertEdgeLabel(t, dag, "extract", "cleanup", "always") +} + +// TestAVerbResultCanGoIntoSeveralVerbs checks that a verb copies the Node it is given instead of +// appending to it. The slice behind the Node that a verb returns can have room for more endpoints. +// A verb that appended to that slice would change the Node that another verb returned. +func TestAVerbResultCanGoIntoSeveralVerbs(t *testing.T) { + dag := Dag("etl") + source := orderedTask(t, dag, "source") + extracted := orderedTask(t, dag, "extract") + loaded := orderedTask(t, dag, "load") + cleaned := orderedTask(t, dag, "cleanup") + audited := orderedTask(t, dag, "audit") + first := orderedTask(t, dag, "first") + second := orderedTask(t, dag, "second") + done := orderedTask(t, dag, "done") + + var none []Node + // shared stands for extract and load, and its slice has room for one more endpoint. + shared := source.Before(extracted, loaded, source.Before(none...)) + fromFirst := first.Before(shared, cleaned) + second.Before(shared, audited) + fromFirst.Before(done) + + assertTasks(t, done.upstreams, extracted, loaded, cleaned) +} + // TestLabelPassesThroughNoVerbTwice covers the same rule for After, which reads the label from // the node it was given rather than from its receiver. func TestLabelPassesThroughNoVerbTwice(t *testing.T) { @@ -384,6 +428,10 @@ func TestLabelOnPanicsOnANilNode(t *testing.T) { "airflow.LabelRef.On: more[1] is a nil airflow.Node", func() { Label("when empty").On(loaded, loaded, nil) }, ) + assert.PanicsWithValue(t, + "airflow.LabelRef.On: more[0] is a nil *airflow.TaskRef", + func() { Label("when empty").On(loaded, (*TaskRef)(nil)) }, + ) } func TestLabelPanicsOnABlankLabel(t *testing.T) { @@ -433,8 +481,8 @@ func TestEdgeVerbsRejectANilNode(t *testing.T) { func() { loaded.After(nil) }, ) assert.PanicsWithValue(t, - "airflow.Node.Before: got a nil *airflow.TaskRef", - func() { loaded.Before((*TaskRef)(nil)) }, + "airflow.Node.Before: nodes[1] is a nil *airflow.TaskRef", + func() { loaded.Before(cleaned, (*TaskRef)(nil)) }, ) } @@ -525,6 +573,10 @@ func TestEdgeVerbsWithNoNodeDeclareNoEdge(t *testing.T) { var none []Node loaded.Before(none...).Before(cleaned) loaded.After(none...) + // When On is given only that empty set, it returns a Node that holds no task or task group, so + // nothing carries the label. A verb called on that Node then has no label to reject, so it does + // not panic. + Label("always").On(loaded.Before(none...)).Before(cleaned) assert.Empty(t, loaded.downstreams) assert.Empty(t, loaded.upstreams) @@ -532,20 +584,30 @@ func TestEdgeVerbsWithNoNodeDeclareNoEdge(t *testing.T) { assert.Empty(t, dag.edgeLabels) } -// TestAVerbThatDeclaresNoEdgeReturnsNoLabel covers a verb whose receiver stands for no task. The -// verb declares no edge. The Node that the verb returns carries none of the labels on the verb's -// arguments, so a chain can go on from that Node. -func TestAVerbThatDeclaresNoEdgeReturnsNoLabel(t *testing.T) { +// TestAVerbWithAnEmptyReceiverRejectsALabel covers a verb whose receiver stands for no task or task +// group, such as the Node that a verb given no node returns. The verb declares no edge, so a label +// that it is given has no edge to go on. +func TestAVerbWithAnEmptyReceiverRejectsALabel(t *testing.T) { dag := Dag("etl") loaded := orderedTask(t, dag, "load") cleaned := orderedTask(t, dag, "cleanup") + done := orderedTask(t, dag, "done") var none []Node - loaded.Before(none...).Before(Label("always").On(cleaned)).Before(done) - + assert.PanicsWithValue(t, + `airflow.Node.Before: Dag "etl": the receiver stands for no task or task group, `+ + `so there is no edge for task "cleanup" labelled "always"`, + func() { loaded.Before(none...).Before(Label("always").On(cleaned)) }, + ) + assert.PanicsWithValue(t, + `airflow.Node.Before: Dag "etl": the receiver stands for no task or task group, `+ + `so there is no edge for task "done" labelled "always"`, + func() { loaded.Before(none...).Before(cleaned, Label("always").On(done)) }, + ) assert.Empty(t, cleaned.upstreams) - assertEdgeLabel(t, dag, "cleanup", "done", "") + assert.Empty(t, done.upstreams) + assert.Empty(t, dag.edgeLabels) } // TestEdgeVerbsCheckTheTasksEvenWithNoNode pins that a verb with no node to point at still checks @@ -562,7 +624,7 @@ func TestEdgeVerbsCheckTheTasksEvenWithNoNode(t *testing.T) { func() { (&TaskRef{}).Before(none...) }, ) assert.PanicsWithValue(t, - "airflow.Node.After: got a nil *airflow.TaskRef", + "airflow.Node.After: the receiver is a nil *airflow.TaskRef", func() { (*TaskRef)(nil).After(none...) }, ) diff --git a/go-sdk/airflow/task_group_test.go b/go-sdk/airflow/task_group_test.go index 0f28b6d5bdb28..8cc431e6116df 100644 --- a/go-sdk/airflow/task_group_test.go +++ b/go-sdk/airflow/task_group_test.go @@ -1263,7 +1263,8 @@ func snapshot(dag *DagRef) dagSnapshot { s.groups = append(s.groups, group.groupID) for _, child := range group.children { s.children[group.groupID] = append( - s.children[group.groupID], endpoints("child", child)[0].id(), + s.children[group.groupID], + appendEndpoints(nil, child, position{"snapshot", "child", -1})[0].id(), ) } }