mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
Go SDK: add airflow.Dag and dag.Task, the native Dag constructors (#73767)
* Go SDK: add airflow.Dag and dag.Task, the native Dag constructors Signed-off-by: PoAn Yang <payang@apache.org> * Make the DagSpec of airflow.Dag optional Signed-off-by: PoAn Yang <payang@apache.org> * Go SDK: Change dagId to dagID and taskIds to taskIDs Signed-off-by: PoAn Yang <payang@apache.org> --------- Signed-off-by: PoAn Yang <payang@apache.org>
This commit is contained in:
PoAn Yang authored and 
GitHub committed
1 parent
5428b3b995
commit
232255e7fa
12 files changed
+891
-31
No files matched your search
@@ -75,7 +75,7 @@ pick := dag.Switch(pickPath) // task_id pickPath, from the function name
|
||||
pick.Case(handleLongRef).
|
||||
Case(handleShortRef)
|
||||
|
||||
dag.Task(airflow.TriggerDagRun(airflow.TriggerDagRunSpec{DagId: "downstream_etl"}), airflow.TaskSpec{TaskId: "trigger_downstream"}).After(gate)
|
||||
dag.Task(airflow.TriggerDagRun(airflow.TriggerDagRunSpec{DagID: "downstream_etl"}), airflow.TaskSpec{TaskID: "trigger_downstream"}).After(gate)
|
||||
```
|
||||
|
||||
A decider has to see the refs it returns, so either they are package-level or it is a closure where the Dag is built.
|
||||
|
||||
@@ -61,11 +61,11 @@ extracted := dag.Task(extract) // native Go: runs on a Go worker
|
||||
|
||||
staged := dag.Task(amazon.S3ToRedshiftOperator{
|
||||
SchemaName: "public", TableName: "events", S3Bucket: "raw", S3Key: "events/{{ ds }}",
|
||||
}, airflow.TaskSpec{TaskId: "stage"}).After(extracted) // DSL only: runs on a Python worker
|
||||
}, airflow.TaskSpec{TaskID: "stage"}).After(extracted) // DSL only: runs on a Python worker
|
||||
|
||||
dag.Task(kubernetes.KubernetesPodOperator{
|
||||
Namespace: "airflow", Image: "report:latest", Name: "report",
|
||||
}, airflow.TaskSpec{TaskId: "report"}).After(staged)
|
||||
}, airflow.TaskSpec{TaskID: "report"}).After(staged)
|
||||
```
|
||||
|
||||
`airflow.TriggerDagRun` ([ADR-0008](0008-control-flow-constructs.md)) follows the same concept but it is the hand-written.
|
||||
|
||||
@@ -27,11 +27,11 @@ Proposed.
|
||||
|
||||
## Decision
|
||||
|
||||
1. **One Dag type, constructed then registered.** `airflow.Dag(dagId, spec)` returns a `*airflow.DagRef` that is complete before `bundle.Register(dag)` takes it — the same verb that registers Mixed Lang task handlers ([ADR 7](0007-mixed-lang-task-handler-interface.md)).
|
||||
1. **One Dag type, constructed then registered.** `airflow.Dag(dagID, spec)` returns a `*airflow.DagRef` that is complete before `bundle.Register(dag)` takes it — the same verb that registers Mixed Lang task handlers ([ADR 7](0007-mixed-lang-task-handler-interface.md)).
|
||||
Naming rule: `airflow.X(...)` constructs, `*airflow.XRef` is the entity; what every Dag must have is a positional parameter for dag_id, and the rest travels in a spec struct.
|
||||
2. **Tasks register through `dag.Task(fn any, opts ...airflow.TaskOption)`**, returning a `*airflow.TaskRef`. `airflow.Inputs(...)` and a bare `airflow.TaskSpec{}` both implement `TaskOption`.
|
||||
3. **At most one `airflow.TaskSpec` per task.** A second one is a registration error rather than something to merge, so a task's attributes are only ever written in one place.
|
||||
4. **task_id is the Go function name by default**, spelled exactly as the function is, and `airflow.TaskSpec{TaskId: ...}` sets it to anything else.
|
||||
4. **task_id is the Go function name by default**, spelled exactly as the function is, and `airflow.TaskSpec{TaskID: ...}` sets it to anything else.
|
||||
5. **`airflow.Inputs(refs...)` declares the data and the edge in one call** for defining graph with TaskFlow syntax.
|
||||
6. **`Before` and `After` are order-only edges on `airflow.Node`**, which both `*airflow.TaskRef` and `*airflow.TaskGroupRef` satisfy. They are the Go pair for `>>` and `<<`, and both return their argument set as one `Node`, so `a.Before(b, c).Before(d)` is Python's `a >> [b, c] >> d`.
|
||||
7. **An edge label wraps the endpoint**: `loaded.Before(airflow.Label(notify, "when empty"))` is Python's `loaded >> Label("when empty") >> notify`. Labelling the endpoint rather than the call lets one fan-out give each edge its own label.
|
||||
@@ -91,7 +91,7 @@ The task_ids are the function names by default: `extract`, `transform`, and `loa
|
||||
loaded := dag.Task(load, airflow.Inputs(transformed))
|
||||
cleaned := dag.Task(cleanup, airflow.TaskSpec{TriggerRule: airflow.TriggerRuleAllDone})
|
||||
notified := dag.Task(notify)
|
||||
emptyNotice := dag.Task(notifyEmpty, airflow.TaskSpec{TaskId: "notify_empty"})
|
||||
emptyNotice := dag.Task(notifyEmpty, airflow.TaskSpec{TaskID: "notify_empty"})
|
||||
staging := dag.TaskGroup("staging") // a group carries edges like a task
|
||||
staging.Task(stageRows) // tasks join a group through the group
|
||||
|
||||
@@ -107,13 +107,13 @@ loaded.Before(airflow.Label(emptyNotice, "when empty")) // loaded >> Label("whe
|
||||
```go
|
||||
package airflow
|
||||
|
||||
func Dag(dagId string, spec DagSpec) *DagRef
|
||||
func Dag(dagID string, spec ...DagSpec) *DagRef
|
||||
|
||||
func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef
|
||||
func (d *DagRef) TaskGroup(groupId string, opts ...TaskGroupOption) *TaskGroupRef
|
||||
func (d *DagRef) TaskGroup(groupID string, opts ...TaskGroupOption) *TaskGroupRef
|
||||
|
||||
func (g *TaskGroupRef) Task(fn any, opts ...TaskOption) *TaskRef
|
||||
func (g *TaskGroupRef) TaskGroup(groupId string, opts ...TaskGroupOption) *TaskGroupRef
|
||||
func (g *TaskGroupRef) TaskGroup(groupID string, opts ...TaskGroupOption) *TaskGroupRef
|
||||
|
||||
// DagSpec and TaskSpec are generated into this package from
|
||||
// airflow-core/src/airflow/serialization/schema.json and committed.
|
||||
@@ -127,7 +127,7 @@ type DagSpec struct {
|
||||
|
||||
// TaskSpec implements TaskOption, so it travels in the same variadic as Inputs.
|
||||
type TaskSpec struct {
|
||||
TaskId string // defaults to the Go function name
|
||||
TaskID string // defaults to the Go function name
|
||||
Retries int
|
||||
TriggerRule TriggerRule
|
||||
// ...
|
||||
@@ -177,7 +177,7 @@ const (
|
||||
Generation needs an exclusion list and a hand-written field or two, the same kind of rule [ADR-0009](../../airflow-core/adr/lang-sdk/0009-provider-operators-as-generated-dsl.md) states for provider operators.
|
||||
- **A data edge is labelled by redeclaring it.** Declaring an edge that already exists is idempotent, so `extracted.Before(airflow.Label(transformed, "rows"))` labels the edge `Inputs` created.
|
||||
- **Renaming a Go function renames the task.** The id is derived, and history, clears, and the UI all key on task_id, so renaming a function whose task carries no `TaskSpec` id is a Dag change.
|
||||
`airflow.TaskSpec{TaskId: ...}` pins an id that has to outlive the function's name, and it is also how a Dag gets snake_case ids, since nothing transforms a Go name.
|
||||
`airflow.TaskSpec{TaskID: ...}` pins an id that has to outlive the function's name, and it is also how a Dag gets snake_case ids, since nothing transforms a Go name.
|
||||
|
||||
## Alternatives
|
||||
|
||||
@@ -187,5 +187,5 @@ const (
|
||||
callable per run, so an edge existing only in execution order cannot be parsed without running the
|
||||
program to completion. `Inputs` keeps the typed outputs that style is reached for.
|
||||
- **Labelling the call**, `loaded.Before(notify).Label("when empty")`. Rejected: one call fans out, so a single label on the call cannot give `a.Before(b, c)` a different label per edge.
|
||||
- **A positional task_id**, `dag.Task(taskId, fn, opts...)`. Rejected: It's more straightforward and native for Go user to define a Task without defining the task_id explicitly. They could still set the task_id other than the function name in the TaskSpec.
|
||||
- **A positional task_id**, `dag.Task(taskID, fn, opts...)`. Rejected: It's more straightforward and native for Go user to define a Task without defining the task_id explicitly. They could still set the task_id other than the function name in the TaskSpec.
|
||||
- **Separate `Dag` and `MixedLangDag` types.** Rejected: Python has one Dag class, and the Mixed Lang case is not a Dag at all ([ADR 7](0007-mixed-lang-task-handler-interface.md)).
|
||||
@@ -26,13 +26,14 @@ import (
|
||||
"github.com/apache/airflow/go-sdk/internal/bundle"
|
||||
)
|
||||
|
||||
// BundleRef holds the task handlers that this executable runs for Airflow.
|
||||
// BundleRef holds the task handlers and the Dags that this executable registers for Airflow.
|
||||
// [Bundle] returns an empty one.
|
||||
type BundleRef struct {
|
||||
// closed ends registration for everything the bundle can hold, so a kind added later
|
||||
// is covered without a flag of its own. Serve sets it; Register reads it.
|
||||
closed atomic.Bool
|
||||
taskHandlers taskHandlerMap
|
||||
dags dagMap
|
||||
}
|
||||
|
||||
// Bundle returns an empty bundle. Register the task handlers on it, then call Serve as the
|
||||
@@ -51,7 +52,8 @@ type BundleRef struct {
|
||||
// }
|
||||
func Bundle() *BundleRef { return &BundleRef{} }
|
||||
|
||||
// Registerable is what [BundleRef.Register] accepts. [TaskHandler] returns one.
|
||||
// Registerable is what [BundleRef.Register] accepts: the value [TaskHandler] returns, or the
|
||||
// [*DagRef] that [Dag] returns.
|
||||
//
|
||||
// Its only method is unexported, so a type outside this package cannot declare it.
|
||||
// A struct that embeds a Registerable still satisfies the interface, and Register panics
|
||||
@@ -70,8 +72,12 @@ type Registerable interface{ registerable() }
|
||||
//
|
||||
// bundle.Register(reports.Handlers()...)
|
||||
//
|
||||
// Add every task to a Dag before registering the Dag. [DagRef.Task] panics once the Dag is
|
||||
// registered.
|
||||
//
|
||||
// Register panics if a task handler with the same dag_id and task_id is already registered,
|
||||
// and if [BundleRef.Serve] has already been called: registration closes when serving starts.
|
||||
// if a Dag with the same dag_id is already registered, or if [BundleRef.Serve] has already
|
||||
// been called: registration closes when serving starts.
|
||||
func (b *BundleRef) Register(items ...Registerable) {
|
||||
if b.closed.Load() {
|
||||
panic(
|
||||
@@ -82,6 +88,11 @@ func (b *BundleRef) Register(items ...Registerable) {
|
||||
switch item := item.(type) {
|
||||
case *taskHandler:
|
||||
b.taskHandlers.add(item.dagId, item.taskId, item.task)
|
||||
case *DagRef:
|
||||
if item == nil {
|
||||
panic("airflow.BundleRef.Register: cannot register a nil *airflow.DagRef")
|
||||
}
|
||||
b.dags.add(item)
|
||||
default:
|
||||
// Either a nil item, or a struct from another package that embeds a Registerable.
|
||||
panic(fmt.Sprintf("airflow.BundleRef.Register: cannot register %T", item))
|
||||
@@ -138,3 +149,23 @@ func (m *taskHandlerMap) ListTaskHandlers() []bundle.TaskHandlerInfo {
|
||||
|
||||
return slices.Clone(m.order)
|
||||
}
|
||||
|
||||
// dagMap holds the registered Dags by dag_id.
|
||||
type dagMap struct {
|
||||
mu sync.Mutex
|
||||
dags map[string]*DagRef
|
||||
}
|
||||
|
||||
func (m *dagMap) add(dag *DagRef) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if _, exists := m.dags[dag.dagID]; exists {
|
||||
panic(fmt.Sprintf("airflow.BundleRef.Register: Dag %q is already registered", dag.dagID))
|
||||
}
|
||||
dag.markRegistered()
|
||||
if m.dags == nil {
|
||||
m.dags = make(map[string]*DagRef)
|
||||
}
|
||||
m.dags[dag.dagID] = dag
|
||||
}
|
||||
@@ -117,16 +117,48 @@ func TestRegisterRejectsDuplicateTask(t *testing.T) {
|
||||
)
|
||||
}
|
||||
|
||||
func TestRegisterTakesTaskHandlersAndDagsTogether(t *testing.T) {
|
||||
etl := Dag("etl")
|
||||
etl.Task(extract)
|
||||
|
||||
b := Bundle()
|
||||
b.Register(TaskHandler("py_etl", "transform", noop), etl)
|
||||
|
||||
_, ok := b.taskHandlers.LookupTask("py_etl", "transform")
|
||||
assert.True(t, ok)
|
||||
assert.Same(t, etl, b.dags.dags["etl"])
|
||||
assert.True(t, etl.registered)
|
||||
}
|
||||
|
||||
func TestRegisterRejectsDuplicateDag(t *testing.T) {
|
||||
etl := Dag("etl")
|
||||
b := Bundle()
|
||||
b.Register(etl)
|
||||
|
||||
want := `airflow.BundleRef.Register: Dag "etl" is already registered`
|
||||
assert.PanicsWithValue(t, want, func() { b.Register(etl) })
|
||||
second := Dag("etl")
|
||||
assert.PanicsWithValue(t, want, func() { b.Register(second) })
|
||||
assert.Same(t, etl, b.dags.dags["etl"])
|
||||
assert.False(t, second.registered, "a Dag that Register rejects can still take tasks")
|
||||
}
|
||||
|
||||
func TestRegisterRejectsNilDag(t *testing.T) {
|
||||
var dag *DagRef
|
||||
assert.PanicsWithValue(t, "airflow.BundleRef.Register: cannot register a nil *airflow.DagRef",
|
||||
func() { Bundle().Register(dag) },
|
||||
)
|
||||
}
|
||||
|
||||
func TestRegisterAfterServePanics(t *testing.T) {
|
||||
b := Bundle()
|
||||
b.Register(TaskHandler("py_etl", "transform", noop))
|
||||
require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard))
|
||||
|
||||
assert.PanicsWithValue(
|
||||
t,
|
||||
"airflow.BundleRef.Register: Serve has already been called; register everything before Serve",
|
||||
func() { b.Register(TaskHandler("py_etl", "load", noop)) },
|
||||
)
|
||||
want := "airflow.BundleRef.Register: Serve has already been called; " +
|
||||
"register everything before Serve"
|
||||
assert.PanicsWithValue(t, want, func() { b.Register(TaskHandler("py_etl", "load", noop)) })
|
||||
assert.PanicsWithValue(t, want, func() { b.Register(Dag("etl")) })
|
||||
}
|
||||
|
||||
// The flag lives on the bundle, not on the task-handler map, so a kind of item added to
|
||||
@@ -160,6 +192,7 @@ func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
|
||||
defer wg.Done()
|
||||
for i := range perWorker {
|
||||
b.Register(TaskHandler("py_etl", fmt.Sprintf("task_%d_%d", worker, i), noop))
|
||||
b.Register(Dag(fmt.Sprintf("dag_%d_%d", worker, i)))
|
||||
b.taskHandlers.LookupTask("py_etl", "task_0_0")
|
||||
b.taskHandlers.ListTaskHandlers()
|
||||
}
|
||||
@@ -168,6 +201,7 @@ func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
|
||||
wg.Wait()
|
||||
|
||||
assert.Len(t, b.taskHandlers.ListTaskHandlers(), workers*perWorker)
|
||||
assert.Len(t, b.dags.dags, workers*perWorker)
|
||||
}
|
||||
|
||||
func TestRegisterRejectsNilItem(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
// 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"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// DagRef is a Dag authored in Go. [Dag] returns a new one.
|
||||
type DagRef struct {
|
||||
dagID string
|
||||
// Dag and Task copy the specs they are given, so a caller cannot change a registered Dag
|
||||
// through a spec it still holds. A slice or map field in DagSpec or TaskSpec would share its
|
||||
// contents with the caller, so Dag and Task would have to copy that field too.
|
||||
spec DagSpec
|
||||
|
||||
mu sync.Mutex
|
||||
registered bool
|
||||
tasks []*TaskRef
|
||||
taskIDs map[string]struct{}
|
||||
}
|
||||
|
||||
// Dag returns an empty Dag with the given dag_id. An optional [DagSpec] holds the rest of the
|
||||
// Dag's attributes, and Dag panics if it gets more than one DagSpec. Add the tasks with
|
||||
// [DagRef.Task], then pass the Dag to [BundleRef.Register]:
|
||||
//
|
||||
// dag := airflow.Dag("etl")
|
||||
// dag.Task(extract)
|
||||
// dag.Task(load, airflow.TaskSpec{TaskID: "load_rows"})
|
||||
//
|
||||
// bundle.Register(dag)
|
||||
//
|
||||
// Add every task before Register. [DagRef.Task] panics once the Dag is registered.
|
||||
//
|
||||
// [BundleRef.Serve] does not yet serve the Dags that Dag returns. It leaves them out of the
|
||||
// --airflow-metadata manifest and cannot run their tasks.
|
||||
func Dag(dagID string, spec ...DagSpec) *DagRef {
|
||||
if len(spec) > 1 {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.Dag: Dag %q got %d airflow.DagSpec values; "+
|
||||
"set all of the Dag's attributes in one DagSpec",
|
||||
dagID, len(spec),
|
||||
))
|
||||
}
|
||||
d := &DagRef{dagID: dagID}
|
||||
if len(spec) == 1 {
|
||||
d.spec = spec[0]
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func (*DagRef) registerable() {}
|
||||
|
||||
// TaskRef is a task that [DagRef.Task] added to a Dag.
|
||||
type TaskRef struct {
|
||||
taskID string
|
||||
spec TaskSpec
|
||||
}
|
||||
|
||||
// Task adds a task that runs fn to the Dag and returns the new task.
|
||||
//
|
||||
// fn takes a [Context] first and returns either error or (result, error), like a function
|
||||
// passed to [TaskHandler].
|
||||
//
|
||||
// The task_id is the name of fn, spelled exactly as it is in Go. dag.Task(extractRows) adds the
|
||||
// task extractRows, and dag.Task(svc.Extract), which passes a method value, adds the task
|
||||
// Extract. A task needs a TaskID when Task cannot read a name from fn. That is the case for a
|
||||
// function literal, for a function that package reflect made, such as one from
|
||||
// reflect.MakeFunc, and for a method value that a generic function takes on a value whose type
|
||||
// uses a type parameter.
|
||||
//
|
||||
// A [TaskSpec] sets the task_id and the other attributes of a task. Pass at most one, so that
|
||||
// each attribute is set in one place:
|
||||
//
|
||||
// dag.Task(extractRows, airflow.TaskSpec{TaskID: "extract_rows"})
|
||||
//
|
||||
// Renaming fn renames a task that takes its task_id from fn. Airflow identifies a task by its
|
||||
// task_id in the run history, in clears and in the UI, so to Airflow the renamed function runs
|
||||
// a new task. Set TaskID when the task_id has to stay the same through such a rename.
|
||||
//
|
||||
// Task panics if:
|
||||
// - fn is not a valid task function
|
||||
// - fn has no name that Task can read and no TaskSpec sets a TaskID
|
||||
// - an option is nil or is not one that package airflow defines
|
||||
// - opts holds more than one TaskSpec
|
||||
// - the Dag already has a task with the same task_id
|
||||
// - the Dag is already registered
|
||||
func (d *DagRef) Task(fn any, opts ...TaskOption) *TaskRef {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
|
||||
if d.registered {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: Dag %q has already been registered; "+
|
||||
"add every task before Register",
|
||||
d.dagID,
|
||||
))
|
||||
}
|
||||
if _, err := newTaskFunction(fn); err != nil {
|
||||
panic(fmt.Sprintf("airflow.DagRef.Task: Dag %q: %v", d.dagID, err))
|
||||
}
|
||||
var cfg taskConfig
|
||||
for i, opt := range opts {
|
||||
switch opt := opt.(type) {
|
||||
case nil:
|
||||
panic(fmt.Sprintf("airflow.DagRef.Task: Dag %q: opts[%d] is nil", d.dagID, i))
|
||||
case *TaskSpec:
|
||||
if opt == nil {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: Dag %q: opts[%d] is a nil *airflow.TaskSpec", d.dagID, i,
|
||||
))
|
||||
}
|
||||
case TaskSpec:
|
||||
default:
|
||||
// Only a struct that embeds a TaskSpec or a TaskOption gets here.
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: Dag %q: opts[%d] has type %T, "+
|
||||
"which is not an option that package airflow defines",
|
||||
d.dagID, i, opt,
|
||||
))
|
||||
}
|
||||
opt.applyTask(&cfg)
|
||||
}
|
||||
if len(cfg.specs) > 1 {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: task %q of Dag %q got %d airflow.TaskSpec values; "+
|
||||
"set all of the task's attributes in one TaskSpec",
|
||||
findTaskName(fn, cfg.specs), d.dagID, len(cfg.specs),
|
||||
))
|
||||
}
|
||||
|
||||
var spec TaskSpec
|
||||
if len(cfg.specs) == 1 {
|
||||
spec = cfg.specs[0]
|
||||
}
|
||||
taskID := spec.TaskID
|
||||
if taskID == "" {
|
||||
var ok bool
|
||||
if taskID, ok = taskIDFromFuncName(funcName(fn)); !ok {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: Dag %q: %s has no name to use as the task_id; "+
|
||||
"set one with airflow.TaskSpec{TaskID: ...}",
|
||||
d.dagID, funcName(fn),
|
||||
))
|
||||
}
|
||||
}
|
||||
if _, exists := d.taskIDs[taskID]; exists {
|
||||
panic(fmt.Sprintf(
|
||||
"airflow.DagRef.Task: Dag %q already has a task %q; "+
|
||||
"set another task_id with airflow.TaskSpec{TaskID: ...}",
|
||||
d.dagID, taskID,
|
||||
))
|
||||
}
|
||||
if d.taskIDs == nil {
|
||||
d.taskIDs = make(map[string]struct{})
|
||||
}
|
||||
d.taskIDs[taskID] = struct{}{}
|
||||
|
||||
task := &TaskRef{taskID: taskID, spec: spec}
|
||||
d.tasks = append(d.tasks, task)
|
||||
return task
|
||||
}
|
||||
|
||||
func (d *DagRef) markRegistered() {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
|
||||
d.registered = true
|
||||
}
|
||||
|
||||
func funcName(fn any) string { return runtime.FuncForPC(reflect.ValueOf(fn).Pointer()).Name() }
|
||||
|
||||
// findTaskName names a task in an error that Task raises before it settles the task_id.
|
||||
func findTaskName(fn any, specs []TaskSpec) string {
|
||||
for _, spec := range specs {
|
||||
if spec.TaskID != "" {
|
||||
return spec.TaskID
|
||||
}
|
||||
}
|
||||
if taskID, ok := taskIDFromFuncName(funcName(fn)); ok {
|
||||
return taskID
|
||||
}
|
||||
return funcName(fn)
|
||||
}
|
||||
|
||||
// taskIDFromFuncName takes the runtime name of a function and returns the name that the
|
||||
// function is declared with. It reports false when the runtime name does not carry one.
|
||||
func taskIDFromFuncName(name string) (string, bool) {
|
||||
// Package reflect runs every function it makes through a stub of its own, such as
|
||||
// reflect.makeFuncStub, so the name is the stub's.
|
||||
if strings.HasPrefix(name, "reflect.") {
|
||||
return "", false
|
||||
}
|
||||
// A method value such as svc.Extract is named after the receiver type of the method, as in
|
||||
// (*Service).Extract-fm, and the last segment is the method name.
|
||||
if method, ok := strings.CutSuffix(name, "-fm"); ok {
|
||||
return method[strings.LastIndex(method, ".")+1:], true
|
||||
}
|
||||
// The runtime prefixes the name with the import path, and it escapes a dot in the last
|
||||
// element of that path as %2e. So the first dot after the last slash ends the package name.
|
||||
name = name[strings.LastIndex(name, "/")+1:]
|
||||
name = name[strings.Index(name, ".")+1:]
|
||||
// The runtime writes the type arguments of a generic function as [...].
|
||||
name = strings.ReplaceAll(name, "[...]", "")
|
||||
// Any dot left comes from a function literal, which the runtime names after the function it
|
||||
// is declared in, as in main.func1. The compiler builds a function literal for a method value
|
||||
// that a generic function takes on a value whose type uses a type parameter, so that case
|
||||
// ends up here too. A method expression such as Service.Extract leaves a dot as well, but
|
||||
// Task rejects most method expressions earlier, because the receiver is their first
|
||||
// parameter.
|
||||
if strings.Contains(name, ".") {
|
||||
return "", false
|
||||
}
|
||||
return name, true
|
||||
}
|
||||
@@ -0,0 +1,461 @@
|
||||
// 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 (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func extract(Context) error { return nil }
|
||||
|
||||
func ExtractRows(Context) error { return nil }
|
||||
|
||||
func extractAny[T any](Context) error { return nil }
|
||||
|
||||
type extractor struct{}
|
||||
|
||||
func (extractor) Extract(Context) error { return nil }
|
||||
|
||||
func (*extractor) ExtractByPointer(Context) error { return nil }
|
||||
|
||||
var packageLevelLiteral = func(Context) error { return nil }
|
||||
|
||||
func addExtract[T interface{ Extract(Context) error }](
|
||||
dag *DagRef,
|
||||
src T,
|
||||
opts ...TaskOption,
|
||||
) *TaskRef {
|
||||
return dag.Task(src.Extract, opts...)
|
||||
}
|
||||
|
||||
// Embedding promotes applyTask, so these structs compile as a TaskOption.
|
||||
type (
|
||||
wrappedSpec struct {
|
||||
TaskSpec
|
||||
MaxRetries int
|
||||
}
|
||||
wrappedOption struct{ TaskOption }
|
||||
)
|
||||
|
||||
func TestTaskIDIsTheFunctionName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fn any
|
||||
want string
|
||||
}{
|
||||
{name: "function", fn: extract, want: "extract"},
|
||||
{name: "spelling kept", fn: ExtractRows, want: "ExtractRows"},
|
||||
{name: "generic function", fn: extractAny[int], want: "extractAny"},
|
||||
{name: "method value", fn: extractor{}.Extract, want: "Extract"},
|
||||
{
|
||||
name: "pointer method value",
|
||||
fn: (&extractor{}).ExtractByPointer,
|
||||
want: "ExtractByPointer",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
task := Dag("etl").Task(tt.fn)
|
||||
assert.Equal(t, tt.want, task.taskID)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The runtime name in each row has the form that runtime.FuncForPC reports for the kind of
|
||||
// function that the row names.
|
||||
func TestTaskIDFromFuncName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
runtimeName string
|
||||
want string
|
||||
ok bool
|
||||
}{
|
||||
{"function in main", "main.extract", "extract", true},
|
||||
{
|
||||
"function in another package",
|
||||
"github.com/acme/etl/tasks.ExtractRows",
|
||||
"ExtractRows",
|
||||
true,
|
||||
},
|
||||
{"dots in the import path", "example.com/a.b/c%2ed.Extract", "Extract", true},
|
||||
{"generic function", "main.process[...]", "process", true},
|
||||
{"pointer-receiver method value", "main.(*Service).Extract-fm", "Extract", true},
|
||||
{"value-receiver method value", "main.Service.Extract-fm", "Extract", true},
|
||||
{"method value on a generic type", "example.com/etl.(*Box[...]).Run-fm", "Run", true},
|
||||
{
|
||||
"method value on an interface literal",
|
||||
"go:interface { Run(github.com/apache/airflow/go-sdk/airflow.Context) error }.Run-fm",
|
||||
"Run",
|
||||
true,
|
||||
},
|
||||
{"function literal", "main.main.func1", "", false},
|
||||
{"function literal in a package variable", "main.init.func1", "", false},
|
||||
{"inlined function literal", "main.main.makeClosure.func2", "", false},
|
||||
{"method value taken on a type parameter", "main.addSource[...].func1", "", false},
|
||||
{"method expression", "main.Service.Extract", "", false},
|
||||
{"reflect.MakeFunc", "reflect.makeFuncStub", "", false},
|
||||
{"method value from reflect", "reflect.methodValueCall", "", false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, ok := taskIDFromFuncName(tt.runtimeName)
|
||||
assert.Equal(t, tt.ok, ok)
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDagTakesAtMostOneDagSpec(t *testing.T) {
|
||||
assert.NotPanics(t, func() { Dag("etl", DagSpec{}) })
|
||||
assert.PanicsWithValue(t,
|
||||
`airflow.Dag: Dag "etl" got 2 airflow.DagSpec values; `+
|
||||
`set all of the Dag's attributes in one DagSpec`,
|
||||
func() { Dag("etl", DagSpec{}, DagSpec{}) },
|
||||
)
|
||||
}
|
||||
|
||||
func TestTaskSpecSetsTheTaskID(t *testing.T) {
|
||||
dag := Dag("etl")
|
||||
|
||||
spec := TaskSpec{TaskID: "extract_rows"}
|
||||
task := dag.Task(ExtractRows, spec)
|
||||
assert.Equal(t, "extract_rows", task.taskID)
|
||||
assert.Equal(t, spec, task.spec)
|
||||
|
||||
task = dag.Task(extract, TaskSpec{})
|
||||
assert.Equal(t, "extract", task.taskID, "an empty TaskID keeps the function name")
|
||||
}
|
||||
|
||||
func TestTaskTakesAPointerToTaskSpec(t *testing.T) {
|
||||
task := Dag("etl").Task(extract, &TaskSpec{TaskID: "extract_rows"})
|
||||
assert.Equal(t, "extract_rows", task.taskID)
|
||||
assert.Equal(t, TaskSpec{TaskID: "extract_rows"}, task.spec)
|
||||
}
|
||||
|
||||
func TestTaskNeedsTaskIDWhenFnHasNoName(t *testing.T) {
|
||||
const pkg = `github\.com/apache/airflow/go-sdk/airflow\.`
|
||||
literal := func(Context) error { return nil }
|
||||
madeByReflect := reflect.MakeFunc(
|
||||
reflect.TypeFor[func(Context) error](),
|
||||
func([]reflect.Value) []reflect.Value {
|
||||
return []reflect.Value{reflect.Zero(reflect.TypeFor[error]())}
|
||||
},
|
||||
).Interface()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
add func(dag *DagRef, opts ...TaskOption) *TaskRef
|
||||
// runtimeName is a regexp for the name the panic message shows.
|
||||
runtimeName string
|
||||
}{
|
||||
{
|
||||
name: "function literal",
|
||||
add: func(dag *DagRef, opts ...TaskOption) *TaskRef {
|
||||
return dag.Task(literal, opts...)
|
||||
},
|
||||
runtimeName: pkg + `TestTaskNeedsTaskIDWhenFnHasNoName\.func\d+`,
|
||||
},
|
||||
{
|
||||
name: "function literal in a package variable",
|
||||
add: func(dag *DagRef, opts ...TaskOption) *TaskRef {
|
||||
return dag.Task(packageLevelLiteral, opts...)
|
||||
},
|
||||
runtimeName: pkg + `init\.func\d+`,
|
||||
},
|
||||
{
|
||||
name: "reflect.MakeFunc",
|
||||
add: func(dag *DagRef, opts ...TaskOption) *TaskRef {
|
||||
return dag.Task(madeByReflect, opts...)
|
||||
},
|
||||
runtimeName: `reflect\.makeFuncStub`,
|
||||
},
|
||||
{
|
||||
name: "method value taken on a type parameter",
|
||||
add: func(dag *DagRef, opts ...TaskOption) *TaskRef {
|
||||
return addExtract(dag, extractor{}, opts...)
|
||||
},
|
||||
runtimeName: pkg + `addExtract\[\.\.\.\]\.func\d+`,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
msg := panicMessage(t, func() { tt.add(Dag("etl")) })
|
||||
assert.Regexp(t,
|
||||
`^airflow\.DagRef\.Task: Dag "etl": `+tt.runtimeName+` has no name to use `+
|
||||
`as the task_id; set one with airflow\.TaskSpec\{TaskID: \.\.\.\}$`,
|
||||
msg,
|
||||
)
|
||||
|
||||
task := tt.add(Dag("etl"), TaskSpec{TaskID: "transform"})
|
||||
assert.Equal(t, "transform", task.taskID)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskRejectsASecondTaskSpec(t *testing.T) {
|
||||
literal := func(Context) error { return nil }
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
fn any
|
||||
opts []TaskOption
|
||||
// task is a regexp for how the panic message names the task.
|
||||
task string
|
||||
}{
|
||||
{
|
||||
name: "TaskID in the first spec",
|
||||
fn: ExtractRows,
|
||||
opts: []TaskOption{TaskSpec{TaskID: "extract_rows"}, TaskSpec{}},
|
||||
task: `extract_rows`,
|
||||
},
|
||||
{
|
||||
name: "TaskID in the second spec",
|
||||
fn: literal,
|
||||
opts: []TaskOption{TaskSpec{}, TaskSpec{TaskID: "given"}},
|
||||
task: `given`,
|
||||
},
|
||||
{
|
||||
name: "no TaskID",
|
||||
fn: ExtractRows,
|
||||
opts: []TaskOption{TaskSpec{}, TaskSpec{}},
|
||||
task: `ExtractRows`,
|
||||
},
|
||||
{
|
||||
name: "no TaskID and no function name",
|
||||
fn: literal,
|
||||
opts: []TaskOption{TaskSpec{}, TaskSpec{}},
|
||||
task: `github\.com/apache/airflow/go-sdk/airflow\.` +
|
||||
`TestTaskRejectsASecondTaskSpec\.func\d+`,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
msg := panicMessage(t, func() { Dag("etl").Task(tt.fn, tt.opts...) })
|
||||
assert.Regexp(t,
|
||||
`^airflow\.DagRef\.Task: task "`+tt.task+`" of Dag "etl" got 2 airflow\.TaskSpec `+
|
||||
`values; set all of the task's attributes in one TaskSpec$`,
|
||||
msg,
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskRejectsADuplicateTaskID(t *testing.T) {
|
||||
dag := Dag("etl")
|
||||
dag.Task(extract)
|
||||
|
||||
want := `airflow.DagRef.Task: Dag "etl" already has a task "extract"; ` +
|
||||
`set another task_id with airflow.TaskSpec{TaskID: ...}`
|
||||
assert.PanicsWithValue(t, want, func() { dag.Task(extract) })
|
||||
assert.PanicsWithValue(t, want, func() { dag.Task(ExtractRows, TaskSpec{TaskID: "extract"}) })
|
||||
assert.NotPanics(t,
|
||||
func() { Dag("reports").Task(extract) },
|
||||
"the same task_id in another Dag is a different task",
|
||||
)
|
||||
}
|
||||
|
||||
func TestTaskPanicsOnBadFunction(t *testing.T) {
|
||||
var unassigned func(Context) error
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
fn any
|
||||
want string
|
||||
}{
|
||||
{name: "not a function", fn: "extract", want: "fn is string, not a function"},
|
||||
{
|
||||
name: "nil function value",
|
||||
fn: unassigned,
|
||||
want: "fn is a nil func(airflow.Context) error",
|
||||
},
|
||||
{
|
||||
name: "no leading Context",
|
||||
fn: func(context.Context) error { return nil },
|
||||
want: "parameter 0 is context.Context, but the first parameter must be airflow.Context",
|
||||
},
|
||||
{
|
||||
// The name of a method expression keeps a dot, as a function literal's name does.
|
||||
// This row checks that the signature check rejects the method expression for its
|
||||
// receiver before Task looks at the name.
|
||||
name: "method expression",
|
||||
fn: extractor.Extract,
|
||||
want: "parameter 0 is airflow.extractor, but the first parameter must be",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
msg := panicMessage(t, func() { Dag("etl").Task(tt.fn) })
|
||||
assert.Contains(t, msg, `airflow.DagRef.Task: Dag "etl": `)
|
||||
assert.Contains(t, msg, tt.want)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskRejectsOptionsItDoesNotDefine(t *testing.T) {
|
||||
var unset TaskOption
|
||||
var nilSpec *TaskSpec
|
||||
const notDefined = ", which is not an option that package airflow defines"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
opts []TaskOption
|
||||
want string
|
||||
}{
|
||||
{name: "nil", opts: []TaskOption{TaskSpec{}, unset}, want: "opts[1] is nil"},
|
||||
{
|
||||
name: "nil *TaskSpec",
|
||||
opts: []TaskOption{nilSpec},
|
||||
want: "opts[0] is a nil *airflow.TaskSpec",
|
||||
},
|
||||
{
|
||||
name: "struct that embeds a TaskSpec",
|
||||
opts: []TaskOption{wrappedSpec{TaskSpec: TaskSpec{TaskID: "x"}, MaxRetries: 5}},
|
||||
want: "opts[0] has type airflow.wrappedSpec" + notDefined,
|
||||
},
|
||||
{
|
||||
name: "struct that embeds a TaskOption",
|
||||
opts: []TaskOption{wrappedOption{TaskSpec{TaskID: "x"}}},
|
||||
want: "opts[0] has type airflow.wrappedOption" + notDefined,
|
||||
},
|
||||
{
|
||||
name: "struct that embeds a nil TaskOption",
|
||||
opts: []TaskOption{wrappedOption{}},
|
||||
want: "opts[0] has type airflow.wrappedOption" + notDefined,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
assert.PanicsWithValue(t, `airflow.DagRef.Task: Dag "etl": `+tt.want, func() {
|
||||
Dag("etl").Task(extract, tt.opts...)
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskAfterRegisterPanics(t *testing.T) {
|
||||
dag := Dag("etl")
|
||||
dag.Task(extract)
|
||||
Bundle().Register(dag)
|
||||
|
||||
assert.PanicsWithValue(t,
|
||||
`airflow.DagRef.Task: Dag "etl" has already been registered; `+
|
||||
`add every task before Register`,
|
||||
func() { dag.Task(ExtractRows) },
|
||||
)
|
||||
assert.Len(t, dag.tasks, 1)
|
||||
}
|
||||
|
||||
func TestTaskIsSafeForConcurrentUse(t *testing.T) {
|
||||
const workers, perWorker = 8, 100
|
||||
|
||||
dag := Dag("etl")
|
||||
var wg sync.WaitGroup
|
||||
for worker := range workers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := range perWorker {
|
||||
dag.Task(extract, TaskSpec{TaskID: fmt.Sprintf("task_%d_%d", worker, i)})
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
assert.Len(t, dag.tasks, workers*perWorker)
|
||||
}
|
||||
|
||||
// Task and Register both take the lock of the Dag, so each task is either added before the Dag
|
||||
// is registered or not added at all. If Register marked the Dag registered without the lock,
|
||||
// only a run with -race would fail.
|
||||
func TestRegisterWhileTasksAreAdded(t *testing.T) {
|
||||
dag := Dag("etl")
|
||||
started := make(chan struct{})
|
||||
registered := make(chan struct{})
|
||||
type outcome struct {
|
||||
added int
|
||||
recovered any
|
||||
}
|
||||
done := make(chan outcome)
|
||||
go func() {
|
||||
var o outcome
|
||||
defer func() {
|
||||
o.recovered = recover()
|
||||
done <- o
|
||||
}()
|
||||
close(started)
|
||||
for ; ; o.added++ {
|
||||
select {
|
||||
case <-registered:
|
||||
// Register has returned, so this Task call has to panic. If the call returns
|
||||
// instead, the goroutine stops after it, and the assertions below fail.
|
||||
dag.Task(extract, TaskSpec{TaskID: "after_register"})
|
||||
return
|
||||
default:
|
||||
dag.Task(extract, TaskSpec{TaskID: fmt.Sprintf("task_%d", o.added)})
|
||||
}
|
||||
}
|
||||
}()
|
||||
// Wait until the goroutine is running before calling Register. If the goroutine only started
|
||||
// after close(registered), that close would order every Task call after Register, and -race
|
||||
// could not catch a Register that skips the lock.
|
||||
<-started
|
||||
Bundle().Register(dag)
|
||||
close(registered)
|
||||
o := <-done
|
||||
|
||||
assert.Equal(t,
|
||||
`airflow.DagRef.Task: Dag "etl" has already been registered; `+
|
||||
`add every task before Register`,
|
||||
o.recovered,
|
||||
)
|
||||
assert.Len(t, dag.tasks, o.added)
|
||||
}
|
||||
|
||||
func TestTaskOptionIsSealed(t *testing.T) {
|
||||
typ := reflect.TypeFor[TaskOption]()
|
||||
require.Equal(t, 1, typ.NumMethod())
|
||||
// reflect reports a package path only for an unexported method.
|
||||
assert.Equal(t, "github.com/apache/airflow/go-sdk/airflow", typ.Method(0).PkgPath)
|
||||
}
|
||||
|
||||
func TestTaskOptionRejectsForeignTypes(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("shells out to `go build`")
|
||||
}
|
||||
if _, err := exec.LookPath("go"); err != nil {
|
||||
t.Skip("go toolchain not on PATH")
|
||||
}
|
||||
|
||||
out, err := exec.Command("go", "build", "-o", os.DevNull, "./testdata/foreigntaskoption").
|
||||
CombinedOutput()
|
||||
|
||||
require.Error(t, err, "a type defined outside package airflow must not compile as an option")
|
||||
// The test checks only the two type names, so that a change in the wording of the compiler
|
||||
// error does not break it.
|
||||
assert.Contains(t, string(out), "foreignOption")
|
||||
assert.Contains(t, string(out), "airflow.TaskOption")
|
||||
}
|
||||
@@ -56,7 +56,7 @@ const (
|
||||
//
|
||||
// The command-line flags of the executable decide what Serve does.
|
||||
// With --airflow-metadata it prints the bundle's manifest and returns, which is how
|
||||
// airflow-go-pack reads the registered Dag and task ids.
|
||||
// airflow-go-pack reads the Dag and task ids of the registered task handlers.
|
||||
// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the
|
||||
// coordinator protocol.
|
||||
//
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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
|
||||
|
||||
// DagSpec holds the attributes of a Dag other than its dag_id. [Dag] takes one.
|
||||
type DagSpec struct{}
|
||||
|
||||
// TaskSpec holds the attributes of a task. [DagRef.Task] takes at most one per task.
|
||||
type TaskSpec struct {
|
||||
// TaskID is the task_id of the task. When TaskID is empty, the task_id is the name of the
|
||||
// Go function that the task runs.
|
||||
TaskID string
|
||||
}
|
||||
@@ -46,18 +46,23 @@ func (*taskHandler) registerable() {}
|
||||
// main calls TaskHandler before Serve, so a handler that fails the check stops the executable as
|
||||
// soon as it starts instead of when the task first runs.
|
||||
func TaskHandler(dagId, taskId string, fn any) Registerable {
|
||||
v := reflect.ValueOf(fn)
|
||||
if v.Kind() != reflect.Func {
|
||||
panic(
|
||||
fmt.Sprintf("airflow.TaskHandler(%q, %q): fn is %T, not a function", dagId, taskId, fn),
|
||||
)
|
||||
}
|
||||
if v.IsNil() {
|
||||
panic(fmt.Sprintf("airflow.TaskHandler(%q, %q): fn is a nil %T", dagId, taskId, fn))
|
||||
}
|
||||
task, err := bundle.NewTaskFunction(fn)
|
||||
task, err := newTaskFunction(fn)
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("airflow.TaskHandler(%q, %q): %v", dagId, taskId, err))
|
||||
}
|
||||
return &taskHandler{dagId: dagId, taskId: taskId, task: task}
|
||||
}
|
||||
|
||||
// newTaskFunction checks fn before bundle.NewTaskFunction does, so that the error names the Go
|
||||
// type of a value that is not a function. It also rejects a nil function, which NewTaskFunction
|
||||
// accepts even though the task would fail once it runs.
|
||||
func newTaskFunction(fn any) (bundle.Task, error) {
|
||||
v := reflect.ValueOf(fn)
|
||||
if v.Kind() != reflect.Func {
|
||||
return nil, fmt.Errorf("fn is %T, not a function", fn)
|
||||
}
|
||||
if v.IsNil() {
|
||||
return nil, fmt.Errorf("fn is a nil %T", fn)
|
||||
}
|
||||
return bundle.NewTaskFunction(fn)
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
// 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
|
||||
|
||||
// TaskOption is an option to [DagRef.Task]. [TaskSpec] is one.
|
||||
//
|
||||
// Its only method is unexported, so a type outside this package cannot declare it.
|
||||
// A struct that embeds a TaskSpec or a TaskOption still satisfies the interface, and Task
|
||||
// panics when it is given one.
|
||||
type TaskOption interface{ applyTask(*taskConfig) }
|
||||
|
||||
type taskConfig struct {
|
||||
// specs keeps every TaskSpec passed to DagRef.Task, so that Task can reject a second one
|
||||
// instead of merging the two.
|
||||
specs []TaskSpec
|
||||
}
|
||||
|
||||
func (s TaskSpec) applyTask(c *taskConfig) { c.specs = append(c.specs, s) }
|
||||
@@ -0,0 +1,34 @@
|
||||
// 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.
|
||||
|
||||
// Command foreigntaskoption tries to pass a task option that package airflow did not define.
|
||||
// It must not compile. TestTaskOptionRejectsForeignTypes builds it and expects that failure.
|
||||
package main
|
||||
|
||||
import "github.com/apache/airflow/go-sdk/airflow"
|
||||
|
||||
type foreignOption struct{}
|
||||
|
||||
// An unexported method name belongs to the package that declares it, so this method is not
|
||||
// the applyTask method of airflow.TaskOption even though it is spelled the same.
|
||||
func (foreignOption) applyTask() {}
|
||||
|
||||
func extract(airflow.Context) error { return nil }
|
||||
|
||||
func main() {
|
||||
airflow.Dag("etl", airflow.DagSpec{}).Task(extract, foreignOption{})
|
||||
}
|
||||
Reference in new issue
Block a user