aaron-y-chen commented on code in PR #75700:
URL: https://github.com/apache/airflow/pull/75700#discussion_r4529413292
##########
go-sdk/pkg/execution/messages.go:
##########
@@ -117,7 +119,109 @@ 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 !mayHoldDecodedAny(v.Type().Elem()) {
+ return
+ }
+ for i := range v.Len() {
+ normalizeDecodedNumbers(v.Index(i))
+ }
+ case reflect.Map:
+ if !mayHoldDecodedAny(v.Type().Elem()) {
+ return
+ }
+ iter := v.MapRange()
+ for iter.Next() {
+ // A map value is not addressable, so normalize a copy
and store it back.
+ elem := reflect.New(v.Type().Elem()).Elem()
+ elem.Set(iter.Value())
+ normalizeDecodedNumbers(elem)
+ v.SetMapIndex(iter.Key(), elem)
+ }
+ }
+}
+
+// mayHoldDecodedAny 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 mayHoldDecodedAny(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 int8:
+ return int64(v)
Review Comment:
Nice catch, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
