Go SDK: embed each native Dag's source file in the bundle

airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md

@@ -95,6 +95,9 @@ packers adopt the same shape — per-Dag source resolution, de-duplication by pa
 always-embedded entrypoint fallback — so a bundle reader treats every language the same way and the
 Code tab behaves identically regardless of which SDK produced the artifact.
 
+Go has no module evaluation, so the Go SDK records the file that calls `airflow.Dag`. A Dag built in
+a factory function maps to the factory's file, not to the file that calls the factory.
+
 ## What a reader returns
 
 | the Task is… | reader returns |
@@ -113,11 +116,11 @@ Code tab behaves identically regardless of which SDK produced the artifact.
   `DagCode` → the Code tab, deferred to [ADR-0010](0010-native-dag-processing.md)'s open question and
   future work) sees one contract.
 - Generated Dags (built in a loop or a factory during module evaluation) are resolved to their
-  generator file like any other Dag — they are fully supported, not a fallback case. Best-effort
-  resolution only bites a Dag built in a detached callback after its module finished, which then
-  shows the entrypoint — a real file in the bundle rather than a wrong one. This mirrors the source
-  view already being best effort for Python factory-function Dags
-  ([ADR-0006](0006-no-lang-sdk-source-display.md), "Why Not" #3).
+  generator file (for Go, the file that calls `airflow.Dag`) like any other Dag — they are fully
+  supported, not a fallback case. Best-effort resolution only bites a Dag built in a detached
+  callback after its module finished, which then shows the entrypoint — a real file in the bundle
+  rather than a wrong one. This mirrors the source view already being best effort for Python
+  factory-function Dags ([ADR-0006](0006-no-lang-sdk-source-display.md), "Why Not" #3).
 
 ## References
 

airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst

@@ -28,7 +28,8 @@ to a compiled Go *bundle* that is launched by
 :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` for each task instance.
 
 Because Go is a compiled language, every task must be compiled ahead of time and registered inside a single,
-self-contained native executable called a **bundle**. The bundle also embeds its Dag source and a metadata
+self-contained native executable called a **bundle**. The bundle also embeds the source file of each Dag built with ``airflow.Dag``
+plus the entrypoint file (the one with ``func main``), and a metadata
 manifest (the ``dag_id`` and ``task_id`` map) in a footer appended to the executable, so the executable *is*
 the bundle: one runnable file to ship, with no separate manifest or archive. The
 :ref:`airflow-go-pack <go-sdk/build>` tool builds and packs that bundle.
@@ -529,7 +530,7 @@ surface as Go values when read back via ``GetXCom``.
 Building and packaging
 ----------------------
 
-A plain ``go build`` produces a runnable binary, but a *deployable* bundle (binary + embedded source +
+A plain ``go build`` produces a runnable binary, but a *deployable* bundle (binary + embedded source files +
 manifest) must be produced with ``airflow-go-pack``. The packer compiles the bundle and appends the embedded
 metadata footer, so the coordinator can read its ``dag_id``\ s without executing the binary, producing a
 single runnable file. The on-disk format the packer emits (the ``AFBNDL01`` footer and the

go-sdk/README.md

@@ -38,7 +38,9 @@ Python tasks are imported and run in-process. Go is compiled, so the model is di
 
 A single binary that bundles one or more Dags' task functions is called a **bundle**. You build one with
 the SDK's packer, `airflow-go-pack`, which compiles your code and appends a metadata footer (the manifest
-of `dag_id`s and `task_id`s, plus the Dag source) to the executable. The result is a **self-contained
+of `dag_id`s and `task_id`s, plus the source files) to the executable. The source files are the
+entrypoint (the file with `func main`) and the file that declares each Dag built with `airflow.Dag`.
+A Dag declared outside your module, such as in a dependency, uses the entrypoint as its source. The result is a **self-contained
 executable bundle**: a single runnable file that *is* the bundle, with no separate manifest or archive to
 ship alongside it.
 

go-sdk/airflow/bundle.go

@@ -217,6 +217,37 @@ func (m *dagMap) add(dag *DagRef) {
 	m.order = append(m.order, dag)
 }
 
+// ListDagSourceFiles returns the file that declared each registered Dag, by dag_id.
+func (m *dagMap) ListDagSourceFiles() map[string]string {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	files := make(map[string]string, len(m.dags))
+	for id, dag := range m.dags {
+		files[id] = dag.file
+	}
+	return files
+}
+
+// lookupTask returns the Go function of a task of a registered Dag. A task from TriggerDagRun has
+// none, because a Python worker runs it.
+func (m *dagMap) lookupTask(dagID, taskID string) (bundle.Task, bool) {
+	m.mu.Lock()
+	dag, ok := m.dags[dagID]
+	m.mu.Unlock()
+	if !ok {
+		return nil, false
+	}
+
+	dag.mu.Lock()
+	defer dag.mu.Unlock()
+	ref, ok := dag.tasksByID[taskID]
+	if !ok || ref.task == nil {
+		return nil, false
+	}
+	return ref.task, true
+}
+
 func (m *dagMap) has(dagID string) bool {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -254,18 +285,32 @@ func serializeRecovering(dag *DagRef, fileloc, relativeFileloc string) (s bundle
 	return s
 }
 
-// coordinatorSource is what Serve hands to execution.Serve: the task handlers and the Dags of one
-// bundle.
+// coordinatorSource is what Serve hands to execution.Serve and execution.DumpAirflowMetadata: the
+// task handlers and the Dags of one bundle.
 type coordinatorSource struct {
 	*taskHandlerMap
 	dags *dagMap
 }
 
 var (
-	_ bundle.Bundle        = coordinatorSource{}
-	_ bundle.DagSerializer = coordinatorSource{}
+	_ bundle.Bundle          = coordinatorSource{}
+	_ bundle.DagSerializer   = coordinatorSource{}
+	_ bundle.DagSourceLister = coordinatorSource{}
 )
 
 func (s coordinatorSource) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag {
 	return s.dags.serialize(fileloc, relativeFileloc)
 }
+
+func (s coordinatorSource) ListDagSourceFiles() map[string]string {
+	return s.dags.ListDagSourceFiles()
+}
+
+// LookupTask finds a task of a Dag from airflow.Dag, then a task handler. A dag_id belongs to only
+// one of the two, as Register checks.
+func (s coordinatorSource) LookupTask(dagID, taskID string) (bundle.Task, bool) {
+	if task, ok := s.dags.lookupTask(dagID, taskID); ok {
+		return task, true
+	}
+	return s.taskHandlerMap.LookupTask(dagID, taskID)
+}

go-sdk/airflow/bundle_test.go

@@ -453,3 +453,30 @@ func TestSerializeDagsLeavesOutADagThatRegisterRejected(t *testing.T) {
 
 	assert.Empty(t, b.dags.serialize("/bundles/go/etl", "etl"))
 }
+
+func TestServeLooksUpTheTasksOfDagsAndTaskHandlers(t *testing.T) {
+	dag := Dag("native_etl")
+	dag.Task(ping, TaskSpec{TaskID: "extract"})
+	dag.Task(
+		TriggerDagRun(TriggerDagRunSpec{DagID: "downstream_etl"}),
+		TaskSpec{TaskID: "trigger"},
+	)
+	b := Bundle()
+	b.Register(dag, TaskHandler("py_etl", "load", noop))
+	source := coordinatorSource{&b.taskHandlers, &b.dags}
+
+	for _, id := range [][2]string{{"native_etl", "extract"}, {"py_etl", "load"}} {
+		task, ok := source.LookupTask(id[0], id[1])
+		assert.True(t, ok, "%s.%s", id[0], id[1])
+		assert.NotNil(t, task)
+	}
+	for _, id := range [][2]string{
+		{"native_etl", "trigger"},
+		{"native_etl", "load"},
+		{"py_etl", "extract"},
+		{"unknown", "extract"},
+	} {
+		_, ok := source.LookupTask(id[0], id[1])
+		assert.False(t, ok, "%s.%s", id[0], id[1])
+	}
+}

go-sdk/airflow/dag.go

@@ -33,6 +33,9 @@ import (
 // DagRef is a Dag authored in Go. [Dag] returns a new one.
 type DagRef struct {
 	dagID string
+	// file is the source file that called Dag, as the compiler recorded it. It is empty when the
+	// runtime cannot report a caller.
+	file string
 	// Dag and Task copy the specs they are given with copySpec, so a caller cannot change a
 	// registered Dag through a spec it still holds.
 	spec DagSpec
@@ -71,8 +74,11 @@ type DagRef struct {
 // [IfRef.Then], [IfRef.Else], [SwitchRef.Case] and the methods of [TaskGroupRef] panic once the Dag
 // is registered.
 //
-// [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.
+// The bundle embeds the source file that calls Dag, so call it from the file that declares the Dag.
+// A Dag built in a factory function belongs to the file of the factory.
+//
+// [BundleRef.Serve] sends the registered Dags to the Dag processor and runs their tasks, but does
+// not list them in the dags of the --airflow-metadata manifest.
 //
 // 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:
@@ -91,6 +97,7 @@ func Dag(dagID string, spec ...DagSpec) *DagRef {
 		))
 	}
 	d := &DagRef{dagID: dagID}
+	_, d.file, _, _ = runtime.Caller(1)
 	if len(spec) == 1 {
 		if err := checkDagSpec(spec[0]); err != nil {
 			panic(fmt.Sprintf("airflow.Dag: Dag %q: %v", dagID, err))

go-sdk/airflow/dag_factory_test.go

@@ -0,0 +1,24 @@
+// 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
+
+// newFactoryDag builds a Dag in a file other than the one that calls it, so a test can check
+// that the Dag belongs to this file.
+func newFactoryDag(dagID string) *DagRef {
+	return Dag(dagID)
+}

go-sdk/airflow/dag_source_test.go

@@ -0,0 +1,65 @@
+// 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 (
+	"bytes"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
+
+	"github.com/apache/airflow/go-sdk/internal/airflowmetadata"
+)
+
+func TestDagRecordsCallerFile(t *testing.T) {
+	assert.Equal(t, "dag_source_test.go", filepath.Base(Dag("direct").file))
+}
+
+func TestDagFromFactoryRecordsFactoryFile(t *testing.T) {
+	assert.Equal(t, "dag_factory_test.go", filepath.Base(newFactoryDag("built").file))
+}
+
+func TestServeListsDagSourceFiles(t *testing.T) {
+	b := Bundle()
+	b.Register(
+		TaskHandler("py_etl", "extract", noop),
+		Dag("direct"),
+		newFactoryDag("built"),
+	)
+
+	var stdout bytes.Buffer
+	require.NoError(t, b.serve([]string{"--airflow-metadata"}, &stdout))
+
+	var got airflowmetadata.Manifest
+	require.NoError(t, yaml.Unmarshal(stdout.Bytes(), &got))
+	require.Len(t, got.DagSourceFiles, 2)
+	assert.Equal(t, "dag_source_test.go", filepath.Base(got.DagSourceFiles["direct"]))
+	assert.Equal(t, "dag_factory_test.go", filepath.Base(got.DagSourceFiles["built"]))
+	assert.NotContains(t, got.DagSourceFiles, "py_etl")
+	assert.NotContains(t, got.Dags, "direct", "native Dags stay out of dags")
+}
+
+func TestServeOmitsDagSourceFilesWithoutNativeDags(t *testing.T) {
+	var stdout bytes.Buffer
+	require.NoError(t, etlBundle().serve([]string{"--airflow-metadata"}, &stdout))
+
+	assert.NotContains(t, stdout.String(), "dag_source_files")
+}

go-sdk/airflow/serve.go

@@ -140,7 +140,11 @@ func (b *BundleRef) serve(args []string, stdout io.Writer) error {
 		if err != nil {
 			return err
 		}
-		return execution.DumpAirflowMetadata(stdout, &b.taskHandlers, format)
+		return execution.DumpAirflowMetadata(
+			stdout,
+			coordinatorSource{&b.taskHandlers, &b.dags},
+			format,
+		)
 	case modeCoordinator:
 		return execution.Serve(coordinatorSource{&b.taskHandlers, &b.dags}, *commAddr, *logsAddr)
 	case modeCoordinatorUsageError:

go-sdk/cmd/airflow-go-pack/inspect.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
+	"gopkg.in/yaml.v3"
 
 	"github.com/apache/airflow/go-sdk/internal/bundlefooter"
 )
@@ -29,7 +30,7 @@ func newInspectCmd() *cobra.Command {
 	var showSource bool
 	cmd := &cobra.Command{
 		Use:   "inspect <bundle>",
-		Short: "Print the manifest (and optionally source) embedded in a bundle",
+		Short: "Print the manifest (and optionally source files) embedded in a bundle",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			source, manifest, err := bundlefooter.Read(args[0])
@@ -38,10 +39,16 @@ func newInspectCmd() *cobra.Command {
 			}
 			out := cmd.OutOrStdout()
 			if showSource {
-				fmt.Fprintln(out, "# --- source ---")
-				out.Write(source)
-				if len(source) > 0 && source[len(source)-1] != '\n' {
-					fmt.Fprintln(out)
+				files, err := embeddedSources(source, manifest)
+				if err != nil {
+					return err
+				}
+				for _, f := range files {
+					fmt.Fprintf(out, "# --- source: %s ---\n", f.path)
+					out.Write(f.data)
+					if len(f.data) > 0 && f.data[len(f.data)-1] != '\n' {
+						fmt.Fprintln(out)
+					}
 				}
 				fmt.Fprintln(out, "# --- manifest ---")
 			}
@@ -52,6 +59,46 @@ func newInspectCmd() *cobra.Command {
 			return nil
 		},
 	}
-	cmd.Flags().BoolVar(&showSource, "source", false, "also print the embedded source file")
+	cmd.Flags().BoolVar(&showSource, "source", false, "also print each embedded source file")
 	return cmd
 }
+
+type embeddedSource struct {
+	path string
+	data []byte
+}
+
+// embeddedSources cuts the source region into the files the manifest's sources index lists.
+func embeddedSources(region, manifest []byte) ([]embeddedSource, error) {
+	var index struct {
+		Sources []struct {
+			Path   string `yaml:"path"`
+			Offset int    `yaml:"offset"`
+			Length int    `yaml:"length"`
+		} `yaml:"sources"`
+	}
+	if err := yaml.Unmarshal(manifest, &index); err != nil {
+		return nil, fmt.Errorf("decoding manifest: %w", err)
+	}
+	if len(index.Sources) == 0 && len(region) > 0 {
+		return nil, fmt.Errorf(
+			"the manifest lists no sources for the %d-byte source region; "+
+				"repack the bundle with this airflow-go-pack",
+			len(region),
+		)
+	}
+	files := make([]embeddedSource, 0, len(index.Sources))
+	for _, src := range index.Sources {
+		if src.Offset < 0 || src.Length < 0 || src.Offset > len(region)-src.Length {
+			return nil, fmt.Errorf(
+				"manifest source %q (offset %d, length %d) does not fit the %d-byte source region",
+				src.Path, src.Offset, src.Length, len(region),
+			)
+		}
+		files = append(
+			files,
+			embeddedSource{path: src.Path, data: region[src.Offset : src.Offset+src.Length]},
+		)
+	}
+	return files, nil
+}