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


##########
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 {

Review Comment:
   Got it, I'll be more careful next time.



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