This is an automated email from the ASF dual-hosted git repository.
jason810496 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 529f429400a Go SDK: reject a task handler and a Dag that share a
dag_id (#73876)
529f429400a is described below
commit 529f429400a3311590217622f86223691c3c4e6f
Author: PoAn Yang <[email protected]>
AuthorDate: Thu Oct 1 20:45:18 2026 +0900
Go SDK: reject a task handler and a Dag that share a dag_id (#73876)
Signed-off-by: PoAn Yang <[email protected]>
---
go-sdk/airflow/bundle.go | 46 +++++++++++++++--
go-sdk/airflow/bundle_test.go | 111 ++++++++++++++++++++++++++++++++++++++++++
2 files changed, 154 insertions(+), 3 deletions(-)
diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go
index 697b6b0ae7d..8f82b071211 100644
--- a/go-sdk/airflow/bundle.go
+++ b/go-sdk/airflow/bundle.go
@@ -31,7 +31,11 @@ import (
type BundleRef struct {
// closed ends registration for everything the bundle can hold, so a
kind added later
// is covered without a flag of its own. Serve sets it; Register reads
it.
- closed atomic.Bool
+ closed atomic.Bool
+ // mu is the lock for writes to taskHandlers and dags. Register holds
it for the whole call,
+ // so two concurrent calls cannot register a task handler and a Dag
with the same dag_id.
+ // Readers such as LookupTask take only the lock of the map they read.
+ mu sync.Mutex
taskHandlers taskHandlerMap
dags dagMap
}
@@ -76,22 +80,42 @@ type Registerable interface{ registerable() }
// registered.
//
// 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, or if
[BundleRef.Serve] has already
-// been called: registration closes when serving starts.
+// 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.
func (b *BundleRef) Register(items ...Registerable) {
if b.closed.Load() {
panic(
"airflow.BundleRef.Register: Serve has already been
called; register everything before Serve",
)
}
+ b.mu.Lock()
+ defer b.mu.Unlock()
for _, item := range items {
switch item := item.(type) {
case *taskHandler:
+ if b.dags.has(item.dagID) {
+ panic(fmt.Sprintf(
+ "airflow.BundleRef.Register: Dag %q is
already registered as a Dag from "+
+ "airflow.Dag, so it cannot also
have task handlers from "+
+ "airflow.TaskHandler",
+ item.dagID,
+ ))
+ }
b.taskHandlers.add(item.dagID, item.taskID, item.task)
case *DagRef:
if item == nil {
panic("airflow.BundleRef.Register: cannot
register a nil *airflow.DagRef")
}
+ if b.taskHandlers.hasDag(item.dagID) {
+ panic(fmt.Sprintf(
+ "airflow.BundleRef.Register: Dag %q
already has task handlers from "+
+ "airflow.TaskHandler, so it
cannot also be registered as a Dag from "+
+ "airflow.Dag",
+ item.dagID,
+ ))
+ }
b.dags.add(item)
default:
// Either a nil item, or a struct from another package
that embeds a Registerable.
@@ -135,6 +159,14 @@ func (m *taskHandlerMap) add(dagID, taskID string, task
bundle.Task) {
m.order = append(m.order, bundle.TaskHandlerInfo{DagID: dagID, TaskID:
taskID})
}
+func (m *taskHandlerMap) hasDag(dagID string) bool {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+
+ _, exists := m.handlers[dagID]
+ return exists
+}
+
func (m *taskHandlerMap) LookupTask(dagID, taskID string) (bundle.Task, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -169,3 +201,11 @@ func (m *dagMap) add(dag *DagRef) {
}
m.dags[dag.dagID] = dag
}
+
+func (m *dagMap) has(dagID string) bool {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ _, exists := m.dags[dagID]
+ return exists
+}
diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go
index d61acbf0454..31a654c453f 100644
--- a/go-sdk/airflow/bundle_test.go
+++ b/go-sdk/airflow/bundle_test.go
@@ -143,6 +143,80 @@ func TestRegisterRejectsDuplicateDag(t *testing.T) {
assert.False(t, second.registered, "a Dag that Register rejects can
still take tasks")
}
+func TestRegisterRejectsADagWithTheDagIDOfATaskHandler(t *testing.T) {
+ tests := []struct {
+ name string
+ register func(b *BundleRef, handler Registerable, dag *DagRef)
+ }{
+ {
+ name: "separate calls",
+ register: func(b *BundleRef, handler Registerable, dag
*DagRef) {
+ b.Register(handler)
+ b.Register(dag)
+ },
+ },
+ {
+ name: "one call",
+ register: func(b *BundleRef, handler Registerable, dag
*DagRef) {
+ b.Register(handler, dag)
+ },
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ b := Bundle()
+ dag := Dag("etl")
+
+ want := `airflow.BundleRef.Register: Dag "etl" already
has task handlers from ` +
+ `airflow.TaskHandler, so it cannot also be
registered as a Dag from airflow.Dag`
+ assert.PanicsWithValue(t, want, func() {
+ tt.register(b, TaskHandler("etl", "transform",
noop), dag)
+ })
+ _, ok := b.taskHandlers.LookupTask("etl", "transform")
+ assert.True(t, ok)
+ assert.NotContains(t, b.dags.dags, "etl")
+ assert.NotPanics(t, func() { dag.Task(extract) },
+ "a Dag that Register rejects can still take
tasks")
+ })
+ }
+}
+
+func TestRegisterRejectsATaskHandlerWithTheDagIDOfADag(t *testing.T) {
+ tests := []struct {
+ name string
+ register func(b *BundleRef, handler Registerable, dag *DagRef)
+ }{
+ {
+ name: "separate calls",
+ register: func(b *BundleRef, handler Registerable, dag
*DagRef) {
+ b.Register(dag)
+ b.Register(handler)
+ },
+ },
+ {
+ name: "one call",
+ register: func(b *BundleRef, handler Registerable, dag
*DagRef) {
+ b.Register(dag, handler)
+ },
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ b := Bundle()
+ dag := Dag("etl")
+
+ want := `airflow.BundleRef.Register: Dag "etl" is
already registered as a Dag from ` +
+ `airflow.Dag, so it cannot also have task
handlers from airflow.TaskHandler`
+ assert.PanicsWithValue(t, want, func() {
+ tt.register(b, TaskHandler("etl", "transform",
noop), dag)
+ })
+ assert.Same(t, dag, b.dags.dags["etl"])
+ assert.NotContains(t, b.taskHandlers.handlers, "etl")
+ assert.Empty(t, b.taskHandlers.ListTaskHandlers())
+ })
+ }
+}
+
func TestRegisterRejectsNilDag(t *testing.T) {
var dag *DagRef
assert.PanicsWithValue(t, "airflow.BundleRef.Register: cannot register
a nil *airflow.DagRef",
@@ -204,6 +278,43 @@ func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
assert.Len(t, b.dags.dags, workers*perWorker)
}
+// For each dag_id, one goroutine registers a task handler and another
registers a Dag at the
+// same time, and exactly one of the two calls must succeed. Without
BundleRef.mu both can
+// succeed. A run with -race does not report that, because every map access
still takes the
+// map's lock.
+func TestRegisterGivesEachDagIDToOneKindUnderConcurrentUse(t *testing.T) {
+ const dagCount = 200
+
+ b := Bundle()
+ start := make(chan struct{})
+ var wg sync.WaitGroup
+ register := func(item Registerable) {
+ defer wg.Done()
+ defer func() { recover() }()
+ <-start
+ b.Register(item)
+ }
+ dags := make([]*DagRef, dagCount)
+ for i := range dags {
+ dags[i] = Dag(fmt.Sprintf("dag_%d", i))
+ wg.Add(2)
+ go register(TaskHandler(dags[i].dagID, "transform", noop))
+ go register(dags[i])
+ }
+ close(start)
+ wg.Wait()
+
+ var wrong []string
+ for _, dag := range dags {
+ _, hasHandler := b.taskHandlers.LookupTask(dag.dagID,
"transform")
+ if hasHandler == dag.registered {
+ wrong = append(wrong, dag.dagID)
+ }
+ }
+ assert.Empty(t, wrong,
+ "each dag_id must end up with a task handler or a Dag, not both
or neither")
+}
+
func TestRegisterRejectsNilItem(t *testing.T) {
var item Registerable
assert.PanicsWithValue(t, "airflow.BundleRef.Register: cannot register
<nil>", func() {
