Skip to content

Commit b1fb8e4

Browse files
chrisochriso
authored
Merge pull request #79 from dispatchrun/well-known-types
Support introspection of built-in protobuf message types
2 parents 1022c64 + 2b61003 commit b1fb8e4

File tree

2 files changed

+278
-35
lines changed

2 files changed

+278
-35
lines changed

‎cli/any.go‎

Lines changed: 119 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,15 @@ package cli
33
import (
44
"fmt"
55
"log/slog"
6+
"strconv"
7+
"strings"
68

79
pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
8-
"google.golang.org/protobuf/proto"
910
"google.golang.org/protobuf/types/known/anypb"
11+
"google.golang.org/protobuf/types/known/durationpb"
12+
"google.golang.org/protobuf/types/known/emptypb"
13+
"google.golang.org/protobuf/types/known/structpb"
14+
"google.golang.org/protobuf/types/known/timestamppb"
1015
"google.golang.org/protobuf/types/known/wrapperspb"
1116
)
1217

@@ -15,51 +20,130 @@ func anyString(any *anypb.Any) string {
1520
return "nil"
1621
}
1722

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-
}
23+
m, err := any.UnmarshalNew()
24+
if err != nil {
25+
return unsupportedAny(any, err)
26+
}
27+
28+
switch mm := m.(type) {
29+
case *wrapperspb.BytesValue:
30+
// The Python SDK originally wrapped pickled values in a
31+
// wrapperspb.BytesValue. Try to unpickle the bytes first,
32+
// and return literal bytes if they cannot be unpickled.
33+
s, err := pythonPickleString(mm.Value)
34+
if err != nil {
35+
s = fmt.Sprintf("bytes(%s)", truncateBytes(mm.Value))
36+
}
37+
return s
38+
39+
case *wrapperspb.Int32Value:
40+
return strconv.FormatInt(int64(mm.Value), 10)
41+
42+
case *wrapperspb.Int64Value:
43+
return strconv.FormatInt(mm.Value, 10)
44+
45+
case *wrapperspb.UInt32Value:
46+
return strconv.FormatUint(uint64(mm.Value), 10)
47+
48+
case *wrapperspb.UInt64Value:
49+
return strconv.FormatUint(mm.Value, 10)
50+
51+
case *wrapperspb.StringValue:
52+
return fmt.Sprintf("%q", mm.Value)
53+
54+
case *wrapperspb.BoolValue:
55+
return strconv.FormatBool(mm.Value)
56+
57+
case *wrapperspb.FloatValue:
58+
return fmt.Sprintf("%v", mm.Value)
59+
60+
case *wrapperspb.DoubleValue:
61+
return fmt.Sprintf("%v", mm.Value)
62+
63+
case *emptypb.Empty:
64+
return "empty()"
65+
66+
case *timestamppb.Timestamp:
67+
return mm.AsTime().String()
68+
69+
case *durationpb.Duration:
70+
return mm.AsDuration().String()
71+
72+
case *structpb.Struct:
73+
return structpbStructString(mm)
74+
75+
case *structpb.ListValue:
76+
return structpbListString(mm)
77+
78+
case *structpb.Value:
79+
return structpbValueString(mm)
80+
81+
case *pythonv1.Pickled:
82+
s, err := pythonPickleString(mm.PickledValue)
83+
if err != nil {
84+
return unsupportedAny(any, fmt.Errorf("pickle error: %w", err))
3085
}
31-
case "type.googleapis.com/google.protobuf.BytesValue":
32-
s, err = anyBytesString(any)
86+
return s
87+
3388
default:
34-
// TODO: support unpacking other types of serialized values
35-
err = fmt.Errorf("not implemented: %s", any.TypeUrl)
89+
return unsupportedAny(any, fmt.Errorf("not implemented: %T", m))
3690
}
37-
if err != nil {
38-
slog.Debug("cannot parse input/output value", "error", err)
39-
return fmt.Sprintf("%s(?)", any.TypeUrl)
91+
}
92+
93+
func structpbStructString(s *structpb.Struct) string {
94+
var b strings.Builder
95+
b.WriteByte('{')
96+
i := 0
97+
for name, value := range s.Fields {
98+
if i > 0 {
99+
b.WriteString(", ")
100+
}
101+
b.WriteString(fmt.Sprintf("%q", name))
102+
b.WriteString(": ")
103+
b.WriteString(structpbValueString(value))
104+
i++
40105
}
41-
return s
106+
b.WriteByte('}')
107+
return b.String()
42108
}
43109

