Go SDK: add airflow.Dag and dag.Task, the native Dag constructors

airflow-core/adr/lang-sdk/0008-control-flow-constructs.md

@@ -75,7 +75,7 @@ pick := dag.Switch(pickPath) // task_id pickPath, from the function name
 pick.Case(handleLongRef).
     Case(handleShortRef)
 
-dag.Task(airflow.TriggerDagRun(airflow.TriggerDagRunSpec{DagId: "downstream_etl"}), airflow.TaskSpec{TaskId: "trigger_downstream"}).After(gate)
+dag.Task(airflow.TriggerDagRun(airflow.TriggerDagRunSpec{DagID: "downstream_etl"}), airflow.TaskSpec{TaskID: "trigger_downstream"}).After(gate)
 ```
 
 A decider has to see the refs it returns, so either they are package-level or it is a closure where the Dag is built.

airflow-core/adr/lang-sdk/0009-provider-operators-as-generated-dsl.md

@@ -61,11 +61,11 @@ extracted := dag.Task(extract) // native Go: runs on a Go worker
 
 staged := dag.Task(amazon.S3ToRedshiftOperator{
     SchemaName: "public", TableName: "events", S3Bucket: "raw", S3Key: "events/{{ ds }}",
-}, airflow.TaskSpec{TaskId: "stage"}).After(extracted) // DSL only: runs on a Python worker
+}, airflow.TaskSpec{TaskID: "stage"}).After(extracted) // DSL only: runs on a Python worker
 
 dag.Task(kubernetes.KubernetesPodOperator{
     Namespace: "airflow", Image: "report:latest", Name: "report",
-}, airflow.TaskSpec{TaskId: "report"}).After(staged)
+}, airflow.TaskSpec{TaskID: "report"}).After(staged)
 ```
 
 `airflow.TriggerDagRun` ([ADR-0008](0008-control-flow-constructs.md)) follows the same concept but it is the hand-written.

go-sdk/adr/0008-native-dag-interface.md

@@ -27,11 +27,11 @@ Proposed.
 
 ## Decision
 
-1. **One Dag type, constructed then registered.** `airflow.Dag(dagId, spec)` returns a `*airflow.DagRef` that is complete before `bundle.Register(dag)` takes it — the same verb that registers Mixed Lang task handlers ([ADR 7](0007-mixed-lang-task-handler-interface.md)).
+1. **One Dag type, constructed then registered.** `airflow.Dag(dagID, spec)` returns a `*airflow.DagRef` that is complete before `bundle.Register(dag)` takes it — the same verb that registers Mixed Lang task handlers ([ADR 7](0007-mixed-lang-task-handler-interface.md)).
    Naming rule: `airflow.X(...)` constructs, `*airflow.XRef` is the entity; what every Dag must have is a positional parameter for dag_id, and the rest travels in a spec struct.
 2. **Tasks register through `dag.Task(fn any, opts ...airflow.TaskOption)`**, returning a `*airflow.TaskRef`. `airflow.Inputs(...)` and a bare `airflow.TaskSpec{}` both implement `TaskOption`.
 3. **At most one `airflow.TaskSpec` per task.** A second one is a registration error rather than something to merge, so a task's attributes are only ever written in one place.
-4. **task_id is the Go function name by default**, spelled exactly as the function is, and `airflow.TaskSpec{TaskId: ...}` sets it to anything else.
+4. **task_id is the Go function name by default**, spelled exactly as the function is, and `airflow.TaskSpec{TaskID: ...}` sets it to anything else.
 5. **`airflow.Inputs(refs...)` declares the data and the edge in one call** for defining graph with TaskFlow syntax.
 6. **`Before` and `After` are order-only edges on `airflow.Node`**, which both `*airflow.TaskRef` and `*airflow.TaskGroupRef` satisfy. They are the Go pair for `>>` and `<<`, and both return their argument set as one `Node`, so `a.Before(b, c).Before(d)` is Python's `a >> [b, c] >> d`.
 7. **An edge label wraps the endpoint**: `loaded.Before(airflow.Label(notify, "when empty"))` is Python's `loaded >> Label("when empty") >> notify`. Labelling the endpoint rather than the call lets one fan-out give each edge its own label.
@@ -91,7 +91,7 @@ The task_ids are the function names by default: `extract`, `transform`, and `loa
 loaded := dag.Task(load, airflow.Inputs(transformed))
 cleaned := dag.Task(cleanup, airflow.TaskSpec{TriggerRule: airflow.TriggerRuleAllDone})
 notified := dag.Task(notify)
-emptyNotice := dag.Task(notifyEmpty, airflow.TaskSpec{TaskId: "notify_empty"})
+emptyNotice := dag.Task(notifyEmpty, airflow.TaskSpec{TaskID: "notify_empty"})
 staging := dag.TaskGroup("staging")  // a group carries edges like a task
 staging.Task(stageRows)              // tasks join a group through the group
 
@@ -107,13 +107,13 @@ loaded.Before(airflow.Label(emptyNotice, "when empty"))  // loaded >> Label("whe
 ```go
 package airflow
 
