aaron-y-chen commented on code in PR #75700:
URL: https://github.com/apache/airflow/pull/75700#discussion_r4531568452


##########
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)
+               }

Review Comment:
   Nice improvement!



-- 
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]

Reply via email to