This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 498724fc13a Go SDK: enumerate every user-facing enum in the airflow
package (#74064)
498724fc13a is described below
commit 498724fc13a41835f5cf68b352799e920eb47cd0
Author: PoAn Yang <[email protected]>
AuthorDate: Fri Oct 2 16:29:50 2026 +0900
Go SDK: enumerate every user-facing enum in the airflow package (#74064)
Signed-off-by: PoAn Yang <[email protected]>
---
go-sdk/README.md | 12 +-
go-sdk/adr/0008-native-dag-interface.md | 5 +
go-sdk/airflow/dag.go | 5 +
go-sdk/airflow/enums.go | 156 ++++++++++++++++++++++
go-sdk/airflow/enums_test.go | 221 ++++++++++++++++++++++++++++++++
go-sdk/airflow/spec.go | 31 -----
go-sdk/airflow/spec_test.go | 71 ----------
go-sdk/airflow/trigger_dag_run.go | 36 ++----
go-sdk/airflow/trigger_dag_run_test.go | 24 ++--
go-sdk/internal/genspec/authoring.go | 6 +-
10 files changed, 423 insertions(+), 144 deletions(-)
diff --git a/go-sdk/README.md b/go-sdk/README.md
index a6dedde0804..3bd9f527502 100644
--- a/go-sdk/README.md
+++ b/go-sdk/README.md
@@ -396,6 +396,11 @@ drift, so a missed bump is caught by `go test`. A snapshot
can also grow a field
struct does not declare — so the `check-go-sdk-generated-drift` prek hook
regenerates them and fails
when the committed files differ.
+[`airflow/enums.go`](./airflow/enums.go) declares the `genmodels.DagRunState`
constants again, so that
+a Dag author does not import `genmodels`. `TestDagRunStateMatchesGenmodels`
fails when regenerating
+the models adds, renames or removes a `DagRunState` constant. The test keeps
failing until
+`airflow/enums.go` declares the same constants as `genmodels`.
+
## Regenerating the Dag and task specs
`airflow.DagSpec` and `airflow.TaskSpec` in
[`airflow/spec.gen.go`](./airflow/spec.gen.go) are
@@ -416,9 +421,10 @@ it into the authoring shape, each entry carrying the
reason it exists:
schema rather than hand-writing the field keeps every field in one struct
declaration, which is
what lets `TaskSpec` implement the sealed `TaskOption`.
-`trigger_rule` is typed as a plain string with no values named, so
`TriggerRule` and its constants
-are hand-written in [`airflow/spec.go`](./airflow/spec.go);
-`TestTriggerRuleConstantsMatchPython` is their tripwire against Airflow's own
enum.
+The schema types `trigger_rule` and `weight_rule` as plain strings and does
not list their values.
+`TriggerRule`, `WeightRule` and their constants are therefore hand-written in
+[`airflow/enums.go`](./airflow/enums.go). `TestRuleConstantsMatchPython`
checks those constants
+against the Python enums in airflow-core.
The `check-go-sdk-generated-drift` prek hook regenerates the file and fails
when the committed one
differs, so a schema change that never reached Go cannot merge.
diff --git a/go-sdk/adr/0008-native-dag-interface.md
b/go-sdk/adr/0008-native-dag-interface.md
index aebb92f5514..974218268a5 100644
--- a/go-sdk/adr/0008-native-dag-interface.md
+++ b/go-sdk/adr/0008-native-dag-interface.md
@@ -175,6 +175,11 @@ const (
Either the schema gains titles or the generate step keeps the map.
- **The schema is the serialized shape, not the authoring shape.** It requires
`fileloc` and `tasks` on a Dag, and `task_type`, `_task_module`, `ui_color`,
`ui_fgcolor`, and `template_fields` on an operator, all of which the SDK fills
in, and it carries a serialized `timetable` object where an author writes a
schedule.
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 state enum generated into `genmodels` is re-exported from `airflow` once
a field that an author can reach holds its values.** Today the only re-exported
enum is `DagRunState`, the type of the values in
`TriggerDagRunSpec.AllowedStates` and `TriggerDagRunSpec.FailedStates`.
`DagRunType` is re-exported once a field such as `allowed_run_types` on
`DagSpec` holds its values. `TaskInstanceState` is re-exported once
`airflow.TaskInstance` gets a state field.
+ `airflow/enums.go` lists every user-facing enum with the Python enum that it
mirrors.
+ `airflow.DagRunState` is a type alias of `genmodels.DagRunState`, not a new
type. `airflow.DagRun` is an alias of `sdk.DagRun`, and package `sdk` cannot
import `airflow`. So a state field that `sdk.DagRun` gains later has a type
declared outside `airflow`. `actx.DagRun().State == airflow.DagRunStateFailed`
compiles only if both sides have the same type.
+ The cost is that `%T` and package reflect report `airflow.DagRunState` as
`genmodels.DagRunState`, and package `airflow` cannot declare a method on that
type.
+- **A value outside an enum is rejected when the task is added.** Go converts
a string literal to an enum type without a cast, so the enum type alone does
not stop `TaskSpec{TriggerRule: "all_sucess"}`. `dag.Task` panics with a
message that names the field and lists the valid values. Python rejects an
unknown trigger rule when the operator is constructed. It rejects an unknown
weight rule only when the Dag is serialized.
- **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.
diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go
index 757ee52f817..5f301676d15 100644
--- a/go-sdk/airflow/dag.go
+++ b/go-sdk/airflow/dag.go
@@ -121,6 +121,8 @@ type TaskRef struct {
// - fn comes from TriggerDagRun and opts holds an Inputs
// - an option is nil or is not one that package airflow defines
// - opts holds more than one TaskSpec or more than one Inputs
+// - the TaskSpec sets TriggerRule to a value that is not a TriggerRule
constant
+// - the TaskSpec sets WeightRule to a value that is not a WeightRule
constant
// - the tasks passed to Inputs do not match the parameters of fn after the
Context
// - the Dag already has a task with the same task_id
// - the Dag is already registered
@@ -209,6 +211,9 @@ func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef {
"airflow.DagRef.Task: task %q of Dag %q: %v", taskID,
d.dagID, triggerErr,
))
}
+ if err := checkTaskSpec(cfg.spec); err != nil {
+ panic(fmt.Sprintf("airflow.DagRef.Task: task %q of Dag %q: %v",
taskID, d.dagID, err))
+ }
if _, exists := d.tasksByID[taskID]; exists {
panic(fmt.Sprintf(
"airflow.DagRef.Task: Dag %q already has a task %q; "+
diff --git a/go-sdk/airflow/enums.go b/go-sdk/airflow/enums.go
new file mode 100644
index 00000000000..4aba07c91c3
--- /dev/null
+++ b/go-sdk/airflow/enums.go
@@ -0,0 +1,156 @@
+// 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 (
+ "fmt"
+ "slices"
+
+ "github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
+)
+
+// This file declares every enum type whose values a Dag author sets. Each
item below gives
+// the Go enum, the Python file that declares the same enum, and the fields
that hold its values:
+//
+// - TriggerRule: airflow-core/src/airflow/task/trigger_rule.py, for
TaskSpec.TriggerRule
+// - WeightRule: airflow-core/src/airflow/task/weight_rule.py, for
TaskSpec.WeightRule
+// - DagRunState: airflow-core/src/airflow/utils/state.py, for
TriggerDagRunSpec.AllowedStates
+// and TriggerDagRunSpec.FailedStates
+//
+// Two more Python enums are not declared here yet, because no field that an
author can reach
+// holds their values:
+//
+// - DagRunType: airflow-core/src/airflow/utils/types.py. DagSpec leaves out
+// allowed_run_types, and DagRun has no run type field.
+// - TaskInstanceState: airflow-core/src/airflow/utils/state.py.
TaskInstance has no state field.
+//
+// Adding such a field also means declaring its enum here, the way DagRunState
is declared.
+//
+// An enum is a named string type, and each constant has the type name as its
prefix, as in
+// TriggerRuleAllDone. Go does not scope a constant to its type, so without
the prefix, a name
+// such as airflow.AllDone would not say which enum the constant belongs to. A
field with a fixed
+// set of values gets its enum type before the field is released. Changing a
released field from
+// string to an enum type breaks every Dag that sets the field from a value of
type string, such
+// as a string variable.
+//
+// DagRef.Task panics on a TriggerRule, WeightRule or DagRunState value that
is not one of the
+// constants of that type. The one exception is an empty TriggerRule or
WeightRule, which leaves
+// that rule unset.
+//
+// The Dag serialization schema types trigger_rule and weight_rule as plain
strings and does not
+// list their values. TriggerRule, WeightRule and their constants are
therefore written by hand in
+// this file, and internal/genspec/authoring.go gives these types to the
generated fields
+// TaskSpec.TriggerRule and TaskSpec.WeightRule.
+
+// TriggerRule decides when a task runs, based on the states of its upstream
tasks. An empty
+// TriggerRule leaves the trigger rule of the task unset.
+type TriggerRule string
+
+const (
+ TriggerRuleAllSuccess TriggerRule = "all_success"
+ TriggerRuleAllFailed TriggerRule = "all_failed"
+ TriggerRuleAllDone TriggerRule = "all_done"
+ TriggerRuleAllDoneMinOneSuccess TriggerRule =
"all_done_min_one_success"
+ TriggerRuleAllDoneSetupSuccess TriggerRule =
"all_done_setup_success"
+ TriggerRuleOneSuccess TriggerRule = "one_success"
+ TriggerRuleOneFailed TriggerRule = "one_failed"
+ TriggerRuleOneDone TriggerRule = "one_done"
+ TriggerRuleNoneFailed TriggerRule = "none_failed"
+ TriggerRuleNoneFailedMinOneSuccess TriggerRule =
"none_failed_min_one_success"
+ TriggerRuleNoneSkipped TriggerRule = "none_skipped"
+ TriggerRuleAllSkipped TriggerRule = "all_skipped"
+ TriggerRuleAlways TriggerRule = "always"
+)
+
+// WeightRule decides how Airflow computes the effective priority weight of a
task.
+// WeightRuleAbsolute uses the TaskSpec.PriorityWeight of the task alone.
WeightRuleDownstream
+// sums the PriorityWeight of the task and of every task downstream of it.
WeightRuleUpstream sums
+// the PriorityWeight of the task and of every task upstream of it. An empty
WeightRule leaves the
+// weight rule of the task unset.
+//
+// Python also accepts the import path of a PriorityWeightStrategy class for
weight_rule. The
+// class can be built into Airflow or registered by a plugin. A Go task can
only name one of the
+// three built-in rules.
+type WeightRule string
+
+const (
+ WeightRuleDownstream WeightRule = "downstream"
+ WeightRuleUpstream WeightRule = "upstream"
+ WeightRuleAbsolute WeightRule = "absolute"
+)
+
+// DagRunState is the state of a Dag run.
+//
+// DagRunState is an alias of genmodels.DagRunState so that an author does not
import genmodels.
+// The types in pkg/execution/genmodels are generated from the supervisor
schema. Because
+// DagRunState is an alias, %T and package reflect report its name as
genmodels.DagRunState.
+type DagRunState = genmodels.DagRunState
+
+const (
+ DagRunStateQueued DagRunState = "queued"
+ DagRunStateRunning DagRunState = "running"
+ DagRunStateSuccess DagRunState = "success"
+ DagRunStateFailed DagRunState = "failed"
+)
+
+// triggerRules, weightRules and dagRunStates list the valid values of
TriggerRule, WeightRule and
+// DagRunState. Tests in this package check these lists and the constants
above against the
+// Python enums for TriggerRule and WeightRule, and against genmodels for
DagRunState.
+var (
+ triggerRules = []TriggerRule{
+ TriggerRuleAllSuccess,
+ TriggerRuleAllFailed,
+ TriggerRuleAllDone,
+ TriggerRuleAllDoneMinOneSuccess,
+ TriggerRuleAllDoneSetupSuccess,
+ TriggerRuleOneSuccess,
+ TriggerRuleOneFailed,
+ TriggerRuleOneDone,
+ TriggerRuleNoneFailed,
+ TriggerRuleNoneFailedMinOneSuccess,
+ TriggerRuleNoneSkipped,
+ TriggerRuleAllSkipped,
+ TriggerRuleAlways,
+ }
+ weightRules = []WeightRule{WeightRuleDownstream, WeightRuleUpstream,
WeightRuleAbsolute}
+ dagRunStates = []DagRunState{
+ DagRunStateQueued,
+ DagRunStateRunning,
+ DagRunStateSuccess,
+ DagRunStateFailed,
+ }
+)
+
+// checkTaskSpec rejects a TaskSpec whose TriggerRule is not a TriggerRule
constant or whose
+// WeightRule is not a WeightRule constant. An empty TriggerRule or WeightRule
is valid and leaves
+// that rule unset.
+func checkTaskSpec(spec TaskSpec) error {
+ if err := checkRule("TriggerRule", "trigger rule", spec.TriggerRule,
triggerRules); err != nil {
+ return err
+ }
+ return checkRule("WeightRule", "weight rule", spec.WeightRule,
weightRules)
+}
+
+func checkRule[T ~string](field, kind string, value T, valid []T) error {
+ if value == "" || slices.Contains(valid, value) {
+ return nil
+ }
+ return fmt.Errorf(
+ "airflow.TaskSpec.%s is %q, which is not a %s; use one of %q",
field, value, kind, valid,
+ )
+}
diff --git a/go-sdk/airflow/enums_test.go b/go-sdk/airflow/enums_test.go
new file mode 100644
index 00000000000..32c8cd7fb09
--- /dev/null
+++ b/go-sdk/airflow/enums_test.go
@@ -0,0 +1,221 @@
+// 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 (
+ "fmt"
+ "maps"
+ "os"
+ "path/filepath"
+ "regexp"
+ "slices"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// enumsPath is the path, relative to this package, of the file that declares
each enum.
+const enumsPath = "enums.go"
+
+func readMatches(t *testing.T, path, pattern string) [][]string {
+ t.Helper()
+
+ body, err := os.ReadFile(filepath.FromSlash(path))
+ require.NoError(t, err)
+ matches :=
regexp.MustCompile(pattern).FindAllStringSubmatch(string(body), -1)
+ require.NotEmpty(t, matches, "%s matched nothing, so the file's shape
has changed", path)
+ return matches
+}
+
+// readPairs maps the first submatch of each match of pattern in the file at
path to the second
+// submatch, such as a constant name to its value.
+func readPairs(t *testing.T, path, pattern string) map[string]string {
+ t.Helper()
+
+ pairs := make(map[string]string)
+ for _, match := range readMatches(t, path, pattern) {
+ pairs[match[1]] = match[2]
+ }
+ return pairs
+}
+
+// readPythonEnum maps the name of each member of the Python enum in the file
at path to the
+// value of the member. It fails the test if a line that assigns to an
upper-case name at the
+// indentation of a member is not written as NAME = "value". Without that
check, a member written
+// another way would go unread instead of failing the test.
+func readPythonEnum(t *testing.T, path string) map[string]string {
+ t.Helper()
+
+ members := readPairs(t, path, `(?m)^ ([A-Z][A-Z0-9_]*) = "([^"]*)"`)
+ var assigned []string
+ for _, match := range readMatches(t, path, `(?m)^
([A-Z][A-Z0-9_]*)\s*[:=]`) {
+ assigned = append(assigned, match[1])
+ }
+ require.ElementsMatch(t, assigned, slices.Collect(maps.Keys(members)),
+ "%s has a member that is not written as NAME = \"value\"", path)
+ return members
+}
+
+// goConstantSuffix turns the name of a Python enum member, such as
ALL_DONE_MIN_ONE_SUCCESS,
+// into the part of a Go constant name after the type name, such as
AllDoneMinOneSuccess.
+func goConstantSuffix(pythonName string) string {
+ var suffix strings.Builder
+ for _, word := range strings.Split(strings.ToLower(pythonName), "_") {
+ if word != "" {
+ suffix.WriteString(strings.ToUpper(word[:1]) + word[1:])
+ }
+ }
+ return suffix.String()
+}
+
+func toStrings[T ~string](values []T) []string {
+ out := make([]string, 0, len(values))
+ for _, value := range values {
+ out = append(out, string(value))
+ }
+ return out
+}
+
+// TestRuleConstantsMatchPython fails when the Python side adds or renames a
trigger rule or a
+// weight rule. The Dag serialization schema types trigger_rule and
weight_rule as plain strings
+// and does not list their values. The TriggerRule and WeightRule constants
are therefore
+// hand-written, outside the generated files that the
check-go-sdk-generated-drift prek hook
+// checks. The test also compares triggerRules and weightRules with the Python
enums.
+// checkTaskSpec accepts only the values in those lists, so an author cannot
set a rule that the
+// lists leave out.
+func TestRuleConstantsMatchPython(t *testing.T) {
+ for _, tt := range []struct {
+ goType string
+ pythonPath string
+ valid []string
+ }{
+ {
+ goType: "TriggerRule",
+ pythonPath:
"../../airflow-core/src/airflow/task/trigger_rule.py",
+ valid: toStrings(triggerRules),
+ },
+ {
+ goType: "WeightRule",
+ pythonPath:
"../../airflow-core/src/airflow/task/weight_rule.py",
+ valid: toStrings(weightRules),
+ },
+ } {
+ t.Run(tt.goType, func(t *testing.T) {
+ python := readPythonEnum(t, tt.pythonPath)
+ want := make(map[string]string, len(python))
+ for name, value := range python {
+ want[goConstantSuffix(name)] = value
+ }
+ spelled := readPairs(
+ t,
+ enumsPath,
+ `(?m)^\t`+tt.goType+`(\w+) +`+tt.goType+` =
"([^"]*)"$`,
+ )
+
+ assert.Equal(t, want, spelled)
+ assert.ElementsMatch(t,
slices.Collect(maps.Values(python)), tt.valid)
+ })
+ }
+}
+
+// TestDagRunStateMatchesGenmodels fails when regenerating genmodels from the
supervisor schema
+// adds, renames or removes a DagRunState constant. The DagRunState alias
brings the type into
+// this package but not the genmodels constants, so enums.go declares each of
those constants
+// again.
+func TestDagRunStateMatchesGenmodels(t *testing.T) {
+ generated := readPairs(
+ t,
+ "../pkg/execution/genmodels/models.gen.go",
+ `(?m)^const DagRunState(\w+) DagRunState = "([^"]*)"$`,
+ )
+ declared := readPairs(t, enumsPath, `(?m)^\tDagRunState(\w+)
+DagRunState = "([^"]*)"$`)
+
+ assert.Equal(t, generated, declared)
+ assert.ElementsMatch(t, slices.Collect(maps.Values(generated)),
toStrings(dagRunStates))
+}
+
+func TestTaskRejectsAValueOutsideAnEnum(t *testing.T) {
+ const (
+ triggerRuleValues = `["all_success" "all_failed" "all_done"
"all_done_min_one_success" ` +
+ `"all_done_setup_success" "one_success" "one_failed"
"one_done" "none_failed" ` +
+ `"none_failed_min_one_success" "none_skipped"
"all_skipped" "always"]`
+ weightRuleValues = `["downstream" "upstream" "absolute"]`
+ )
+ tests := []struct {
+ name string
+ fn any
+ spec TaskSpec
+ taskID string
+ want string
+ }{
+ {
+ name: "TriggerRule of a task named after its
function",
+ fn: extract,
+ spec: TaskSpec{TriggerRule: "all_sucess"},
+ taskID: "extract",
+ want: `airflow.TaskSpec.TriggerRule is "all_sucess",
which is not a trigger rule; ` +
+ `use one of ` + triggerRuleValues,
+ },
+ {
+ name: "WeightRule next to a valid TriggerRule",
+ fn: extract,
+ spec: TaskSpec{
+ TaskID: "extract_rows", TriggerRule:
TriggerRuleAllDone, WeightRule: "Absolute",
+ },
+ taskID: "extract_rows",
+ want: `airflow.TaskSpec.WeightRule is "Absolute", which
is not a weight rule; ` +
+ `use one of ` + weightRuleValues,
+ },
+ {
+ name: "WeightRule of a task from TriggerDagRun",
+ fn: TriggerDagRun(TriggerDagRunSpec{DagID:
"downstream_etl"}),
+ spec: TaskSpec{TaskID: "trigger_downstream",
WeightRule: "Absolute"},
+ taskID: "trigger_downstream",
+ want: `airflow.TaskSpec.WeightRule is "Absolute", which
is not a weight rule; ` +
+ `use one of ` + weightRuleValues,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ dag := Dag("etl")
+ assert.PanicsWithValue(t,
+ fmt.Sprintf(`airflow.DagRef.Task: task %q of
Dag "etl": %s`, tt.taskID, tt.want),
+ func() { dag.Task(tt.fn, tt.spec) },
+ )
+ assert.NotPanics(t,
+ func() { dag.Task(extract, TaskSpec{TaskID:
tt.taskID}) },
+ "a rejected task does not take its task_id",
+ )
+ })
+ }
+}
+
+func TestTaskAcceptsEveryValueOfAnEnum(t *testing.T) {
+ dag := Dag("etl")
+ dag.Task(extract, TaskSpec{TaskID: "no_rules"})
+ for _, rule := range triggerRules {
+ dag.Task(extract, TaskSpec{TaskID: "trigger_" + string(rule),
TriggerRule: rule})
+ }
+ for _, rule := range weightRules {
+ dag.Task(extract, TaskSpec{TaskID: "weight_" + string(rule),
WeightRule: rule})
+ }
+
+ assert.Len(t, dag.tasks, 1+len(triggerRules)+len(weightRules))
+}
diff --git a/go-sdk/airflow/spec.go b/go-sdk/airflow/spec.go
index 4b928999d4e..af477b795fc 100644
--- a/go-sdk/airflow/spec.go
+++ b/go-sdk/airflow/spec.go
@@ -39,37 +39,6 @@ import "reflect"
//go:generate go run github.com/atombender/[email protected] --only-models
--struct-name-from-title --tags "" --capitalization ID --capitalization JSON
--capitalization MD --capitalization XCom -p airflow -o spec.gen.go
../../.build/go-sdk/spec.schema.json
//go:generate go run ../internal/genspec -license spec.gen.go
-// TriggerRule is when a task runs, given the state of the tasks upstream of
it, and
-// WeightRule is how a task's priority weight is worked out from the graph
around it.
-// The serialization schema types both as a plain string and names none of the
values
-// either takes, so the constants are written here and the generated fields
are given
-// these types by internal/genspec/authoring.go.
-type TriggerRule string
-
-const (
- TriggerRuleAllSuccess TriggerRule = "all_success"
- TriggerRuleAllFailed TriggerRule = "all_failed"
- TriggerRuleAllDone TriggerRule = "all_done"
- TriggerRuleAllDoneMinOneSuccess TriggerRule =
"all_done_min_one_success"
- TriggerRuleAllDoneSetupSuccess TriggerRule =
"all_done_setup_success"
- TriggerRuleOneSuccess TriggerRule = "one_success"
- TriggerRuleOneFailed TriggerRule = "one_failed"
- TriggerRuleOneDone TriggerRule = "one_done"
- TriggerRuleNoneFailed TriggerRule = "none_failed"
- TriggerRuleNoneFailedMinOneSuccess TriggerRule =
"none_failed_min_one_success"
- TriggerRuleNoneSkipped TriggerRule = "none_skipped"
- TriggerRuleAllSkipped TriggerRule = "all_skipped"
- TriggerRuleAlways TriggerRule = "always"
-)
-
-type WeightRule string
-
-const (
- WeightRuleDownstream WeightRule = "downstream"
- WeightRuleUpstream WeightRule = "upstream"
- WeightRuleAbsolute WeightRule = "absolute"
-)
-
// copySpec returns a copy of spec that shares no slice, map or pointer with
it, so that a
// caller changing what it still holds cannot change a registered Dag.
Assigning a spec
// copies a slice header or a pointer and not what it points at, which is the
sharing this
diff --git a/go-sdk/airflow/spec_test.go b/go-sdk/airflow/spec_test.go
deleted file mode 100644
index 2221bd56e91..00000000000
--- a/go-sdk/airflow/spec_test.go
+++ /dev/null
@@ -1,71 +0,0 @@
-// 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 (
- "os"
- "path/filepath"
- "regexp"
- "testing"
-
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
-)
-
-// specPath is where the Go half of each value set is declared, reached from
this package.
-const specPath = "spec.go"
-
-func readLines(t *testing.T, path, pattern string) []string {
- t.Helper()
-
- body, err := os.ReadFile(filepath.FromSlash(path))
- require.NoError(t, err)
- matches :=
regexp.MustCompile(pattern).FindAllStringSubmatch(string(body), -1)
- require.NotEmpty(t, matches, "%s matched nothing, so the file's shape
has changed", path)
- values := make([]string, 0, len(matches))
- for _, match := range matches {
- values = append(values, match[1])
- }
- return values
-}
-
-// TestRuleConstantsMatchPython is the tripwire for a trigger rule or a weight
rule
-// added or renamed on the Python side. The serialization schema types both
fields as a
-// plain string and names none of their values, so the constants here are
hand-written
-// and the spec drift check cannot see them; Airflow rejects a rule it does
not know, so
-// a missing constant is a Dag that fails to register.
-func TestRuleConstantsMatchPython(t *testing.T) {
- for _, tt := range []struct {
- goType string
- pythonPath string
- }{
- {goType: "TriggerRule", pythonPath:
"../../airflow-core/src/airflow/task/trigger_rule.py"},
- {goType: "WeightRule", pythonPath:
"../../airflow-core/src/airflow/task/weight_rule.py"},
- } {
- t.Run(tt.goType, func(t *testing.T) {
- python := readLines(t, tt.pythonPath, `(?m)^ [A-Z_]+
= "([a-z_]+)"$`)
- spelled := readLines(
- t,
- specPath,
- `(?m)^\t`+tt.goType+`\w+ +`+tt.goType+` =
"([a-z_]+)"$`,
- )
-
- assert.ElementsMatch(t, python, spelled)
- })
- }
-}
diff --git a/go-sdk/airflow/trigger_dag_run.go
b/go-sdk/airflow/trigger_dag_run.go
index dcec75a9cca..d6ede6f724c 100644
--- a/go-sdk/airflow/trigger_dag_run.go
+++ b/go-sdk/airflow/trigger_dag_run.go
@@ -24,8 +24,6 @@ import (
"fmt"
"slices"
"time"
-
- "github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
)
// TriggerDagRunSpec holds the options of a task that triggers a Dag run.
[TriggerDagRun] takes a
@@ -61,16 +59,15 @@ type TriggerDagRunSpec struct {
// PokeInterval is how often a task that waits checks the state of the
new Dag run. It must be
// a whole number of seconds. When PokeInterval is nil, the task checks
every 60 seconds.
PokeInterval *time.Duration
- // AllowedStates are the states of the new Dag run in which a task that
waits succeeds. Each
- // state is one of queued, running, success and failed. When
AllowedStates is empty, the task
- // succeeds in the success state.
- AllowedStates []string
- // FailedStates are the states of the new Dag run in which a task that
waits fails. Each state
- // is one of queued, running, success and failed. A nil FailedStates
fails the task in the
- // failed state. A FailedStates that is empty but not nil means that no
state fails the task.
- // The task then succeeds once the new Dag run is in a state that
AllowedStates lists. In any
- // other state, including failed, the task keeps waiting.
- FailedStates []string
+ // AllowedStates are the states of the new Dag run in which a task that
waits succeeds. When
+ // AllowedStates is empty, the task succeeds in the success state of
the new Dag run.
+ AllowedStates []DagRunState
+ // FailedStates are the states of the new Dag run in which a task that
waits fails. When
+ // FailedStates is nil, the task fails in the failed state of the new
Dag run. When
+ // FailedStates is empty but not nil, no Dag run state fails the task.
The task then succeeds
+ // once the new Dag run is in a state that AllowedStates lists. In any
other state, including
+ // failed, the task keeps waiting.
+ FailedStates []DagRunState
// SkipWhenAlreadyExists marks the task skipped if the Dag run already
exists.
SkipWhenAlreadyExists bool
// FailWhenDagIsPaused fails the task when the Dag to trigger is paused.
@@ -108,15 +105,6 @@ func TriggerDagRun(spec TriggerDagRunSpec)
TriggerDagRunTask {
return TriggerDagRunTask{spec: spec}
}
-// validDagRunStates are the values that TriggerDagRunOperator accepts in
allowed_states and
-// failed_states.
-var validDagRunStates = []string{
- string(genmodels.DagRunStateQueued),
- string(genmodels.DagRunStateRunning),
- string(genmodels.DagRunStateSuccess),
- string(genmodels.DagRunStateFailed),
-}
-
// copyTriggerDagRunSpec checks spec and returns a deep copy of it, so that
nothing the caller
// still holds, such as Conf, a state slice or a pointer field, can change the
task that
// DagRef.Task added.
@@ -133,14 +121,14 @@ func copyTriggerDagRunSpec(spec TriggerDagRunSpec)
(TriggerDagRunSpec, error) {
}
for _, field := range []struct {
name string
- states []string
+ states []DagRunState
}{{"AllowedStates", spec.AllowedStates}, {"FailedStates",
spec.FailedStates}} {
for _, state := range field.states {
- if !slices.Contains(validDagRunStates, state) {
+ if !slices.Contains(dagRunStates, state) {
return TriggerDagRunSpec{}, fmt.Errorf(
"airflow.TriggerDagRunSpec.%s has %q,
which is not a Dag run state; "+
"use one of %q",
- field.name, state, validDagRunStates,
+ field.name, state, dagRunStates,
)
}
}
diff --git a/go-sdk/airflow/trigger_dag_run_test.go
b/go-sdk/airflow/trigger_dag_run_test.go
index b812f146cec..d7cf45b92ec 100644
--- a/go-sdk/airflow/trigger_dag_run_test.go
+++ b/go-sdk/airflow/trigger_dag_run_test.go
@@ -40,8 +40,8 @@ func TestTriggerDagRunIsATask(t *testing.T) {
ResetDagRun: true,
WaitForCompletion: true,
PokeInterval: ptr(30 * time.Second),
- AllowedStates: []string{"success", "failed"},
- FailedStates: []string{"queued", "running"},
+ AllowedStates: []DagRunState{DagRunStateSuccess,
DagRunStateFailed},
+ FailedStates: []DagRunState{DagRunStateQueued,
DagRunStateRunning},
SkipWhenAlreadyExists: true,
FailWhenDagIsPaused: true,
Note: "triggered by etl",
@@ -71,8 +71,8 @@ func TestTriggerDagRunKeepsNilApartFromZero(t *testing.T) {
DagID: "downstream_etl",
Conf: map[string]any{},
PokeInterval: ptr(time.Duration(0)),
- AllowedStates: []string{},
- FailedStates: []string{},
+ AllowedStates: []DagRunState{},
+ FailedStates: []DagRunState{},
Deferrable: ptr(false),
}
task = dag.Task(TriggerDagRun(zero), TaskSpec{TaskID: "zero"})
@@ -130,7 +130,7 @@ func TestTriggerDagRunRejectsAnInvalidSpec(t *testing.T) {
{
name: "unknown state in AllowedStates",
trigger: TriggerDagRun(TriggerDagRunSpec{
- DagID: "downstream_etl", AllowedStates:
[]string{"success", "SUCCESS"},
+ DagID: "downstream_etl", AllowedStates:
[]DagRunState{DagRunStateSuccess, "SUCCESS"},
}),
want: `airflow.TriggerDagRunSpec.AllowedStates has
"SUCCESS", which is not a Dag ` +
`run state; use one of ["queued" "running"
"success" "failed"]`,
@@ -138,7 +138,7 @@ func TestTriggerDagRunRejectsAnInvalidSpec(t *testing.T) {
{
name: "unknown state in FailedStates",
trigger: TriggerDagRun(TriggerDagRunSpec{
- DagID: "downstream_etl", FailedStates:
[]string{"skipped"},
+ DagID: "downstream_etl", FailedStates:
[]DagRunState{"skipped"},
}),
want: `airflow.TriggerDagRunSpec.FailedStates has
"skipped", which is not a Dag ` +
`run state; use one of ["queued" "running"
"success" "failed"]`,
@@ -175,8 +175,8 @@ func TestTriggerDagRunRejectsAnInvalidSpec(t *testing.T) {
func TestTriggerDagRunCopiesTheSpec(t *testing.T) {
nested := map[string]any{"table": "rows"}
- allowed := []string{"success"}
- failed := []string{"failed"}
+ allowed := []DagRunState{DagRunStateSuccess}
+ failed := []DagRunState{DagRunStateFailed}
poke := 30 * time.Second
deferrable := true
spec := TriggerDagRunSpec{
@@ -192,8 +192,8 @@ func TestTriggerDagRunCopiesTheSpec(t *testing.T) {
nested["table"] = "changed"
spec.Conf["added"] = true
- allowed[0] = "running"
- failed[0] = "queued"
+ allowed[0] = DagRunStateRunning
+ failed[0] = DagRunStateQueued
poke = time.Minute
deferrable = false
@@ -205,8 +205,8 @@ func TestTriggerDagRunCopiesTheSpec(t *testing.T) {
},
stored.Conf,
)
- assert.Equal(t, []string{"success"}, stored.AllowedStates)
- assert.Equal(t, []string{"failed"}, stored.FailedStates)
+ assert.Equal(t, []DagRunState{DagRunStateSuccess}, stored.AllowedStates)
+ assert.Equal(t, []DagRunState{DagRunStateFailed}, stored.FailedStates)
assert.Equal(t, 30*time.Second, *stored.PokeInterval)
assert.True(t, *stored.Deferrable)
}
diff --git a/go-sdk/internal/genspec/authoring.go
b/go-sdk/internal/genspec/authoring.go
index 4eac1c5ccc1..d515065a5d5 100644
--- a/go-sdk/internal/genspec/authoring.go
+++ b/go-sdk/internal/genspec/authoring.go
@@ -81,7 +81,7 @@ var dagShape = authoringShape{
"dag_dependencies": "derived from the edges and the
assets a Dag declares",
"timezone": "carried by the time.Time an
author sets on StartDate",
"timetable": "the serialized form of Schedule,
which is injected instead",
- "allowed_run_types": "a union the author expresses by
setting Schedule",
+ "allowed_run_types": "no Go authoring type yet: a list
of DagRunType values",
"_concurrency": "the pre-2.2 spelling of
MaxActiveTasks",
"has_on_success_callback": "derived from whether a callback
is registered",
"has_on_failure_callback": "derived from whether a callback
is registered",
@@ -171,8 +171,8 @@ var taskShape = authoringShape{
"max_active_tis_per_dag": {goType: "int"},
"max_active_tis_per_dagrun": {goType: "int"},
// TriggerRule, WeightRule and their constants are hand-written
in the airflow
- // package: the schema types both fields as a plain string and
names none of
- // the values either one takes.
+ // package, because the schema types trigger_rule and
weight_rule as plain
+ // strings and does not list their values.
"trigger_rule": {goType: "TriggerRule"},
"weight_rule": {goType: "WeightRule"},
// A multiplier, not a switch: 0 keeps the delay constant, 2.0
doubles it each