-func Dag(dagId string, spec DagSpec) *DagRef
+func Dag(dagID string, spec ...DagSpec) *DagRef
 
 func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef
-func (d *DagRef) TaskGroup(groupId string, opts ...TaskGroupOption) *TaskGroupRef
+func (d *DagRef) TaskGroup(groupID string, opts ...TaskGroupOption) *TaskGroupRef
 
 func (g *TaskGroupRef) Task(fn any, opts ...TaskOption) *TaskRef
-func (g *TaskGroupRef) TaskGroup(groupId string, opts ...TaskGroupOption) *TaskGroupRef
+func (g *TaskGroupRef) TaskGroup(groupID string, opts ...TaskGroupOption) *TaskGroupRef
 
 // DagSpec and TaskSpec are generated into this package from
 // airflow-core/src/airflow/serialization/schema.json and committed.
@@ -127,7 +127,7 @@ type DagSpec struct {
 
 // TaskSpec implements TaskOption, so it travels in the same variadic as Inputs.
 type TaskSpec struct {
-    TaskId      string // defaults to the Go function name
+    TaskID      string // defaults to the Go function name
     Retries     int
     TriggerRule TriggerRule
     // ...
@@ -177,7 +177,7 @@ const (
   Generation needs an exclusion list and a hand-written field or two, the same kind of rule [ADR-0009](../../airflow-core/adr/lang-sdk/0009-provider-operators-as-generated-dsl.md) states for provider operators.
 - **A data edge is labelled by redeclaring it.** Declaring an edge that already exists is idempotent, so `extracted.Before(airflow.Label(transformed, "rows"))` labels the edge `Inputs` created.
 - **Renaming a Go function renames the task.** The id is derived, and history, clears, and the UI all key on task_id, so renaming a function whose task carries no `TaskSpec` id is a Dag change.
-  `airflow.TaskSpec{TaskId: ...}` pins an id that has to outlive the function's name, and it is also how a Dag gets snake_case ids, since nothing transforms a Go name.
+  `airflow.TaskSpec{TaskID: ...}` pins an id that has to outlive the function's name, and it is also how a Dag gets snake_case ids, since nothing transforms a Go name.
 
 ## Alternatives
 
@@ -187,5 +187,5 @@ const (
   callable per run, so an edge existing only in execution order cannot be parsed without running the
   program to completion. `Inputs` keeps the typed outputs that style is reached for.
 - **Labelling the call**, `loaded.Before(notify).Label("when empty")`. Rejected: one call fans out, so a single label on the call cannot give `a.Before(b, c)` a different label per edge.
-- **A positional task_id**, `dag.Task(taskId, fn, opts...)`. Rejected: It's more straightforward and native for Go user to define a Task without defining the task_id explicitly. They could still set the task_id other than the function name in the TaskSpec.
+- **A positional task_id**, `dag.Task(taskID, fn, opts...)`. Rejected: It's more straightforward and native for Go user to define a Task without defining the task_id explicitly. They could still set the task_id other than the function name in the TaskSpec.
 - **Separate `Dag` and `MixedLangDag` types.** Rejected: Python has one Dag class, and the Mixed Lang case is not a Dag at all ([ADR 7](0007-mixed-lang-task-handler-interface.md)).

go-sdk/airflow/bundle.go

@@ -26,13 +26,14 @@ import (
 	"github.com/apache/airflow/go-sdk/internal/bundle"
 )
 
-// BundleRef holds the task handlers that this executable runs for Airflow.
+// BundleRef holds the task handlers and the Dags that this executable registers for Airflow.
 // [Bundle] returns an empty one.
 type BundleRef struct {
 	// closed ends registration for everything the bundle can hold, so a kind added later
 	// is covered without a flag of its own. Serve sets it; Register reads it.
 	closed       atomic.Bool
 	taskHandlers taskHandlerMap
+	dags         dagMap
 }
 
 // Bundle returns an empty bundle. Register the task handlers on it, then call Serve as the
@@ -51,7 +52,8 @@ type BundleRef struct {
 //	}
 func Bundle() *BundleRef { return &BundleRef{} }
 
-// Registerable is what [BundleRef.Register] accepts. [TaskHandler] returns one.
+// Registerable is what [BundleRef.Register] accepts: the value [TaskHandler] returns, or the
+// [*DagRef] that [Dag] returns.
 //
 // Its only method is unexported, so a type outside this package cannot declare it.
 // A struct that embeds a Registerable still satisfies the interface, and Register panics
@@ -70,8 +72,12 @@ type Registerable interface{ registerable() }
 //
 //	bundle.Register(reports.Handlers()...)
 //
+// Add every task to a Dag before registering the Dag. [DagRef.Task] panics once the Dag is
+// registered.
+//
 // Register panics if a task handler with the same dag_id and task_id is already registered,
-// and if [BundleRef.Serve] has already been called: registration closes when serving starts.
+// if a Dag with the same dag_id is already registered, or if [BundleRef.Serve] has already
+// been called: registration closes when serving starts.
 func (b *BundleRef) Register(items ...Registerable) {
 	if b.closed.Load() {
 		panic(
@@ -82,6 +88,11 @@ func (b *BundleRef) Register(items ...Registerable) {
 		switch item := item.(type) {
 		case *taskHandler:
 			b.taskHandlers.add(item.dagId, item.taskId, item.task)
+		case *DagRef:
+			if item == nil {
+				panic("airflow.BundleRef.Register: cannot register a nil *airflow.DagRef")
+			}
+			b.dags.add(item)
 		default:
 			// Either a nil item, or a struct from another package that embeds a Registerable.
 			panic(fmt.Sprintf("airflow.BundleRef.Register: cannot register %T", item))
@@ -138,3 +149,23 @@ func (m *taskHandlerMap) ListTaskHandlers() []bundle.TaskHandlerInfo {
 
 	return slices.Clone(m.order)
 }
+
+// dagMap holds the registered Dags by dag_id.
+type dagMap struct {
+	mu   sync.Mutex
+	dags map[string]*DagRef
+}
+
+func (m *dagMap) add(dag *DagRef) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if _, exists := m.dags[dag.dagID]; exists {
+		panic(fmt.Sprintf("airflow.BundleRef.Register: Dag %q is already registered", dag.dagID))
+	}
+	dag.markRegistered()
+	if m.dags == nil {
+		m.dags = make(map[string]*DagRef)
+	}
+	m.dags[dag.dagID] = dag
+}

go-sdk/airflow/bundle_test.go

@@ -117,16 +117,48 @@ func TestRegisterRejectsDuplicateTask(t *testing.T) {
 	)
 }
 
+func TestRegisterTakesTaskHandlersAndDagsTogether(t *testing.T) {
+	etl := Dag("etl")
+	etl.Task(extract)
+
+	b := Bundle()
+	b.Register(TaskHandler("py_etl", "transform", noop), etl)
+
+	_, ok := b.taskHandlers.LookupTask("py_etl", "transform")
+	assert.True(t, ok)
+	assert.Same(t, etl, b.dags.dags["etl"])
+	assert.True(t, etl.registered)
+}
+
+func TestRegisterRejectsDuplicateDag(t *testing.T) {
+	etl := Dag("etl")
+	b := Bundle()
+	b.Register(etl)
+
+	want := `airflow.BundleRef.Register: Dag "etl" is already registered`
+	assert.PanicsWithValue(t, want, func() { b.Register(etl) })
+	second := Dag("etl")
+	assert.PanicsWithValue(t, want, func() { b.Register(second) })
+	assert.Same(t, etl, b.dags.dags["etl"])
+	assert.False(t, second.registered, "a Dag that Register rejects can still take tasks")
+}
+
+func TestRegisterRejectsNilDag(t *testing.T) {
+	var dag *DagRef
+	assert.PanicsWithValue(t, "airflow.BundleRef.Register: cannot register a nil *airflow.DagRef",
+		func() { Bundle().Register(dag) },
+	)
+}
+
 func TestRegisterAfterServePanics(t *testing.T) {
 	b := Bundle()
 	b.Register(TaskHandler("py_etl", "transform", noop))
 	require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard))
 
-	assert.PanicsWithValue(
-		t,
-		"airflow.BundleRef.Register: Serve has already been called; register everything before Serve",
-		func() { b.Register(TaskHandler("py_etl", "load", noop)) },
-	)
+	want := "airflow.BundleRef.Register: Serve has already been called; " +
+		"register everything before Serve"
+	assert.PanicsWithValue(t, want, func() { b.Register(TaskHandler("py_etl", "load", noop)) })
+	assert.PanicsWithValue(t, want, func() { b.Register(Dag("etl")) })
 }
 
 // The flag lives on the bundle, not on the task-handler map, so a kind of item added to
@@ -160,6 +192,7 @@ func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
 			defer wg.Done()
 			for i := range perWorker {
 				b.Register(TaskHandler("py_etl", fmt.Sprintf("task_%d_%d", worker, i), noop))
+				b.Register(Dag(fmt.Sprintf("dag_%d_%d", worker, i)))
 				b.taskHandlers.LookupTask("py_etl", "task_0_0")
 				b.taskHandlers.ListTaskHandlers()
 			}
@@ -168,6 +201,7 @@ func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
 	wg.Wait()
 
 	assert.Len(t, b.taskHandlers.ListTaskHandlers(), workers*perWorker)
+	assert.Len(t, b.dags.dags, workers*perWorker)
 }
 
 func TestRegisterRejectsNilItem(t *testing.T) {