Skip to content

Commit dbce3be

Browse files
chrisochriso
authored
Merge pull request #76 from dispatchrun/pickled-value-container
Pickled value container support
2 parents 7ca508c + 775e933 commit dbce3be

File tree

5 files changed

+85
-24
lines changed

5 files changed

+85
-24
lines changed

‎Makefile‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,5 @@ push: image
3737
$(DOCKER) push $(IMAGE)
3838

3939
update:
40-
buf mod update ./proto
4140
for ref in $$(yq -r '.deps[] | .remote + "/gen/go/" + .owner + "/" + .repository + "/protocolbuffers/go@" + .commit' proto/buf.lock); do go get $$ref; done
4241
go mod tidy

‎cli/any.go‎

Lines changed: 39 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
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

2944
func 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
}

‎cli/tui.go‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -552,7 +552,20 @@ func (t *TUI) detailView(id DispatchID) string {
552552
add("Input", rt.request.input)
553553

554554
case *sdkv1.RunRequest_PollResult:
555-
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.PollResult.CoroutineState))))
555+
switch s := d.PollResult.State.(type) {
556+
case *sdkv1.PollResult_CoroutineState:
557+
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
558+
case *sdkv1.PollResult_TypedCoroutineState:
559+
if any := s.TypedCoroutineState; any != nil {
560+
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
561+
} else {
562+
add("Input", detailLowPriorityStyle.Render("<no state>"))
563+
}
564+
case nil:
565+
add("Input", detailLowPriorityStyle.Render("<no state>"))
566+
default:
567+
add("Input", detailLowPriorityStyle.Render("<unknown state>"))
568+
}
556569
// TODO: show call results
557570
// TODO: show poll error
558571
}
@@ -593,7 +606,21 @@ func (t *TUI) detailView(id DispatchID) string {
593606

594607
case *sdkv1.RunResponse_Poll:
595608
add("Status", suspendedStyle.Render("Suspended"))
596-
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.Poll.CoroutineState))))
609+
610+
switch s := d.Poll.State.(type) {
611+
case *sdkv1.Poll_CoroutineState:
612+
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
613+
case *sdkv1.Poll_TypedCoroutineState:
614+
if any := s.TypedCoroutineState; any != nil {
615+
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
616+
} else {
617+
add("Output", detailLowPriorityStyle.Render("<no state>"))
618+
}
619+
case nil:
620+
add("Output", detailLowPriorityStyle.Render("<no state>"))
621+
default:
622+
add("Output", detailLowPriorityStyle.Render("<unknown state>"))
623+
}
597624

598625
if len(d.Poll.Calls) > 0 {
599626
var calls strings.Builder
@@ -1033,3 +1060,11 @@ func terminalHTTPStatusCode(code int) bool {
10331060
return true
10341061
}
10351062
}
1063+
1064+
func typeName(typeUrl string) string {
1065+
i := strings.LastIndexByte(typeUrl, '/')
1066+
if i < 0 {
1067+
return typeUrl
1068+
}
1069+
return typeUrl[i+1:]
1070+
}

‎go.mod‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module github.com/dispatchrun/dispatch
33
go 1.22.0
44

55
require (
6-
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1
6+
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2
77
github.com/charmbracelet/bubbles v0.18.0
88
github.com/charmbracelet/bubbletea v0.25.0
99
github.com/charmbracelet/lipgloss v0.9.1
@@ -14,11 +14,11 @@ require (
1414
github.com/spf13/cobra v1.8.0
1515
github.com/stretchr/testify v1.9.0
1616
golang.org/x/term v0.19.0
17-
google.golang.org/protobuf v1.33.0
17+
google.golang.org/protobuf v1.34.2
1818
)
1919

2020
require (
21-
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 // indirect
21+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 // indirect
2222
github.com/atotto/clipboard v0.1.4 // indirect
2323
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
2424
github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 // indirect

‎go.sum‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 h1:SbKXkoZduR4jQ+aCrX71I+dAcHKZPk1SiriRs/XYX9s=
2-
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1/go.mod h1:Tgn5bgL220vkFOI0KPStlcClPeOJzAv4uT+V8JXGUnw=
3-
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1 h1:C8qLXgA5Y2iK8VkXsNHi8E4qU3aUsZQ0mFqwgxIH7KI=
4-
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1/go.mod h1:yJei/TBJwZBJ8ZUWCKVKceUHk/8gniSGs812SZA9TEE=
1+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 h1:ilkjxsfdhrdeLL582R4XH+IQLrV2Y31lHB8jb3AQSJA=
2+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2/go.mod h1:ylS4c28ACSI59oJrOdW4pHS4n0Hw4TgSPHn8rpHl4Yw=
3+
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2 h1:GfZ4I61ZZiXk6zXLZaMNz1rSW0MxlBRRc5n5GOUbIg4=
4+
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2/go.mod h1:Z1Y+G5dzXnHxfdt7l3S9kLBWV6iO2Zw3z2DvY9yKpaw=
55
github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
66
github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
77
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
@@ -18,7 +18,6 @@ github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
1818
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
1919
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
2020
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
21-
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
2221
github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
2322
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
2423
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
@@ -78,9 +77,8 @@ golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
7877
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
7978
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
8079
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
81-
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
82-
google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
83-
google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
80+
google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
81+
google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
8482
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
8583
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
8684
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

0 commit comments

Comments
 (0)