Skip to content

Commit 219b268

Browse files
chrisochriso
committed
Support introspection of built-in protobuf message types
1 parent 1022c64 commit 219b268

File tree

2 files changed

+185
-39
lines changed

2 files changed

+185
-39
lines changed

‎cli/any.go

+63-39
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@ package cli
33
import (
44
"fmt"
55
"log/slog"
6+
"strconv"
67

78
pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
8-
"google.golang.org/protobuf/proto"
99
"google.golang.org/protobuf/types/known/anypb"
10+
"google.golang.org/protobuf/types/known/durationpb"
11+
"google.golang.org/protobuf/types/known/emptypb"
12+
"google.golang.org/protobuf/types/known/timestamppb"
1013
"google.golang.org/protobuf/types/known/wrapperspb"
1114
)
1215

@@ -15,51 +18,72 @@ func anyString(any *anypb.Any) string {
1518
return "nil"
1619
}
1720

18-
var s string
19-
var err error
20-
switch any.TypeUrl {
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-
}
30-
}
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-
}
21+
m, err := any.UnmarshalNew()
3722
if err != nil {
38-
slog.Debug("cannot parse input/output value", "error", err)
39-
return fmt.Sprintf("%s(?)", any.TypeUrl)
23+
return unsupportedAny(any, err)
4024
}
41-
return s
42-
}
4325

44-
func anyBytesString(any *anypb.Any) (string, error) {
45-
m, err := anypb.UnmarshalNew(any, proto.UnmarshalOptions{})
46-
if err != nil {
47-
return "", err
48-
}
49-
bv, ok := m.(*wrapperspb.BytesValue)
50-
if !ok {
51-
return "", fmt.Errorf("invalid bytes value: %T", m)
26+
switch mm := m.(type) {
27+
case *wrapperspb.BytesValue:
28+
// The Python SDK originally wrapped pickled values in a
29+
// wrapperspb.BytesValue. Try to unpickle the bytes first,
30+
// and return literal bytes if they cannot be unpickled.
31+
s, err := pythonPickleString(mm.Value)
32+
if err != nil {
33+
s = fmt.Sprintf("bytes(%s)", truncateBytes(mm.Value))
34+
}
35+
return s
36+
37+
case *wrapperspb.Int32Value:
38+
return strconv.FormatInt(int64(mm.Value), 10)
39+
40+
case *wrapperspb.Int64Value:
41+
return strconv.FormatInt(mm.Value, 10)
42+
43+
case *wrapperspb.UInt32Value:
44+
return strconv.FormatUint(uint64(mm.Value), 10)
45+
46+
case *wrapperspb.UInt64Value:
47+
return strconv.FormatUint(mm.Value, 10)
48+
49+
case *wrapperspb.StringValue:
50+
return fmt.Sprintf("%q", mm.Value)
51+
52+
case *wrapperspb.BoolValue:
53+
return strconv.FormatBool(mm.Value)
54+
55+
case *wrapperspb.FloatValue:
56+
return fmt.Sprintf("%v", mm.Value)
57+
58+
case *wrapperspb.DoubleValue:
59+
return fmt.Sprintf("%v", mm.Value)
60+
61+
case *emptypb.Empty:
62+
return "empty()"
63+
64+
case *timestamppb.Timestamp:
65+
return mm.AsTime().String()
66+
67+
case *durationpb.Duration:
68+
return mm.AsDuration().String()
69+
70+
case *pythonv1.Pickled:
71+
s, err := pythonPickleString(mm.PickledValue)
72+
if err != nil {
73+
return unsupportedAny(any, fmt.Errorf("pickle error: %w", err))
74+
}
75+
return s
76+
77+
default:
78+
return unsupportedAny(any, fmt.Errorf("not implemented: %T", m))
5279
}
53-
b := bv.Value
80+
}
5481

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)
82+
func unsupportedAny(any *anypb.Any, err error) string {
5983
if err != nil {
60-
s = string(truncateBytes(b))
84+
slog.Debug("cannot parse input/output value", "error", err)
6185
}
62-
return s, nil
86+
return fmt.Sprintf("%s(?)", any.TypeUrl)
6387
}
6488

