diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index b13b74ec76733..14522f47af984 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -21,6 +21,7 @@ import ( "fmt" "slices" "sync" + "sync/atomic" "github.com/apache/airflow/go-sdk/internal/bundle" ) @@ -28,6 +29,9 @@ import ( // BundleRef holds the task handlers that this executable runs for Airflow. // [Bundle] returns an empty one. 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 taskHandlers taskHandlerMap } @@ -66,8 +70,14 @@ type Registerable interface{ registerable() } // // bundle.Register(reports.Handlers()...) // -// Register panics if a task handler with the same dag_id and task_id is already registered. +// Register panics if a task handler with the same dag_id and task_id is already registered, +// and if [BundleRef.Serve] has already been called: registration closes when serving starts. func (b *BundleRef) Register(items ...Registerable) { + if b.closed.Load() { + panic( + "airflow.BundleRef.Register: Serve has already been called; register everything before Serve", + ) + } for _, item := range items { switch item := item.(type) { case *taskHandler: diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index 9ad5217b30b6a..fb9503e9a3268 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -19,6 +19,7 @@ package airflow import ( "fmt" + "io" "os" "os/exec" "reflect" @@ -116,6 +117,38 @@ func TestRegisterRejectsDuplicateTask(t *testing.T) { ) } +func TestRegisterAfterServePanics(t *testing.T) { + b := Bundle() + b.Register(TaskHandler("py_etl", "transform", noop)) + require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard)) + + assert.PanicsWithValue( + t, + "airflow.BundleRef.Register: Serve has already been called; register everything before Serve", + func() { b.Register(TaskHandler("py_etl", "load", noop)) }, + ) +} + +// The flag lives on the bundle, not on the task-handler map, so a kind of item added to +// Register later is covered without a flag of its own. +func TestRegisterAfterServeRejectsEveryKindOfItem(t *testing.T) { + b := Bundle() + require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard)) + + var nilItem Registerable + assert.Panics(t, func() { b.Register(nilItem) }, + "the closed check runs before Register looks at what the item is") +} + +// Serve closes registration whatever the run does, so a bundle that only printed its usage +// still refuses a late Register. +func TestRegisterAfterAFailedServePanics(t *testing.T) { + b := Bundle() + require.NoError(t, b.serve([]string{"--help"}, io.Discard)) + + assert.Panics(t, func() { b.Register(TaskHandler("py_etl", "transform", noop)) }) +} + func TestRegisterIsSafeForConcurrentUse(t *testing.T) { const workers, perWorker = 8, 100 diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index eb61002a25e71..f16cc6963f679 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -51,6 +51,9 @@ const ( // Serve runs the bundle. Call it as the last statement of main. // +// Serve closes registration: a [BundleRef.Register] call that reaches the bundle afterwards +// panics rather than changing what the running bundle answers for. +// // The command-line flags of the executable decide what Serve does. // With --airflow-metadata it prints the bundle's manifest and returns, which is how // airflow-go-pack reads the registered Dag and task ids. @@ -68,6 +71,10 @@ func (b *BundleRef) Serve() error { } func (b *BundleRef) serve(args []string, stdout io.Writer) error { + // Registration closes here, whatever this run turns out to do, so that a Register left + // below Serve in main is reported as the mistake it is rather than racing the runtime. + b.closed.Store(true) + // The flags go on their own FlagSet. On pflag.CommandLine, every program that imports this // package would get them, and one that defines its own --format there would panic. flags := flag.NewFlagSet(os.Args[0], flag.ContinueOnError) diff --git a/go-sdk/airflow/task_handler_test.go b/go-sdk/airflow/task_handler_test.go index 5df87af0e62c0..0112335935c56 100644 --- a/go-sdk/airflow/task_handler_test.go +++ b/go-sdk/airflow/task_handler_test.go @@ -120,6 +120,11 @@ func TestTaskHandlerPanicsOnBadHandler(t *testing.T) { fn: func(Context, chan int) error { return nil }, want: "type chan int cannot receive a task argument", }, + { + name: "variadic", + fn: func(Context, ...string) error { return nil }, + want: "is variadic; a task argument cannot fill a ... parameter", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index c012b1a80aa14..c8febc84bee9c 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -156,6 +156,15 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error { return fmt.Errorf("expected a func as input but was %s", fnType.Kind()) } + // Execute calls the function with Call, which passes a variadic parameter its slice + // rather than spreading it, so a variadic task function would panic at execution. + if fnType.IsVariadic() { + return fmt.Errorf( + "task function %s is variadic; a task argument cannot fill a ... parameter", + f.fullName, + ) + } + // Return values // `, error`, or just `error` if fnType.NumOut() < 1 || fnType.NumOut() > 2 {