Repository: yarpc/yab
Pull request #314 (2021-01-08 09:02:02)

gRPC streams support - 1
## Aim
Support gRPC streams in Yab

## Changes
* Exposes a new interface in serializer for creating stream request and multiple stream message body reader
* Refactors `getRequestInput` to provider a reader interface which would be used for reading multiple stream message body

Next Stack: #315 


Patch:
diff --git a/encoding/encoding.go b/encoding/encoding.go
index b6b1c75..145bb9a 100644
--- a/encoding/encoding.go
+++ b/encoding/encoding.go
@@ -24,6 +24,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/yarpc/yab/thrift"
@@ -51,6 +52,20 @@ type Serializer interface {
 	CheckSuccess(body *transport.Response) error
 }
 
+// StreamRequestReader interface exposes method to read multiple request body
+type StreamRequestReader interface {
+	// NextBody returns the encoded request body if available, and if not, returns an
+	// io.EOF to indicate end of requests and caller must not call it again
+	NextBody() ([]byte, error)
+}
+
+// StreamSerializer serializes and deserializes data for a stream requests
+type StreamSerializer interface {
+	// StreamRequest creates a root stream request, a stream request reader using
+	// body reader provided
+	StreamRequest(body io.Reader) (*transport.StreamRequest, StreamRequestReader, error)
+}
+
 // The list of supported encodings.
 const (
 	UnspecifiedEncoding Encoding = ""
diff --git a/encoding/inputdecoder/decoder.go b/encoding/inputdecoder/decoder.go
new file mode 100644
index 0000000..4087dbc
--- /dev/null
+++ b/encoding/inputdecoder/decoder.go
@@ -0,0 +1,96 @@
+// Copyright (c) 2021 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package inputdecoder
+
+import (
+	"bufio"
+	"encoding/json"
+	"io"
+
+	gyaml "github.com/ghodss/yaml"
+	"gopkg.in/yaml.v2"
+)
+
+// Decoder interface exposes method for reading multiple input requests
+type Decoder interface {
+	// NextJSONBytes returns JSON marshaled bytes of the next request body
+	// It returns io.EOF when there are no more requests
+	NextJSONBytes() ([]byte, error)
+}
+
+// jsonInputDecoder parses multiple JSON objects from given reader
+// JSON objects can be delimited by space or newline
+type jsonInputDecoder struct{ dec *json.Decoder }
+
+func (r *jsonInputDecoder) NextJSONBytes() ([]byte, error) {
+	if !r.dec.More() {
+		return nil, io.EOF
+	}
+
+	var v json.RawMessage
+	err := r.dec.Decode(&v)
+	return []byte(v), err
+}
+
+// yamlInputDecoder parses multiple YAML objects from the given reader
+// consecutive YAML objects must be delimited by `---`
+type yamlInputDecoder struct{ dec *yaml.Decoder }
+
+func (r *yamlInputDecoder) NextJSONBytes() ([]byte, error) {
+	var v interface{}
+	if err := r.dec.Decode(&v); err != nil {
+		return nil, err
+	}
+
+	bytes, err := yaml.Marshal(v)
+	if err != nil {
+		return nil, err
+	}
+
+	return gyaml.YAMLToJSON(bytes)
+}
+
+// isJSONInput assumes the input is JSON compatible if the initial
+// byte read is `{`
+func isJSONInput(r *bufio.Reader) (bool, error) {
+	b, err := r.ReadByte()
+	if err != nil {
+		return false, err
+	}
+
+	return b == '{', r.UnreadByte()
+}
+
+// New detects the input encoding type, returns either
+// json or yaml request decoder
+func New(reader io.Reader) (Decoder, error) {
+	bufReader := bufio.NewReader(reader)
+	isJSON, err := isJSONInput(bufReader)
+	if err != nil && err != io.EOF {
+		return nil, err
+	}
+
+	if isJSON {
+		return &jsonInputDecoder{json.NewDecoder(bufReader)}, nil
+	}
+
+	return &yamlInputDecoder{yaml.NewDecoder(bufReader)}, nil
+}
diff --git a/encoding/protobuf.go b/encoding/protobuf.go
index a552c2e..f7ee72d 100644
--- a/encoding/protobuf.go
+++ b/encoding/protobuf.go
@@ -3,9 +3,11 @@ package encoding
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/yarpc/yab/encoding/encodingerror"
+	"github.com/yarpc/yab/encoding/inputdecoder"
 	"github.com/yarpc/yab/protobuf"
 	"github.com/yarpc/yab/transport"
 
@@ -109,14 +111,9 @@ func (p protoSerializer) Request(body []byte) (*transport.Request, error) {
 	if err != nil {
 		return nil, err
 	}
-
-	req := dynamic.NewMessage(p.method.GetInputType())
-	if err := req.UnmarshalJSON(jsonContent); err != nil {
-		return nil, fmt.Errorf("could not parse given request body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
-	}
-	bytes, err := proto.Marshal(req)
+	bytes, err := p.encode(jsonContent)
 	if err != nil {
-		return nil, fmt.Errorf("could marshal message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
+		return nil, err
 	}
 	return &transport.Request{
 		Method: procedure.ToName(p.serviceName, p.methodName),
@@ -144,11 +141,59 @@ func (p protoSerializer) Response(body *transport.Response) (interface{}, error)
 	return unmarshaledJSON, nil
 }
 
+func (p protoSerializer) StreamRequest(body io.Reader) (*transport.StreamRequest, StreamRequestReader, error) {
+	decoder, err := inputdecoder.New(body)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	streamReq := &transport.StreamRequest{
+		Request: &transport.Request{
+			Method: procedure.ToName(p.serviceName, p.methodName),
+		},
+	}
+
+	reader := protoStreamRequestReader{
+		decoder: decoder,
+		proto:   p,
+	}
+
+	return streamReq, reader, nil
+}
+
 func (p protoSerializer) CheckSuccess(body *transport.Response) error {
 	_, err := p.Response(body)
 	return err
 }
 
+func (p protoSerializer) encode(jsonBytes []byte) ([]byte, error) {
+	req := dynamic.NewMessage(p.method.GetInputType())
+	if err := req.UnmarshalJSON(jsonBytes); err != nil {
+		return nil, fmt.Errorf("could not parse given request body as message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
+	}
+
+	bytes, err := proto.Marshal(req)
+	if err != nil {
+		return nil, fmt.Errorf("could marshal message of type %q: %v", p.method.GetInputType().GetFullyQualifiedName(), err)
+	}
+
+	return bytes, nil
+}
+
+type protoStreamRequestReader struct {
+	decoder inputdecoder.Decoder
+	proto   protoSerializer
+}
+
+func (p protoStreamRequestReader) NextBody() ([]byte, error) {
+	body, err := p.decoder.NextJSONBytes()
+	if err != nil {
+		return nil, err
+	}
+
+	return p.proto.encode(body)
+}
+
 func splitMethod(fullMethod string) (svc, method string, err error) {
 	parts := strings.Split(fullMethod, "/")
 	switch len(parts) {
diff --git a/main.go b/main.go
index 0041fa6..fc86234 100644
--- a/main.go
+++ b/main.go
@@ -25,6 +25,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 	"os"
 	"regexp"
@@ -287,10 +288,11 @@ func runWithOptions(opts Options, out output, logger *zap.Logger) {
 		return
 	}
 
-	reqInput, err := getRequestInput(opts.ROpts.RequestJSON, opts.ROpts.RequestFile)
+	reqReader, err := getRequestInput(opts.ROpts.RequestJSON, opts.ROpts.RequestFile)
 	if err != nil {
-		out.Fatalf("Failed while loading body input: %v\n", err)
+		out.Fatalf("Failed while creating request reader: %v\n", err)
 	}
+	defer reqReader.Close()
 
 	headers, err := getHeaders(opts.ROpts.HeadersJSON, opts.ROpts.HeadersFile, opts.ROpts.Headers)
 	if err != nil {
@@ -338,6 +340,11 @@ func runWithOptions(opts Options, out output, logger *zap.Logger) {
 		out.Fatalf("Failed while parsing options: %v\n", err)
 	}
 
+	reqInput, err := ioutil.ReadAll(reqReader)
+	if err != nil {
+		out.Fatalf("Failed while reading body input: %v\n", err)
+	}
+
 	// req is the transport.Request that will be used to make a call.
 	req, err := serializer.Request(reqInput)
 	if err != nil {
diff --git a/request.go b/request.go
index d8d26d9..5ab7905 100644
--- a/request.go
+++ b/request.go
@@ -21,9 +21,11 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -42,29 +44,30 @@ var (
 )
 
 // getRequestInput gets the byte body passed in by the user via flags or through a file.
-func getRequestInput(inline, file string) ([]byte, error) {
+func getRequestInput(inline, file string) (io.ReadCloser, error) {
 	if file == "-" || inline == "-" {
-		return ioutil.ReadAll(os.Stdin)
+		return os.Stdin, nil
 	}
 
 	if file != "" {
-		bs, err := ioutil.ReadFile(file)
+		f, err := os.Open(file)
 		if err != nil {
 			return nil, fmt.Errorf("failed to open request file: %v", err)
 		}
-		return bs, nil
+		return f, nil
 	}
 
-	if inline != "" {
-		return []byte(inline), nil
-	}
-
-	// It is valid to have an empty body.
-	return nil, nil
+	return ioutil.NopCloser(bytes.NewReader([]byte(inline))), nil
 }
 
 func getHeaders(inline, file string, override map[string]string) (map[string]string, error) {
-	contents, err := getRequestInput(inline, file)
+	r, err := getRequestInput(inline, file)
+	if err != nil {
+		return nil, err
+	}
+	defer r.Close()
+
+	contents, err := ioutil.ReadAll(r)
 	if err != nil {
 		return nil, err
 	}
diff --git a/transport/interface.go b/transport/interface.go
index b90d3cd..ed8eb24 100644
--- a/transport/interface.go
+++ b/transport/interface.go
@@ -40,6 +40,11 @@ type Request struct {
 	Body             []byte
 }
 
+// StreamRequest is a wrapper of Request, to be used for streaming RPC
+type StreamRequest struct {
+	Request *Request
+}
+
 // Response represents the result of an RPC.
 type Response struct {
 	Headers map[string]string


Test patch:
diff --git a/encoding/inputdecoder/decoder_test.go b/encoding/inputdecoder/decoder_test.go
new file mode 100644
index 0000000..cb997aa
--- /dev/null
+++ b/encoding/inputdecoder/decoder_test.go
@@ -0,0 +1,133 @@
+// Copyright (c) 2021 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package inputdecoder
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsJSONInput(t *testing.T) {
+	t.Run("valid json", func(t *testing.T) {
+		buf := bufio.NewReader(bytes.NewReader([]byte(`{}`)))
+		isJSON, err := isJSONInput(buf)
+		assert.True(t, isJSON, "expected true")
+		assert.NoError(t, err, "unexpected error")
+	})
+	t.Run("invalid json", func(t *testing.T) {
+		buf := bufio.NewReader(bytes.NewReader([]byte(`test: json`)))
+		isJSON, err := isJSONInput(buf)
+		assert.False(t, isJSON, "expected false")
+		assert.NoError(t, err, "unexpected error")
+	})
+}
+
+func TestRequestReader(t *testing.T) {
+	t.Run("Use json decoder", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`{}{}{}`))
+		reader, err := New(buf)
+		assert.IsTypef(t, &jsonInputDecoder{}, reader, "expected json decoder")
+		assert.NoError(t, err, "unexpected error")
+	})
+	t.Run("Use yaml decoder", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`a:b`))
+		reader, err := New(buf)
+		assert.IsTypef(t, &yamlInputDecoder{}, reader, "expected yaml decoder")
+		assert.NoError(t, err, "unexpected error")
+	})
+	t.Run("Parse json request", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`{"test":1} {"test":2}`))
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(`{"test":1}`), body)
+		assert.NoError(t, err)
+		body, err = reader.NextJSONBytes()
+		assert.Equal(t, []byte(`{"test":2}`), body)
+		assert.NoError(t, err)
+		body, err = reader.NextJSONBytes()
+		assert.Nil(t, body)
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+	t.Run("Parse yaml request", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`test: 1
+---
+test: 2`))
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(`{"test":1}`), body)
+		assert.NoError(t, err)
+		body, err = reader.NextJSONBytes()
+		assert.Equal(t, []byte(`{"test":2}`), body)
+		assert.NoError(t, err)
+		body, err = reader.NextJSONBytes()
+		assert.Nil(t, body)
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+	t.Run("error parsing second non-json input", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`{"test": 1}
+---
+test: 2`))
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(`{"test": 1}`), body)
+		assert.NoError(t, err)
+		body, err = reader.NextJSONBytes()
+		assert.Error(t, err)
+	})
+	t.Run("empty input", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(``))
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(nil), body)
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+	t.Run("nil input", func(t *testing.T) {
+		buf := bytes.NewReader(nil)
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(nil), body)
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+	t.Run("invalid yaml", func(t *testing.T) {
+		buf := bytes.NewReader([]byte(`a:
+		b`))
+		reader, err := New(buf)
+		assert.NotNil(t, reader, "expected non-nil reader")
+		assert.NoError(t, err)
+		body, err := reader.NextJSONBytes()
+		assert.Equal(t, []byte(nil), body)
+		assert.EqualError(t, err, `yaml: line 2: found character that cannot start any token`)
+	})
+}
diff --git a/encoding/protobuf_test.go b/encoding/protobuf_test.go
index 0d492ac..a4a05e0 100644
--- a/encoding/protobuf_test.go
+++ b/encoding/protobuf_test.go
@@ -1,9 +1,11 @@
 package encoding
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"testing"
 
 	"github.com/yarpc/yab/protobuf"
