44 "fmt"
55 "log/slog"
66
7+ pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
78 "google.golang.org/protobuf/proto"
89 "google.golang.org/protobuf/types/known/anypb"
910 "google.golang.org/protobuf/types/known/wrapperspb"
@@ -13,17 +14,31 @@ func anyString(any *anypb.Any) string {
1314 if any == nil {
1415 return "nil"
1516 }
17+
18+ var s string
19+ var err error
1620 switch any .TypeUrl {
17- case "type.googleapis.com/google.protobuf.BytesValue" :
18- s , err := anyBytesString (any )
19- if err != nil {
20- slog .Debug ("cannot parse input/output value" , "error" , err )
21- // fallthrough
22- } else {
23- return s
21+ case "buf.build/stealthrocket/dispatch-proto/dispatch.sdk.python.v1.Pickled" :
22+ var pickled proto.Message
23+ pickled , err = any .UnmarshalNew ()
24+ if err == nil {
25+ if p , ok := pickled .(* pythonv1.Pickled ); ok {
26+ s , err = pythonPickleString (p .PickledValue )
27+ } else {
28+ err = fmt .Errorf ("invalid pickled message: %T" , p )
29+ }
2430 }
31+ case "type.googleapis.com/google.protobuf.BytesValue" :
32+ s , err = anyBytesString (any )
33+ default :
34+ // TODO: support unpacking other types of serialized values
35+ err = fmt .Errorf ("not implemented: %s" , any .TypeUrl )
36+ }
37+ if err != nil {
38+ slog .Debug ("cannot parse input/output value" , "error" , err )
39+ return fmt .Sprintf ("%s(?)" , any .TypeUrl )
2540 }
26- return fmt . Sprintf ( "%s(?)" , any . TypeUrl )
41+ return s
2742}
2843
2944func anyBytesString (any * anypb.Any ) (string , error ) {
@@ -37,6 +52,20 @@ func anyBytesString(any *anypb.Any) (string, error) {
3752 }
3853 b := bv .Value
3954
40- // TODO: support unpacking other types of serialized values
41- return pythonPickleString (b )
55+ // The Python SDK originally wrapped pickled values in a
56+ // wrapperspb.BytesValue. Try to unpickle the bytes first,
57+ // and return literal bytes if they cannot be unpickled.
58+ s , err := pythonPickleString (b )
59+ if err != nil {
60+ s = string (truncateBytes (b ))
61+ }
62+ return s , nil
63+ }
64+
65+ func truncateBytes (b []byte ) []byte {
66+ const n = 4
67+ if len (b ) < n {
68+ return b
69+ }
70+ return append (b [:n :n ], "..." ... )
4271}
0 commit comments