6589
func truncateBytes(b []byte) []byte {

‎cli/any_test.go

+122
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
package cli
2+
3+
import (
4+
"testing"
5+
"time"
6+
7+
pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
8+
"google.golang.org/protobuf/proto"
9+
"google.golang.org/protobuf/types/known/anypb"
10+
"google.golang.org/protobuf/types/known/durationpb"
11+
"google.golang.org/protobuf/types/known/emptypb"
12+
"google.golang.org/protobuf/types/known/timestamppb"
13+
"google.golang.org/protobuf/types/known/wrapperspb"
14+
)
15+
16+
func TestAnyString(t *testing.T) {
17+
for _, test := range []struct {
18+
input *anypb.Any
19+
want string
20+
}{
21+
{
22+
input: asAny(wrapperspb.Bool(true)),
23+
want: "true",
24+
},
25+
{
26+
input: asAny(wrapperspb.Int32(-1)),
27+
want: "-1",
28+
},
29+
{
30+
input: asAny(wrapperspb.Int64(2)),
31+
want: "2",
32+
},
33+
{
34+
input: asAny(wrapperspb.UInt32(3)),
35+
want: "3",
36+
},
37+
{
38+
input: asAny(wrapperspb.UInt64(4)),
39+
want: "4",
40+
},
41+
{
42+
input: asAny(wrapperspb.Float(1.25)),
43+
want: "1.25",
44+
},
45+
{
46+
input: asAny(wrapperspb.Double(3.14)),
47+
want: "3.14",
48+
},
49+
{
50+
input: asAny(wrapperspb.String("foo")),
51+
want: `"foo"`,
52+
},
53+
{
54+
input: asAny(wrapperspb.Bytes([]byte("foobar"))),
55+
want: "bytes(foob...)",
56+
},
57+
{
58+
input: asAny(timestamppb.New(time.Date(2024, time.June, 25, 10, 56, 11, 1234, time.UTC))),
59+
want: "2024-06-25 10:56:11.000001234 +0000 UTC",
60+
},
61+
{
62+
input: asAny(durationpb.New(1 * time.Second)),
63+
want: "1s",
64+
},
65+
{
66+
// $ python3 -c 'import pickle; print(pickle.dumps(1))'
67+
// b'\x80\x04K\x01.'
68+
input: pickled([]byte("\x80\x04K\x01.")),
69+
want: "1",
70+
},
71+
{
72+
// Legacy way that the Python SDK wrapped pickled values:
73+
input: asAny(wrapperspb.Bytes([]byte("\x80\x04K\x01."))),
74+
want: "1",
75+
},
76+
{
77+
// $ python3 -c 'import pickle; print(pickle.dumps("bar"))'
78+
// b'\x80\x04\x95\x07\x00\x00\x00\x00\x00\x00\x00\x8c\x03foo\x94.'
79+
input: pickled([]byte("\x80\x04\x95\x07\x00\x00\x00\x00\x00\x00\x00\x8c\x03bar\x94.")),
80+
want: `"bar"`,
81+
},
82+
{
83+
input: pickled([]byte("!!!invalid!!!")),
84+
want: "buf.build/stealthrocket/dispatch-proto/dispatch.sdk.python.v1.Pickled(?)",
85+
},
86+
{
87+
input: &anypb.Any{TypeUrl: "com.example/some.Message"},
88+
want: "com.example/some.Message(?)",
89+
},
90+
{
91+
input: asAny(&emptypb.Empty{}),
92+
want: "empty()",
93+
},
94+
} {
95+
t.Run(test.want, func(*testing.T) {
96+
got := anyString(test.input)
97+
if got != test.want {
98+
t.Errorf("unexpected string: got %v, want %v", got, test.want)
99+
}
100+
})
101+
}
102+
}
103+
104+
func asAny(m proto.Message) *anypb.Any {
105+
any, err := anypb.New(m)
106+
if err != nil {
107+
panic(err)
108+
}
109+
return any
110+
}
111+
112+
func pickled(b []byte) *anypb.Any {
113+
m := &pythonv1.Pickled{PickledValue: b}
114+
mb, err := proto.Marshal(m)
115+
if err != nil {
116+
panic(err)
117+
}
118+
return &anypb.Any{
119+
TypeUrl: "buf.build/stealthrocket/dispatch-proto/" + string(m.ProtoReflect().Descriptor().FullName()),
120+
Value: mb,
121+
}
122+
}

0 commit comments

Comments
 (0)