jason810496 commented on code in PR #73426: URL: https://github.com/apache/airflow/pull/73426#discussion_r4059688158
########## go-sdk/airflow/bundle.go: ########## @@ -0,0 +1,134 @@ +// 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" + "sync" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +// BundleRef holds the task handlers that this executable runs for Airflow. +// [Bundle] returns an empty one. +type BundleRef struct { + tasks taskMap Review Comment: I wonder should we store the `taskHandlerMap` and the `dagMap` to separate the purpose, since they will be completely different entity eventually. What do you think? ########## go-sdk/internal/bundlev1/doc.go: ########## @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -// Package bundlev1server runs a bundle through Airflow's coordinator protocol. -// Bundle entry points call [Serve] so the Python ExecutableCoordinator can -// connect over the local comm and logs channels. -package bundlev1server +// Package bundlev1 defines what the coordinator runtime needs from a bundle: +// the tasks it looks up and runs, and the Dag and task ids it lists in the manifest. +// +// Package airflow builds both from the task handlers a bundle registers. +package bundlev1 Review Comment: Would `internal/bundle` (without v1) be better package namespace? ########## go-sdk/airflow/task_handler.go: ########## @@ -0,0 +1,60 @@ +// 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" + "reflect" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +type taskHandler struct { + dagId, taskId string + task bundlev1.Task +} + +func (*taskHandler) registraterable() {} + +// TaskHandler makes fn the Go body of a task that a Python Dag declares with @task.stub. +// Pass what it returns to [BundleRef.Register]. +// +// dagId is the dag_id of that Python Dag, and taskId is the task_id of the stub task. +// +// fn takes a [Context] first, as the package documentation describes. +// Every parameter after the Context is data, filled from the arguments of the Python stub's +// TaskFlow call. Review Comment: ```suggestion // Every parameter after the Context is data, filled from the arguments of the Python stub's TaskFlow call. ``` ########## go-sdk/airflow/bundle.go: ########## @@ -0,0 +1,134 @@ +// 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" + "sync" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +// BundleRef holds the task handlers that this executable runs for Airflow. +// [Bundle] returns an empty one. +type BundleRef struct { + tasks taskMap +} + +// Bundle returns an empty bundle. Register the task handlers on it, then call Serve as the +// last statement of main: +// +// func main() { +// bundle := airflow.Bundle() +// +// bundle.Register( +// airflow.TaskHandler("py_etl", "transform", transform), +// ) +// +// if err := bundle.Serve(); err != nil { +// log.Fatal(err) +// } +// } +func Bundle() *BundleRef { return &BundleRef{} } + +// Registraterable is what [BundleRef.Register] accepts. [TaskHandler] returns one. +// +// Its only method is unexported, so a type outside this package cannot declare it. +// A struct that embeds a Registraterable still satisfies the interface, and Register panics +// when it is given one. +type Registraterable interface{ registraterable() } + +// Register adds items to the bundle. +// +// A package that defines task handlers can return them as a []Registraterable, +// and main passes that slice as bundle.Register(pkg.Handlers()...). Review Comment: Not sure would this be more clear for the users / agents when invoking the method? ```suggestion // and main passes that slice as bundle.Register(airflow.TaskHandler(...), airflow.TaskHandler(...)). ``` ########## go-sdk/internal/contexttest/contexttest.go: ########## @@ -0,0 +1,60 @@ +// 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 contexttest provides a stand-in for airflow.Context. +// It is for the tests of the packages that airflow imports. +// +// Those tests cannot import airflow back, because that would be an import cycle, so they cannot +// declare a task function that takes an airflow.Context. +// They pass [New] to binding.RegisterTaskContext and declare [Context] parameters instead. +package contexttest + +import ( + "context" + "log/slog" + + "github.com/apache/airflow/go-sdk/sdk" +) + +// Context has the same methods as airflow.Context. +type Context struct { + context.Context + + logger *slog.Logger + client sdk.Client + ti sdk.TaskInstance + dagRun sdk.DagRun +} + +// New has the same signature as airflow.NewContext. +func New( + ctx context.Context, + logger *slog.Logger, + client sdk.Client, + ti sdk.TaskInstance, + dagRun sdk.DagRun, +) Context { + return Context{Context: ctx, logger: logger, client: client, ti: ti, dagRun: dagRun} +} + Review Comment: Not necessary in this PR, if we really need the the `context` definition here for the `internal/bundle` test, but we can have a prek hook to make sure the signature and the implementation of `context` here is same as the `airflow.Context` as follow-up. Another direction (but I'm not sure if it's feasible), could we define the `airflow.Context` at `internal/context:Context` then export it as `airflow.Context` so users can't tell the difference but it's better for codebase level testing, etc. ########## go-sdk/internal/bundlev1/task.go: ########## @@ -31,18 +31,38 @@ import ( ) // Task is one registered task that the coordinator runtime can execute. Bundle -// authors do not implement this directly; Dag.AddTask wraps a plain Go +// authors do not implement this directly. airflow.TaskHandler wraps a plain Go // function into a Task. type Task interface { Execute(ctx context.Context, logger *slog.Logger, args []binding.Arg) error } -// Bundle is the execution-time view of a registry. It looks up a task by -// dag_id and task_id. +// Bundle looks up a registered task by dag_id and task_id. The coordinator +// runtime uses Bundle to find the task the supervisor asked for. type Bundle interface { LookupTask(dagId, taskId string) (Task, bool) } +// TaskInfo describes a registered task by its user-visible id. +type TaskInfo struct { + ID string +} + +// DagInfo describes a registered dag together with its tasks in +// registration order. +type DagInfo struct { + DagID string + Tasks []TaskInfo +} + +// EnumerableBundle lists the registered Dags and their tasks in registration +// order. DumpAirflowMetadata in pkg/execution builds the --airflow-metadata +// manifest from that list, which is how airflow-go-pack reads a bundle's Dag and +// task ids without running a task. +type EnumerableBundle interface { + OrderedDags() []DagInfo Review Comment: ```suggestion OrderedTaskHandler() []SomethingElse ``` Or `ListTaskHandlers() []SomethingElse`? ########## go-sdk/airflow/task_handler.go: ########## @@ -0,0 +1,60 @@ +// 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" + "reflect" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +type taskHandler struct { + dagId, taskId string + task bundlev1.Task +} + +func (*taskHandler) registraterable() {} + +// TaskHandler makes fn the Go body of a task that a Python Dag declares with @task.stub. +// Pass what it returns to [BundleRef.Register]. +// +// dagId is the dag_id of that Python Dag, and taskId is the task_id of the stub task. +// +// fn takes a [Context] first, as the package documentation describes. +// Every parameter after the Context is data, filled from the arguments of the Python stub's +// TaskFlow call. +// fn returns either error or (result, error). +// A non-nil error fails the task, and a non-nil result is pushed as the task's return-value XCom. +// +// TaskHandler checks the signature of fn and panics if the check fails, for example when fn is +// not a function, does not take a Context first, or does not return an error. +// main calls TaskHandler before Serve, so a handler that fails the check stops the executable as +// soon as it starts instead of when the task first runs. +func TaskHandler(dagId, taskId string, fn any) Registraterable { + if reflect.ValueOf(fn).Kind() != reflect.Func { Review Comment: `Kind() != reflect.Func` does not catch a typed-nil func value — `Kind()` reflects the static type, not nilness. A `var h func(airflow.Context) error` left unassigned would pass this check and only panic ("reflect: call of nil function") when the task actually runs, which contradicts the doc comment's promise that `TaskHandler` "stops the executable as soon as it starts instead of when the task first runs." ```suggestion if v := reflect.ValueOf(fn); v.Kind() != reflect.Func || v.IsNil() { ``` ########## go-sdk/airflow/bundle.go: ########## @@ -0,0 +1,134 @@ +// 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" + "sync" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +// BundleRef holds the task handlers that this executable runs for Airflow. +// [Bundle] returns an empty one. +type BundleRef struct { + tasks taskMap +} + +// Bundle returns an empty bundle. Register the task handlers on it, then call Serve as the +// last statement of main: +// +// func main() { +// bundle := airflow.Bundle() +// +// bundle.Register( +// airflow.TaskHandler("py_etl", "transform", transform), +// ) +// +// if err := bundle.Serve(); err != nil { +// log.Fatal(err) +// } +// } +func Bundle() *BundleRef { return &BundleRef{} } + +// Registraterable is what [BundleRef.Register] accepts. [TaskHandler] returns one. +// +// Its only method is unexported, so a type outside this package cannot declare it. +// A struct that embeds a Registraterable still satisfies the interface, and Register panics +// when it is given one. +type Registraterable interface{ registraterable() } + +// Register adds items to the bundle. +// +// A package that defines task handlers can return them as a []Registraterable, +// and main passes that slice as bundle.Register(pkg.Handlers()...). +// +// Register panics if a task handler with the same dag_id and task_id is already registered. +func (b *BundleRef) Register(items ...Registraterable) { + for _, item := range items { + switch item := item.(type) { + case *taskHandler: + b.tasks.add(item.dagId, item.taskId, item.task) + default: + // Either a nil item, or a struct from another package that embeds a Registraterable. + panic(fmt.Sprintf("airflow.BundleRef.Register: cannot register %T", item)) + } + } +} + +// taskMap holds the registered tasks by dag_id and task_id. +// It also keeps registration order. The --airflow-metadata manifest lists the tasks of each Dag +// in that order. +type taskMap struct { + mu sync.RWMutex + tasks map[string]map[string]bundlev1.Task Review Comment: Ditto. How about having the "TaskHandler" term instead of the "Task" term. ########## go-sdk/airflow/bundle.go: ########## @@ -0,0 +1,134 @@ +// 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" + "sync" + + "github.com/apache/airflow/go-sdk/internal/bundlev1" +) + +// BundleRef holds the task handlers that this executable runs for Airflow. +// [Bundle] returns an empty one. +type BundleRef struct { + tasks taskMap +} + +// Bundle returns an empty bundle. Register the task handlers on it, then call Serve as the +// last statement of main: +// +// func main() { +// bundle := airflow.Bundle() +// +// bundle.Register( +// airflow.TaskHandler("py_etl", "transform", transform), +// ) +// +// if err := bundle.Serve(); err != nil { +// log.Fatal(err) +// } +// } +func Bundle() *BundleRef { return &BundleRef{} } + +// Registraterable is what [BundleRef.Register] accepts. [TaskHandler] returns one. +// +// Its only method is unexported, so a type outside this package cannot declare it. +// A struct that embeds a Registraterable still satisfies the interface, and Register panics +// when it is given one. +type Registraterable interface{ registraterable() } Review Comment: Typo: "Registraterable" / "registraterable" (should be "Registerable" / "registerable") caught by Claude Code. ```suggestion type Registerable interface{ registerable() } ``` ########## go-sdk/airflow/serve.go: ########## @@ -15,113 +15,97 @@ // specific language governing permissions and limitations // under the License. -package bundlev1server +package airflow import ( "errors" + "io" + "os" flag "github.com/spf13/pflag" - "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/pkg/execution" ) -// ErrCoordinatorFlagsRequired is returned by [Serve] unless both --comm and +// errCoordinatorFlagsRequired is returned by Serve unless both --comm and // --logs are supplied. Bundle execution always uses the coordinator protocol. -var ErrCoordinatorFlagsRequired = errors.New( +var errCoordinatorFlagsRequired = errors.New( "--comm and --logs are required for bundle execution", ) -// ErrFormatRequiresMetadata is returned by [Serve] when --format is supplied +// errFormatRequiresMetadata is returned by Serve when --format is supplied // without --airflow-metadata, the only mode whose encoding it selects. -var ErrFormatRequiresMetadata = errors.New( +var errFormatRequiresMetadata = errors.New( "--format is only valid together with --airflow-metadata", ) -// CLI Flags, all read by Serve to choose a server mode below. -// --airflow-metadata prints the bundle's manifest and exits (airflow-go-pack -// consumes it to build the embedded airflow-metadata.yaml); --format selects -// its encoding. --comm and --logs select coordinator mode. -var ( - printMetadata = flag.Bool( +// serveMode tags the protocol the binary will speak this run. +type serveMode int + +const ( + modeAirflowMetadata serveMode = iota // --airflow-metadata: print the manifest JSON (ADR 0002/0004) + modeCoordinator // --comm/--logs: msgpack-over-IPC (ADR 0003) + modeCoordinatorUsageError // missing coordinator flags +) + +// Serve runs the bundle. Call it as the last statement of main. +// +// 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. +// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the +// coordinator protocol. +// +// main must exit with a non-zero status when Serve returns an error, because the exit status +// is how the supervisor learns that the task failed: +// +// if err := bundle.Serve(); err != nil { +// log.Fatal(err) +// } +func (b *BundleRef) Serve() error { + return b.serve(os.Args[1:], os.Stdout) +} + +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) + printMetadata := flags.Bool( "airflow-metadata", false, "print the bundle's airflow-metadata manifest and exit", ) - metadataFormat = flag.String( + metadataFormat := flags.String( "format", string(execution.MetadataFormatYAML), "encoding for --airflow-metadata: yaml (default) or json; only valid with --airflow-metadata", ) - commAddr = flag.String( + commAddr := flags.String( "comm", "", "host:port of the supervisor's coordinator comm channel (selects coordinator mode)", ) - logsAddr = flag.String( + logsAddr := flags.String( "logs", "", "host:port of the supervisor's coordinator logs channel (selects coordinator mode)", ) -) - -// ServeOpt is an interface for defining options that can be passed to the -// Serve function. Each implementation modifies the ServeConfig being -// generated. A slice of ServeOpts then, cumulatively applied, render a full -// ServeConfig. -type ServeOpt interface { - ApplyServeOpt(*ServerConfig) error -} - -type serveConfigFunc func(*ServerConfig) error - -func (s serveConfigFunc) ApplyServeOpt(in *ServerConfig) error { - return s(in) -} - -// ServerConfig holds settings that ServeOpt values apply before the bundle -// server starts. It is currently empty; it exists so options can be added later -// without changing Serve's signature. -type ServerConfig struct{} - -// serveMode tags the protocol the binary will speak this run. -type serveMode int - -const ( - modeAirflowMetadata serveMode = iota // --airflow-metadata: print the manifest JSON (ADR 0002/0004) - modeCoordinator // --comm/--logs: msgpack-over-IPC (ADR 0003) - modeCoordinatorUsageError // missing coordinator flags -) - -// Serve is the entrypoint for a bundle executed by Airflow's coordinator. -// -// The mode is decided from CLI flags. Callers should -// surface the returned error so misuse (e.g. only one of --comm/--logs -// supplied) produces a non-zero exit: -// -// func main() { -// if err := bundlev1server.Serve(&myBundle{}); err != nil { -// log.Fatal(err) -// } -// } -// -// Zero or more options to configure the server may also be passed. There are -// no options yet; the parameter exists to allow future additions without -// breaking compatibility. -func Serve(bundle bundlev1.BundleProvider, opts ...ServeOpt) error { - flag.Parse() - - serveConfig := &ServerConfig{} - for _, c := range opts { - c.ApplyServeOpt(serveConfig) + // A bundle may define flags of its own on pflag.CommandLine. Serve parses the whole command + // line, so it has to accept those too. + flags.AddFlagSet(flag.CommandLine) Review Comment: Comments from Claude Code: `flags.AddFlagSet(flag.CommandLine)` runs after the four SDK flags (`airflow-metadata`, `format`, `comm`, `logs`) are already defined on `flags`, and pflag silently skips merging any flag from `flag.CommandLine` that shares one of those names — a bundle author who defines their own `--format` on `pflag.CommandLine` gets it silently dropped instead of an error, and `serve` proceeds using the SDK's own `--format` under a name the bundle thought was its own. ```suggestion for _, name := range []string{"airflow-metadata", "format", "comm", "logs"} { if flag.CommandLine.Lookup(name) != nil { return fmt.Errorf("bundle defines a flag named %q, which collides with a flag Serve reserves", name) } } flags.AddFlagSet(flag.CommandLine) ``` (needs `"fmt"` added to the import block) ########## go-sdk/internal/bundlev1/task.go: ########## @@ -31,18 +31,38 @@ import ( ) // Task is one registered task that the coordinator runtime can execute. Bundle -// authors do not implement this directly; Dag.AddTask wraps a plain Go +// authors do not implement this directly. airflow.TaskHandler wraps a plain Go // function into a Task. type Task interface { Execute(ctx context.Context, logger *slog.Logger, args []binding.Arg) error } -// Bundle is the execution-time view of a registry. It looks up a task by -// dag_id and task_id. +// Bundle looks up a registered task by dag_id and task_id. The coordinator +// runtime uses Bundle to find the task the supervisor asked for. type Bundle interface { LookupTask(dagId, taskId string) (Task, bool) } +// TaskInfo describes a registered task by its user-visible id. +type TaskInfo struct { + ID string +} + +// DagInfo describes a registered dag together with its tasks in +// registration order. +type DagInfo struct { + DagID string + Tasks []TaskInfo +} Review Comment: So we probably need to remove the `DagInfo` at this moment. ```suggestion ``` ########## go-sdk/pkg/execution/metadata.go: ########## @@ -62,39 +62,25 @@ func ParseMetadataFormat(s string) (MetadataFormat, error) { } } -// DumpAirflowMetadata writes the bundle's airflow-metadata manifest to stdout -// (YAML by default, JSON when format is MetadataFormatJSON). It runs -// RegisterDags against an in-memory recorder only — no task execution, no external -// services. airflow-go-pack execs the binary with --airflow-metadata and -// decodes this output to build the embedded manifest. -func DumpAirflowMetadata(bundle bundlev1.BundleProvider, format MetadataFormat) error { - meta, err := collectManifest(bundle) +// DumpAirflowMetadata writes the bundle's airflow-metadata manifest to w Review Comment: It seems the comment was cut. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
