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 53134ff6a69 Go SDK: Return XCom and task state store integers as int64
(#75700)
53134ff6a69 is described below
commit 53134ff6a697fe50f6813e9295dc3eb3462a3615
Author: Aaron Chen <[email protected]>
AuthorDate: Tue Nov 10 21:40:17 2026 -0800
Go SDK: Return XCom and task state store integers as int64 (#75700)
* Go SDK: Return XCom and task state store integers as int64
* Normalize handling of decoded integers and improve test coverage for
number normalization
---
.../authoring-and-scheduling/language-sdks/go.rst | 13 +-
go-sdk/pkg/execution/messages.go | 116 ++++++++++++-
go-sdk/pkg/execution/messages_test.go | 185 +++++++++++++++++++++
go-sdk/pkg/execution/task_runner.go | 55 +-----
go-sdk/pkg/execution/task_runner_test.go | 89 ----------
5 files changed, 304 insertions(+), 154 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 03dcb72d056..66738f562d7 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
@@ -602,7 +602,7 @@ XCom type mapping
-----------------
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``:
+applies to a value read with ``TaskStateStore().Get`` and to the values in
``actx.DagRun().Conf``:
.. list-table::
:header-rows: 1
@@ -613,7 +613,7 @@ applies to the values in ``actx.DagRun().Conf``:
- Go type
* - ``int``
- number (integer)
- - an integer type (see note)
+ - ``int64`` (see note)
* - ``float``
- number (decimal)
- ``float64``
@@ -635,11 +635,10 @@ applies to the values in ``actx.DagRun().Conf``:
.. note::
- 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.
+ An integer arrives as ``int64``, or as ``uint64`` when it is above
``math.MaxInt64``. This also applies to
+ values in nested maps and arrays. To get a typed Go value, round-trip the
value through ``json.Marshal`` and
+ ``json.Unmarshal``. 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/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go
index e494cf36edb..8fe334311d6 100644
--- a/go-sdk/pkg/execution/messages.go
+++ b/go-sdk/pkg/execution/messages.go
@@ -19,6 +19,8 @@ package execution
import (
"fmt"
+ "math"
+ "reflect"
"github.com/vmihailenco/msgpack/v5"
@@ -67,7 +69,7 @@ func decodeIncomingBody(raw msgpack.RawMessage) (any, error) {
switch typ {
case genmodels.TypeStartupDetails:
var msg genmodels.StartupDetails
- if err := msgpack.Unmarshal(raw, &msg); err != nil {
+ if err := unmarshalNormalized(raw, &msg); err != nil {
return nil, fmt.Errorf("decoding StartupDetails: %w",
err)
}
if err := validateStartupTI(msg.TI); err != nil {
@@ -76,13 +78,13 @@ func decodeIncomingBody(raw msgpack.RawMessage) (any,
error) {
return &msg, nil
case genmodels.TypeDagFileParseRequest:
var msg genmodels.DagFileParseRequest
- if err := msgpack.Unmarshal(raw, &msg); err != nil {
+ if err := unmarshalNormalized(raw, &msg); err != nil {
return nil, fmt.Errorf("decoding DagFileParseRequest:
%w", err)
}
return &msg, nil
case genmodels.TypeErrorResponse:
var msg genmodels.ErrorResponse
- if err := msgpack.Unmarshal(raw, &msg); err != nil {
+ if err := unmarshalNormalized(raw, &msg); err != nil {
return nil, fmt.Errorf("decoding ErrorResponse: %w",
err)
}
return &msg, nil
@@ -116,7 +118,113 @@ func decodeBody(raw msgpack.RawMessage, dst any) error {
if isNilRaw(raw) {
return fmt.Errorf("empty response body")
}
- return msgpack.Unmarshal(raw, dst)
+ return unmarshalNormalized(raw, dst)
+}
+
+// unmarshalNormalized decodes raw into dst, then normalizes the numbers that
+// msgpack decoded into any-typed values, wherever they sit in dst.
+func unmarshalNormalized(raw []byte, dst any) error {
+ if err := msgpack.Unmarshal(raw, dst); err != nil {
+ return err
+ }
+ normalizeDecodedNumbers(reflect.ValueOf(dst))
+ return nil
+}
+
+// normalizeDecodedNumbers applies normalizeDecodedValue, in place, to every
+// any-typed struct field, map value and slice element reachable from v.
+func normalizeDecodedNumbers(v reflect.Value) {
+ switch v.Kind() {
+ case reflect.Pointer:
+ if !v.IsNil() {
+ normalizeDecodedNumbers(v.Elem())
+ }
+ case reflect.Interface:
+ if !v.IsNil() && v.CanSet() {
+
v.Set(reflect.ValueOf(normalizeDecodedValue(v.Interface())))
+ }
+ case reflect.Struct:
+ for i := range v.NumField() {
+ if v.Type().Field(i).IsExported() {
+ normalizeDecodedNumbers(v.Field(i))
+ }
+ }
+ case reflect.Slice, reflect.Array:
+ if !canHoldDecodedAny(v.Type().Elem()) {
+ return
+ }
+ for i := range v.Len() {
+ normalizeDecodedNumbers(v.Index(i))
+ }
+ case reflect.Map:
+ if !canHoldDecodedAny(v.Type().Elem()) {
+ return
+ }
+ elem := reflect.New(v.Type().Elem()).Elem()
+ iter := v.MapRange()
+ for iter.Next() {
+ // A map value is not addressable, so normalize a copy
and store it back.
+ elem.Set(iter.Value())
+ normalizeDecodedNumbers(elem)
+ v.SetMapIndex(iter.Key(), elem)
+ }
+ }
+}
+
+// canHoldDecodedAny reports whether a value of type t can contain an any-typed
+// value, so that normalizeDecodedNumbers skips the elements of a []byte or a
+// map[string]string.
+func canHoldDecodedAny(t reflect.Type) bool {
+ switch t.Kind() {
+ case reflect.Interface, reflect.Pointer, reflect.Struct,
+ reflect.Slice, reflect.Array, reflect.Map:
+ return true
+ default:
+ return false
+ }
+}
+
+// normalizeDecodedValue returns value with its integers as int64, or uint64
+// above math.MaxInt64, and its floats as float64, including the values nested
+// in its maps and slices, which it updates in place. msgpack decodes a number
+// into the narrowest Go type that its wire encoding names, and the supervisor
+// packs each integer in the fewest bytes, so 7, 200 and 70000 would otherwise
+// arrive as int8, uint8 and uint32. The Java SDK decodes every integer to a
+// Long in the same way.
+func normalizeDecodedValue(value any) any {
+ switch v := value.(type) {
+ case map[string]any:
+ for key, item := range v {
+ v[key] = normalizeDecodedValue(item)
+ }
+ case []any:
+ for i, item := range v {
+ v[i] = normalizeDecodedValue(item)
+ }
+ case int:
+ return int64(v)
+ case uint:
+ return normalizeDecodedValue(uint64(v))
+ case int8:
+ return int64(v)
+ case int16:
+ return int64(v)
+ case int32:
+ return int64(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)
+ }
+ case float32:
+ return float64(v)
+ }
+ return value
}
// apiErrorFromFrame returns the supervisor error carried by a frame, or nil if
diff --git a/go-sdk/pkg/execution/messages_test.go
b/go-sdk/pkg/execution/messages_test.go
index b94e26b14fc..0cb773bd1bd 100644
--- a/go-sdk/pkg/execution/messages_test.go
+++ b/go-sdk/pkg/execution/messages_test.go
@@ -18,8 +18,11 @@
package execution
import (
+ "bytes"
"encoding/json"
+ "math"
"os"
+ "reflect"
"testing"
"time"
@@ -540,6 +543,188 @@ func TestAPIErrorFromFrame(t *testing.T) {
})
}
+// marshalCompactBody encodes body with each integer in the fewest bytes, as
the
+// supervisor's msgpack packer does.
+func marshalCompactBody(t *testing.T, body any) msgpack.RawMessage {
+ t.Helper()
+ var buf bytes.Buffer
+ enc := msgpack.NewEncoder(&buf)
+ enc.UseCompactInts(true)
+ require.NoError(t, enc.Encode(body))
+ return buf.Bytes()
+}
+
+func TestDecodeBodyNormalizesNumbers(t *testing.T) {
+ cases := []struct {
+ name string
+ value any
+ want any
+ }{
+ {name: "positive fixint", value: 7, want: int64(7)},
+ {name: "uint8", value: 200, want: int64(200)},
+ {name: "uint16", value: 500, want: int64(500)},
+ {name: "uint32", value: 70000, want: int64(70000)},
+ {name: "uint64 up to max int64", value: uint64(math.MaxInt64),
want: int64(math.MaxInt64)},
+ {
+ name: "uint64 above max int64",
+ value: uint64(math.MaxInt64) + 1,
+ want: uint64(math.MaxInt64) + 1,
+ },
+ {name: "negative fixint", value: -1, want: int64(-1)},
+ {name: "int8", value: -100, want: int64(-100)},
+ {name: "int16", value: -129, want: int64(-129)},
+ {name: "int32", value: -32769, want: int64(-32769)},
+ {name: "float32", value: float32(1.25), want: float64(1.25)},
+ {
+ name: "nested in maps and arrays",
+ value: map[string]any{
+ "rows": 7,
+ "batches": []any{200, "eu", nil,
map[string]any{"ratio": float32(0.5)}},
+ },
+ want: map[string]any{
+ "rows": int64(7),
+ "batches": []any{int64(200), "eu", nil,
map[string]any{"ratio": float64(0.5)}},
+ },
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ raw := marshalCompactBody(t, map[string]any{
+ "type": "XComResult",
+ "key": "return_value",
+ "value": tc.value,
+ })
+ var result genmodels.XComResult
+ require.NoError(t, decodeBody(raw, &result))
+ assert.Equal(t, tc.want, result.Value)
+ })
+ }
+}
+
+func TestDecodeIncomingBodyNormalizesNumbers(t *testing.T) {
+ cases := []struct {
+ name string
+ body map[string]any
+ got func(msg any) any
+ want any
+ }{
+ {
+ name: "StartupDetails dag run conf",
+ body: map[string]any{
+ "type": "StartupDetails",
+ "ti": map[string]any{
+ "id":
"550e8400-e29b-41d4-a716-446655440000",
+ "task_id": "extract",
+ "dag_id": "tutorial_dag",
+ "run_id": "manual__2024-01-15",
+ "try_number": 1,
+ },
+ "bundle_info": map[string]any{"name":
"example_dags", "version": "1.0.0"},
+ "ti_context": map[string]any{
+ "dag_run": map[string]any{
+ "conf": map[string]any{"rows":
200, "window": map[string]any{"days": 7}},
+ },
+ },
+ },
+ got: func(msg any) any {
+ return
*msg.(*genmodels.StartupDetails).TIContext.DagRun.Conf
+ },
+ want: genmodels.Conf{"rows": int64(200), "window":
map[string]any{"days": int64(7)}},
+ },
+ {
+ name: "DagFileParseRequest callback requests",
+ body: map[string]any{
+ "type": "DagFileParseRequest",
+ "file": "/files/dags/tutorial.go",
+ "bundle_name": "example_dags",
+ "bundle_path": "/files/dags",
+ "callback_requests": []any{map[string]any{
+ "type": "TaskCallbackRequest",
+ "ti": map[string]any{"task_id":
"extract", "try_number": 2, "map_index": -1},
+ }},
+ },
+ got: func(msg any) any {
+ return
msg.(*genmodels.DagFileParseRequest).CallbackRequests
+ },
+ want: []any{map[string]any{
+ "type": "TaskCallbackRequest",
+ "ti": map[string]any{
+ "task_id": "extract",
+ "try_number": int64(2),
+ "map_index": int64(-1),
+ },
+ }},
+ },
+ {
+ name: "ErrorResponse detail",
+ body: map[string]any{
+ "type": "ErrorResponse",
+ "error": "API_SERVER_ERROR",
+ "detail": map[string]any{"status_code": 400,
"message": "Unhandled request"},
+ },
+ got: func(msg any) any {
+ return *msg.(*genmodels.ErrorResponse).Detail
+ },
+ want: genmodels.Detail{"status_code": int64(400),
"message": "Unhandled request"},
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ msg, err := decodeIncomingBody(marshalCompactBody(t,
tc.body))
+ require.NoError(t, err)
+ assert.Equal(t, tc.want, tc.got(msg))
+ })
+ }
+}
+
+func TestUnmarshalNormalizedMapStructValues(t *testing.T) {
+ // No supervisor message has a map with struct values yet.
+ type entry struct {
+ Value any `msgpack:"value"`
+ }
+ var dst struct {
+ Entries map[string]entry `msgpack:"entries"`
+ }
+ raw := marshalCompactBody(t, map[string]any{
+ "entries": map[string]any{
+ "rows": map[string]any{"value": 200},
+ "ratio": map[string]any{"value": float32(0.5)},
+ },
+ })
+ require.NoError(t, unmarshalNormalized(raw, &dst))
+ assert.Equal(
+ t,
+ map[string]entry{"rows": {Value: int64(200)}, "ratio": {Value:
float64(0.5)}},
+ dst.Entries,
+ )
+}
+
+func TestCanHoldDecodedAny(t *testing.T) {
+ for _, tc := range []struct {
+ container any
+ want bool
+ }{
+ {container: []byte(nil), want: false},
+ {container: map[string]string(nil), want: false},
+ {container: []any(nil), want: true},
+ {container: []*genmodels.XComResult(nil), want: true},
+ {container: map[string]genmodels.XComResult(nil), want: true},
+ {container: [][]any(nil), want: true},
+ {container: [][2]any(nil), want: true},
+ {container: []map[string]any(nil), want: true},
+ } {
+ elem := reflect.TypeOf(tc.container).Elem()
+ assert.Equal(t, tc.want, canHoldDecodedAny(elem), "%T",
tc.container)
+ }
+}
+
+func TestNormalizeDecodedValueGoInts(t *testing.T) {
+ // msgpack never decodes a number into int or uint, so a decoded body
cannot
+ // reach these cases.
+ assert.Equal(t, int64(-7), normalizeDecodedValue(int(-7)))
+ assert.Equal(t, int64(500), normalizeDecodedValue(uint(500)))
+}
+
func TestIfaceInt(t *testing.T) {
for _, v := range []any{
int(42), int8(42), int16(42), int32(42), int64(42),
diff --git a/go-sdk/pkg/execution/task_runner.go
b/go-sdk/pkg/execution/task_runner.go
index 3981e1ec799..792aff79a73 100644
--- a/go-sdk/pkg/execution/task_runner.go
+++ b/go-sdk/pkg/execution/task_runner.go
@@ -21,7 +21,6 @@ import (
"context"
"fmt"
"log/slog"
- "math"
"runtime/debug"
"time"
@@ -77,7 +76,7 @@ func RunTask(
dagRun := details.TIContext.DagRun
var conf map[string]any
if dagRun.Conf != nil {
- conf = normalizeDagRunConf(*dagRun.Conf)
+ conf = *dagRun.Conf
}
runtimeContext := sdk.NewTIRunContext(
context.Background(),
@@ -133,58 +132,6 @@ 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
deleted file mode 100644
index ecc06a30d75..00000000000
--- a/go-sdk/pkg/execution/task_runner_test.go
+++ /dev/null
@@ -1,89 +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 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])
-}