44-
func anyBytesString(any *anypb.Any) (string, error) {
45-
m, err := anypb.UnmarshalNew(any, proto.UnmarshalOptions{})
46-
if err != nil {
47-
return "", err
110+
func structpbListString(s *structpb.ListValue) string {
111+
var b strings.Builder
112+
b.WriteByte('[')
113+
for i, value := range s.Values {
114+
if i > 0 {
115+
b.WriteString(", ")
116+
}
117+
b.WriteString(structpbValueString(value))
48118
}
49-
bv, ok := m.(*wrapperspb.BytesValue)
50-
if !ok {
51-
return "", fmt.Errorf("invalid bytes value: %T", m)
119+
b.WriteByte(']')
120+
return b.String()
121+
}
122+
123+
func structpbValueString(s *structpb.Value) string {
124+
switch v := s.Kind.(type) {
125+
case *structpb.Value_StructValue:
126+
return structpbStructString(v.StructValue)
127+
case *structpb.Value_ListValue:
128+
return structpbListString(v.ListValue)
129+
case *structpb.Value_BoolValue:
130+
return strconv.FormatBool(v.BoolValue)
131+
case *structpb.Value_NumberValue:
132+
return fmt.Sprintf("%v", v.NumberValue)
133+
case *structpb.Value_StringValue:
134+
return fmt.Sprintf("%q", v.StringValue)
135+
case *structpb.Value_NullValue:
136+
return "null"
137+
default:
138+
panic("unreachable")
52139
}
53-
b := bv.Value
140+
}
54141

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)
142+
func unsupportedAny(any *anypb.Any, err error) string {
59143
if err != nil {
60-
s = string(truncateBytes(b))
144+
slog.Debug("cannot parse input/output value", "error", err)
61145
}
62-
return s, nil
146+
return fmt.Sprintf("%s(?)", any.TypeUrl)
63147
}
64148

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

‎cli/any_test.go‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
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/structpb"
13+
"google.golang.org/protobuf/types/known/timestamppb"
14+
"google.golang.org/protobuf/types/known/wrapperspb"
15+
)
16+
17+
func TestAnyString(t *testing.T) {
18+
for _, test := range []struct {
19+
input *anypb.Any
20+
want string
21+
}{
22+
{
23+
input: asAny(wrapperspb.Bool(true)),
24+
want: "true",
25+
},
26+
{
27+
input: asAny(wrapperspb.Int32(-1)),
28+
want: "-1",
29+
},
30+
{
31+
input: asAny(wrapperspb.Int64(2)),
32+
want: "2",
33+
},
34+
{
35+
input: asAny(wrapperspb.UInt32(3)),
36+
want: "3",
37+
},
38+
{
39+
input: asAny(wrapperspb.UInt64(4)),
40+
want: "4",
41+
},
42+
{
43+
input: asAny(wrapperspb.Float(1.25)),
44+
want: "1.25",
45+
},
46+
{
47+
input: asAny(wrapperspb.Double(3.14)),
48+
want: "3.14",
49+
},
50+
{
51+
input: asAny(wrapperspb.String("foo")),
52+
want: `"foo"`,
53+
},
54+
{
55+
input: asAny(wrapperspb.Bytes([]byte("foobar"))),
56+
want: "bytes(foob...)",
57+
},
58+
{
59+
input: asAny(timestamppb.New(time.Date(2024, time.June, 25, 10, 56, 11, 1234, time.UTC))),
60+
want: "2024-06-25 10:56:11.000001234 +0000 UTC",
61+
},
62+
{
63+
input: asAny(durationpb.New(1 * time.Second)),
64+
want: "1s",
65+
},
66+
{
67+
// $ python3 -c 'import pickle; print(pickle.dumps(1))'
68+
// b'\x80\x04K\x01.'
69+
input: pickled([]byte("\x80\x04K\x01.")),
70+
want: "1",
71+
},
72+
{
73+
// Legacy way that the Python SDK wrapped pickled values:
74+
input: asAny(wrapperspb.Bytes([]byte("\x80\x04K\x01."))),
75+
want: "1",
76+
},
77+
{
78+
// $ python3 -c 'import pickle; print(pickle.dumps("bar"))'
79+
// b'\x80\x04\x95\x07\x00\x00\x00\x00\x00\x00\x00\x8c\x03foo\x94.'
80+
input: pickled([]byte("\x80\x04\x95\x07\x00\x00\x00\x00\x00\x00\x00\x8c\x03bar\x94.")),
81+
want: `"bar"`,
82+
},
83+
{
84+
input: pickled([]byte("!!!invalid!!!")),
85+
want: "buf.build/stealthrocket/dispatch-proto/dispatch.sdk.python.v1.Pickled(?)",
86+
},
87+
{
88+
input: &anypb.Any{TypeUrl: "com.example/some.Message"},
89+
want: "com.example/some.Message(?)",
90+
},
91+
{
92+
input: asAny(&emptypb.Empty{}),
93+
want: "empty()",
94+
},
95+
{
96+
input: asAny(structpb.NewNullValue()),
97+
want: "null",
98+
},
99+
{
100+
input: asAny(structpb.NewBoolValue(false)),
101+
want: "false",
102+
},
103+
{
104+
input: asAny(structpb.NewNumberValue(1111)),
105+
want: "1111",
106+
},
107+
{
108+
input: asAny(structpb.NewNumberValue(3.14)),
109+
want: "3.14",
110+
},
111+
{
112+
input: asAny(structpb.NewStringValue("foobar")),
113+
want: `"foobar"`,
114+
},
115+
{
116+
input: asStructValue([]any{1, true, "abc", nil, map[string]any{}, []any{}}),
117+
want: `[1, true, "abc", null, {}, []]`,
118+
},
119+
{
120+
input: asStructValue(map[string]any{"foo": []any{"bar", "baz"}}),
121+
want: `{"foo": ["bar", "baz"]}`,
122+
},
123+
} {
124+
t.Run(test.want, func(*testing.T) {
125+
got := anyString(test.input)
126+
if got != test.want {
127+
t.Errorf("unexpected string: got %v, want %v", got, test.want)
128+
}
129+
})
130+
}
131+
}
132+
133+
func asAny(m proto.Message) *anypb.Any {
134+
any, err := anypb.New(m)
135+
if err != nil {
136+
panic(err)
137+
}
138+
return any
139+
}
140+
141+
func asStructValue(v any) *anypb.Any {
142+
m, err := structpb.NewValue(v)
143+
if err != nil {
144+
panic(err)
145+
}
146+
return asAny(m)
147+
}
148+
149+
func pickled(b []byte) *anypb.Any {
150+
m := &pythonv1.Pickled{PickledValue: b}
151+
mb, err := proto.Marshal(m)
152+
if err != nil {
153+
panic(err)
154+
}
155+
return &anypb.Any{
156+
TypeUrl: "buf.build/stealthrocket/dispatch-proto/" + string(m.ProtoReflect().Descriptor().FullName()),
157+
Value: mb,
158+
}
159+
}

0 commit comments

Comments
 (0)