diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index f16cc6963f679..3a22b27ca4f2b 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -78,6 +78,12 @@ func (b *BundleRef) serve(args []string, stdout io.Writer) error { // 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) + // --help is output the caller asked for, so it goes to stdout. Anything else pflag prints, + // such as a deprecation warning for a flag the bundle defines, stays on stderr where it + // cannot land in the middle of the --airflow-metadata manifest. + flags.Usage = func() { + fmt.Fprintf(stdout, "Usage of %s:\n%s", flags.Name(), flags.FlagUsages()) + } printMetadata := flags.Bool( "airflow-metadata", false, diff --git a/go-sdk/airflow/serve_test.go b/go-sdk/airflow/serve_test.go index 937df061a608e..93de70af0231c 100644 --- a/go-sdk/airflow/serve_test.go +++ b/go-sdk/airflow/serve_test.go @@ -144,6 +144,24 @@ func TestServeRejectsBundleFlagWithReservedName(t *testing.T) { } } +// pflag prints a deprecation warning while it parses, so it must not reach the writer that +// carries the --airflow-metadata manifest. +func TestServeKeepsDeprecationWarningOutOfTheManifest(t *testing.T) { + saved := flag.CommandLine + t.Cleanup(func() { flag.CommandLine = saved }) + flag.CommandLine = flag.NewFlagSet("bundle", flag.ContinueOnError) + flag.String("region", "", "a flag the bundle author defined") + require.NoError(t, flag.CommandLine.MarkDeprecated("region", "use --zone")) + + var stdout bytes.Buffer + args := []string{"--region", "us", "--airflow-metadata", "--format", "json"} + require.NoError(t, etlBundle().serve(args, &stdout)) + + assert.NotContains(t, stdout.String(), "deprecated") + var got manifest + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) +} + func TestServeRejectsBadFlags(t *testing.T) { tests := []struct { name string @@ -195,7 +213,10 @@ func TestServeRejectsBadFlags(t *testing.T) { } func TestServeHelpIsNotAnError(t *testing.T) { - assert.NoError(t, etlBundle().serve([]string{"--help"}, io.Discard)) + var stdout bytes.Buffer + require.NoError(t, etlBundle().serve([]string{"--help"}, &stdout)) + + assert.Contains(t, stdout.String(), "--airflow-metadata") } // A fake supervisor sends StartupDetails over the comm socket, as the Python diff --git a/go-sdk/airflow/task_handler_test.go b/go-sdk/airflow/task_handler_test.go index 0112335935c56..4f0787c198298 100644 --- a/go-sdk/airflow/task_handler_test.go +++ b/go-sdk/airflow/task_handler_test.go @@ -123,7 +123,7 @@ func TestTaskHandlerPanicsOnBadHandler(t *testing.T) { { name: "variadic", fn: func(Context, ...string) error { return nil }, - want: "is variadic; a task argument cannot fill a ... parameter", + want: "is variadic; declare the last parameter as []T instead of ...T", }, } for _, tt := range tests { diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index c8febc84bee9c..a2399fd9200fd 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -67,9 +67,12 @@ var _ Task = (*taskFunction)(nil) // NewTaskFunction validates and wraps a Go function as a Task. func NewTaskFunction(fn any) (Task, error) { + // The kind comes first: Value.Pointer panics on an int, and Value.Type on an untyped nil. v := reflect.ValueOf(fn) - fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{fn: v, fullName: fullName} + if v.Kind() != reflect.Func { + return nil, fmt.Errorf("expected a func as input but was %s", v.Kind()) + } + f := &taskFunction{fn: v, fullName: runtime.FuncForPC(v.Pointer()).Name()} if err := f.validateFn(v.Type()); err != nil { return nil, err } @@ -152,15 +155,12 @@ func (f *taskFunction) sendXcom( } func (f *taskFunction) validateFn(fnType reflect.Type) error { - if fnType.Kind() != reflect.Func { - 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. + // binding.Analyze turns a ... tail into one []T parameter, so Execute would have to call + // the function with CallSlice rather than Call to fill it. That is only worth doing for a + // signature []T cannot already express, and ...T is not one: both take a single argument. if fnType.IsVariadic() { return fmt.Errorf( - "task function %s is variadic; a task argument cannot fill a ... parameter", + "task function %s is variadic; declare the last parameter as []T instead of ...T", f.fullName, ) } diff --git a/go-sdk/internal/bundle/task_test.go b/go-sdk/internal/bundle/task_test.go index a7ce5eae37c83..d3633e9449c71 100644 --- a/go-sdk/internal/bundle/task_test.go +++ b/go-sdk/internal/bundle/task_test.go @@ -78,6 +78,24 @@ func (s *TaskSuite) TestReturnValidation() { } } +func (s *TaskSuite) TestRejectsNonFunc() { + cases := map[string]struct { + fn any + wantErr string + }{ + "int": {3, "expected a func as input but was int"}, + "nil": {nil, "expected a func as input but was invalid"}, + "pointer": {new(int), "expected a func as input but was ptr"}, + } + + for name, tt := range cases { + s.Run(name, func() { + _, err := NewTaskFunction(tt.fn) + s.Assert().EqualError(err, tt.wantErr) + }) + } +} + // probeKey is an unexported context key used to confirm the live task context // (not a freshly built one) backs the airflow.Context a task receives. type probeKeyType struct{}