Go SDK: answer the Dag-parsing request from Serve

go-sdk/README.md

@@ -389,6 +389,10 @@ Python supervisor / task runner
   protocol on the comm socket, with structured JSON-line logs on the logs socket.
 - The Python runtime is the worker. It proxies every `GetConnection` / `GetVariable` / `GetXCom` /
   `SetXCom` call through to the Execution API. The Go binary just runs the task function.
+- When the first frame on the comm socket is a `DagFileParseRequest` from the Dag processor, the
+  binary runs no task. It answers with one `DagFileParsingResult` that holds the Dags from
+  `airflow.Dag` that the binary registered. The Dags are serialized as
+  [Serializing a native Dag](#serializing-a-native-dag) describes. The binary then exits.
 
 The Go side of the protocol is implemented in `pkg/execution/`. On the Python side it is the
 `ExecutableCoordinator` in `task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py`.

go-sdk/airflow/bundle.go

@@ -19,6 +19,8 @@ package airflow
 
 import (
 	"fmt"
+	"log/slog"
+	"runtime/debug"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -193,10 +195,11 @@ func (m *taskHandlerMap) ListTaskHandlers() []bundle.TaskHandlerInfo {
 	return slices.Clone(m.order)
 }
 
-// dagMap holds the registered Dags by dag_id.
+// dagMap holds the registered Dags by dag_id, in registration order.
 type dagMap struct {
-	mu   sync.Mutex
-	dags map[string]*DagRef
+	mu    sync.Mutex
+	dags  map[string]*DagRef
+	order []*DagRef
 }
 
 func (m *dagMap) add(dag *DagRef) {
@@ -211,6 +214,7 @@ func (m *dagMap) add(dag *DagRef) {
 		m.dags = make(map[string]*DagRef)
 	}
 	m.dags[dag.dagID] = dag
+	m.order = append(m.order, dag)
 }
 
 func (m *dagMap) has(dagID string) bool {
@@ -220,3 +224,48 @@ func (m *dagMap) has(dagID string) bool {
 	_, exists := m.dags[dagID]
 	return exists
 }
+
+// serialize serializes the Dags in registration order. If a Dag panics, the panic becomes that
+// Dag's Err.
+func (m *dagMap) serialize(fileloc, relativeFileloc string) []bundle.SerializedDag {
+	m.mu.Lock()
+	dags := slices.Clone(m.order)
+	m.mu.Unlock()
+
+	serialized := make([]bundle.SerializedDag, len(dags))
+	for i, dag := range dags {
+		serialized[i] = serializeRecovering(dag, fileloc, relativeFileloc)
+	}
+	return serialized
+}
+
+func serializeRecovering(dag *DagRef, fileloc, relativeFileloc string) (s bundle.SerializedDag) {
+	s.DagID = dag.dagID
+	defer func() {
+		if r := recover(); r != nil {
+			slog.Error(
+				"Dag serialization panicked",
+				"dag_id", dag.dagID, "panic", r, "stack", string(debug.Stack()),
+			)
+			s.Data, s.Err = nil, fmt.Errorf("%v", r)
+		}
+	}()
+	s.Data = dag.serialize(fileloc, relativeFileloc)
+	return s
+}
+
+// coordinatorSource is what Serve hands to execution.Serve: the task handlers and the Dags of one
+// bundle.
+type coordinatorSource struct {
+	*taskHandlerMap
+	dags *dagMap
+}
+
+var (
+	_ bundle.Bundle        = coordinatorSource{}
+	_ bundle.DagSerializer = coordinatorSource{}
+)
+
+func (s coordinatorSource) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag {
+	return s.dags.serialize(fileloc, relativeFileloc)
+}

go-sdk/airflow/bundle_test.go

@@ -412,3 +412,44 @@ func TestRegisterableRejectsForeignTypes(t *testing.T) {
 	assert.Contains(t, string(out), "foreignItem does not implement airflow.Registerable")
 	assert.Contains(t, string(out), "unexported method registerable")
 }
+
+func TestSerializeDagsKeepsTheOtherDagsWhenADagCannotBeSerialized(t *testing.T) {
+	b := Bundle()
+	etl := Dag("etl")
+	etl.Task(noop)
+	b.Register(etl)
+	// A Dag that skipped Register stands in for one the serializer fails on.
+	broken := Dag("broken")
+	b.dags.dags["broken"] = broken
+	b.dags.order = append(b.dags.order, broken)
+	reports := Dag("reports")
+	reports.Task(noop)
+	b.Register(reports)
+
+	serialized := b.dags.serialize("/bundles/go/etl", "etl")
+
+	require.Len(t, serialized, 3)
+	assert.Equal(t, bundle.SerializedDag{
+		DagID: "etl",
+		Data:  etl.serialize("/bundles/go/etl", "etl"),
+	}, serialized[0])
+	assert.Equal(t, "broken", serialized[1].DagID)
+	assert.Nil(t, serialized[1].Data)
+	assert.ErrorContains(t, serialized[1].Err, `Dag "broken" is not registered`)
+	assert.Equal(t, bundle.SerializedDag{
+		DagID: "reports",
+		Data:  reports.serialize("/bundles/go/etl", "etl"),
+	}, serialized[2])
+}
+
+func TestSerializeDagsLeavesOutADagThatRegisterRejected(t *testing.T) {
+	b := Bundle()
+	cyclic := Dag("cyclic")
+	extracted := orderedTask(t, cyclic, "extract")
+	loaded := orderedTask(t, cyclic, "load")
+	extracted.Before(loaded)
+	loaded.Before(extracted)
+	require.Panics(t, func() { b.Register(cyclic) })
+
+	assert.Empty(t, b.dags.serialize("/bundles/go/etl", "etl"))
+}

go-sdk/airflow/dag.go

@@ -71,8 +71,8 @@ type DagRef struct {
 // [IfRef.Then], [IfRef.Else], [SwitchRef.Case] and the methods of [TaskGroupRef] panic once the Dag
 // is registered.
 //
-// [BundleRef.Serve] does not yet serve the Dags that Dag returns. It leaves them out of the
-// --airflow-metadata manifest and cannot run their tasks.
+// [BundleRef.Serve] sends the registered Dags to the Dag processor, but does not yet list them in
+// the --airflow-metadata manifest or run their tasks.
 //
 // Dag panics if it gets more than one DagSpec, or if the DagSpec has a value that Python rejects
 // when it builds or validates a Dag:

go-sdk/airflow/serve.go

@@ -57,8 +57,9 @@ const (
 // 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 Dag and task ids of the registered task handlers.
-// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the
-// coordinator protocol.
+// Airflow starts the executable with --comm and --logs. Serve then speaks the coordinator
+// protocol and either runs one task or answers the Dag processor's parse request with the Dags
+// from [Dag].
 //
 // 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:
@@ -141,7 +142,7 @@ func (b *BundleRef) serve(args []string, stdout io.Writer) error {
 		}
 		return execution.DumpAirflowMetadata(stdout, &b.taskHandlers, format)
 	case modeCoordinator:
-		return execution.Serve(&b.taskHandlers, *commAddr, *logsAddr)
+		return execution.Serve(coordinatorSource{&b.taskHandlers, &b.dags}, *commAddr, *logsAddr)
 	case modeCoordinatorUsageError:
 		return errCoordinatorFlagsRequired
 	}

go-sdk/airflow/serve_test.go

@@ -33,6 +33,7 @@ import (
 	"gopkg.in/yaml.v3"
 
 	"github.com/apache/airflow/go-sdk/pkg/execution"
+	"github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
 )
 
 func TestDecideMode(t *testing.T) {
@@ -219,28 +220,22 @@ func TestServeHelpIsNotAnError(t *testing.T) {
 	assert.Contains(t, stdout.String(), "--airflow-metadata")
 }
 
-// A fake supervisor sends StartupDetails over the comm socket, as the Python
-// ExecutableCoordinator does after it starts the bundle with --comm and --logs.
-func TestServeRunsTaskForSupervisor(t *testing.T) {
+// serveForSupervisor runs b.serve with --comm and --logs and returns the supervisor's end of the
+// comm socket and the channel that gets what serve returns.
+func serveForSupervisor(t *testing.T, b *BundleRef) (*execution.CoordinatorComm, <-chan error) {
+	t.Helper()
 	commLn, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
-	defer commLn.Close()
+	t.Cleanup(func() { commLn.Close() })
 	logsLn, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
-	defer logsLn.Close()
+	t.Cleanup(func() { logsLn.Close() })
 	// Without a deadline, a Serve that never dials would leave Accept blocked until the test
 	// binary times out.
 	deadline := time.Now().Add(10 * time.Second)
 	require.NoError(t, commLn.(*net.TCPListener).SetDeadline(deadline))
 	require.NoError(t, logsLn.(*net.TCPListener).SetDeadline(deadline))
 
-	ran := false
-	b := Bundle()
-	b.Register(TaskHandler("py_etl", "transform", func(Context) error {
-		ran = true
-		return nil
-	}))
-
 	done := make(chan error, 1)
 	go func() {
 		done <- b.serve(
@@ -251,13 +246,26 @@ func TestServeRunsTaskForSupervisor(t *testing.T) {
 
 	commConn, err := commLn.Accept()
 	require.NoError(t, err)
-	defer commConn.Close()
+	t.Cleanup(func() { commConn.Close() })
 	logsConn, err := logsLn.Accept()
 	require.NoError(t, err)
-	defer logsConn.Close()
+	t.Cleanup(func() { logsConn.Close() })
 	require.NoError(t, commConn.SetDeadline(deadline))
 
-	supervisor := execution.NewCoordinatorComm(commConn, commConn, discardLogger())
+	return execution.NewCoordinatorComm(commConn, commConn, discardLogger()), done
+}
+
+// A fake supervisor sends StartupDetails over the comm socket, as the Python
+// ExecutableCoordinator does after it starts the bundle with --comm and --logs.
+func TestServeRunsTaskForSupervisor(t *testing.T) {
+	ran := false
+	b := Bundle()
+	b.Register(TaskHandler("py_etl", "transform", func(Context) error {
+		ran = true
+		return nil
+	}))
+
+	supervisor, done := serveForSupervisor(t, b)
 	require.NoError(t, supervisor.SendRequest(0, map[string]any{
 		"type": "StartupDetails",
 		"ti": map[string]any{
@@ -284,3 +292,67 @@ func TestServeRunsTaskForSupervisor(t *testing.T) {
 	}
 	assert.True(t, ran)
 }
+
+// The Dag processor sends a DagFileParseRequest instead of StartupDetails.
+func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) {
+	var ran []string
+	record := func(name string) func(Context) error {
+		return func(Context) error {
+			ran = append(ran, name)
+			return nil
+		}
+	}
+	etl := Dag("etl")
+	extracted := etl.Task(func(Context) (int, error) {
+		ran = append(ran, "extract")
+		return 3, nil
+	}, TaskSpec{TaskID: "extract"})
+	loaded := etl.Task(record("load"), TaskSpec{TaskID: "load"})
+	etl.If(func(_ Context, rows int) (bool, error) {
+		ran = append(ran, "has_rows")
+		return rows > 0, nil
+	}, Inputs(extracted), TaskSpec{TaskID: "has_rows"}).Then(loaded)
+	reports := Dag("reports")
+	reports.Task(record("publish"), TaskSpec{TaskID: "publish"})
+
+	b := Bundle()
+	b.Register(etl, TaskHandler("py_etl", "transform", record("transform")), reports)
+
+	supervisor, done := serveForSupervisor(t, b)
+	const requestID = 7
+	require.NoError(t, supervisor.SendRequest(requestID, map[string]any{
+		"type":        "DagFileParseRequest",
+		"file":        "/bundles/go/etl",
+		"bundle_path": "/bundles/go",
+		"bundle_name": "go",
+	}))
+
+	frame, err := supervisor.ReadMessage()
+	require.NoError(t, err)
+	assert.EqualValues(t, requestID, frame.ID)
+	var result genmodels.DagFileParsingResult
+	require.NoError(t, msgpack.Unmarshal(frame.Body, &result))
+	assert.Equal(t, "DagFileParsingResult", result.Type)
+	assert.Equal(t, "/bundles/go/etl", result.Fileloc)
+	assert.Nil(t, result.ImportErrors)
+	require.Len(t, result.SerializedDags, 2)
+	for i, dag := range []*DagRef{etl, reports} {
+		want := jsonOf(t, dag.serialize("/bundles/go/etl", "etl"))
+		assertJSON(t, want, result.SerializedDags[i].Data)
+	}
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("Serve did not return after it answered the Dag parse request")
+	}
+	assert.Empty(t, ran)
+}
+
+func jsonOf(t *testing.T, value any) string {
+	t.Helper()
+	raw, err := json.Marshal(value)
+	require.NoError(t, err)
+	return string(raw)
+}

go-sdk/internal/bundle/doc.go

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-// Package bundle 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 bundle defines what the coordinator runtime needs from a bundle: the
+// tasks it looks up and runs, the Dag and task ids it lists in the manifest, and
+// the serialized Dags it sends to the Dag processor.
 //
-// Package airflow builds both from the task handlers a bundle registers.
+// Package airflow builds the tasks and ids from the task handlers a bundle
+// registers, and the serialized Dags from its Dags.
 package bundle

go-sdk/internal/bundle/task.go

@@ -57,6 +57,20 @@ type EnumerableBundle interface {
 	ListTaskHandlers() []TaskHandlerInfo
 }
 
+// SerializedDag is one Dag from airflow.Dag, serialized for the Dag processor. Data is nil when
+// Err is set.
+type SerializedDag struct {
+	DagID string
+	Data  map[string]any
+	Err   error
+}
+
+// DagSerializer serializes the Dags from airflow.Dag that a bundle registered, in registration
+// order.
+type DagSerializer interface {
+	SerializeDags(fileloc, relativeFileloc string) []SerializedDag
+}
+
 type taskFunction struct {
 	fn       reflect.Value
 	fullName string