@@ -286,6 +288,69 @@ func Test_anyResolver_Resolve(t *testing.T) {
 	}
 }
 
+type errorReader struct {
+	err error
+}
+
+func (e errorReader) Read(b []byte) (int, error) {
+	return 0, e.err
+}
+
+func TestProtobufStreamReader(t *testing.T) {
+	source, err := protobuf.NewDescriptorProviderFileDescriptorSetBins("../testdata/protobuf/simple/simple.proto.bin")
+	assert.NoError(t, err)
+
+	t.Run("pass", func(t *testing.T) {
+		serializer, err := NewProtobuf("Bar/Baz", source)
+		assert.NoError(t, err)
+		streamSerializer, ok := serializer.(StreamSerializer)
+		assert.True(t, ok)
+		req, reader, err := streamSerializer.StreamRequest(bytes.NewReader([]byte(`{"test": 10}`)))
+		assert.NoError(t, err)
+		expectedReq := &transport.StreamRequest{
+			Request: &transport.Request{Method: "Bar::Baz"},
+		}
+		assert.Equal(t, expectedReq, req)
+		body, err := reader.NextBody()
+		assert.NoError(t, err)
+		assert.Equal(t, []byte{0x8, 0xa}, body)
+		_, err = reader.NextBody()
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+
+	t.Run("invalid input", func(t *testing.T) {
+		serializer, err := NewProtobuf("Bar/Baz", source)
+		assert.NoError(t, err)
+		streamSerializer := serializer.(StreamSerializer)
+		_, reader, err := streamSerializer.StreamRequest(bytes.NewReader([]byte(`{`)))
+		assert.NoError(t, err)
+		body, err := reader.NextBody()
+		assert.EqualError(t, err, "unexpected EOF")
+		assert.Nil(t, body)
+	})
+
+	t.Run("reader error", func(t *testing.T) {
+		serializer, err := NewProtobuf("Bar/Baz", source)
+		assert.NoError(t, err)
+		streamSerializer := serializer.(StreamSerializer)
+		reader := errorReader{err: errors.New("test error")}
+		_, _, err = streamSerializer.StreamRequest(reader)
+		assert.EqualError(t, err, "test error")
+	})
+
+	t.Run("reader error", func(t *testing.T) {
+		serializer, err := NewProtobuf("Bar/Baz", source)
+		assert.NoError(t, err)
+		streamSerializer := serializer.(StreamSerializer)
+		reader := &errorReader{err: io.EOF}
+		_, streamReqReader, err := streamSerializer.StreamRequest(reader)
+		assert.NoError(t, err)
+		reader.err = errors.New("test error")
+		_, err = streamReqReader.NextBody()
+		assert.EqualError(t, err, "yaml: input error: test error")
+	})
+}
+
 func getAnyType(t *testing.T, typeURL string, value proto.Message) []byte {
 	valueContent, err := proto.Marshal(value)
 	require.NoError(t, err)
diff --git a/main_test.go b/main_test.go
index aa0991c..260de56 100644
--- a/main_test.go
+++ b/main_test.go
@@ -209,6 +209,22 @@ func TestRunWithOptions(t *testing.T) {
 			},
 			errMsg: "Disallowed caller name: testBlockedCallerName",
 		},
+		{
+			desc: "Fail due to incorrect file",
+			opts: Options{
+				ROpts: RequestOptions{
+					ThriftFile:  validThrift,
+					Procedure:   fooMethod,
+					Timeout:     timeMillisFlag(time.Nanosecond),
+					RequestFile: "invalid",
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+					Peers:       []string{echoServer(t, fooMethod, nil)},
+				},
+			},
+			errMsg: "Failed while creating request reader: failed to open request file: open invalid: no such file or directory",
+		},
 	}
 
 	var errBuf bytes.Buffer
diff --git a/request_test.go b/request_test.go
index 911c24d..58c2b25 100644
--- a/request_test.go
+++ b/request_test.go
@@ -60,7 +60,7 @@ func TestGetRequestInput(t *testing.T) {
 		want   []byte
 	}{
 		{
-			want: nil,
+			want: []byte{},
 		},
 		{
 			file:   "/fake/file",
@@ -114,7 +114,9 @@ func TestGetRequestInput(t *testing.T) {
 		}
 
 		if assert.NoError(t, err, "getRequestInput(%v, %v) should not fail", tt.inline, tt.file) {
-			assert.Equal(t, tt.want, got, "getRequestInput(%v, %v) mismatch", tt.inline, tt.file)
+			contents, err := ioutil.ReadAll(got)
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, contents, "getRequestInput(%v, %v) mismatch", tt.inline, tt.file)
 		}
 	}
 }