Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Pickled value container support #76

Merged
merged 4 commits into from
Jun 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,5 @@ push: image
$(DOCKER) push $(IMAGE)

update:
buf mod update ./proto
for ref in $$(yq -r '.deps[] | .remote + "/gen/go/" + .owner + "/" + .repository + "/protocolbuffers/go@" + .commit' proto/buf.lock); do go get $$ref; done
go mod tidy
49 changes: 39 additions & 10 deletions cli/any.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"log/slog"

pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/wrapperspb"
Expand All @@ -13,17 +14,31 @@ func anyString(any *anypb.Any) string {
if any == nil {
return "nil"
}

var s string
var err error
switch any.TypeUrl {
case "type.googleapis.com/google.protobuf.BytesValue":
s, err := anyBytesString(any)
if err != nil {
slog.Debug("cannot parse input/output value", "error", err)
// fallthrough
} else {
return s
case "buf.build/stealthrocket/dispatch-proto/dispatch.sdk.python.v1.Pickled":
var pickled proto.Message
pickled, err = any.UnmarshalNew()
if err == nil {
if p, ok := pickled.(*pythonv1.Pickled); ok {
s, err = pythonPickleString(p.PickledValue)
} else {
err = fmt.Errorf("invalid pickled message: %T", p)
}
}
case "type.googleapis.com/google.protobuf.BytesValue":
s, err = anyBytesString(any)
default:
// TODO: support unpacking other types of serialized values
err = fmt.Errorf("not implemented: %s", any.TypeUrl)
}
if err != nil {
slog.Debug("cannot parse input/output value", "error", err)
return fmt.Sprintf("%s(?)", any.TypeUrl)
}
return fmt.Sprintf("%s(?)", any.TypeUrl)
return s
}

func anyBytesString(any *anypb.Any) (string, error) {
Expand All @@ -37,6 +52,20 @@ func anyBytesString(any *anypb.Any) (string, error) {
}
b := bv.Value

// TODO: support unpacking other types of serialized values
return pythonPickleString(b)
// The Python SDK originally wrapped pickled values in a
// wrapperspb.BytesValue. Try to unpickle the bytes first,
// and return literal bytes if they cannot be unpickled.
s, err := pythonPickleString(b)
if err != nil {
s = string(truncateBytes(b))
}
return s, nil
}

func truncateBytes(b []byte) []byte {
const n = 4
if len(b) < n {
return b
}
return append(b[:n:n], "..."...)
}
39 changes: 37 additions & 2 deletions cli/tui.go
Original file line number Diff line number Diff line change
Expand Up @@ -552,7 +552,20 @@ func (t *TUI) detailView(id DispatchID) string {
add("Input", rt.request.input)

case *sdkv1.RunRequest_PollResult:
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.PollResult.CoroutineState))))
switch s := d.PollResult.State.(type) {
case *sdkv1.PollResult_CoroutineState:
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
case *sdkv1.PollResult_TypedCoroutineState:
if any := s.TypedCoroutineState; any != nil {
add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
} else {
add("Input", detailLowPriorityStyle.Render("<no state>"))
}
case nil:
add("Input", detailLowPriorityStyle.Render("<no state>"))
default:
add("Input", detailLowPriorityStyle.Render("<unknown state>"))
}
// TODO: show call results
// TODO: show poll error
}
Expand Down Expand Up @@ -593,7 +606,21 @@ func (t *TUI) detailView(id DispatchID) string {

case *sdkv1.RunResponse_Poll:
add("Status", suspendedStyle.Render("Suspended"))
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.Poll.CoroutineState))))

switch s := d.Poll.State.(type) {
case *sdkv1.Poll_CoroutineState:
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
case *sdkv1.Poll_TypedCoroutineState:
if any := s.TypedCoroutineState; any != nil {
add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
} else {
add("Output", detailLowPriorityStyle.Render("<no state>"))
}
case nil:
add("Output", detailLowPriorityStyle.Render("<no state>"))
default:
add("Output", detailLowPriorityStyle.Render("<unknown state>"))
}

if len(d.Poll.Calls) > 0 {
var calls strings.Builder
Expand Down Expand Up @@ -1033,3 +1060,11 @@ func terminalHTTPStatusCode(code int) bool {
return true
}
}

func typeName(typeUrl string) string {
i := strings.LastIndexByte(typeUrl, '/')
if i < 0 {
return typeUrl
}
return typeUrl[i+1:]
}
6 changes: 3 additions & 3 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ module github.com/dispatchrun/dispatch
go 1.22.0

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

require (
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 // indirect
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 // indirect
github.com/atotto/clipboard v0.1.4 // indirect
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 // indirect
Expand Down
14 changes: 6 additions & 8 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 h1:SbKXkoZduR4jQ+aCrX71I+dAcHKZPk1SiriRs/XYX9s=
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1/go.mod h1:Tgn5bgL220vkFOI0KPStlcClPeOJzAv4uT+V8JXGUnw=
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1 h1:C8qLXgA5Y2iK8VkXsNHi8E4qU3aUsZQ0mFqwgxIH7KI=
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1/go.mod h1:yJei/TBJwZBJ8ZUWCKVKceUHk/8gniSGs812SZA9TEE=
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 h1:ilkjxsfdhrdeLL582R4XH+IQLrV2Y31lHB8jb3AQSJA=
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2/go.mod h1:ylS4c28ACSI59oJrOdW4pHS4n0Hw4TgSPHn8rpHl4Yw=
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2 h1:GfZ4I61ZZiXk6zXLZaMNz1rSW0MxlBRRc5n5GOUbIg4=
buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2/go.mod h1:Z1Y+G5dzXnHxfdt7l3S9kLBWV6iO2Zw3z2DvY9yKpaw=
github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
Expand All @@ -18,7 +18,6 @@ github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
Expand Down Expand Up @@ -78,9 +77,8 @@ golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
Expand Down
Loading