diff --git a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json index 17df5644bf1e9..031f73bebd1a7 100644 --- a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json +++ b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json @@ -1621,7 +1621,7 @@ "ratio": 0.5, "rows": 2 }, - "logical_date": "{{ ds }}", + "logical_date": "2026-09-30 00:00:00+00:00", "skip_when_already_exists": true, "task_id": "trigger_all", "task_type": "TriggerDagRunOperator", @@ -1637,7 +1637,7 @@ "conf": "py" }, "trigger_dag_id": "conformance_minimal", - "trigger_run_id": "triggered_{{ run_id }}", + "trigger_run_id": "triggered_by_conformance", "ui_color": "#ffefeb", "wait_for_completion": true } diff --git a/go-sdk/README.md b/go-sdk/README.md index a865e52f24bef..5ca2568faaa18 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -337,8 +337,8 @@ prek hook regenerate it. | state: `success` | MUST | ✓ | 3.3 | | | state: `failed` | MUST | ✓ | 3.3 | | | state: `up_for_retry` | MUST | ✓ | 3.3 | RetryTask | -| state: `skipped` | SHOULD | ✗ | – | runtime does not emit TaskState skipped yet | -| state: `deferred` | MAY | ✗ | – | runtime does not emit DeferTask yet | +| state: `skipped` | SHOULD | ✓ | 3.4 | airflow.TriggerDagRun with SkipWhenAlreadyExists only | +| state: `deferred` | MAY | ✓ | 3.4 | airflow.TriggerDagRun with Deferrable only; resumes in the Go runtime | | state: `up_for_reschedule` | MAY | ✗ | – | runtime does not emit RescheduleTask yet | | state: `awaiting_input` | MAY | ✗ | – | runtime does not emit AwaitInputTask yet | | state: `removed` | MAY | ✓ | 3.3 | | diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 61bcd77aa8fd5..8349b5d4a9027 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -229,8 +229,9 @@ func (m *dagMap) ListDagSourceFiles() map[string]string { 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. +// lookupTask returns the task of a registered Dag. It is the Go function of a task from DagRef.Task, +// or a bundle.TriggerTask with the options of a task from TriggerDagRun, which the runtime runs +// itself. func (m *dagMap) lookupTask(dagID, taskID string) (bundle.Task, bool) { m.mu.Lock() dag, ok := m.dags[dagID] @@ -242,7 +243,13 @@ func (m *dagMap) lookupTask(dagID, taskID string) (bundle.Task, bool) { dag.mu.Lock() defer dag.mu.Unlock() ref, ok := dag.tasksByID[taskID] - if !ok || ref.task == nil { + if !ok { + return nil, false + } + if ref.triggerDagRun != nil { + return &bundle.TriggerTask{Spec: triggerSpec(*ref.triggerDagRun)}, true + } + if ref.task == nil { return nil, false } return ref.task, true diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index 1ddc3b2b1097f..95b24e12b1c3e 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -25,6 +25,7 @@ import ( "reflect" "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -471,7 +472,6 @@ func TestServeLooksUpTheTasksOfDagsAndTaskHandlers(t *testing.T) { assert.NotNil(t, task) } for _, id := range [][2]string{ - {"native_etl", "trigger"}, {"native_etl", "load"}, {"py_etl", "extract"}, {"unknown", "extract"}, @@ -480,3 +480,60 @@ func TestServeLooksUpTheTasksOfDagsAndTaskHandlers(t *testing.T) { assert.False(t, ok, "%s.%s", id[0], id[1]) } } + +func TestServeLooksUpATriggerDagRunTaskWithACopyOfItsSpec(t *testing.T) { + poke := 30 * time.Second + logicalDate := time.Date(2026, 9, 30, 0, 0, 0, 0, time.UTC) + dag := Dag("native_etl") + dag.Task(TriggerDagRun(TriggerDagRunSpec{ + DagID: "downstream_etl", + RunID: "etl_downstream", + Note: "from etl", + Conf: map[string]any{"tables": []any{"rows"}}, + LogicalDate: logicalDate, + ResetDagRun: true, + WaitForCompletion: true, + SkipWhenAlreadyExists: true, + FailWhenDagIsPaused: true, + PokeInterval: &poke, + AllowedStates: []DagRunState{DagRunStateSuccess, DagRunStateQueued}, + FailedStates: []DagRunState{}, + Deferrable: ptr(true), + }), TaskSpec{TaskID: "trigger"}) + dag.Task(TriggerDagRun(TriggerDagRunSpec{DagID: "bare"}), TaskSpec{TaskID: "bare"}) + b := Bundle() + b.Register(dag) + source := coordinatorSource{&b.taskHandlers, &b.dags} + + task, ok := source.LookupTask("native_etl", "trigger") + require.True(t, ok) + trigger, ok := task.(*bundle.TriggerTask) + require.True(t, ok, "got %T", task) + assert.Equal(t, bundle.TriggerSpec{ + DagID: "downstream_etl", + RunID: "etl_downstream", + Note: "from etl", + Conf: map[string]any{"tables": []any{"rows"}}, + LogicalDate: logicalDate, + ResetDagRun: true, + WaitForCompletion: true, + SkipWhenAlreadyExists: true, + FailWhenDagIsPaused: true, + PokeInterval: &poke, + AllowedStates: []string{"success", "queued"}, + FailedStates: []string{}, + Deferrable: ptr(true), + }, trigger.Spec) + + // The runtime gets a copy, so that nothing it does changes the task. + trigger.Spec.Conf["tables"].([]any)[0] = "changed" + again, _ := source.LookupTask("native_etl", "trigger") + assert.Equal(t, []any{"rows"}, again.(*bundle.TriggerTask).Spec.Conf["tables"]) + + task, ok = source.LookupTask("native_etl", "bare") + require.True(t, ok) + spec := task.(*bundle.TriggerTask).Spec + assert.Nil(t, spec.AllowedStates) + assert.Nil(t, spec.FailedStates, "an unset FailedStates stays nil, which means the default") + assert.Nil(t, spec.Conf) +} diff --git a/go-sdk/airflow/conformance_test.go b/go-sdk/airflow/conformance_test.go index b8887a230e51f..149ad6aed4000 100644 --- a/go-sdk/airflow/conformance_test.go +++ b/go-sdk/airflow/conformance_test.go @@ -289,6 +289,13 @@ func conformanceTriggerDagRun( key, ) } + if moment, ok := target.(*time.Time); ok { + // The value is spelled as Python's str(datetime), such as "2026-09-30 00:00:00+00:00". + parsed, err := time.Parse("2006-01-02 15:04:05Z07:00", value.Value) + require.NoError(t, err, "%s: %s", label, key) + *moment = parsed + continue + } require.NoError(t, value.Decode(target), "%s: %s", label, key) } return spec diff --git a/go-sdk/airflow/serialize.go b/go-sdk/airflow/serialize.go index cfbb7d3d8b7ca..bd294b85b4f8e 100644 --- a/go-sdk/airflow/serialize.go +++ b/go-sdk/airflow/serialize.go @@ -210,7 +210,7 @@ var dagSpecRules = specRules{ fields: dagSpecFields, set: map[string]bool{"Tags": true}, skip: map[string]bool{ - // serializeTaskLocked writes Queue as the queue of each Go task whose TaskSpec sets no Queue. + // serializeTaskLocked writes Queue as the queue of each task whose TaskSpec sets no Queue. "Queue": true, // serializeTimetable writes Schedule as the timetable. "Schedule": true, @@ -330,6 +330,17 @@ func jsonNumber(value any) (float64, bool) { // to the microsecond. A Python datetime holds nothing finer than a microsecond. func encodeTime(t time.Time) float64 { return float64(t.UnixMicro()) / 1e6 } +// pythonStr returns t as Python's str(datetime) writes it in UTC, such as +// "2026-09-30 00:00:00+00:00". It writes the microseconds, six digits, only when they are not zero. +func pythonStr(t time.Time) string { + t = t.UTC() + layout := "2006-01-02 15:04:05" + if t.Nanosecond()/1000 != 0 { + layout += ".000000" + } + return t.Format(layout) + "+00:00" +} + // checkTime returns an error for a time that a Python datetime cannot hold, which is a time whose // year in UTC is not from 1 to 9999. Airflow would reject a serialized Dag with such a time when it // loads the Dag. field names the field that holds t. The zero Time passes, because it means that @@ -425,9 +436,10 @@ func (d *DagRef) serializeTaskLocked(task *TaskRef) map[string]any { // condition. data["_can_skip_downstream"] = true } - if spec.Queue == "" { - spec.Queue = d.spec.Queue - } + } + // A task from TriggerDagRun takes the queue of the Dag too, because the Go runtime runs it. + if spec.Queue == "" { + spec.Queue = d.spec.Queue } writeSpecFields(data, spec, taskSpecRules) if len(task.downstreams) > 0 { @@ -461,12 +473,12 @@ func serializeArgBindings(inputs []taskInput) []any { } // writeTriggerDagRun writes a task from TriggerDagRun as Python's serializer writes a -// TriggerDagRunOperator, because a Python worker runs the task. +// TriggerDagRunOperator, so that the Airflow UI and the Dag dependencies treat it as one. // // For a Python Dag, Python's serializer writes the template fields of the operator and leaves out -// its other parameters. A Python worker gets those other parameters by parsing the Python Dag file -// again. A Go Dag has no Python Dag file, so writeTriggerDagRun also writes each other parameter -// that spec sets, under the name that TriggerDagRunOperator gives the parameter. +// its other parameters. A Go Dag has no Python Dag file, so writeTriggerDagRun also writes each +// other parameter that spec sets, under the name that TriggerDagRunOperator gives the parameter. +// The Go runtime runs the task from its own copy of spec and reads none of this. func writeTriggerDagRun(data map[string]any, spec TriggerDagRunSpec) { data["task_type"] = triggerDagRunTaskType data["_task_module"] = triggerDagRunTaskModule @@ -481,8 +493,12 @@ func writeTriggerDagRun(data map[string]any, spec TriggerDagRunSpec) { } // Python leaves out a template field that is None, but the default logical_date is NOTSET, a // sentinel that lets TriggerDagRunOperator pick the logical date. Python writes the sentinel as - // the string NOTSET, and Airflow loads the string as it is. - data["logical_date"] = cmp.Or(spec.LogicalDate, "NOTSET") + // the string NOTSET, and Airflow loads the string as it is. Python writes a datetime in a + // template field as str(datetime), not as a {"__type": "datetime"} object. + data["logical_date"] = "NOTSET" + if !spec.LogicalDate.IsZero() { + data["logical_date"] = pythonStr(spec.LogicalDate) + } if spec.Conf != nil { data["conf"] = copyJSON(spec.Conf) } diff --git a/go-sdk/airflow/serialize_test.go b/go-sdk/airflow/serialize_test.go index 7f2337106e16f..174791840c84c 100644 --- a/go-sdk/airflow/serialize_test.go +++ b/go-sdk/airflow/serialize_test.go @@ -355,7 +355,7 @@ func TestSerializeWritesTimesAndDurationsToTheMicrosecond(t *testing.T) { assert.NotContains(t, got, "end_date") } -func TestSerializeGivesEachGoTaskTheQueueOfTheDag(t *testing.T) { +func TestSerializeGivesEachTaskTheQueueOfTheDag(t *testing.T) { dag := Dag("etl", DagSpec{Queue: "golang"}) dag.Task(extract) dag.Task(ping, TaskSpec{Queue: "golang_large"}) @@ -371,9 +371,13 @@ func TestSerializeGivesEachGoTaskTheQueueOfTheDag(t *testing.T) { assert.Equal(t, "golang", serializedTask(t, got, "extract")["queue"]) assert.Equal(t, "golang_large", serializedTask(t, got, "ping")["queue"]) assert.Equal(t, "golang", serializedTask(t, got, "hasRows")["queue"]) - // A task from TriggerDagRun runs on a Python worker. - assert.NotContains(t, serializedTask(t, got, "trigger"), "queue") + // The Go runtime runs a task from TriggerDagRun, so it takes the queue as a Go task does. + assert.Equal(t, "golang", serializedTask(t, got, "trigger")["queue"]) assert.Equal(t, "python", serializedTask(t, got, "trigger_on_python")["queue"]) + + noQueue := Dag("etl") + noQueue.Task(TriggerDagRun(TriggerDagRunSpec{DagID: "reports"}), TaskSpec{TaskID: "trigger"}) + assert.NotContains(t, serializedTask(t, serializedDag(t, noQueue), "trigger"), "queue") } func TestSerializeWritesTheTimetableOfTheSchedule(t *testing.T) { @@ -901,9 +905,9 @@ func TestSerializeWritesATriggerDagRunAsATriggerDagRunOperator(t *testing.T) { dag := Dag("etl", DagSpec{Queue: "golang"}) dag.Task(TriggerDagRun(TriggerDagRunSpec{ DagID: "downstream_etl", - RunID: "{{ run_id }}", + RunID: "etl_downstream", Conf: map[string]any{"rows": 2, "nested": map[string]any{"ok": true}}, - LogicalDate: "{{ ds }}", + LogicalDate: time.Date(2026, 1, 31, 8, 30, 0, 500000000, time.UTC), RunAfter: time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC), ResetDagRun: true, WaitForCompletion: true, @@ -934,8 +938,8 @@ func TestSerializeWritesATriggerDagRunAsATriggerDagRunOperator(t *testing.T) { "template_fields_renderers": {"conf": "py"}, "_operator_extra_links": {"Triggered DAG": "_link_TriggerDagRunLink"}, "trigger_dag_id": "downstream_etl", - "trigger_run_id": "{{ run_id }}", - "logical_date": "{{ ds }}", + "trigger_run_id": "etl_downstream", + "logical_date": "2026-01-31 08:30:00.500000+00:00", "conf": {"rows": 2, "nested": {"ok": true}}, "wait_for_completion": true, "skip_when_already_exists": true, @@ -947,10 +951,48 @@ func TestSerializeWritesATriggerDagRunAsATriggerDagRunOperator(t *testing.T) { "fail_when_dag_is_paused": true, "note": "from etl", "deferrable": false, + "queue": "golang", "retries": 1 }`, got) } +func TestSerializeWritesTheLogicalDateOfATriggerDagRunAsPythonStr(t *testing.T) { + tests := []struct { + name string + date time.Time + want string + }{ + { + "whole seconds", + time.Date(2026, 9, 30, 0, 0, 0, 0, time.UTC), + "2026-09-30 00:00:00+00:00", + }, + { + "microseconds are not trimmed", + time.Date(2026, 9, 30, 1, 2, 3, 120000000, time.UTC), + "2026-09-30 01:02:03.120000+00:00", + }, + { + "another zone is written in UTC", + time.Date(2026, 9, 30, 9, 0, 0, 0, time.FixedZone("UTC+9", 9*3600)), + "2026-09-30 00:00:00+00:00", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + dag := Dag("etl") + dag.Task( + TriggerDagRun(TriggerDagRunSpec{DagID: "reports", LogicalDate: tc.date}), + TaskSpec{TaskID: "trigger"}, + ) + + got := serializedTask(t, serializedDag(t, dag), "trigger") + + assert.Equal(t, tc.want, got["logical_date"]) + }) + } +} + func TestSerializeWritesOnlyTheTriggerDagRunOptionsThatAreSet(t *testing.T) { dag := Dag("etl") dag.Task(TriggerDagRun(TriggerDagRunSpec{DagID: "reports"}), TaskSpec{TaskID: "trigger"}) diff --git a/go-sdk/airflow/spec.gen.go b/go-sdk/airflow/spec.gen.go index 001b7ac7995d8..f13b1fa32cfff 100644 --- a/go-sdk/airflow/spec.gen.go +++ b/go-sdk/airflow/spec.gen.go @@ -64,8 +64,7 @@ type DagSpec struct { // Queue is the queue that each task of the Dag runs on, unless the TaskSpec of // the task sets a Queue. The queue_to_coordinator option in the [sdk] section of // the Airflow configuration maps the queue to the coordinator that runs Go code. - // A task from TriggerDagRun runs on a Python worker, so it does not take this - // queue. + // A task from TriggerDagRun takes the queue too, because the Go runtime runs it. Queue string // RenderTemplateAsNativeObj corresponds to the JSON schema field diff --git a/go-sdk/airflow/trigger_dag_run.go b/go-sdk/airflow/trigger_dag_run.go index b5c3c098c9e9e..ddc7da062185b 100644 --- a/go-sdk/airflow/trigger_dag_run.go +++ b/go-sdk/airflow/trigger_dag_run.go @@ -27,6 +27,8 @@ import ( "strconv" "strings" "time" + + "github.com/apache/airflow/go-sdk/internal/bundle" ) // TriggerDagRunSpec holds the options of a task that triggers a Dag run. [TriggerDagRun] takes a @@ -36,21 +38,24 @@ import ( // leaves its parameter at the Python default. PokeInterval and Deferrable are pointers so that a // pointer to 0 or false can set that value instead of leaving the default. // -// DagID, RunID, LogicalDate and the values in Conf are templated. They can hold Jinja such as -// "{{ ds }}", which Airflow renders when the task runs. +// The task does not render templates. Note and the values in Conf are sent as they are, so a +// string such as "{{ ds }}" reaches the new Dag run unchanged. DagID and RunID must not contain +// "{{", since Airflow rejects such a dag_id or run_id. The task does not offer +// openlineage_inject_parent_info, so it adds nothing to Conf. type TriggerDagRunSpec struct { // DagID is the dag_id of the Dag to trigger. It is required. DagID string - // RunID is the run_id of the new Dag run. When RunID is empty, Airflow generates one. + // RunID is the run_id of the new Dag run. When RunID is empty, the task names the run after + // the time it can start, such as "manual__2026-09-30T00:00:00+00:00", and adds an underscore + // and eight random characters when the run has no logical date. RunID string // Conf is the conf of the new Dag run. Each value must marshal to JSON, and each integer in // Conf must fit in 64 bits. Conf map[string]any - // LogicalDate is the logical date of the new Dag run, as an ISO 8601 string such as - // "2026-09-30T00:00:00+00:00" or a template such as "{{ ds }}". When LogicalDate is empty and - // RunAfter is the zero Time, the logical date is the time the task runs. When LogicalDate is - // empty and RunAfter is set, the new Dag run has no logical date. - LogicalDate string + // LogicalDate is the logical date of the new Dag run. When LogicalDate is the zero Time and + // RunAfter is also the zero Time, the logical date is the time the task runs. When LogicalDate + // is the zero Time and RunAfter is set, the new Dag run has no logical date. + LogicalDate time.Time // RunAfter is the earliest time at which the new Dag run can start. When RunAfter is the zero // Time, the new Dag run can start as soon as the task triggers it. RunAfter time.Time @@ -77,9 +82,11 @@ type TriggerDagRunSpec struct { FailWhenDagIsPaused bool // Note is the note of the new Dag run. Note string - // Deferrable makes a task that waits defer instead of holding a worker slot. When Deferrable - // is nil, the task follows the default_deferrable option in the operators section of the - // Airflow configuration. + // Deferrable makes a task that waits defer instead of holding a worker slot. The wait then + // runs in the Airflow triggerer, which needs the standard provider installed, and the task + // resumes in the Go runtime when the new Dag run finishes. Deferrable has no effect on a task + // that does not wait. When Deferrable is nil, the task follows the default_deferrable option + // in the operators section of the Airflow configuration. Deferrable *bool } @@ -97,10 +104,12 @@ type TriggerDagRunTask struct { // airflow.TaskSpec{TaskID: "trigger_downstream"}, // ) // -// The task runs no Go code. Once [BundleRef.Serve] serves the Dags from [Dag], Airflow will run -// the task as TriggerDagRunOperator on a Python worker. So the task does not take the Queue of the -// [DagSpec], which routes the tasks that run Go code. A [TaskSpec] can set a Queue for the task, -// to pick the Python workers that run it. +// The task runs no Go code. Once [BundleRef.Serve] serves the Dags from [Dag], the Go runtime runs +// the task as TriggerDagRunOperator does, so a Dag from [Dag] needs no Python worker for it. The +// task takes the Queue of the [DagSpec], unless its [TaskSpec] sets one. The runtime reads the +// settings it needs from its environment: [api] base_url for the link to the new Dag run, and +// [operators] default_deferrable for a spec that leaves Deferrable nil. Set Deferrable to defer the +// wait, which hands it to the Python triggerer as the DagStateTrigger of the standard provider. // // Because the task has no Go function to take a task_id from, DagRef.Task needs a [TaskSpec] // that sets TaskID. For the same reason, the task takes no [Inputs], and it returns no result @@ -117,6 +126,16 @@ func copyTriggerDagRunSpec(spec TriggerDagRunSpec) (TriggerDagRunSpec, error) { if spec.DagID == "" { return TriggerDagRunSpec{}, errors.New("airflow.TriggerDagRunSpec has no DagID") } + for _, field := range []struct{ name, value string }{ + {"DagID", spec.DagID}, {"RunID", spec.RunID}, + } { + if strings.Contains(field.value, "{{") { + return TriggerDagRunSpec{}, fmt.Errorf( + "airflow.TriggerDagRunSpec.%s is %q; the task does not render templates", + field.name, field.value, + ) + } + } if poke := spec.PokeInterval; poke != nil && (*poke < 0 || *poke%time.Second != 0) { return TriggerDagRunSpec{}, fmt.Errorf( "airflow.TriggerDagRunSpec.PokeInterval is %v; "+ @@ -141,6 +160,9 @@ func copyTriggerDagRunSpec(spec TriggerDagRunSpec) (TriggerDagRunSpec, error) { if err := checkTime("airflow.TriggerDagRunSpec.RunAfter", spec.RunAfter); err != nil { return TriggerDagRunSpec{}, err } + if err := checkTime("airflow.TriggerDagRunSpec.LogicalDate", spec.LogicalDate); err != nil { + return TriggerDagRunSpec{}, err + } conf, err := copyConf(spec.Conf) if err != nil { return TriggerDagRunSpec{}, fmt.Errorf("airflow.TriggerDagRunSpec.Conf: %w", err) @@ -153,6 +175,36 @@ func copyTriggerDagRunSpec(spec TriggerDagRunSpec) (TriggerDagRunSpec, error) { return copied, nil } +// triggerSpec returns the options of spec for the runtime, which cannot import this package. The +// checked spec is a deep copy, but the runtime gets one of its own, so that nothing the runtime +// does can change the task. +func triggerSpec(spec TriggerDagRunSpec) bundle.TriggerSpec { + copied := copySpec(spec) + out := bundle.TriggerSpec{ + DagID: copied.DagID, + RunID: copied.RunID, + Note: copied.Note, + LogicalDate: copied.LogicalDate, + RunAfter: copied.RunAfter, + ResetDagRun: copied.ResetDagRun, + WaitForCompletion: copied.WaitForCompletion, + SkipWhenAlreadyExists: copied.SkipWhenAlreadyExists, + FailWhenDagIsPaused: copied.FailWhenDagIsPaused, + PokeInterval: copied.PokeInterval, + Deferrable: copied.Deferrable, + } + if copied.Conf != nil { + out.Conf = copyJSON(copied.Conf).(map[string]any) + } + if len(copied.AllowedStates) > 0 { + out.AllowedStates = dagRunStateNames(copied.AllowedStates) + } + if copied.FailedStates != nil { + out.FailedStates = dagRunStateNames(copied.FailedStates) + } + return out +} + // copyConf copies conf by way of JSON, so it also rejects a conf that JSON cannot hold. // // Each number in the copy is an int64, a uint64 or a float64, which encoding/json and msgpack both diff --git a/go-sdk/airflow/trigger_dag_run_test.go b/go-sdk/airflow/trigger_dag_run_test.go index ab2f8fea5a81a..fdfe813c64522 100644 --- a/go-sdk/airflow/trigger_dag_run_test.go +++ b/go-sdk/airflow/trigger_dag_run_test.go @@ -33,9 +33,9 @@ func TestTriggerDagRunIsATask(t *testing.T) { gate := dag.Task(extract) spec := TriggerDagRunSpec{ DagID: "downstream_etl", - RunID: "{{ run_id }}_downstream", + RunID: "etl_downstream", Conf: map[string]any{"source": "etl"}, - LogicalDate: "{{ ds }}", + LogicalDate: time.Date(2026, 9, 29, 0, 0, 0, 0, time.UTC), RunAfter: time.Date(2026, 9, 30, 0, 0, 0, 0, time.UTC), ResetDagRun: true, WaitForCompletion: true, @@ -111,6 +111,20 @@ func TestTriggerDagRunRejectsAnInvalidSpec(t *testing.T) { trigger: TriggerDagRun(TriggerDagRunSpec{WaitForCompletion: true}), want: "airflow.TriggerDagRunSpec has no DagID", }, + { + name: "templated DagID", + trigger: TriggerDagRun(TriggerDagRunSpec{DagID: "{{ params.target }}"}), + want: `airflow.TriggerDagRunSpec.DagID is "{{ params.target }}"; ` + + "the task does not render templates", + }, + { + name: "templated RunID", + trigger: TriggerDagRun( + TriggerDagRunSpec{DagID: "downstream_etl", RunID: "triggered_{{ run_id }}"}, + ), + want: `airflow.TriggerDagRunSpec.RunID is "triggered_{{ run_id }}"; ` + + "the task does not render templates", + }, { name: "negative PokeInterval", trigger: TriggerDagRun( @@ -158,6 +172,14 @@ func TestTriggerDagRunRejectsAnInvalidSpec(t *testing.T) { want: "airflow.TriggerDagRunSpec.RunAfter is 10000-01-01T00:00:00Z in UTC; " + "Airflow takes a time only from year 1 to year 9999", }, + { + name: "LogicalDate before year 1 in UTC", + trigger: TriggerDagRun(TriggerDagRunSpec{ + DagID: "downstream_etl", LogicalDate: time.Date(0, 12, 31, 0, 0, 0, 0, time.UTC), + }), + want: "airflow.TriggerDagRunSpec.LogicalDate is 0000-12-31T00:00:00Z in UTC; " + + "Airflow takes a time only from year 1 to year 9999", + }, { name: "Conf with an integer that does not fit in 64 bits", trigger: TriggerDagRun(TriggerDagRunSpec{ diff --git a/go-sdk/capabilities.yaml b/go-sdk/capabilities.yaml index cabf5d4587151..4527ffd1da1d8 100644 --- a/go-sdk/capabilities.yaml +++ b/go-sdk/capabilities.yaml @@ -23,8 +23,9 @@ min_airflow_version: "3.4" # bundle metadata. The render hook fails if the two disagree. supervisor_schema_version: "2026-10-30" -# The task runner terminates a task with SucceedTask, RetryTask, or TaskState (failed/removed); it -# does not yet emit skipped, DeferTask, RescheduleTask, or AwaitInputTask. +# The task runner terminates a task with SucceedTask, RetryTask, or TaskState (failed/removed); only +# an airflow.TriggerDagRun task emits TaskState skipped or DeferTask. It does not emit RescheduleTask +# or AwaitInputTask. states: success: supported: true @@ -37,11 +38,13 @@ states: since: "3.3" note: "RetryTask" skipped: - supported: false - note: "runtime does not emit TaskState skipped yet" + supported: true + since: "3.4" + note: "airflow.TriggerDagRun with SkipWhenAlreadyExists only" deferred: - supported: false - note: "runtime does not emit DeferTask yet" + supported: true + since: "3.4" + note: "airflow.TriggerDagRun with Deferrable only; resumes in the Go runtime" up_for_reschedule: supported: false note: "runtime does not emit RescheduleTask yet" diff --git a/go-sdk/internal/bundle/trigger.go b/go-sdk/internal/bundle/trigger.go new file mode 100644 index 0000000000000..2b45418f5310a --- /dev/null +++ b/go-sdk/internal/bundle/trigger.go @@ -0,0 +1,64 @@ +// 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 bundle + +import ( + "context" + "errors" + "log/slog" + "time" + + "github.com/apache/airflow/go-sdk/pkg/binding" +) + +// TriggerSpec is what the runtime needs to run a task from airflow.TriggerDagRun. It holds the +// options of the task as the author wrote them, so the runtime applies the defaults. The states +// are the names of Dag run states. +type TriggerSpec struct { + DagID string + RunID string + Note string + Conf map[string]any + LogicalDate time.Time + RunAfter time.Time + ResetDagRun bool + WaitForCompletion bool + SkipWhenAlreadyExists bool + FailWhenDagIsPaused bool + // PokeInterval is nil when the author left it unset. + PokeInterval *time.Duration + // AllowedStates is empty when the author left it unset. + AllowedStates []string + // FailedStates is nil when the author left it unset, and empty when no state fails the task. + FailedStates []string + // Deferrable is nil when the author left it unset. + Deferrable *bool +} + +// TriggerTask is the Task of a task from airflow.TriggerDagRun. It runs no Go code, so the runtime +// recognizes it and runs the task itself. +type TriggerTask struct { + Spec TriggerSpec +} + +var _ Task = (*TriggerTask)(nil) + +// Execute always fails, because the runtime runs a TriggerTask without calling it. +func (*TriggerTask) Execute(context.Context, *slog.Logger, []binding.Arg) error { + return errors.New("a TriggerDagRun task is run by the runtime, so Execute is never called") +} diff --git a/go-sdk/internal/genspec/authoring.go b/go-sdk/internal/genspec/authoring.go index fd6c407cd8fe5..5481af5915f46 100644 --- a/go-sdk/internal/genspec/authoring.go +++ b/go-sdk/internal/genspec/authoring.go @@ -123,7 +123,7 @@ var dagShape = authoringShape{ // them there. "queue": { "type": "string", - "description": "Queue is the queue that each task of the Dag runs on, unless the TaskSpec of the task sets a Queue. The queue_to_coordinator option in the [sdk] section of the Airflow configuration maps the queue to the coordinator that runs Go code. A task from TriggerDagRun runs on a Python worker, so it does not take this queue.", + "description": "Queue is the queue that each task of the Dag runs on, unless the TaskSpec of the task sets a Queue. The queue_to_coordinator option in the [sdk] section of the Airflow configuration maps the queue to the coordinator that runs Go code. A task from TriggerDagRun takes the queue too, because the Go runtime runs it.", }, }, } diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go index 610a5c53a5ca4..8fda6b0864bea 100644 --- a/go-sdk/pkg/execution/client.go +++ b/go-sdk/pkg/execution/client.go @@ -43,6 +43,7 @@ const ( errCodeConnectionNotFound = "CONNECTION_NOT_FOUND" errCodeXComNotFound = "XCOM_NOT_FOUND" errCodeTaskStoreNotFound = "TASK_STORE_NOT_FOUND" + errCodeDagRunExists = string(genmodels.ErrorTypeDAGRUNALREADYEXISTS) ) // A language SDK runtime cannot read Airflow config, so the supervisor passes @@ -361,6 +362,51 @@ func (c *CoordinatorClient) skipDownstreamTasks(ctx context.Context, taskIDs []s return err } +// triggerDagRun asks the supervisor to create a Dag run. It reports alreadyExists, not an error, +// when a Dag run with the same run_id exists, because a trigger task decides what to do then. +func (c *CoordinatorClient) triggerDagRun( + ctx context.Context, + msg genmodels.TriggerDagRun, +) (alreadyExists bool, err error) { + if _, err := c.comm.Communicate(ctx, msg); err != nil { + var apiErr *APIError + if errors.As(err, &apiErr) && apiErr.Err == errCodeDagRunExists { + return true, nil + } + return false, err + } + return false, nil +} + +// getDagRunState asks the supervisor for the state of a Dag run. +func (c *CoordinatorClient) getDagRunState( + ctx context.Context, + dagID, runID string, +) (string, error) { + resp, err := c.comm.Communicate(ctx, genmodels.GetDagRunState{DagID: dagID, RunID: runID}) + if err != nil { + return "", err + } + var result genmodels.DagRunStateResult + if err := decodeBody(resp, &result); err != nil { + return "", fmt.Errorf("decoding Dag run state result: %w", err) + } + return string(result.State), nil +} + +// isDagPaused asks the supervisor whether a Dag is paused. +func (c *CoordinatorClient) isDagPaused(ctx context.Context, dagID string) (bool, error) { + resp, err := c.comm.Communicate(ctx, genmodels.GetDag{DagID: dagID}) + if err != nil { + return false, err + } + var result genmodels.DagResult + if err := decodeBody(resp, &result); err != nil { + return false, fmt.Errorf("decoding Dag result: %w", err) + } + return result.IsPaused, nil +} + // TaskStateStore returns the task state store scoped to this task instance. func (c *CoordinatorClient) TaskStateStore() sdk.TaskStateStore { return taskStateStore{client: c} diff --git a/go-sdk/pkg/execution/client_test.go b/go-sdk/pkg/execution/client_test.go index bd8dfd3f6e7f0..ec61b474ec8ec 100644 --- a/go-sdk/pkg/execution/client_test.go +++ b/go-sdk/pkg/execution/client_test.go @@ -25,12 +25,14 @@ import ( "log/slog" "math" "os" + "strconv" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" "github.com/apache/airflow/go-sdk/sdk" ) @@ -1132,3 +1134,137 @@ func TestCoordinatorClientDeleteXCom(t *testing.T) { }) } } + +// clientWithReply returns a client whose supervisor answers its one request with body and errBody, +// and the buffer that holds the request it sends. +func clientWithReply( + t *testing.T, + body, errBody map[string]any, +) (*CoordinatorClient, *bytes.Buffer) { + t.Helper() + var responseBuf bytes.Buffer + require.NoError(t, writeFrame(&responseBuf, encodeResponseFrame(t, 0, body, errBody))) + var requestBuf bytes.Buffer + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(&responseBuf, &requestBuf, logger) + return NewCoordinatorClient(comm, testTIID), &requestBuf +} + +// TestCoordinatorClientTriggerDagRun verifies the TriggerDagRun frame, that a Dag run which already +// exists is reported instead of returned as an error, and that any other supervisor error is +// returned. +func TestCoordinatorClientTriggerDagRun(t *testing.T) { + logicalDate := time.Date(2026, 9, 30, 1, 2, 3, 0, time.UTC) + tests := []struct { + name string + errBody map[string]any + wantAlreadyExists bool + wantErrCode string + }{ + {name: "triggered"}, + { + name: "already exists", + errBody: map[string]any{ + "type": "ErrorResponse", + "error": "DAGRUN_ALREADY_EXISTS", + "detail": map[string]any{"message": "exists"}, + }, + wantAlreadyExists: true, + }, + { + name: "other error", + errBody: map[string]any{ + "type": "ErrorResponse", + "error": "API_SERVER_ERROR", + "detail": map[string]any{"status_code": 500}, + }, + wantErrCode: "API_SERVER_ERROR", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + client, requestBuf := clientWithReply(t, nil, tc.errBody) + + alreadyExists, err := client.triggerDagRun( + context.Background(), + genmodels.TriggerDagRun{ + DagID: "downstream", + RunID: "run_1", + LogicalDate: logicalDate, + Conf: &genmodels.Conf{"k": "v"}, + }, + ) + assert.Equal(t, tc.wantAlreadyExists, alreadyExists) + if tc.wantErrCode != "" { + var apiErr *APIError + require.ErrorAs(t, err, &apiErr) + assert.Equal(t, tc.wantErrCode, apiErr.Err) + } else { + require.NoError(t, err) + } + + sent, err := readFrame(requestBuf) + require.NoError(t, err) + got := rawToMap(t, sent.Body) + assert.Equal(t, "TriggerDagRun", got["type"]) + assert.Equal(t, "downstream", got["dag_id"]) + assert.Equal(t, "run_1", got["run_id"]) + assert.Equal(t, map[string]any{"k": "v"}, got["conf"]) + assert.NotContains(t, got, "run_after", "an unset run_after must stay off the wire") + }) + } +} + +func TestCoordinatorClientGetDagRunState(t *testing.T) { + client, requestBuf := clientWithReply( + t, map[string]any{"type": "DagRunStateResult", "state": "running"}, nil, + ) + + state, err := client.getDagRunState(context.Background(), "downstream", "run_1") + require.NoError(t, err) + assert.Equal(t, "running", state) + + sent, err := readFrame(requestBuf) + require.NoError(t, err) + assert.Equal(t, map[string]any{ + "type": "GetDagRunState", + "dag_id": "downstream", + "run_id": "run_1", + }, rawToMap(t, sent.Body)) +} + +func TestCoordinatorClientGetDagRunStateError(t *testing.T) { + client, _ := clientWithReply(t, nil, map[string]any{ + "type": "ErrorResponse", + "error": "DAGRUN_NOT_FOUND", + "detail": map[string]any{"status_code": 404}, + }) + + _, err := client.getDagRunState(context.Background(), "downstream", "run_1") + var apiErr *APIError + require.ErrorAs(t, err, &apiErr) + assert.Equal(t, "DAGRUN_NOT_FOUND", apiErr.Err) +} + +func TestCoordinatorClientIsDagPaused(t *testing.T) { + for _, paused := range []bool{true, false} { + t.Run(strconv.FormatBool(paused), func(t *testing.T) { + client, requestBuf := clientWithReply(t, map[string]any{ + "type": "DagResult", + "dag_id": "downstream", + "is_paused": paused, + }, nil) + + got, err := client.isDagPaused(context.Background(), "downstream") + require.NoError(t, err) + assert.Equal(t, paused, got) + + sent, err := readFrame(requestBuf) + require.NoError(t, err) + assert.Equal(t, map[string]any{ + "type": "GetDag", + "dag_id": "downstream", + }, rawToMap(t, sent.Body)) + }) + } +} diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index 45bd228ff5ad3..16791640890d4 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -36,7 +36,8 @@ import ( // It looks up the task in the bundle, creates a CoordinatorClient for SDK // calls, deletes the XComs that ti_context.xcom_keys_to_clear lists, executes // the task, and returns the terminal body to ship as the final response frame: -// one of genmodels.SucceedTask, TaskState, or RetryTask. When RunTask cannot +// one of genmodels.SucceedTask, TaskState, RetryTask, or DeferTask. A task from airflow.TriggerDagRun +// is the only one that defers, and RunTask runs it without a Go function. When RunTask cannot // delete an XCom, or arg_bindings is invalid, the task fails without running. // // The supervisor owns the Execution-API state transitions, so the runtime only @@ -71,13 +72,7 @@ func RunTask( // context is a placeholder, because binding reads only the task instance // and Dag run from this value and builds the airflow.Context around the // live task context. - ti := sdk.TaskInstance{ - DagID: details.TI.DagID, - RunID: details.TI.RunID, - TaskID: details.TI.TaskID, - MapIndex: mapIndexPtr(details.TI.MapIndex), - TryNumber: details.TI.TryNumber, - } + ti := taskInstanceOf(details) dagRun := details.TIContext.DagRun runtimeContext := sdk.NewTIRunContext( context.Background(), @@ -112,6 +107,11 @@ func RunTask( } } + // The runtime runs a TriggerDagRun task itself. It has no Go function to bind arguments to. + if trigger, ok := task.(*bundle.TriggerTask); ok { + return runTriggerDagRun(ctx, details, trigger.Spec, client, logger, defaultTriggerEnv) + } + args, err := convertArgBindings(details.TIContext.ArgBindings) if err != nil { logger.Error("Invalid arg_bindings spec from supervisor", @@ -125,6 +125,16 @@ func RunTask( return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) } +func taskInstanceOf(details *genmodels.StartupDetails) sdk.TaskInstance { + return sdk.TaskInstance{ + DagID: details.TI.DagID, + RunID: details.TI.RunID, + TaskID: details.TI.TaskID, + MapIndex: mapIndexPtr(details.TI.MapIndex), + TryNumber: details.TI.TryNumber, + } +} + // failTask returns the terminal body for a task that fails before it runs: RetryTask when // ti_context.should_retry is set, otherwise a FAILED TaskState. func failTask(shouldRetry bool, err error) any { @@ -282,10 +292,14 @@ func executeTask( } } - // task_outlets / outlet_events must be sent as empty lists, not omitted: - // the supervisor's SucceedTask validation rejects a null/absent value - // ("Input should be a valid list"). A task that emits no asset events - // reports empty collections rather than None. + return succeedTask() +} + +// succeedTask returns the terminal body of a task that succeeded. task_outlets / outlet_events must +// be sent as empty lists, not omitted: the supervisor's SucceedTask validation rejects a +// null/absent value ("Input should be a valid list"). A task that emits no asset events reports +// empty collections rather than None. +func succeedTask() genmodels.SucceedTask { return genmodels.SucceedTask{ EndDate: time.Now().UTC(), TaskOutlets: &genmodels.TaskOutlets{}, diff --git a/go-sdk/pkg/execution/trigger_runner.go b/go-sdk/pkg/execution/trigger_runner.go new file mode 100644 index 0000000000000..8c85ace9a27b0 --- /dev/null +++ b/go-sdk/pkg/execution/trigger_runner.go @@ -0,0 +1,418 @@ +// 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 execution + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math/big" + "os" + "runtime/debug" + "slices" + "strconv" + "strings" + "time" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// This file runs a task from airflow.TriggerDagRun the way TriggerDagRunOperator.execute and +// execute_complete in the standard provider do, with the difference that nothing is Jinja-rendered. + +const ( + dagStateTriggerClasspath = "airflow.providers.standard.triggers.external_task.DagStateTrigger" + // triggerDagRunLinkXComKey is the XCom that the "Triggered DAG" extra link reads. + triggerDagRunLinkXComKey = "_link_TriggerDagRunLink" + triggerRunIDXComKey = "trigger_run_id" + // triggerFailMethod is the next_method of a task whose trigger failed or timed out. + triggerFailMethod = "__fail__" + executeCompleteMethod = "execute_complete" + + defaultTriggerPokeInterval = 60 * time.Second +) + +// The runtime cannot read airflow.cfg, so the coordinator exports these settings. +const ( + apiBaseURLEnv = "AIRFLOW__API__BASE_URL" + defaultDeferrableEnv = "AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE" + triggererQueuesEnabledEnv = "AIRFLOW__TRIGGERER__QUEUES_ENABLED" +) + +const runIDSuffixChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + +// triggerEnv holds what a trigger task takes from its surroundings, so a test can replace it. +type triggerEnv struct { + now func() time.Time + getenv func(string) string + sleep func(ctx context.Context, d time.Duration) error + randomSuffix func() string +} + +var defaultTriggerEnv = triggerEnv{ + now: time.Now, + getenv: os.Getenv, + sleep: func(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } + }, + randomSuffix: func() string { + out := make([]byte, 8) + for i := range out { + n, err := rand.Int(rand.Reader, big.NewInt(int64(len(runIDSuffixChars)))) + if err != nil { + panic(err) + } + out[i] = runIDSuffixChars[n.Int64()] + } + return string(out) + }, +} + +// runTriggerDagRun runs a trigger task and returns the terminal body for the supervisor. It +// returns a genmodels.DeferTask for a task that waits and is deferrable. +func runTriggerDagRun( + ctx context.Context, + details *genmodels.StartupDetails, + spec bundle.TriggerSpec, + client *CoordinatorClient, + logger *slog.Logger, + env triggerEnv, +) (result any) { + shouldRetry := details.TIContext.ShouldRetry + fail := func(message string) any { + logger.ErrorContext(ctx, "Task failed", "error", message) + return failTask(shouldRetry, errors.New(message)) + } + defer func() { + if r := recover(); r != nil { + logger.ErrorContext(ctx, "Recovered panic in task", + "error", r, "stack", string(debug.Stack()), + ) + result = failTask(shouldRetry, fmt.Errorf("panic: %v", r)) + } + }() + allowed := spec.AllowedStates + if len(allowed) == 0 { + allowed = []string{string(genmodels.DagRunStateSuccess)} + } + failed := spec.FailedStates + if failed == nil { + failed = []string{string(genmodels.DagRunStateFailed)} + } + + if nextMethod := ifaceString(details.TIContext.NextMethod); nextMethod != "" { + return resumeTriggerDagRun( + nextMethod, details.TIContext.NextKwargs, spec.DagID, allowed, failed, logger, fail, + ) + } + + logicalDate := triggerLogicalDate(spec, env.now().UTC()) + runID := spec.RunID + if runID == "" { + // Python names the run after run_after, which is the logical date when the task sets none. + runAfter := logicalDate + if !spec.RunAfter.IsZero() { + runAfter = spec.RunAfter + } + runID = "manual__" + pythonIsoformat(runAfter) + if logicalDate.IsZero() { + runID += "_" + env.randomSuffix() + } + } + + if spec.FailWhenDagIsPaused { + paused, err := client.isDagPaused(ctx, spec.DagID) + if err != nil { + return fail(err.Error()) + } + if paused { + return fail(fmt.Sprintf("Dag %s is paused", spec.DagID)) + } + } + + ti := taskInstanceOf(details) + link := dagRunURL(env.getenv(apiBaseURLEnv), spec.DagID, runID) + if err := client.PushXCom(ctx, ti, triggerDagRunLinkXComKey, link); err != nil { + return fail(err.Error()) + } + + logger.InfoContext(ctx, "Triggering Dag Run.", "trigger_dag_id", spec.DagID) + msg := genmodels.TriggerDagRun{ + DagID: spec.DagID, + RunID: runID, + ResetDagRun: spec.ResetDagRun, + } + if !logicalDate.IsZero() { + msg.LogicalDate = logicalDate + } + if !spec.RunAfter.IsZero() { + msg.RunAfter = spec.RunAfter.UTC() + } + if spec.Conf != nil { + conf := genmodels.Conf(spec.Conf) + msg.Conf = &conf + } + if spec.Note != "" { + msg.Note = spec.Note + } + alreadyExists, err := client.triggerDagRun(ctx, msg) + if err != nil { + return fail(err.Error()) + } + if alreadyExists { + if spec.SkipWhenAlreadyExists { + logger.InfoContext(ctx, + "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", + "dag_id", spec.DagID, + ) + return genmodels.TaskState{ + State: genmodels.TaskStateStateSkipped, + EndDate: time.Now().UTC(), + } + } + logger.ErrorContext(ctx, "Dag Run already exists, marking task as failed.", + "dag_id", spec.DagID, + ) + return genmodels.TaskState{State: genmodels.TaskStateStateFailed, EndDate: time.Now().UTC()} + } + logger.InfoContext(ctx, "Dag Run triggered successfully.", "trigger_dag_id", spec.DagID) + if err := client.PushXCom(ctx, ti, triggerRunIDXComKey, runID); err != nil { + return fail(err.Error()) + } + + deferrable, err := triggerIsDeferrable(spec, env) + if err != nil { + return fail(err.Error()) + } + if !spec.WaitForCompletion { + if deferrable { + logger.InfoContext(ctx, + "Ignoring deferrable=True because wait_for_completion=False. "+ + "Task will complete immediately without waiting for the triggered DAG run.", + "trigger_dag_id", spec.DagID, + ) + } + return succeedTask() + } + + pokeInterval := defaultTriggerPokeInterval + if spec.PokeInterval != nil { + pokeInterval = *spec.PokeInterval + } + if deferrable { + logger.InfoContext(ctx, "Pausing task as DEFERRED.", + "trigger_dag_id", spec.DagID, "run_id", runID, + ) + queuesEnabled, err := envBool(env.getenv, triggererQueuesEnabledEnv) + if err != nil { + return fail(err.Error()) + } + deferred := genmodels.DeferTask{ + State: "deferred", + Classpath: dagStateTriggerClasspath, + // The keys are those of DagStateTrigger.serialize(). + TriggerKwargs: &genmodels.TriggerKwargs{ + "dag_id": spec.DagID, + "states": slices.Concat(allowed, failed), + "poll_interval": int(pokeInterval / time.Second), + "run_ids": []string{runID}, + "execution_dates": nil, + }, + NextMethod: executeCompleteMethod, + NextKwargs: &genmodels.NextKwargs{}, + } + // The trigger runs on a triggerer that serves the queue of the task only when + // [triggerer] queues_enabled is set. + if queuesEnabled && details.TI.Queue != "" { + deferred.Queue = details.TI.Queue + } + return deferred + } + + for { + logger.InfoContext(ctx, "Waiting for dag run to complete execution in allowed state.", + "dag_id", spec.DagID, "run_id", runID, "allowed_state", allowed, + ) + if err := env.sleep(ctx, pokeInterval); err != nil { + return fail(err.Error()) + } + state, err := client.getDagRunState(ctx, spec.DagID, runID) + if err != nil { + return fail(err.Error()) + } + if slices.Contains(failed, state) { + logger.ErrorContext(ctx, "DagRun finished with failed state.", + "dag_id", spec.DagID, "state", state, + ) + return fail(fmt.Sprintf("%s failed with failed state %s", spec.DagID, state)) + } + if slices.Contains(allowed, state) { + logger.InfoContext(ctx, "DagRun finished with allowed state.", + "dag_id", spec.DagID, "state", state, + ) + return succeedTask() + } + logger.DebugContext(ctx, "DagRun not yet in allowed or failed state.", + "dag_id", spec.DagID, "state", state, + ) + } +} + +// triggerLogicalDate returns the logical date of the Dag run to trigger, or the zero Time for a Dag +// run with none, as TriggerDagRunOperator.execute picks it. It is now only when the task sets +// neither LogicalDate nor RunAfter. +func triggerLogicalDate(spec bundle.TriggerSpec, now time.Time) time.Time { + switch { + case !spec.LogicalDate.IsZero(): + return spec.LogicalDate.UTC() + case spec.RunAfter.IsZero(): + return now + } + return time.Time{} +} + +// triggerIsDeferrable returns whether the task defers, which is spec.Deferrable or else +// [operators] default_deferrable. +func triggerIsDeferrable(spec bundle.TriggerSpec, env triggerEnv) (bool, error) { + if spec.Deferrable != nil { + return *spec.Deferrable, nil + } + return envBool(env.getenv, defaultDeferrableEnv) +} + +// envBool reads a setting that the coordinator exports as True or False. An unset setting is +// false, which is Python's fallback for the settings of a trigger task. +func envBool(getenv func(string) string, name string) (bool, error) { + raw := strings.TrimSpace(getenv(name)) + if raw == "" { + return false, nil + } + value, err := strconv.ParseBool(raw) + if err != nil { + return false, fmt.Errorf("%s is %q, which is not a boolean; use True or False", name, raw) + } + return value, nil +} + +// resumeTriggerDagRun finishes a task that came back from a deferral, as BaseOperator.resume_execution +// does: __fail__ fails the task, and execute_complete checks the states of the Dag runs. +func resumeTriggerDagRun( + nextMethod string, + nextKwargs *genmodels.NextKwargs, + dagID string, + allowed, failed []string, + logger *slog.Logger, + fail func(string) any, +) any { + var kwargs genmodels.NextKwargs + if nextKwargs != nil { + kwargs = *nextKwargs + } + switch nextMethod { + case triggerFailMethod: + if traceback, ok := kwargs["traceback"].([]any); ok { + lines := make([]string, len(traceback)) + for i, line := range traceback { + lines[i] = fmt.Sprint(line) + } + logger.Error("Trigger failed:\n" + strings.Join(lines, "\n")) + } + message := "Unknown" + if value, ok := kwargs["error"]; ok && value != nil { + message = fmt.Sprint(value) + } + return fail(message) + case executeCompleteMethod: + default: + return fail(fmt.Sprintf("Task cannot resume with next_method %q", nextMethod)) + } + + event, ok := decodeTriggerEvent(kwargs["event"]) + runIDs, hasRunIDs := event["run_ids"].([]any) + if !ok || !hasRunIDs { + encoded, _ := json.Marshal(kwargs["event"]) + return fail("Task resumed with an event it cannot read: " + string(encoded)) + } + var failedRunIDs []string + for _, id := range runIDs { + runID := fmt.Sprint(id) + state, _ := event[runID].(string) + if slices.Contains(failed, state) { + failedRunIDs = append(failedRunIDs, runID) + } else if slices.Contains(allowed, state) { + logger.Info("Triggered Dag run finished with allowed state.", + "dag_id", dagID, "state", state, "run_id", runID, + ) + } + } + if len(failedRunIDs) > 0 { + states, _ := json.Marshal(failed) + runs, _ := json.Marshal(failedRunIDs) + return fail(fmt.Sprintf( + "%s failed with failed states %s for run_ids %s", dagID, states, runs, + )) + } + return succeedTask() +} + +// decodeTriggerEvent returns the data of the event that DagStateTrigger fired. The event is the +// pair (classpath, data). The triggerer stores it with serde, which writes a tuple as +// {"__classname__": "builtins.tuple", "__data__": [...]}. +func decodeTriggerEvent(event any) (map[string]any, bool) { + if wrapped, ok := event.(map[string]any); ok && wrapped["__classname__"] == "builtins.tuple" { + event = wrapped["__data__"] + } + pair, ok := event.([]any) + if !ok || len(pair) != 2 { + return nil, false + } + data, ok := pair[1].(map[string]any) + return data, ok +} + +// pythonIsoformat returns t in UTC as Python's datetime.isoformat() writes it, such as +// "2026-09-30T01:02:03.500000+00:00". It writes the microseconds, six digits, only when they are +// not zero. +func pythonIsoformat(t time.Time) string { + t = t.UTC() + layout := "2006-01-02T15:04:05" + if t.Nanosecond()/1000 != 0 { + layout += ".000000" + } + return t.Format(layout) + "+00:00" +} + +// dagRunURL returns the URL of a Dag run in the UI, as build_airflow_dagrun_url does on +// [api] base_url. An unset base_url is "/". +func dagRunURL(baseURL, dagID, runID string) string { + if baseURL == "" { + baseURL = "/" + } + return strings.TrimRight(baseURL, "/") + "/dags/" + dagID + "/runs/" + runID +} diff --git a/go-sdk/pkg/execution/trigger_runner_test.go b/go-sdk/pkg/execution/trigger_runner_test.go new file mode 100644 index 0000000000000..d8cf537012d07 --- /dev/null +++ b/go-sdk/pkg/execution/trigger_runner_test.go @@ -0,0 +1,648 @@ +// 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 execution + +import ( + "bytes" + "context" + "io" + "log/slog" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +var triggerNow = time.Date(2026, 9, 30, 1, 2, 3, 500000000, time.UTC) + +const triggerNowRunID = "manual__2026-09-30T01:02:03.500000+00:00" + +// triggerSupervisor answers the requests of a trigger task and keeps them in order. +type triggerSupervisor struct { + mu sync.Mutex + requests []map[string]any + + paused bool + // triggerErr is the error code that TriggerDagRun answers with, if any. + triggerErr string + // states are the answers to GetDagRunState in turn. The last one repeats. + states []string +} + +func (s *triggerSupervisor) answer(req map[string]any) (body, errBody map[string]any) { + s.mu.Lock() + defer s.mu.Unlock() + s.requests = append(s.requests, req) + switch req["type"] { + case "GetDag": + return map[string]any{ + "type": "DagResult", "dag_id": req["dag_id"], "is_paused": s.paused, + }, nil + case "TriggerDagRun": + if s.triggerErr != "" { + return nil, map[string]any{ + "type": "ErrorResponse", "error": s.triggerErr, "detail": map[string]any{}, + } + } + case "GetDagRunState": + state := s.states[0] + if len(s.states) > 1 { + s.states = s.states[1:] + } + return map[string]any{"type": "DagRunStateResult", "state": state}, nil + } + return nil, nil +} + +func (s *triggerSupervisor) sent() []map[string]any { + s.mu.Lock() + defer s.mu.Unlock() + return append([]map[string]any(nil), s.requests...) +} + +// types returns the type of each request that the task sent, with the key of an XCom. +func (s *triggerSupervisor) types() []string { + var types []string + for _, req := range s.sent() { + kind := req["type"].(string) + if kind == "SetXCom" { + kind += ":" + req["key"].(string) + } + types = append(types, kind) + } + return types +} + +func (s *triggerSupervisor) request(t *testing.T, kind string) map[string]any { + t.Helper() + for _, req := range s.sent() { + if req["type"] == kind { + return req + } + } + require.Failf(t, "request not sent", "the task sent no %s: %v", kind, s.types()) + return nil +} + +func (s *triggerSupervisor) comm(t *testing.T) *CoordinatorComm { + t.Helper() + reqR, reqW := io.Pipe() + respR, respW := io.Pipe() + t.Cleanup(func() { + reqR.Close() + reqW.Close() + respR.Close() + respW.Close() + }) + go func() { + for { + frame, err := readFrame(reqR) + if err != nil { + return + } + body, errBody := s.answer(rawToMap(t, frame.Body)) + if err := writeFrame(respW, encodeResponseFrame(t, frame.ID, body, errBody)); err != nil { + return + } + } + }() + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + return NewCoordinatorComm(respR, reqW, logger) +} + +// triggerTestEnv replaces the surroundings of a trigger task for the test. It returns the +// durations that the task slept. +func triggerTestEnv(t *testing.T, settings map[string]string) *[]time.Duration { + t.Helper() + var slept []time.Duration + saved := defaultTriggerEnv + defaultTriggerEnv = triggerEnv{ + now: func() time.Time { return triggerNow }, + getenv: func(name string) string { return settings[name] }, + sleep: func(ctx context.Context, d time.Duration) error { + slept = append(slept, d) + return ctx.Err() + }, + randomSuffix: func() string { return "AbCd1234" }, + } + t.Cleanup(func() { defaultTriggerEnv = saved }) + return &slept +} + +func runTriggerTask( + t *testing.T, + spec bundle.TriggerSpec, + supervisor *triggerSupervisor, + adjust func(*genmodels.StartupDetails), +) any { + t.Helper() + b := testBundle{"test_dag": testDag{"trigger": &bundle.TriggerTask{Spec: spec}}} + details := newStartupDetails("trigger") + details.TI.Queue = "golang" + if adjust != nil { + adjust(details) + } + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + return RunTask(context.Background(), b, details, supervisor.comm(t), logger) +} + +func shouldRetry(details *genmodels.StartupDetails) { details.TIContext.ShouldRetry = true } + +func TestTriggerDagRunFireAndForget(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + + result := runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream"}, supervisor, nil) + + assertSucceedTask(t, result) + assert.Equal(t, + []string{"SetXCom:_link_TriggerDagRunLink", "TriggerDagRun", "SetXCom:trigger_run_id"}, + supervisor.types(), + ) + sent := supervisor.sent() + assert.Equal(t, "/dags/downstream/runs/"+triggerNowRunID, sent[0]["value"]) + assert.Equal(t, "run1", sent[0]["run_id"]) + assert.Equal(t, "trigger", sent[0]["task_id"]) + assert.Equal(t, triggerNowRunID, sent[1]["run_id"]) + assert.Equal(t, triggerNowRunID, sent[2]["value"]) +} + +func TestTriggerDagRunLinkUsesTheBaseURL(t *testing.T) { + triggerTestEnv(t, map[string]string{apiBaseURLEnv: "https://airflow.example.com//"}) + supervisor := &triggerSupervisor{} + + runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream", RunID: "r1"}, supervisor, nil) + + assert.Equal(t, + "https://airflow.example.com/dags/downstream/runs/r1", + supervisor.sent()[0]["value"], + ) +} + +func TestTriggerDagRunSendsTheOptionsAsWritten(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + spec := bundle.TriggerSpec{ + DagID: "downstream", + RunID: "downstream_x", + Note: "{{ ds }}", + Conf: map[string]any{"day": "{{ ds }}", "rows": int64(2)}, + ResetDagRun: true, + } + + result := runTriggerTask(t, spec, supervisor, nil) + + assertSucceedTask(t, result) + sent := supervisor.request(t, "TriggerDagRun") + assert.Equal(t, "downstream_x", sent["run_id"], "an explicit run_id is not generated") + assert.Equal(t, "{{ ds }}", sent["note"]) + assert.Equal(t, map[string]any{"day": "{{ ds }}", "rows": int8(2)}, sent["conf"]) + assert.Equal(t, true, sent["reset_dag_run"]) + assert.Equal(t, "downstream_x", supervisor.sent()[2]["value"]) +} + +func TestTriggerDagRunDatesAndRunID(t *testing.T) { + custom := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + runAfter := time.Date(2026, 2, 3, 4, 5, 6, 0, time.UTC) + tests := []struct { + name string + spec bundle.TriggerSpec + wantLogicalDate any + wantRunAfter any + wantRunID string + }{ + { + name: "neither set", + spec: bundle.TriggerSpec{DagID: "d"}, + wantLogicalDate: triggerNow, + wantRunID: triggerNowRunID, + }, + { + name: "only RunAfter set", + spec: bundle.TriggerSpec{DagID: "d", RunAfter: runAfter}, + wantRunAfter: runAfter, + wantRunID: "manual__2026-02-03T04:05:06+00:00_AbCd1234", + }, + { + name: "only LogicalDate set", + spec: bundle.TriggerSpec{DagID: "d", LogicalDate: custom}, + wantLogicalDate: custom, + wantRunID: "manual__2026-01-02T03:04:05+00:00", + }, + { + name: "both set", + spec: bundle.TriggerSpec{ + DagID: "d", + LogicalDate: custom, + RunAfter: runAfter, + }, + wantLogicalDate: custom, + wantRunAfter: runAfter, + wantRunID: "manual__2026-02-03T04:05:06+00:00", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + + runTriggerTask(t, tc.spec, supervisor, nil) + + sent := supervisor.request(t, "TriggerDagRun") + assert.Equal(t, tc.wantRunID, sent["run_id"]) + for key, want := range map[string]any{ + "logical_date": tc.wantLogicalDate, "run_after": tc.wantRunAfter, + } { + if want == nil { + assert.NotContains(t, sent, key) + continue + } + got, ok := sent[key].(time.Time) + require.True(t, ok, "%s is %T", key, sent[key]) + assert.True(t, want.(time.Time).Equal(got), "%s is %v, want %v", key, got, want) + } + }) + } +} + +func TestTriggerDagRunAlreadyExists(t *testing.T) { + tests := []struct { + name string + skip bool + want genmodels.TaskStateState + }{ + {"skipped", true, genmodels.TaskStateStateSkipped}, + {"failed", false, genmodels.TaskStateStateFailed}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{triggerErr: "DAGRUN_ALREADY_EXISTS"} + spec := bundle.TriggerSpec{DagID: "downstream", SkipWhenAlreadyExists: tc.skip} + + // A task that may retry still ends in this state, as in Python. + result := runTriggerTask(t, spec, supervisor, shouldRetry) + + assertTaskState(t, result, tc.want) + assert.Equal(t, + []string{"SetXCom:_link_TriggerDagRunLink", "TriggerDagRun"}, + supervisor.types(), "no trigger_run_id is pushed for a run that was not created", + ) + }) + } +} + +func TestTriggerDagRunErrorRetries(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{triggerErr: "API_SERVER_ERROR"} + + result := runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream"}, supervisor, shouldRetry) + assertRetryTask(t, result, "API_SERVER_ERROR") + + result = runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream"}, supervisor, nil) + assertTaskState(t, result, genmodels.TaskStateStateFailed) +} + +func TestTriggerDagRunFailsWhenTheDagIsPaused(t *testing.T) { + triggerTestEnv(t, nil) + spec := bundle.TriggerSpec{DagID: "downstream", FailWhenDagIsPaused: true} + + paused := &triggerSupervisor{paused: true} + assertRetryTask(t, runTriggerTask(t, spec, paused, shouldRetry), "Dag downstream is paused") + assert.Equal(t, []string{"GetDag"}, paused.types(), "nothing is triggered") + + running := &triggerSupervisor{} + assertSucceedTask(t, runTriggerTask(t, spec, running, nil)) + assert.Equal(t, "GetDag", running.types()[0]) + assert.Contains(t, running.types(), "TriggerDagRun") + + ignored := &triggerSupervisor{paused: true} + spec.FailWhenDagIsPaused = false + assertSucceedTask(t, runTriggerTask(t, spec, ignored, nil)) + assert.NotContains(t, ignored.types(), "GetDag") +} + +func TestTriggerDagRunPollsUntilAnAllowedState(t *testing.T) { + slept := triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{states: []string{"queued", "running", "success"}} + spec := bundle.TriggerSpec{DagID: "downstream", WaitForCompletion: true} + + result := runTriggerTask(t, spec, supervisor, nil) + + assertSucceedTask(t, result) + assert.Equal(t, []time.Duration{60 * time.Second, 60 * time.Second, 60 * time.Second}, *slept) + assert.Equal(t, triggerNowRunID, supervisor.request(t, "GetDagRunState")["run_id"]) +} + +func TestTriggerDagRunPollFailsOnAFailedState(t *testing.T) { + slept := triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{states: []string{"running", "failed"}} + poke := 5 * time.Second + spec := bundle.TriggerSpec{DagID: "downstream", WaitForCompletion: true, PokeInterval: &poke} + + result := runTriggerTask(t, spec, supervisor, shouldRetry) + + assertRetryTask(t, result, "downstream failed with failed state failed") + assert.Equal(t, []time.Duration{poke, poke}, *slept) + + supervisor = &triggerSupervisor{states: []string{"failed"}} + assertTaskState(t, runTriggerTask(t, spec, supervisor, nil), genmodels.TaskStateStateFailed) +} + +func TestTriggerDagRunPollHonorsTheStatesOfTheTask(t *testing.T) { + triggerTestEnv(t, nil) + // No state fails the task, so the failed Dag run is not an allowed state to wait out. + supervisor := &triggerSupervisor{states: []string{"failed", "queued", "failed", "running"}} + spec := bundle.TriggerSpec{ + DagID: "downstream", + WaitForCompletion: true, + AllowedStates: []string{"running"}, + FailedStates: []string{}, + } + + result := runTriggerTask(t, spec, supervisor, nil) + + assertSucceedTask(t, result) + assert.Equal(t, 4, len(supervisor.types())-2-1, "four polls after the link, trigger and run_id") +} + +func TestTriggerDagRunPollStopsWhenTheTaskIsCancelled(t *testing.T) { + triggerTestEnv(t, nil) + defaultTriggerEnv.sleep = func(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } + } + supervisor := &triggerSupervisor{states: []string{"running"}} + poke := time.Hour + spec := bundle.TriggerSpec{DagID: "downstream", WaitForCompletion: true, PokeInterval: &poke} + b := testBundle{"test_dag": testDag{"trigger": &bundle.TriggerTask{Spec: spec}}} + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(50*time.Millisecond, cancel) + + done := make(chan any, 1) + go func() { + done <- RunTask(ctx, b, newStartupDetails("trigger"), supervisor.comm(t), logger) + }() + + select { + case result := <-done: + assertTaskState(t, result, genmodels.TaskStateStateFailed) + case <-time.After(5 * time.Second): + t.Fatal("the task kept sleeping after its context was cancelled") + } + assert.NotContains(t, supervisor.types(), "GetDagRunState") +} + +func TestTriggerDagRunDefers(t *testing.T) { + poke := 90 * time.Second + tests := []struct { + name string + settings map[string]string + spec bundle.TriggerSpec + wantQueue any + wantStates []any + wantPoll any + }{ + { + name: "defaults, queues off", + spec: bundle.TriggerSpec{Deferrable: ptr(true)}, + wantStates: []any{"success", "failed"}, + wantPoll: int8(60), + }, + { + name: "queues on", + settings: map[string]string{triggererQueuesEnabledEnv: "True"}, + spec: bundle.TriggerSpec{Deferrable: ptr(true)}, + wantQueue: "golang", + wantStates: []any{"success", "failed"}, + wantPoll: int8(60), + }, + { + name: "default_deferrable, own states and interval", + settings: map[string]string{ + defaultDeferrableEnv: "True", + triggererQueuesEnabledEnv: "False", + }, + spec: bundle.TriggerSpec{ + PokeInterval: &poke, + AllowedStates: []string{"success", "queued"}, + FailedStates: []string{"failed"}, + }, + wantStates: []any{"success", "queued", "failed"}, + wantPoll: int8(90), + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + slept := triggerTestEnv(t, tc.settings) + supervisor := &triggerSupervisor{} + tc.spec.DagID = "downstream" + tc.spec.WaitForCompletion = true + + result := runTriggerTask(t, tc.spec, supervisor, nil) + + assert.Empty(t, *slept, "a deferred task does not poll") + assert.NotContains(t, supervisor.types(), "GetDagRunState") + deferred, ok := result.(genmodels.DeferTask) + require.True(t, ok, "expected DeferTask, got %T", result) + wire, err := encodeRequest(0, deferred) + require.NoError(t, err) + var framed bytes.Buffer + require.NoError(t, writeFrame(&framed, wire)) + frame, err := readFrame(&framed) + require.NoError(t, err) + want := map[string]any{ + "type": "DeferTask", + "state": "deferred", + "classpath": "airflow.providers.standard.triggers.external_task.DagStateTrigger", + "next_method": "execute_complete", + "next_kwargs": map[string]any{}, + "trigger_kwargs": map[string]any{ + "dag_id": "downstream", + "states": tc.wantStates, + "poll_interval": tc.wantPoll, + "run_ids": []any{triggerNowRunID}, + "execution_dates": nil, + }, + } + if tc.wantQueue != nil { + want["queue"] = tc.wantQueue + } + assert.Equal(t, want, rawToMap(t, frame.Body)) + }) + } +} + +func TestTriggerDagRunDeferrableFalseOverridesTheConfig(t *testing.T) { + slept := triggerTestEnv(t, map[string]string{defaultDeferrableEnv: "True"}) + supervisor := &triggerSupervisor{states: []string{"success"}} + spec := bundle.TriggerSpec{DagID: "downstream", WaitForCompletion: true, Deferrable: ptr(false)} + + result := runTriggerTask(t, spec, supervisor, nil) + + assertSucceedTask(t, result) + assert.Len(t, *slept, 1, "the task polled instead of deferring") +} + +func TestTriggerDagRunIgnoresDeferrableWithoutWaiting(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + + result := runTriggerTask(t, + bundle.TriggerSpec{DagID: "downstream", Deferrable: ptr(true)}, supervisor, nil, + ) + + assertSucceedTask(t, result) +} + +func TestTriggerDagRunRejectsAnInvalidSetting(t *testing.T) { + triggerTestEnv(t, map[string]string{defaultDeferrableEnv: "maybe"}) + supervisor := &triggerSupervisor{} + + result := runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream"}, supervisor, shouldRetry) + + assertRetryTask(t, result, defaultDeferrableEnv) +} + +func TestTriggerDagRunResume(t *testing.T) { + tuple := func(data map[string]any) map[string]any { + return map[string]any{ + "__classname__": "builtins.tuple", + "__data__": []any{dagStateTriggerClasspath, data}, + } + } + tests := []struct { + name string + nextMethod string + kwargs map[string]any + // wantRetry is a substring of the retry reason, or empty when the task succeeds. + wantRetry string + }{ + { + name: "allowed state", + nextMethod: "execute_complete", + kwargs: map[string]any{ + "event": tuple(map[string]any{"run_ids": []any{"r1"}, "r1": "success"}), + }, + }, + { + name: "failed state", + nextMethod: "execute_complete", + kwargs: map[string]any{ + "event": tuple( + map[string]any{"run_ids": []any{"r1", "r2"}, "r1": "success", "r2": "failed"}, + ), + }, + wantRetry: `downstream failed with failed states ["failed"] for run_ids ["r2"]`, + }, + { + name: "event without the tuple wrapper", + nextMethod: "execute_complete", + kwargs: map[string]any{ + "event": []any{ + dagStateTriggerClasspath, + map[string]any{"run_ids": []any{"r1"}, "r1": "failed"}, + }, + }, + wantRetry: "for run_ids", + }, + { + name: "event that cannot be read", + nextMethod: "execute_complete", + kwargs: map[string]any{"event": "oops"}, + wantRetry: `Task resumed with an event it cannot read: "oops"`, + }, + { + name: "no event", + nextMethod: "execute_complete", + wantRetry: "Task resumed with an event it cannot read: null", + }, + { + name: "trigger failed", + nextMethod: "__fail__", + kwargs: map[string]any{"error": "trigger crashed", "traceback": []any{"line 1"}}, + wantRetry: "trigger crashed", + }, + { + name: "trigger failed without an error", + nextMethod: "__fail__", + wantRetry: "Unknown", + }, + { + name: "unknown method", + nextMethod: "other", + wantRetry: `Task cannot resume with next_method "other"`, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + resume := func(details *genmodels.StartupDetails) { + details.TIContext.ShouldRetry = true + details.TIContext.NextMethod = tc.nextMethod + if tc.kwargs != nil { + kwargs := genmodels.NextKwargs(tc.kwargs) + details.TIContext.NextKwargs = &kwargs + } + } + + result := runTriggerTask(t, bundle.TriggerSpec{DagID: "downstream"}, supervisor, resume) + + if tc.wantRetry == "" { + assertSucceedTask(t, result) + } else { + assertRetryTask(t, result, tc.wantRetry) + } + assert.Empty(t, supervisor.types(), "a resumed task triggers nothing") + }) + } +} + +func TestTriggerDagRunResumeUsesTheStatesOfTheTask(t *testing.T) { + triggerTestEnv(t, nil) + supervisor := &triggerSupervisor{} + spec := bundle.TriggerSpec{DagID: "downstream", FailedStates: []string{"queued"}} + kwargs := genmodels.NextKwargs{"event": []any{ + dagStateTriggerClasspath, map[string]any{"run_ids": []any{"r1"}, "r1": "failed"}, + }} + + result := runTriggerTask(t, spec, supervisor, func(details *genmodels.StartupDetails) { + details.TIContext.NextMethod = "execute_complete" + details.TIContext.NextKwargs = &kwargs + }) + + assertSucceedTask(t, result) +} + +func TestTriggerTaskExecuteIsNeverCalled(t *testing.T) { + err := (&bundle.TriggerTask{}).Execute(context.Background(), nil, nil) + assert.ErrorContains(t, err, "run by the runtime") +} diff --git a/scripts/ci/lang_sdk_serialization/test_dags.yaml b/scripts/ci/lang_sdk_serialization/test_dags.yaml index 174accfdd4f33..5d3fbfad8bc7a 100644 --- a/scripts/ci/lang_sdk_serialization/test_dags.yaml +++ b/scripts/ci/lang_sdk_serialization/test_dags.yaml @@ -306,7 +306,7 @@ dags: - task_id: trigger_all trigger_dag_run: trigger_dag_id: conformance_minimal - trigger_run_id: "triggered_{{ run_id }}" + trigger_run_id: "triggered_by_conformance" conf: {rows: 2, ratio: 0.5, nested: {tags: [a, b]}} logical_date: "2026-09-30 00:00:00+00:00" wait_for_completion: true