This is an automated email from the ASF dual-hosted git repository.
henry3260 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f070997ba63 Go SDK: order tasks with Before, After and Label (#74119)
f070997ba63 is described below
commit f070997ba63e7d3c88eb035bf9398025567e864f
Author: Henry Chen <[email protected]>
AuthorDate: Sun Oct 4 01:15:51 2026 +0800
Go SDK: order tasks with Before, After and Label (#74119)
* Go SDK: order tasks with Before, After and Label
A Dag authored in Go can now order tasks that exchange no data, the way
Python's >> and << do, and label the edge that reaches a task.
loaded.Before(notified, cleaned) // loaded >> [notify, cleanup]
cleaned.After(extracted) // cleanup << extracted
loaded.Before(airflow.Label(emptyNotice, "when empty"))
Both verbs are variadic, so one call fans out, and both return their
argument set as one Node, so a.Before(b, c).Before(d) is a >> [b, c] >> d.
airflow.Label wraps the endpoint rather than the call, so each edge of a
fan-out can carry a label of its own. Declaring an edge the Dag already
has only applies the label, which is how an airflow.Inputs edge gets one.
* Go SDK: check an edge verb's tasks even when it points at none
Before and After with no node declare no edge, but they now still reject a
Dag that has been registered and a *TaskRef that DagRef.Task did not
return, so a mistake in the receiver is not hidden by an empty fan-out.
The label a verb puts on an edge is now settled in one place, mergeLabel,
which both the check over the call's pairs and the recording step use. An
edge verb enumerates its pairs once and records what it checked, so the
checks and the recording cannot drift apart, and addEdgeLocked keeps only
the cycle check, which cannot run up front.
An edge given two labels in one call now says so rather than reporting the
first as already on the edge, which no earlier declaration had put there.
* Go SDK: let a second edge label replace the first, as Python does
Declaring a label on an edge that already carries one now replaces it,
which is what Python's DAG.set_edge_info does: "this will overwrite,
rather than merge with, existing info". It used to panic, a rule the Go
SDK had of its own.
A label on the receiver of an edge verb is still dropped, since Label
marks the edge that reaches a node and the receiver is the node an edge
leaves from. Python is silent there too: Label("x") >> b, with nothing
upstream of the label, sets no label either. Tests now pin both, and the
nesting that Label(Label(x, "inner"), "outer") reads as.
* Go SDK: cover the cycle an Inputs edge can close
DagRef.Task records the edges Inputs declares so that the cycle check sees
them, which is the case ADR-0008 names: b := dag.Task(B, Inputs(a)) then
b.Before(a) is a genuine cycle in accepted syntax. Only the upstream and
downstream lists were asserted, so deleting that recording failed no test
over the cycle itself.
* Go SDK: reject a Dag's cycles at Register, over the whole graph
Before, After and Inputs each walked the graph to see whether the edge they
were recording closed a cycle, which made building a Dag cost one walk per
edge: a 8,000-task chain declared in reverse took seconds. Register now
walks the graph once, after the Dag is whole, as the Java SDK's
Bundle.register does and the TypeScript SDK plans to. The same chain builds
in tens of milliseconds.
Recording an edge no longer panics, so a fan-out that is rejected leaves the
Dag exactly as it was, and the caveat about a cycle keeping the earlier
edges of its call is gone with it. A Dag whose dependencies contain a cycle
stays unregistered, so it can still be corrected.
A task ordered against itself is still rejected where it is declared: that
needs no walk.
* Go SDK: order a condition's Then and Else tasks after it
IfRef.Then and IfRef.Else stored the task on the condition without recording
an edge, so the ordering their doc states, that the two run after the
condition, was in no Dag the SDK builds: the task had no upstream, and a
cycle running through a condition passed Register.
read := dag.Task(readRows)
loaded := dag.Task(load)
dag.If(hasRows, airflow.Inputs(read)).Then(loaded)
loaded.Before(read) // read -> hasRows -> load -> read
Naming a task now records the edge from the condition to it, which is also
the upstream the serialized Dag needs. Declaring that edge again, as an
author who writes it out does, changes nothing.
---
go-sdk/airflow/bundle.go | 12 +-
go-sdk/airflow/bundle_test.go | 54 +++
go-sdk/airflow/dag.go | 29 +-
go-sdk/airflow/if.go | 7 +
go-sdk/airflow/if_test.go | 50 +++
go-sdk/airflow/node.go | 383 ++++++++++++++++++++
go-sdk/airflow/node_test.go | 538 ++++++++++++++++++++++++++++
go-sdk/airflow/testdata/foreignnode/main.go | 39 ++
8 files changed, 1105 insertions(+), 7 deletions(-)
diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go
index 91b2773fd75..d100231d832 100644
--- a/go-sdk/airflow/bundle.go
+++ b/go-sdk/airflow/bundle.go
@@ -79,12 +79,16 @@ type Registerable interface{ registerable() }
// Add every task to a Dag before registering the Dag. [DagRef.Task],
[DagRef.If], [IfRef.Then]
// and [IfRef.Else] panic once the Dag is registered.
//
+// Register is where a Dag's task dependencies are checked for a cycle, over
the whole graph at
+// once: [TaskRef.Before], [TaskRef.After] and [Inputs] each record an edge
without walking the
+// graph, so building a Dag stays linear in its edges however many a task has.
+//
// Register panics if a task handler with the same dag_id and task_id is
already registered,
// if a Dag with the same dag_id is already registered, if a task handler and
a Dag have the
-// same dag_id, or if [BundleRef.Serve] has already been called: registration
closes when
-// serving starts. A task handler runs a task of a Python Dag, so its dag_id
cannot also belong
-// to a Dag authored in Go. Register also panics if a Dag has a condition from
[DagRef.If]
-// without a task from [IfRef.Then].
+// same dag_id, if the task dependencies of a Dag contain a cycle, or if
[BundleRef.Serve] has
+// already been called: registration closes when serving starts. A task
handler runs a task of a
+// Python Dag, so its dag_id cannot also belong to a Dag authored in Go.
Register also panics if a
+// Dag has a condition from [DagRef.If] without a task from [IfRef.Then].
func (b *BundleRef) Register(items ...Registerable) {
if b.closed.Load() {
panic(
diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go
index 31a654c453f..7a5874e29b4 100644
--- a/go-sdk/airflow/bundle_test.go
+++ b/go-sdk/airflow/bundle_test.go
@@ -217,6 +217,60 @@ func TestRegisterRejectsATaskHandlerWithTheDagIDOfADag(t
*testing.T) {
}
}
+// TestRegisterRejectsACycle covers the check that reads a Dag's edges as a
whole. Before and
+// After record an edge without walking the graph, so a cycle between other
tasks is Register's to
+// find, and every edge of the call that closed it is recorded until then.
+func TestRegisterRejectsACycle(t *testing.T) {
+ dag := Dag("etl")
+ extracted := orderedTask(t, dag, "extract")
+ loaded := orderedTask(t, dag, "load")
+ notified := orderedTask(t, dag, "notify")
+ extracted.Before(loaded)
+ loaded.Before(notified, extracted)
+
+ assert.PanicsWithValue(t,
+ `airflow.BundleRef.Register: the task dependencies of Dag "etl"
contain a cycle: `+
+ `extract -> load -> extract`,
+ func() { Bundle().Register(dag) },
+ )
+ // The verbs recorded what they were given, and the Dag can still be
corrected.
+ assertTasks(t, loaded.downstreams, notified, extracted)
+ assert.False(t, dag.registered)
+}
+
+// TestRegisterRejectsACycleThroughAnInputsEdge pins that the check sees the
edges Inputs
+// declared, which is what DagRef.Task recording them is for. It is the case
ADR-0008 names:
+// b := dag.Task(B, Inputs(a)) followed by b.Before(a) is a genuine cycle in
accepted syntax.
+func TestRegisterRejectsACycleThroughAnInputsEdge(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ counted := dag.Task(countRows, Inputs(read))
+ notified := orderedTask(t, dag, "notify")
+ counted.Before(notified)
+ counted.Before(read)
+
+ assert.PanicsWithValue(t,
+ `airflow.BundleRef.Register: the task dependencies of Dag "etl"
contain a cycle: `+
+ `readRows -> countRows -> readRows`,
+ func() { Bundle().Register(dag) },
+ )
+}
+
+// TestRegisterTakesADagWhoseTasksShareADownstream pins that the walk follows
a diamond, where a
+// task is reached twice without any cycle.
+func TestRegisterTakesADagWhoseTasksShareADownstream(t *testing.T) {
+ dag := Dag("etl")
+ extracted := orderedTask(t, dag, "extract")
+ notified := orderedTask(t, dag, "notify")
+ cleaned := orderedTask(t, dag, "cleanup")
+ done := orderedTask(t, dag, "done")
+ extracted.Before(notified, cleaned).Before(done)
+
+ Bundle().Register(dag)
+
+ assert.True(t, dag.registered)
+}
+
func TestRegisterRejectsNilDag(t *testing.T) {
var dag *DagRef
assert.PanicsWithValue(t, "airflow.BundleRef.Register: cannot register
a nil *airflow.DagRef",
diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go
index bbfbf022dc6..8a45c867eb4 100644
--- a/go-sdk/airflow/dag.go
+++ b/go-sdk/airflow/dag.go
@@ -38,6 +38,10 @@ type DagRef struct {
registered bool
tasks []*TaskRef
tasksByID map[string]*TaskRef
+ // edgeLabels holds every edge 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 map[edgeKey]string
}
// Dag returns an empty Dag with the given dag_id. An optional [DagSpec] holds
the rest of the
@@ -74,7 +78,8 @@ func (*DagRef) registerable() {}
// TaskRef is a task that [DagRef.Task] added to a Dag. Pass it to [Inputs] to
give its result
// to a task that DagRef.Task or [DagRef.If] adds later. Pass it to
[IfRef.Then] or [IfRef.Else]
-// to run it on one side of a condition.
+// to run it on one side of a condition. A TaskRef is a [Node], so
[TaskRef.Before] and
+// [TaskRef.After] order it against another task.
type TaskRef struct {
dag *DagRef
taskID string
@@ -85,7 +90,12 @@ type TaskRef struct {
// inputs holds the tasks that Inputs passed, in the order of the
parameters they fill. Each
// of them is an upstream task of this one.
inputs []*TaskRef
- task bundle.Task
+ // upstreams and downstreams hold the edges of the task, in the order
they were declared and
+ // without a repeat, so that an edge is recorded in both directions.
Inputs, Before and After
+ // all record an edge here.
+ upstreams []*TaskRef
+ downstreams []*TaskRef
+ task bundle.Task
// triggerDagRun is the checked copy of the TriggerDagRunSpec of a task
from TriggerDagRun.
// It is nil for a task that runs a Go function. A task from
TriggerDagRun runs no Go
// function, so its resultType, inputs and task are nil.
@@ -270,11 +280,18 @@ func (d *DagRef) addTask(method string, fn any, opts
[]TaskOption, ifRef *IfRef)
}
d.tasksByID[taskID] = task
d.tasks = append(d.tasks, task)
+ // Inputs passes a task once per parameter it fills, so the same task
can arrive twice. The
+ // edge is one either way, and the task is new, so no edge to it
carries a label to settle.
+ for _, upstream := range upstreams {
+ d.addEdgeLocked(upstream, task, "")
+ }
return task
}
// markRegistered marks d as registered, which stops any further change to d.
It panics instead
-// when a condition from If has no task from Then.
+// when a condition from If has no task from Then, or when the edges of d
close a cycle. The Dag
+// is whole by then, so one walk of the graph answers for every edge its tasks
declared, and a Dag
+// that fails a check stays unregistered and can still be corrected.
func (d *DagRef) markRegistered() {
d.mu.Lock()
defer d.mu.Unlock()
@@ -288,6 +305,12 @@ func (d *DagRef) markRegistered() {
))
}
}
+ if cycle := d.cycleLocked(); cycle != nil {
+ panic(fmt.Sprintf(
+ "airflow.BundleRef.Register: the task dependencies of
Dag %q contain a cycle: %s",
+ d.dagID, strings.Join(cycle, " -> "),
+ ))
+ }
d.registered = true
}
diff --git a/go-sdk/airflow/if.go b/go-sdk/airflow/if.go
index 2e49f7409d6..2d5b79d7912 100644
--- a/go-sdk/airflow/if.go
+++ b/go-sdk/airflow/if.go
@@ -64,6 +64,9 @@ type IfRef struct {
// - false skips the task from Then
// - an error fails the task, which then skips nothing
//
+// A task from Then or Else runs after the condition, so naming it records an
edge from the
+// condition to it, as [TaskRef.Before] would. Declaring that edge as well
changes nothing.
+//
// The condition skips only that one task. In Python, a branch skips every
task directly after it
// that it does not follow, and a short circuit skips every task after it.
After the skip, the
// trigger rule of each task after the skipped task decides whether that task
runs. A task that
@@ -172,6 +175,10 @@ func (g *IfRef) setTask(side string, task *TaskRef) {
))
}
*slot = task
+ // The task runs after the condition, so it is a downstream task of it.
Recording the edge is
+ // what puts the condition in the serialized Dag as the task's
upstream, and what lets
+ // registration see a cycle that runs through a condition.
+ d.addEdgeLocked(g.task, task, "")
}
// wrapCondition wraps fn as the task of g. The task skips the side of g that
the result of fn
diff --git a/go-sdk/airflow/if_test.go b/go-sdk/airflow/if_test.go
index cf51c330b1e..2f5730547c3 100644
--- a/go-sdk/airflow/if_test.go
+++ b/go-sdk/airflow/if_test.go
@@ -259,6 +259,39 @@ func TestThenAndElseNameTheTasksOfTheCondition(t
*testing.T) {
assert.Same(t, reported, reversed.elseTask)
}
+// TestThenAndElseRecordTheEdgeFromTheCondition pins that naming a task orders
it after the
+// condition, which is what puts the condition in the serialized Dag as the
task's upstream.
+func TestThenAndElseRecordTheEdgeFromTheCondition(t *testing.T) {
+ dag := Dag("etl")
+ loaded := dag.Task(load)
+ reported := dag.Task(reportEmpty)
+
+ gate := dag.If(isReady)
+ gate.Then(loaded)
+ gate.Else(reported)
+
+ assertTasks(t, gate.task.downstreams, loaded, reported)
+ assertTasks(t, loaded.upstreams, gate.task)
+ assertTasks(t, reported.upstreams, gate.task)
+ assertEdgeLabel(t, dag, "isReady", "load", "")
+ assertEdgeLabel(t, dag, "isReady", "reportEmpty", "")
+}
+
+// TestDeclaringTheEdgeOfAConditionAgainChangesNothing covers an author who
also writes the edge
+// that Then records, which is as idempotent as declaring any edge twice.
+func TestDeclaringTheEdgeOfAConditionAgainChangesNothing(t *testing.T) {
+ dag := Dag("etl")
+ loaded := dag.Task(load)
+
+ gate := dag.If(isReady)
+ gate.Then(loaded)
+ gate.task.Before(Label(loaded, "when ready"))
+
+ assertTasks(t, gate.task.downstreams, loaded)
+ assertTasks(t, loaded.upstreams, gate.task)
+ assertEdgeLabel(t, dag, "isReady", "load", "when ready")
+}
+
func TestThenAndElseRejectATaskOutsideTheDag(t *testing.T) {
dag := Dag("etl")
loaded := dag.Task(load)
@@ -485,6 +518,23 @@ func runCondition(gate *IfRef, results, earlier
map[string]any) (*conditionClien
return client, gate.task.task.Execute(ctx, discardLogger(), args)
}
+// TestRegisterRejectsACycleThroughACondition pins that the edge Then records
is one the cycle
+// check walks, so a cycle that runs through a condition is rejected like any
other.
+func TestRegisterRejectsACycleThroughACondition(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ loaded := dag.Task(load)
+ dag.If(hasRows, Inputs(read)).Then(loaded)
+ loaded.Before(read)
+
+ assert.PanicsWithValue(t,
+ `airflow.BundleRef.Register: the task dependencies of Dag "etl"
contain a cycle: `+
+ `readRows -> hasRows -> load -> readRows`,
+ func() { Bundle().Register(dag) },
+ )
+ assert.False(t, dag.registered)
+}
+
func TestConditionSkipsTheSideThatItDoesNotTake(t *testing.T) {
tests := []struct {
name string
diff --git a/go-sdk/airflow/node.go b/go-sdk/airflow/node.go
new file mode 100644
index 00000000000..6b0cefab74d
--- /dev/null
+++ b/go-sdk/airflow/node.go
@@ -0,0 +1,383 @@
+// 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 airflow
+
+import (
+ "fmt"
+ "slices"
+)
+
+// Node is what an edge between tasks connects. A [TaskRef] is one, so the
task that
+// [DagRef.Task] returns is an edge endpoint.
+//
+// Before and After declare the order-only edges that Python writes with >>
and <<, for tasks
+// that have to run in an order but pass no data. [Inputs] declares the edge
that carries a
+// result.
+//
+// Both verbs are variadic, so one call fans out, and both return their
argument set as one Node
+// rather than the receiver. That is what makes a chain run the way it reads:
+//
+// loaded.Before(notified, cleaned) // loaded >> [notify, cleanup]
+// cleaned.After(extracted) // cleanup << extracted
+// extracted.Before(loaded).Before(done) // extract >> load >> done
+//
+// The second Before above starts from load, not from extract.
+//
+// Node is the Go counterpart of Python's DAGNode, where set_upstream and
set_downstream live.
+// Its node method is unexported, so a type declared outside package airflow
can be a Node only by
+// embedding one, and Before and After reject such a type as an argument.
+type Node interface {
+ // Before makes the receiver an upstream task of every node, as
Python's >> does.
+ Before(nodes ...Node) Node
+ // After makes the receiver a downstream task 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.
+ node()
+}
+
+// nodeEndpoint is one task that a Node stands for, with the label that
[Label] carries into the
+// edge verb the Node is passed to.
+type nodeEndpoint struct {
+ task *TaskRef
+ label 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.
+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 tasks 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}
+ }
+ 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) }
+
+func (*TaskRef) node() {}
+
+// endpoints returns the tasks 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 {
+ switch node := node.(type) {
+ case *TaskRef:
+ return []nodeEndpoint{{task: node}}
+ case nodeSet:
+ return node
+ default:
+ panic(fmt.Sprintf(
+ "%s has type %T, which is not a node that package
airflow defines", where, node,
+ ))
+ }
+}
+
+// Before makes the task an upstream task of every node, which is Python's
+// loaded >> [notify, cleanup]:
+//
+// loaded.Before(notified, cleaned)
+//
+// The edge carries no data, so a task that only comes after another takes no
parameter for it.
+// Use [Inputs] for an edge that passes a result.
+//
+// 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.
+//
+// Declaring an edge that the Dag already has changes nothing, other than to
apply a [Label].
+//
+// Before panics if:
+// - a node is nil, or is a task that [DagRef.Task] did not return
+// - a node belongs to another Dag
+// - the edge would make a task depend on itself
+// - the Dag is already registered
+//
+// Every check runs before the call records any edge, so a fan-out that panics
leaves the Dag as
+// it was. Edges that close a cycle between other tasks are
[BundleRef.Register]'s to reject,
+// over the whole graph at once.
+func (t *TaskRef) Before(nodes ...Node) Node { return declareEdges(t, nodes,
dirBefore) }
+
+// After makes the task a downstream task of every node, which is Python's
cleanup << extracted:
+//
+// cleaned.After(extracted)
+//
+// It is Before with the direction reversed, and it panics for the same
reasons. Like Before, it
+// returns the nodes it was given as one Node, so
cleaned.After(extracted).After(started) is
+// 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:
+//
+// loaded.Before(Label(emptyNotice, "when empty"))
+//
+// 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(processed, "rows found"), Label(emptyNotice, "no
rows"))
+//
+// 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.
+//
+// A label declared on an edge that already carries one replaces it, as
Python's DAG.set_edge_info
+// does.
+//
+// An [Inputs] edge is labelled by declaring it again, which is idempotent:
+//
+// transformed := dag.Task(transform, Inputs(extracted))
+// extracted.Before(Label(transformed, "rows"))
+//
+// 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")
+ }
+ labelling := endpoints("airflow.Label: node", node)
+ labelled := make(nodeSet, len(labelling))
+ for i, endpoint := range labelling {
+ labelled[i] = nodeEndpoint{task: endpoint.task, label: text}
+ }
+ return labelled
+}
+
+// edgeKey identifies an edge of a Dag. The task_ids of a Dag are unique, so
they name the ends.
+type edgeKey struct{ upstream, downstream string }
+
+// edgeDir is which way an edge verb points: [TaskRef.Before] from its
receiver, and
+// [TaskRef.After] at it.
+type edgeDir int
+
+const (
+ dirBefore edgeDir = iota
+ dirAfter
+)
+
+func (dir edgeDir) String() string {
+ if dir == dirAfter {
+ return "After"
+ }
+ return "Before"
+}
+
+// order returns the two ends of the edge between an endpoint of the receiver
and one of the nodes
+// the verb was given.
+func (dir edgeDir) order(recv, arg *TaskRef) (upstream, downstream *TaskRef) {
+ if dir == dirAfter {
+ return arg, recv
+ }
+ return recv, arg
+}
+
+// pendingEdge is an edge that declareEdges has checked and is about to record.
+type pendingEdge struct {
+ upstream, downstream *TaskRef
+ key edgeKey
+ label string
+}
+
+// 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))
+ 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)...)
+ }
+ ends := endpoints(where+": the receiver", receiver)
+ all := slices.Concat(ends, args)
+ if len(all) == 0 {
+ return nodeSet(nil)
+ }
+
+ dag := edgeDag(where, all)
+
+ dag.mu.Lock()
+ defer dag.mu.Unlock()
+
+ if dag.registered {
+ panic(fmt.Sprintf(
+ "%s: Dag %q has already been registered; declare every
edge before Register",
+ where, dag.dagID,
+ ))
+ }
+ for _, endpoint := range all {
+ // A zero TaskRef and a copy of a TaskRef get here.
+ if dag.tasksByID[endpoint.task.taskID] != endpoint.task {
+ panic(fmt.Sprintf(
+ "%s: Dag %q got a *airflow.TaskRef that
DagRef.Task did not return",
+ where, dag.dagID,
+ ))
+ }
+ }
+ // 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.
+ if len(ends) == 0 || len(args) == 0 {
+ return args.unlabelled()
+ }
+
+ // Check and merge every pair the call declares before it records any
of them.
+ pending := make([]pendingEdge, 0, len(ends)*len(args))
+ at := make(map[edgeKey]int, len(ends)*len(args))
+ for _, end := range ends {
+ for _, arg := range args {
+ upstream, downstream := dir.order(end.task, arg.task)
+ if upstream == downstream {
+ panic(fmt.Sprintf(
+ "%s: Dag %q: task %q cannot depend on
itself",
+ where, dag.dagID, upstream.taskID,
+ ))
+ }
+ key := edgeKey{upstream: upstream.taskID, downstream:
downstream.taskID}
+ i, declared := at[key]
+ if !declared {
+ i = len(pending)
+ at[key] = i
+ pending = append(pending, pendingEdge{
+ upstream: upstream, downstream:
downstream, key: key,
+ label: dag.edgeLabels[key],
+ })
+ }
+ // The label belongs to the node the verb was given, in
either direction.
+ pending[i].label = mergeLabel(pending[i].label,
arg.label)
+ }
+ }
+ for _, edge := range pending {
+ dag.addEdgeLocked(edge.upstream, edge.downstream, edge.label)
+ }
+ return args.unlabelled()
+}
+
+// mergeLabel returns the label an edge carries once label is declared on it.
A declaration that
+// carries no label leaves the edge's own label alone, which is what makes
redeclaring an edge
+// idempotent, and one that carries a label overwrites it, as Python's
DAG.set_edge_info does.
+func mergeLabel(declared, label string) string {
+ if label == "" {
+ return declared
+ }
+ return label
+}
+
+// edgeDag returns the Dag that every end of an edge belongs to. It panics
unless each end is a
+// task of that one Dag.
+func edgeDag(where string, ends []nodeEndpoint) *DagRef {
+ var first *TaskRef
+ for _, end := range ends {
+ task := end.task
+ switch {
+ case task == nil:
+ panic(fmt.Sprintf("%s: got a nil *airflow.TaskRef",
where))
+ case task.dag == nil:
+ panic(fmt.Sprintf(
+ "%s: got a *airflow.TaskRef that DagRef.Task
did not return", where,
+ ))
+ case first == nil:
+ first = task
+ case task.dag != first.dag:
+ panic(fmt.Sprintf(
+ "%s: cannot declare an edge between task %q of
Dag %q and task %q of Dag %q; "+
+ "an edge connects tasks of one Dag",
+ where, first.taskID, first.dag.dagID,
task.taskID, task.dag.dagID,
+ ))
+ }
+ }
+ return first.dag
+}
+
+// addEdgeLocked records one edge of d, which the caller holds d.mu for. An
edge d already has is
+// recorded once, and label is what it carries from here on, so a caller that
declares an edge
+// again settles the label with [mergeLabel] first. Whether the edges of a Dag
close a cycle is
+// [BundleRef.Register]'s to answer, over the whole graph at once.
+func (d *DagRef) addEdgeLocked(upstream, downstream *TaskRef, label string) {
+ if d.edgeLabels == nil {
+ d.edgeLabels = make(map[edgeKey]string)
+ }
+ key := edgeKey{upstream: upstream.taskID, downstream: downstream.taskID}
+ if _, exists := d.edgeLabels[key]; !exists {
+ upstream.downstreams = append(upstream.downstreams, downstream)
+ downstream.upstreams = append(downstream.upstreams, upstream)
+ }
+ d.edgeLabels[key] = label
+}
+
+// cycleLocked returns the task_ids on a cycle of the Dag, closed by the task
it starts from
+// again, and nil when the Dag is acyclic. The caller holds d.mu. Registration
calls it once, so
+// building a Dag walks the graph once rather than once per edge.
+func (d *DagRef) cycleLocked() []string {
+ const (
+ unvisited = iota
+ onPath
+ settled
+ )
+ state := make(map[*TaskRef]int, len(d.tasks))
+ var path []*TaskRef
+ // walk is (non-tailrec-eligible) recursive, so a Dag whose
dependencies nest deeper than the
+ // stack takes would overflow it. Airflow's own Dag serialization
recurses over a Dag too.
+ var walk func(task *TaskRef) []string
+ walk = func(task *TaskRef) []string {
+ state[task] = onPath
+ path = append(path, task)
+ for _, downstream := range task.downstreams {
+ switch state[downstream] {
+ case onPath:
+ // The cycle is the path from downstream
onwards, closed by downstream again.
+ return append(taskIDs(path[slices.Index(path,
downstream):]), downstream.taskID)
+ case unvisited:
+ if cycle := walk(downstream); cycle != nil {
+ return cycle
+ }
+ }
+ }
+ path = path[:len(path)-1]
+ state[task] = settled
+ return nil
+ }
+ for _, task := range d.tasks {
+ if state[task] == unvisited {
+ if cycle := walk(task); cycle != nil {
+ return cycle
+ }
+ }
+ }
+ return nil
+}
+
+func taskIDs(tasks []*TaskRef) []string {
+ ids := make([]string, len(tasks))
+ for i, task := range tasks {
+ ids[i] = task.taskID
+ }
+ return ids
+}
diff --git a/go-sdk/airflow/node_test.go b/go-sdk/airflow/node_test.go
new file mode 100644
index 00000000000..e284a9f092d
--- /dev/null
+++ b/go-sdk/airflow/node_test.go
@@ -0,0 +1,538 @@
+// 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 airflow
+
+import (
+ "fmt"
+ "os"
+ "os/exec"
+ "reflect"
+ "sync"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// Embedding promotes node, so this struct compiles as a Node.
+type wrappedNode struct{ Node }
+
+// orderedTask adds a task that passes no data, which is what an order-only
edge connects.
+func orderedTask(t *testing.T, dag *DagRef, taskID string) *TaskRef {
+ t.Helper()
+ return dag.Task(ping, TaskSpec{TaskID: taskID})
+}
+
+func assertTasks(t *testing.T, got []*TaskRef, want ...*TaskRef) {
+ t.Helper()
+ require.Len(t, got, len(want))
+ for i := range want {
+ assert.Same(t, want[i], got[i], "task %d", i)
+ }
+}
+
+func assertEdgeLabel(t *testing.T, dag *DagRef, upstream, downstream, want
string) {
+ t.Helper()
+ label, declared := dag.edgeLabels[edgeKey{upstream: upstream,
downstream: downstream}]
+ require.True(t, declared, "the edge %s -> %s was not declared",
upstream, downstream)
+ assert.Equal(t, want, label)
+}
+
+func TestBeforeDeclaresTheEdgeInBothDirections(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(cleaned)
+
+ assertTasks(t, loaded.downstreams, cleaned)
+ assertTasks(t, cleaned.upstreams, loaded)
+ assert.Empty(t, loaded.upstreams)
+ assert.Empty(t, cleaned.downstreams)
+ assertEdgeLabel(t, dag, "load", "cleanup", "")
+}
+
+func TestAfterDeclaresTheEdgeBeforeWouldDeclare(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ cleaned.After(loaded)
+
+ assertTasks(t, loaded.downstreams, cleaned)
+ assertTasks(t, cleaned.upstreams, loaded)
+}
+
+func TestEdgeVerbsFanOut(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ notified := orderedTask(t, dag, "notify")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(notified, cleaned)
+
+ assertTasks(t, loaded.downstreams, notified, cleaned)
+ assertTasks(t, notified.upstreams, loaded)
+ assertTasks(t, cleaned.upstreams, loaded)
+}
+
+// TestBeforeReturnsItsArgumentSet pins what makes a chain work. Before
returns the nodes it was
+// given, so a.Before(b, c).Before(d) is a >> [b, c] >> d. Returning the
receiver would read like
+// a chain and mean a second fan-out from a.
+func TestBeforeReturnsItsArgumentSet(t *testing.T) {
+ dag := Dag("etl")
+ extracted := orderedTask(t, dag, "extract")
+ notified := orderedTask(t, dag, "notify")
+ cleaned := orderedTask(t, dag, "cleanup")
+ done := orderedTask(t, dag, "done")
+
+ extracted.Before(notified, cleaned).Before(done)
+
+ assertTasks(t, extracted.downstreams, notified, cleaned)
+ assertTasks(t, done.upstreams, notified, cleaned)
+ assertTasks(t, notified.downstreams, done)
+ assertTasks(t, cleaned.downstreams, done)
+}
+
+func TestAfterReturnsItsArgumentSet(t *testing.T) {
+ dag := Dag("etl")
+ started := orderedTask(t, dag, "start")
+ notified := orderedTask(t, dag, "notify")
+ cleaned := orderedTask(t, dag, "cleanup")
+ done := orderedTask(t, dag, "done")
+
+ done.After(notified, cleaned).After(started)
+
+ assertTasks(t, done.upstreams, notified, cleaned)
+ assertTasks(t, started.downstreams, notified, cleaned)
+ assertTasks(t, notified.upstreams, started)
+}
+
+func TestEdgeVerbsReturnTheWholeArgumentSetOfEveryNode(t *testing.T) {
+ dag := Dag("etl")
+ extracted := orderedTask(t, dag, "extract")
+ notified := orderedTask(t, dag, "notify")
+ cleaned := orderedTask(t, dag, "cleanup")
+ done := orderedTask(t, dag, "done")
+
+ // The set that the first verb returns is one node of the second.
+ extracted.Before(extracted.Before(notified, cleaned), done)
+
+ assertTasks(t, extracted.downstreams, notified, cleaned, done)
+}
+
+func TestInputsDeclaresAnEdgeInBothDirections(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ counted := dag.Task(countRows, Inputs(read))
+
+ assertTasks(t, read.downstreams, counted)
+ assertTasks(t, counted.upstreams, read)
+ assertEdgeLabel(t, dag, "readRows", "countRows", "")
+}
+
+func TestInputsDeclaresOneEdgePerTaskItRepeats(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ compared := dag.Task(compareRows, Inputs(read, read))
+
+ assertTasks(t, read.downstreams, compared)
+ assertTasks(t, compared.upstreams, read)
+}
+
+func TestRedeclaringAnEdgeIsIdempotent(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ counted := dag.Task(countRows, Inputs(read))
+
+ read.Before(counted)
+ counted.After(read)
+
+ assertTasks(t, read.downstreams, counted)
+ assertTasks(t, counted.upstreams, read)
+}
+
+func TestLabelLabelsTheEdgeToTheNodeItWraps(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ emptyNotice := orderedTask(t, dag, "notify_empty")
+
+ loaded.Before(Label(emptyNotice, "when empty"))
+
+ assertTasks(t, loaded.downstreams, emptyNotice)
+ assertEdgeLabel(t, dag, "load", "notify_empty", "when empty")
+}
+
+func TestLabelLabelsEachEdgeOfAFanOut(t *testing.T) {
+ dag := Dag("etl")
+ checked := orderedTask(t, dag, "check")
+ processed := orderedTask(t, dag, "process")
+ emptyNotice := orderedTask(t, dag, "notify_empty")
+
+ checked.Before(Label(processed, "rows found"), Label(emptyNotice, "no
rows"))
+
+ assertEdgeLabel(t, dag, "check", "process", "rows found")
+ assertEdgeLabel(t, dag, "check", "notify_empty", "no rows")
+}
+
+func TestAfterLabelsTheEdgeToTheNodeItWasGiven(t *testing.T) {
+ dag := Dag("etl")
+ extracted := orderedTask(t, dag, "extract")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ cleaned.After(Label(extracted, "always"))
+
+ assertEdgeLabel(t, dag, "extract", "cleanup", "always")
+}
+
+// TestRedeclaringAnInputsEdgeLabelsIt pins how a data edge gets a label:
Inputs takes none, so
+// the edge it declared is declared again, which only applies the label.
+func TestRedeclaringAnInputsEdgeLabelsIt(t *testing.T) {
+ dag := Dag("etl")
+ read := dag.Task(readRows)
+ counted := dag.Task(countRows, Inputs(read))
+
+ read.Before(Label(counted, "rows"))
+
+ assertTasks(t, read.downstreams, counted)
+ assertTasks(t, counted.upstreams, read)
+ assertEdgeLabel(t, dag, "readRows", "countRows", "rows")
+}
+
+func TestRedeclaringALabelledEdgeKeepsTheLabel(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(Label(cleaned, "always"))
+ loaded.Before(cleaned)
+ cleaned.After(Label(loaded, "always"))
+
+ assertEdgeLabel(t, dag, "load", "cleanup", "always")
+}
+
+// TestASecondLabelOnAnEdgeReplacesTheFirst follows Python, where
DAG.set_edge_info overwrites
+// rather than merges, so the last label declared on an edge is the one it
carries.
+func TestASecondLabelOnAnEdgeReplacesTheFirst(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(Label(cleaned, "always"))
+ loaded.Before(Label(cleaned, "when empty"))
+
+ assertEdgeLabel(t, dag, "load", "cleanup", "when empty")
+ assertTasks(t, loaded.downstreams, cleaned)
+}
+
+// TestTheLastLabelOfOneCallWins covers the same rule within a single fan-out,
where both labels
+// reach the edge through one call.
+func TestTheLastLabelOfOneCallWins(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(Label(cleaned, "always"), Label(cleaned, "when empty"))
+
+ 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) {
+ 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", "")
+}
+
+// TestTheOuterLabelOfANestedLabelWins pins what nesting means: the outer
Label relabels the
+// endpoints the inner one marked, rather than the two combining.
+func TestTheOuterLabelOfANestedLabelWins(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ loaded.Before(Label(Label(cleaned, "inner"), "outer"))
+
+ assertEdgeLabel(t, dag, "load", "cleanup", "outer")
+}
+
+// TestLabelBelongsToTheVerbItIsPassedTo pins that a label is on the edge of
the call it appears
+// in. The Node that a verb returns stands for the tasks it pointed at, so it
carries no label on
+// either side of the next verb: as its receiver, and as one of its nodes.
+func TestLabelBelongsToTheVerbItIsPassedTo(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+ done := orderedTask(t, dag, "done")
+ extracted := orderedTask(t, dag, "extract")
+
+ labelled := loaded.Before(Label(cleaned, "always"))
+ labelled.Before(done)
+ extracted.Before(labelled)
+
+ assertEdgeLabel(t, dag, "load", "cleanup", "always")
+ assertEdgeLabel(t, dag, "cleanup", "done", "")
+ assertEdgeLabel(t, dag, "extract", "cleanup", "")
+}
+
+// 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) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+ notified := orderedTask(t, dag, "notify")
+
+ notified.After(cleaned.After(Label(loaded, "always")))
+
+ assertEdgeLabel(t, dag, "load", "cleanup", "always")
+ assertEdgeLabel(t, dag, "load", "notify", "")
+}
+
+func TestLabelPanicsOnANilNode(t *testing.T) {
+ assert.PanicsWithValue(t,
+ "airflow.Label: got a nil airflow.Node",
+ func() { Label(nil, "when empty") },
+ )
+}
+
+func TestLabelPanicsOnAnEmptyLabel(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, "") },
+ )
+}
+
+func TestEdgeVerbsRejectANilNode(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ assert.PanicsWithValue(t,
+ "airflow.Node.Before: nodes[1] is a nil airflow.Node",
+ func() { loaded.Before(cleaned, nil) },
+ )
+ assert.PanicsWithValue(t,
+ "airflow.Node.After: nodes[0] is a nil airflow.Node",
+ func() { loaded.After(nil) },
+ )
+ assert.PanicsWithValue(t,
+ "airflow.Node.Before: got a nil *airflow.TaskRef",
+ func() { loaded.Before((*TaskRef)(nil)) },
+ )
+}
+
+func TestEdgeVerbsRejectATaskThatTaskDidNotReturn(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ copied := *loaded
+
+ assert.PanicsWithValue(t,
+ "airflow.Node.Before: got a *airflow.TaskRef that DagRef.Task
did not return",
+ func() { (&TaskRef{}).Before(loaded) },
+ )
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: Dag "etl" got a *airflow.TaskRef that
DagRef.Task did not return`,
+ func() { loaded.Before(&copied) },
+ )
+}
+
+func TestEdgeVerbsRejectATaskOfAnotherDag(t *testing.T) {
+ loaded := orderedTask(t, Dag("etl"), "load")
+ cleaned := orderedTask(t, Dag("reporting"), "cleanup")
+
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: cannot declare an edge between task
"load" of Dag "etl" and `+
+ `task "cleanup" of Dag "reporting"; an edge connects
tasks of one Dag`,
+ func() { loaded.Before(cleaned) },
+ )
+}
+
+func TestEdgeVerbsRejectATaskBeforeItself(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: Dag "etl": task "load" cannot depend on
itself`,
+ func() { loaded.Before(loaded) },
+ )
+ assert.Empty(t, loaded.downstreams)
+}
+
+// TestEdgeVerbsRecordNoEdgeWhenAPairIsRejected pins that the checks run over
the whole fan-out
+// before any of it is recorded, the way DagRef.Task settles every check
before it writes a task.
+func TestEdgeVerbsRecordNoEdgeWhenAPairIsRejected(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+ notified := orderedTask(t, dag, "notify")
+
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: Dag "etl": task "load" cannot depend on
itself`,
+ func() { loaded.Before(cleaned, loaded) },
+ )
+ assert.PanicsWithValue(t,
+ `airflow.Node.After: Dag "etl": task "load" cannot depend on
itself`,
+ func() { loaded.After(notified, loaded) },
+ )
+
+ assert.Empty(t, loaded.downstreams)
+ assert.Empty(t, loaded.upstreams)
+ assert.Empty(t, cleaned.upstreams)
+ assert.Empty(t, notified.downstreams)
+ assert.Empty(t, dag.edgeLabels)
+}
+
+func TestEdgeVerbsAfterRegisterPanic(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+ Bundle().Register(dag)
+
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: Dag "etl" has already been registered; `+
+ `declare every edge before Register`,
+ func() { loaded.Before(cleaned) },
+ )
+ assert.Empty(t, loaded.downstreams)
+}
+
+func TestEdgeVerbsWithNoNodeDeclareNoEdge(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ // A spread of an empty slice reaches a verb with no node, and the
empty set it returns
+ // carries a chain no further.
+ var none []Node
+ loaded.Before(none...).Before(cleaned)
+ loaded.After(none...)
+
+ assert.Empty(t, loaded.downstreams)
+ assert.Empty(t, loaded.upstreams)
+ assert.Empty(t, cleaned.upstreams)
+ assert.Empty(t, dag.edgeLabels)
+}
+
+// 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.
+func TestEdgeVerbsCheckTheTasksEvenWithNoNode(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ var none []Node
+
+ assert.PanicsWithValue(t,
+ "airflow.Node.Before: got a *airflow.TaskRef that DagRef.Task
did not return",
+ func() { (&TaskRef{}).Before(none...) },
+ )
+ assert.PanicsWithValue(t,
+ "airflow.Node.After: got a nil *airflow.TaskRef",
+ func() { (*TaskRef)(nil).After(none...) },
+ )
+
+ Bundle().Register(dag)
+ assert.PanicsWithValue(t,
+ `airflow.Node.Before: Dag "etl" has already been registered; `+
+ `declare every edge before Register`,
+ func() { loaded.Before(none...) },
+ )
+}
+
+func TestEdgeVerbsAreSafeForConcurrentUse(t *testing.T) {
+ const workers = 8
+
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ tasks := make([]*TaskRef, workers)
+ for worker := range workers {
+ tasks[worker] = orderedTask(t, dag, fmt.Sprintf("notify_%d",
worker))
+ }
+
+ var wg sync.WaitGroup
+ for worker := range workers {
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ loaded.Before(tasks[worker])
+ tasks[worker].After(loaded)
+ }()
+ }
+ wg.Wait()
+
+ assert.Len(t, loaded.downstreams, workers)
+ assert.Len(t, dag.edgeLabels, workers)
+}
+
+// TestEdgeVerbsRejectANodeTheyDoNotDefine covers the one gap in the seal:
embedding promotes the
+// node method, so a struct that embeds a Node compiles as one.
+func TestEdgeVerbsRejectANodeTheyDoNotDefine(t *testing.T) {
+ dag := Dag("etl")
+ loaded := orderedTask(t, dag, "load")
+ cleaned := orderedTask(t, dag, "cleanup")
+
+ assert.PanicsWithValue(t,
+ "airflow.Node.Before: nodes[0] has type airflow.wrappedNode, "+
+ "which is not a node that package airflow defines",
+ func() { loaded.Before(wrappedNode{cleaned}) },
+ )
+ assert.PanicsWithValue(t,
+ "airflow.Label: node has type airflow.wrappedNode, "+
+ "which is not a node that package airflow defines",
+ func() { Label(wrappedNode{cleaned}, "when empty") },
+ )
+ assert.Empty(t, loaded.downstreams)
+}
+
+func TestNodeIsSealed(t *testing.T) {
+ typ := reflect.TypeFor[Node]()
+ require.Equal(t, 3, typ.NumMethod())
+ // reflect reports a package path only for an unexported method.
+ method, ok := typ.MethodByName("node")
+ require.True(t, ok, "Node has no unexported node method to seal it")
+ assert.Equal(t, "github.com/apache/airflow/go-sdk/airflow",
method.PkgPath)
+}
+
+func TestNodeRejectsForeignTypes(t *testing.T) {
+ if testing.Short() {
+ t.Skip("shells out to `go build`")
+ }
+ if _, err := exec.LookPath("go"); err != nil {
+ t.Skip("go toolchain not on PATH")
+ }
+
+ out, err := exec.Command("go", "build", "-o", os.DevNull,
"./testdata/foreignnode").
+ CombinedOutput()
+
+ require.Error(t, err, "a type defined outside package airflow must not
compile as a Node")
+ // The test checks only the two type names, so that a change in the
wording of the compiler
+ // error does not break it.
+ assert.Contains(t, string(out), "foreignNode")
+ assert.Contains(t, string(out), "airflow.Node")
+}
diff --git a/go-sdk/airflow/testdata/foreignnode/main.go
b/go-sdk/airflow/testdata/foreignnode/main.go
new file mode 100644
index 00000000000..74af5d1abd0
--- /dev/null
+++ b/go-sdk/airflow/testdata/foreignnode/main.go
@@ -0,0 +1,39 @@
+// 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 foreignnode tries to use a node that package airflow did not define
as an edge
+// endpoint. It must not compile. TestNodeRejectsForeignTypes builds it and
expects that failure.
+package main
+
+import "github.com/apache/airflow/go-sdk/airflow"
+
+type foreignNode struct{}
+
+func (foreignNode) Before(...airflow.Node) airflow.Node { return nil }
+
+func (foreignNode) After(...airflow.Node) airflow.Node { return nil }
+
+// An unexported method name belongs to the package that declares it, so this
method is not the
+// node method of airflow.Node even though it is spelled the same.
+func (foreignNode) node() {}
+
+func extract(airflow.Context) error { return nil }
+
+func main() {
+ dag := airflow.Dag("etl")
+ dag.Task(extract).Before(foreignNode{})
+}
