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 804eee944b4 Expose Dag run configuration and metadata to Go tasks
(#74674)
804eee944b4 is described below
commit 804eee944b4e84beaa3ef4e6d572c9f76493ab44
Author: Aaron Chen <[email protected]>
AuthorDate: Sat Nov 7 04:59:36 2026 -0800
Expose Dag run configuration and metadata to Go tasks (#74674)
* Expose Dag run configuration and metadata to Go tasks
* Refine documentation for DagRun configuration and XCom type mapping in Go
SDK
* Refactor DagRunType to use sdk package and improve type consistency in
context
* Adjust DagRun RunAfter handling to ensure UTC consistency in task
execution
---
.../authoring-and-scheduling/language-sdks/go.rst | 20 +++--
go-sdk/DEVELOPMENT.md | 7 +-
go-sdk/adr/0008-native-dag-interface.md | 2 +-
go-sdk/airflow/context.go | 2 +-
go-sdk/airflow/enums.go | 29 +++++--
go-sdk/airflow/enums_test.go | 14 ++++
go-sdk/airflow/serve_test.go | 89 ++++++++++++++++++++++
go-sdk/example/bundle/main.go | 3 +
go-sdk/pkg/execution/integration_test.go | 11 +++
go-sdk/pkg/execution/task_runner.go | 60 +++++++++++++++
go-sdk/pkg/execution/task_runner_test.go | 89 ++++++++++++++++++++++
go-sdk/sdk/context.go | 18 ++++-
12 files changed, 323 insertions(+), 21 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
index 19bc36d9c2e..03dcb72d056 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
@@ -467,14 +467,19 @@ Its first parameter is an ``airflow.Context``, which
gives it what Airflow provi
- ``DagID``, ``RunID``, ``TaskID`` (including any task group prefix),
``MapIndex`` (``nil`` for an unmapped
task) and ``TryNumber``.
* - ``actx.DagRun()``
- - ``DagID``, ``RunID``, and the ``*time.Time`` fields ``LogicalDate``,
``DataIntervalStart`` and
- ``DataIntervalEnd`` (``nil`` when the run has no such value, such as a
manual run).
+ - ``DagID``, ``RunID``, ``RunAfter``, ``RunType`` (an
``airflow.DagRunType`` such as
+ ``airflow.DagRunTypeManual``), ``Conf`` (empty when no configuration
was supplied), and
+ the ``*time.Time`` fields ``LogicalDate``, ``DataIntervalStart`` and
``DataIntervalEnd`` (``nil`` when the
+ run has no such value, such as a manual run).
``airflow.Context`` is a ``context.Context``, so pass it to a client call or
to ``http.NewRequestWithContext``. Its
``Done()`` fires when Airflow stops the task, so long-running work can stop
cleanly. Cleanup that must outlive
that runs under ``context.WithoutCancel(actx)``. A helper that takes a plain
``context.Context`` gets the
``airflow.Context`` back with ``airflow.FromContext``.
+``Conf`` is a ``map[string]any`` whose values follow the :ref:`XCom type
mapping <go-sdk/types>`. Use a map lookup
+or ``len`` to check for configuration entries.
+
A value the task returns becomes its ``return_value`` XCom. A non-nil
``error``, or a panic, fails the task, and
retries it if the task has retries.
@@ -596,7 +601,8 @@ A namespace level applies to the longest matching dotted
prefix, so ``example=DE
XCom type mapping
-----------------
-XComs are stored as JSON, so a value read with ``GetXCom`` arrives as the
matching Go type:
+XComs are stored as JSON, so a value read with ``GetXCom`` arrives as the
matching Go type. The same mapping
+applies to the values in ``actx.DagRun().Conf``:
.. list-table::
:header-rows: 1
@@ -629,9 +635,11 @@ XComs are stored as JSON, so a value read with ``GetXCom``
arrives as the matchi
.. note::
- An integer arrives as a Go integer type whose width depends on the value.
Type-switch over the integer types
- you expect, or round-trip the value through ``json.Marshal`` and
``json.Unmarshal`` into a typed Go value. A
- task input from ``airflow.Inputs`` or a TaskFlow argument is already decoded
into its parameter type.
+ An integer arrives as a Go integer type whose width depends on the value. In
``actx.DagRun().Conf``, it
+ arrives as ``int64`` instead, or as ``uint64`` when it is above
``math.MaxInt64``, including in nested maps
+ and arrays. Type-switch over the integer types you expect, or round-trip the
value through ``json.Marshal``
+ and ``json.Unmarshal`` into a typed Go value. A task input from
``airflow.Inputs`` or a TaskFlow argument is
+ already decoded into its parameter type.
.. _go-sdk/build:
diff --git a/go-sdk/DEVELOPMENT.md b/go-sdk/DEVELOPMENT.md
index f90a25aa3c4..7419e3605c7 100644
--- a/go-sdk/DEVELOPMENT.md
+++ b/go-sdk/DEVELOPMENT.md
@@ -102,9 +102,10 @@ 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`](./airflow/enums.go) declares the `genmodels.DagRunState`
and
+`genmodels.DagRunType` constants again, so that a Dag author does not import
`genmodels`.
+`TestDagRunStateMatchesGenmodels` and `TestDagRunTypeMatchesGenmodels` fail
when regenerating
+the models adds, renames or removes a constant from either enum. The tests
keep failing until
`airflow/enums.go` declares the same constants as `genmodels`.
## Regenerating the Dag, task and task group specs
diff --git a/go-sdk/adr/0008-native-dag-interface.md
b/go-sdk/adr/0008-native-dag-interface.md
index 806e430644a..79b90b74ca5 100644
--- a/go-sdk/adr/0008-native-dag-interface.md
+++ b/go-sdk/adr/0008-native-dag-interface.md
@@ -184,7 +184,7 @@ 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. The SDK writes each of
those or leaves it at its default instead of taking it from an author, and the
schema 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.
+- **An enum generated into `genmodels` is re-exported from `airflow` once a
field that an author can reach holds its values.** `DagRunState` is the type of
the values in `TriggerDagRunSpec.AllowedStates` and
`TriggerDagRunSpec.FailedStates`; `DagRunType` is the type of `DagRun.RunType`.
`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.
diff --git a/go-sdk/airflow/context.go b/go-sdk/airflow/context.go
index 1f2d0fb438e..7bb949a327e 100644
--- a/go-sdk/airflow/context.go
+++ b/go-sdk/airflow/context.go
@@ -30,7 +30,7 @@ type (
TaskInstance = sdk.TaskInstance
// DagRun identifies the Dag run the task instance belongs to,
- // and carries its scheduling timestamps.
+ // and carries its scheduling timestamps, run type, and configuration.
DagRun = sdk.DagRun
)
diff --git a/go-sdk/airflow/enums.go b/go-sdk/airflow/enums.go
index 49f8a490f0d..709b581a36d 100644
--- a/go-sdk/airflow/enums.go
+++ b/go-sdk/airflow/enums.go
@@ -22,21 +22,22 @@ import (
"slices"
"github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
+ "github.com/apache/airflow/go-sdk/sdk"
)
-// 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:
+// This file declares every enum type whose values a Dag author sets or reads.
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
+// - DagRunType: airflow-core/src/airflow/utils/types.py, for DagRun.RunType
//
-// Two more Python enums are not declared here yet, because no field that an
author can reach
-// holds their values:
+// One more Python enum is not declared here yet, because no field that an
author can reach
+// holds its 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.
@@ -108,9 +109,23 @@ const (
DagRunStateFailed DagRunState = "failed"
)
+// DagRunType says how a Dag run was created. DagRun.RunType holds it.
+//
+// Like DagRunState, %T and package reflect report its name as
genmodels.DagRunType.
+type DagRunType = sdk.DagRunType
+
+const (
+ DagRunTypeBackfill DagRunType = "backfill"
+ DagRunTypeScheduled DagRunType = "scheduled"
+ DagRunTypeManual DagRunType = "manual"
+ DagRunTypeOperatorTriggered DagRunType = "operator_triggered"
+ DagRunTypeAssetTriggered DagRunType = "asset_triggered"
+ DagRunTypeAssetMaterialization DagRunType = "asset_materialization"
+)
+
// 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.
+// Python enums for TriggerRule and WeightRule, and against genmodels for
DagRunState and DagRunType.
var (
triggerRules = []TriggerRule{
TriggerRuleAllSuccess,
diff --git a/go-sdk/airflow/enums_test.go b/go-sdk/airflow/enums_test.go
index 32c8cd7fb09..2cde5a791f4 100644
--- a/go-sdk/airflow/enums_test.go
+++ b/go-sdk/airflow/enums_test.go
@@ -151,6 +151,20 @@ func TestDagRunStateMatchesGenmodels(t *testing.T) {
assert.ElementsMatch(t, slices.Collect(maps.Values(generated)),
toStrings(dagRunStates))
}
+// TestDagRunTypeMatchesGenmodels fails when regenerating genmodels adds,
renames or removes a
+// DagRunType constant. Like DagRunState, DagRunType is an alias, so enums.go
declares each of the
+// genmodels constants again.
+func TestDagRunTypeMatchesGenmodels(t *testing.T) {
+ generated := readPairs(
+ t,
+ "../pkg/execution/genmodels/models.gen.go",
+ `(?m)^const DagRunType(\w+) DagRunType = "([^"]*)"$`,
+ )
+ declared := readPairs(t, enumsPath, `(?m)^\tDagRunType(\w+) +DagRunType
= "([^"]*)"$`)
+
+ assert.Equal(t, generated, declared)
+}
+
func TestTaskRejectsAValueOutsideAnEnum(t *testing.T) {
const (
triggerRuleValues = `["all_success" "all_failed" "all_done"
"all_done_min_one_success" ` +
diff --git a/go-sdk/airflow/serve_test.go b/go-sdk/airflow/serve_test.go
index bca506bc14c..e797e327251 100644
--- a/go-sdk/airflow/serve_test.go
+++ b/go-sdk/airflow/serve_test.go
@@ -22,6 +22,7 @@ import (
"encoding/json"
"errors"
"io"
+ "math"
"net"
"testing"
"time"
@@ -257,6 +258,94 @@ func serveForSupervisor(t *testing.T, b *BundleRef)
(*execution.CoordinatorComm,
// A fake supervisor sends StartupDetails over the comm socket, as the Python
// ExecutableCoordinator does after it starts the bundle with --comm and
--logs.
+func TestServeInjectsDagRunMetadata(t *testing.T) {
+ // msgpack decodes a timestamp in time.Local, so a worker outside UTC
+ // would otherwise get RunAfter in its own zone.
+ local := time.Local
+ time.Local = time.FixedZone("UTC+8", 8*60*60)
+ t.Cleanup(func() { time.Local = local })
+ runAfter := time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
+ cases := []struct {
+ name string
+ sendConf bool
+ conf map[string]any
+ want map[string]any
+ }{
+ {name: "absent"},
+ {name: "null", sendConf: true},
+ {name: "empty", sendConf: true, conf: map[string]any{}, want:
map[string]any{}},
+ {
+ name: "nested numbers",
+ sendConf: true,
+ conf: map[string]any{
+ "days": 7,
+ "batch_size": 500,
+ "big": 70000,
+ "precise": int64(9007199254740993),
+ "max": uint64(math.MaxInt64),
+ "above_max": uint64(math.MaxInt64) + 1,
+ "nested": map[string]any{"attempts": -129,
"enabled": true, "ratio": 1.25},
+ "list": []any{-1, -32769,
int64(-2147483649), "us", nil},
+ },
+ want: map[string]any{
+ "days": int64(7),
+ "batch_size": int64(500),
+ "big": int64(70000),
+ "precise": int64(9007199254740993),
+ "max": int64(math.MaxInt64),
+ "above_max": uint64(math.MaxInt64) + 1,
+ "nested": map[string]any{
+ "attempts": int64(-129),
+ "enabled": true,
+ "ratio": 1.25,
+ },
+ "list": []any{int64(-1), int64(-32769),
int64(-2147483649), "us", nil},
+ },
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ var got DagRun
+ b := Bundle()
+ b.Register(TaskHandler("py_etl", "transform", func(actx
Context) error {
+ got = actx.DagRun()
+ return nil
+ }))
+ supervisor, done := serveForSupervisor(t, b)
+ dagRun := map[string]any{"run_after": runAfter,
"run_type": "manual"}
+ if tc.sendConf {
+ dagRun["conf"] = tc.conf
+ }
+ require.NoError(t, supervisor.SendRequest(0,
map[string]any{
+ "type": "StartupDetails",
+ "ti": map[string]any{
+ "id":
"550e8400-e29b-41d4-a716-446655440000",
+ "dag_id": "py_etl",
+ "task_id": "transform",
+ "run_id": "run1",
+ "try_number": 1,
+ },
+ "ti_context": map[string]any{"dag_run":
dagRun},
+ "bundle_info": map[string]any{"name": "test",
"version": "1.0"},
+ }))
+ frame, err := supervisor.ReadMessage()
+ require.NoError(t, err)
+ var terminal map[string]any
+ require.NoError(t, msgpack.Unmarshal(frame.Body,
&terminal))
+ require.Equal(t, "SucceedTask", terminal["type"])
+ select {
+ case err := <-done:
+ require.NoError(t, err)
+ case <-time.After(2 * time.Second):
+ t.Fatal("Serve did not return after the task
finished")
+ }
+ assert.Equal(t, DagRunTypeManual, got.RunType)
+ assert.Equal(t, runAfter, got.RunAfter)
+ assert.Equal(t, tc.want, got.Conf)
+ })
+ }
+}
+
func TestServeRunsTaskForSupervisor(t *testing.T) {
ran := false
b := Bundle()
diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go
index 5bca037f7f1..08ba77e3490 100644
--- a/go-sdk/example/bundle/main.go
+++ b/go-sdk/example/bundle/main.go
@@ -124,6 +124,9 @@ func extract(actx airflow.Context) (any, error) {
"logical_date", dagRun.LogicalDate,
"data_interval_start", dagRun.DataIntervalStart,
"data_interval_end", dagRun.DataIntervalEnd,
+ "run_after", dagRun.RunAfter,
+ "run_type", dagRun.RunType,
+ "conf", dagRun.Conf,
),
),
)
diff --git a/go-sdk/pkg/execution/integration_test.go
b/go-sdk/pkg/execution/integration_test.go
index 8e4b24358c8..a93ff793a31 100644
--- a/go-sdk/pkg/execution/integration_test.go
+++ b/go-sdk/pkg/execution/integration_test.go
@@ -496,6 +496,11 @@ func TestRunTaskInjectsAirflowContext(t *testing.T) {
logical := time.Date(2026, 6, 9, 12, 0, 0, 0, time.UTC)
start := logical.Add(-time.Hour)
end := logical.Add(time.Hour)
+ runAfter := end.Add(30 * time.Minute)
+ conf := genmodels.Conf{
+ "export_date": "2026-06-09",
+ "output": map[string]any{"bucket": "reports"},
+ }
var got contexttest.Context
bundle := buildBundle(t, func(r testBundle) {
@@ -523,6 +528,9 @@ func TestRunTaskInjectsAirflowContext(t *testing.T) {
LogicalDate: logical,
DataIntervalStart: start,
DataIntervalEnd: end,
+ RunAfter: runAfter,
+ RunType: genmodels.DagRunTypeManual,
+ Conf: &conf,
},
},
}
@@ -552,6 +560,9 @@ func TestRunTaskInjectsAirflowContext(t *testing.T) {
assert.Equal(t, start, *dagRun.DataIntervalStart)
require.NotNil(t, dagRun.DataIntervalEnd)
assert.Equal(t, end, *dagRun.DataIntervalEnd)
+ assert.Equal(t, runAfter, dagRun.RunAfter)
+ assert.Equal(t, genmodels.DagRunTypeManual, dagRun.RunType)
+ assert.Equal(t, map[string]any(conf), dagRun.Conf)
}
// Guards against the runtime binding the task state store to an empty task
instance id.
diff --git a/go-sdk/pkg/execution/task_runner.go
b/go-sdk/pkg/execution/task_runner.go
index 16791640890..d01b770e836 100644
--- a/go-sdk/pkg/execution/task_runner.go
+++ b/go-sdk/pkg/execution/task_runner.go
@@ -21,6 +21,7 @@ import (
"context"
"fmt"
"log/slog"
+ "math"
"runtime/debug"
"time"
@@ -74,6 +75,10 @@ func RunTask(
// live task context.
ti := taskInstanceOf(details)
dagRun := details.TIContext.DagRun
+ var conf map[string]any
+ if dagRun.Conf != nil {
+ conf = normalizeDagRunConf(*dagRun.Conf)
+ }
runtimeContext := sdk.NewTIRunContext(
context.Background(),
ti,
@@ -83,6 +88,9 @@ func RunTask(
LogicalDate: ifaceTimePtr(dagRun.LogicalDate),
DataIntervalStart:
ifaceTimePtr(dagRun.DataIntervalStart),
DataIntervalEnd: ifaceTimePtr(dagRun.DataIntervalEnd),
+ RunAfter: dagRun.RunAfter.UTC(), // msgpack
decodes it in the local zone.
+ RunType: dagRun.RunType,
+ Conf: conf,
},
)
@@ -125,6 +133,58 @@ func RunTask(
return executeTask(ctx, task, args, details.TIContext.ShouldRetry,
logger)
}
+func normalizeDagRunConf(conf map[string]any) map[string]any {
+ if conf == nil {
+ return nil
+ }
+ normalized := make(map[string]any, len(conf))
+ for key, value := range conf {
+ normalized[key] = normalizeDagRunConfValue(value)
+ }
+ return normalized
+}
+
+func normalizeDagRunConfValue(value any) any {
+ switch v := value.(type) {
+ case map[string]any:
+ return normalizeDagRunConf(v)
+ case []any:
+ if v == nil {
+ return v
+ }
+ normalized := make([]any, len(v))
+ for i, item := range v {
+ normalized[i] = normalizeDagRunConfValue(item)
+ }
+ return normalized
+ case int:
+ return int64(v)
+ case int8:
+ return int64(v)
+ case int16:
+ return int64(v)
+ case int32:
+ return int64(v)
+ case uint:
+ return normalizeDagRunConfValue(uint64(v))
+ case uint8:
+ return int64(v)
+ case uint16:
+ return int64(v)
+ case uint32:
+ return int64(v)
+ case uint64:
+ if v <= math.MaxInt64 {
+ return int64(v)
+ }
+ return v
+ case float32:
+ return float64(v)
+ default:
+ return value
+ }
+}
+
func taskInstanceOf(details *genmodels.StartupDetails) sdk.TaskInstance {
return sdk.TaskInstance{
DagID: details.TI.DagID,
diff --git a/go-sdk/pkg/execution/task_runner_test.go
b/go-sdk/pkg/execution/task_runner_test.go
new file mode 100644
index 00000000000..ecc06a30d75
--- /dev/null
+++ b/go-sdk/pkg/execution/task_runner_test.go
@@ -0,0 +1,89 @@
+// 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 (
+ "math"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestNormalizeDagRunConfNumbers(t *testing.T) {
+ cases := []struct {
+ name string
+ value any
+ want any
+ }{
+ {name: "int", value: int(7), want: int64(7)},
+ {name: "int8", value: int8(-1), want: int64(-1)},
+ {name: "int16", value: int16(-129), want: int64(-129)},
+ {name: "int32", value: int32(-32769), want: int64(-32769)},
+ {name: "uint", value: uint(500), want: int64(500)},
+ {name: "uint8", value: uint8(200), want: int64(200)},
+ {name: "uint16", value: uint16(500), want: int64(500)},
+ {name: "uint32", value: uint32(70000), want: int64(70000)},
+ {
+ name: "large exact integer",
+ value: uint64(9007199254740993),
+ want: int64(9007199254740993),
+ },
+ {name: "max int64", value: uint64(math.MaxInt64), want:
int64(math.MaxInt64)},
+ {
+ name: "above max int64",
+ value: uint64(math.MaxInt64) + 1,
+ want: uint64(math.MaxInt64) + 1,
+ },
+ {name: "float32", value: float32(1.25), want: float64(1.25)},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ original := map[string]any{"value": tc.value}
+ got := normalizeDagRunConf(original)
+ assert.Equal(t, tc.want, got["value"])
+ assert.Equal(t, tc.value, original["value"])
+ })
+ }
+}
+
+func TestNormalizeDagRunConfCopiesContainers(t *testing.T) {
+ original := map[string]any{
+ "label": "original",
+ "nested": map[string]any{"bucket": "reports"},
+ "list": []any{map[string]any{"region": "us"},
[]any{"eu"}},
+ "nil_map": map[string]any(nil),
+ "empty_map": map[string]any{},
+ "nil_slice": []any(nil),
+ "empty_slice": []any{},
+ "flag": true,
+ "null": nil,
+ }
+ normalized := normalizeDagRunConf(original)
+ require.Equal(t, original, normalized)
+
+ normalized["label"] = "changed"
+ normalized["nested"].(map[string]any)["bucket"] = "changed"
+ normalized["list"].([]any)[0].(map[string]any)["region"] = "changed"
+ normalized["list"].([]any)[1].([]any)[0] = "changed"
+
+ assert.Equal(t, "original", original["label"])
+ assert.Equal(t, "reports",
original["nested"].(map[string]any)["bucket"])
+ assert.Equal(t, "us",
original["list"].([]any)[0].(map[string]any)["region"])
+ assert.Equal(t, "eu", original["list"].([]any)[1].([]any)[0])
+}
diff --git a/go-sdk/sdk/context.go b/go-sdk/sdk/context.go
index c6ca1dfe8c3..eb0263a10d6 100644
--- a/go-sdk/sdk/context.go
+++ b/go-sdk/sdk/context.go
@@ -20,6 +20,8 @@ package sdk
import (
"context"
"time"
+
+ "github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
)
// TIRunContext is a context.Context that also exposes the identifiers and
@@ -84,14 +86,24 @@ type TaskInstance struct {
TryNumber int
}
+// DagRunType identifies how a Dag run was created.
+type DagRunType = genmodels.DagRunType
+
// DagRun identifies the Dag run the current task instance belongs to and
-// carries its scheduling timestamps. The *time.Time fields are nil when the
-// supervisor did not provide a value (for example, a manually triggered run
-// without a logical date).
+// carries its scheduling timestamps, how it was created, and its conf. The
+// *time.Time fields are nil when the supervisor did not provide a value (for
+// example, a manually triggered run without a logical date).
type DagRun struct {
DagID string
RunID string
LogicalDate *time.Time
DataIntervalStart *time.Time
DataIntervalEnd *time.Time
+ RunAfter time.Time
+ // RunType says how the run was created, such as
+ // [github.com/apache/airflow/go-sdk/airflow.DagRunTypeManual].
+ RunType DagRunType
+ // Conf is the run configuration, empty when none was supplied.
Integers are
+ // int64, or uint64 above math.MaxInt64; floating-point values are
float64.
+ Conf map[string]any
}
