Validate Dag and task IDs in the Go SDK

go-sdk/bundle/bundlev1/registry.go

@@ -20,13 +20,34 @@ package bundlev1
 import (
 	"fmt"
 	"reflect"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync"
+	"unicode/utf8"
 
 	"github.com/apache/airflow/go-sdk/pkg/worker"
 )
 
+const maxKeyLength = 250
+
+var keyRegex = regexp.MustCompile(`^[\p{L}\p{N}_.-]+$`)
+
+func validateKey(name, value string) {
+	if length := utf8.RuneCountInString(value); length > maxKeyLength {
+		panic(fmt.Errorf("%s must be less than %d characters, not %d", name, maxKeyLength, length))
+	}
+	if !keyRegex.MatchString(value) {
+		panic(
+			fmt.Errorf(
+				"%s %q must be made of alphanumeric characters, dashes, dots, and underscores",
+				name,
+				value,
+			),
+		)
+	}
+}
+
 type (
 	// Task is one registered task: something the runtime can Execute. Bundle
 	// authors do not implement this directly; Dag.AddTask wraps a plain Go
@@ -60,7 +81,9 @@ type (
 		// AddTaskWithName is like AddTask but sets task_id explicitly instead of
 		// deriving it from the function name. Use it when the Go function name
 		// cannot match the Python @task.stub id, for example for an anonymous
-		// function or a differing name.
+		// function or a differing name. The task_id must be non-empty, at most
+		// 250 characters, and made of alphanumeric characters, dashes, dots,
+		// and underscores; an invalid id panics.
 		AddTaskWithName(taskId string, fn any)
 	}
 
@@ -70,8 +93,10 @@ type (
 	Registry interface {
 		Bundle
 		// AddDag registers a dag by its dag_id (matching the Python stub dag)
-		// and returns a Dag handle for attaching tasks. Registering the same
-		// dag_id twice panics.
+		// and returns a Dag handle for attaching tasks. The dag_id must be
+		// non-empty, at most 250 characters, and made of alphanumeric
+		// characters, dashes, dots, and underscores. An invalid or duplicate
+		// dag_id panics.
 		AddDag(dagId string) Dag
 	}
 
@@ -135,6 +160,8 @@ func getFnName(fn reflect.Value) string {
 }
 
 func (r *registry) AddDag(dagId string) Dag {
+	validateKey("Dag ID", dagId)
+
 	r.Lock()
 	defer r.Unlock()
 
@@ -159,6 +186,8 @@ func (r *registry) registerTask(dagId string, fn any) {
 }
 
 func (r *registry) registerTaskWithName(dagId, taskId string, fn any) {
+	validateKey("Task ID", taskId)
+
 	task, err := NewTaskFunction(fn)
 	if err != nil {
 		panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err))

go-sdk/bundle/bundlev1/registry_test.go

@@ -20,7 +20,9 @@ package bundlev1
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log/slog"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -66,6 +68,55 @@ func (s *RegistrySuite) TestAddDag_DuplicatePanics() {
 	})
 }
 
+func (s *RegistrySuite) TestAddDag_AcceptsValidIds() {
+	reg := New()
+	for _, id := range []string{"simple", "with-dash", "with.dot", "with_underscore", "0numeric", "café_dag", "任務", strings.Repeat("a", 250), strings.Repeat("任", 250)} {
+		s.NotPanics(func() { reg.AddDag(id) })
+	}
+}
+
+func (s *RegistrySuite) TestAddDag_InvalidCharsPanic() {
+	for _, id := range []string{"", "with space", "with/slash", "with:colon", "with\ttab"} {
+		s.PanicsWithError(
+			fmt.Sprintf(
+				"Dag ID %q must be made of alphanumeric characters, dashes, dots, and underscores",
+				id,
+			),
+			func() { New().AddDag(id) },
+		)
+	}
+}
+
+func (s *RegistrySuite) TestAddDag_TooLongPanics() {
+	s.PanicsWithError("Dag ID must be less than 250 characters, not 251", func() {
+		New().AddDag(strings.Repeat("a", 251))
+	})
+}
+
+func (s *RegistrySuite) TestAddTask_InvalidCharsPanic() {
+	for _, id := range []string{"", "with space", "with/slash", "with:colon", "with\ttab"} {
+		s.PanicsWithError(
+			fmt.Sprintf(
+				"Task ID %q must be made of alphanumeric characters, dashes, dots, and underscores",
+				id,
+			),
+			func() { s.dag.AddTaskWithName(id, myTask) },
+		)
+	}
+}
+
+func (s *RegistrySuite) TestAddTask_AcceptsValidIds() {
+	for _, id := range []string{"simple", "with-dash", "with.dot", "with_underscore", "0numeric", "café_dag", "任務", strings.Repeat("a", 250), strings.Repeat("任", 250)} {
+		s.NotPanics(func() { s.dag.AddTaskWithName(id, myTask) })
+	}
+}
+
+func (s *RegistrySuite) TestAddTask_TooLongPanics() {
+	s.PanicsWithError("Task ID must be less than 250 characters, not 251", func() {
+		s.dag.AddTaskWithName(strings.Repeat("a", 251), myTask)
+	})
+}
+
 func (s *RegistrySuite) TestAddTask_RegistersAndFindsTask() {
 	s.dag.AddTask(myTask)
 	task, exists := s.reg.LookupTask("dag1", "myTask")