Go-SDK: Implement coordinator-mode runtime entry point and task runner

go-sdk/bundle/bundlev1/bundlev1server/server.go

@@ -19,6 +19,7 @@ package bundlev1server
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log/slog"
 	"os"
@@ -32,9 +33,34 @@ import (
 	"github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server/impl"
 	"github.com/apache/airflow/go-sdk/pkg/bundles/shared"
 	"github.com/apache/airflow/go-sdk/pkg/config"
+	"github.com/apache/airflow/go-sdk/pkg/execution"
 )
 
-var versionInfo *bool = flag.Bool("bundle-metadata", false, "show the embedded bundle info")
+// ErrCoordinatorFlagsIncomplete is returned by [Serve] when exactly one of
+// --comm or --logs is supplied. Both flags select coordinator mode and must
+// be set together; callers (typically main) can check for this sentinel to
+// print usage before exiting non-zero.
+var ErrCoordinatorFlagsIncomplete = errors.New(
+	"--comm and --logs must be supplied together",
+)
+
+// CLI Flags.
+// The --bundle-metadata flag is used for showing the embedded bundle info in airflow-metadata.yaml spec format.
+// The --comm and --logs select the coordinator-mode protocol
+// All three are read by Serve to choose a server mode below.
+var (
+	versionInfo = flag.Bool("bundle-metadata", false, "show the embedded bundle info")
+	commAddr    = flag.String(
+		"comm",
+		"",
+		"host:port of the supervisor's coordinator comm channel (selects coordinator mode)",
+	)
+	logsAddr = flag.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
@@ -52,41 +78,99 @@ func (s serveConfigFunc) ApplyServeOpt(in *ServerConfig) error {
 
 type ServerConfig struct{}
 
-// Serve is the entrypoint for your bundle, and sets it up ready for Airflow's Go Worker to use
+// serveMode tags the protocol the binary will speak this run.
+type serveMode int
+
+const (
+	modePlugin                serveMode = iota // go-plugin gRPC (existing Edge Worker path)
+	modeMetadataDump                           // --bundle-metadata: print BundleInfo JSON
+	modeCoordinator                            // --comm/--logs: msgpack-over-IPC (ADR 0003)
+	modeCoordinatorUsageError                  // misuse: print usage and exit non-zero
+)
+
+// Serve is the entrypoint for your bundle, and sets it up ready for Airflow's
+// Go Worker (go-plugin) or Python supervisor (coordinator protocol) to use.
 //
-// Zero or more options to configure the server may also be passed. There are no options yet, this is to allow
-// future changes without breaking compatibility
+// The mode is decided from CLI flags and process environment. 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 {
 	config.SetupViper("")
 
-	hcLogger := hclog.New(&hclog.LoggerOptions{
-		Level:                    hclog.Trace,
-		Output:                   os.Stderr,
-		JSONFormat:               true,
-		IncludeLocation:          true,
-		AdditionalLocationOffset: 3,
-	})
-
-	log := slog.New(hclogslog.Adapt(hcLogger))
-	slog.SetDefault(log)
-
 	flag.Parse()
 
 	serveConfig := &ServerConfig{}
 	for _, c := range opts {
 		c.ApplyServeOpt(serveConfig)
 	}
 
+	switch decideMode() {
+	case modeMetadataDump:
+		return dumpBundleMetadata(bundle)
+	case modeCoordinator:
+		// In coordinator mode the supervisor reads the logs channel for
+		// structured records, so configuring the hclog/stderr default
+		// logger here is unnecessary — execution.Serve installs its own
+		// slog handler against the logs socket before any user code runs.
+		return execution.Serve(bundle, *commAddr, *logsAddr)
+	case modePlugin:
+		installPluginLogger()
+		return servePlugin(bundle)
+	case modeCoordinatorUsageError:
+		return ErrCoordinatorFlagsIncomplete
+	}
+	return nil
+}
+
+func decideMode() serveMode {
 	if *versionInfo {
-		meta := bundle.GetBundleVersion()
-		data, err := json.MarshalIndent(meta, "", "    ")
-		if err != nil {
-			return err
-		}
-		fmt.Println(string(data))
-		return nil
+		return modeMetadataDump
+	}
+	commSet := *commAddr != ""
+	logsSet := *logsAddr != ""
+	if commSet && logsSet {
+		return modeCoordinator
+	}
+	if commSet || logsSet {
+		// Partial use is a hard error, both flags are required
+		return modeCoordinatorUsageError
+	}
+	return modePlugin
+}
+
+func dumpBundleMetadata(bundle bundlev1.BundleProvider) error {
+	meta := bundle.GetBundleVersion()
+	data, err := json.MarshalIndent(meta, "", "    ")
+	if err != nil {
+		return err
 	}
+	fmt.Println(string(data))
+	return nil
+}
+
+func installPluginLogger() {
+	hcLogger := hclog.New(&hclog.LoggerOptions{
+		Level:                    hclog.Trace,
+		Output:                   os.Stderr,
+		JSONFormat:               true,
+		IncludeLocation:          true,
+		AdditionalLocationOffset: 3,
+	})
+	log := slog.New(hclogslog.Adapt(hcLogger))
+	slog.SetDefault(log)
+}
 
+func servePlugin(bundle bundlev1.BundleProvider) error {
 	pluginConfig := &plugin.ServeConfig{
 		HandshakeConfig: shared.Handshake,
 		Plugins: plugin.PluginSet{

go-sdk/example/bundle/main.go

@@ -20,6 +20,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"log"
 	"log/slog"
 	"runtime"
 	"time"
@@ -54,7 +55,9 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error {
 }
 
 func main() {
-	bundlev1server.Serve(&myBundle{})
+	if err := bundlev1server.Serve(&myBundle{}); err != nil {
+		log.Fatal(err)
+	}
 }
 
 func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, error) {
@@ -63,7 +66,14 @@ func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, err
 	if err != nil {
 		log.ErrorContext(ctx, "unable to get conn", "error", err)
 	} else {
-		log.InfoContext(ctx, "got conn", "conn", conn)
+		// Log only non-sensitive fields; conn.Password and any secrets in
+		// conn.Extra must never reach the log stream.
+		log.InfoContext(ctx, "got conn",
+			"conn_id", conn.ID,
+			"conn_type", conn.Type,
+			"host", conn.Host,
+			"port", conn.Port,
+		)
 	}
 	for range 10 {