Repository: googleapis/gapic-generator-go
Pull request #933 (2022-03-16 21:20:02)

feat(gengapic): regapic server-streaming
This adds support for REGAPIC Server Stream RPCs. Since server-stream RPCs in GAPICs must return the gRPC-defined, RPC-specific "stream client" interface, we must implement that interface around the HTTP-JSON stream (the Response Body) opened by the HTTP-JSON request. An RPC-specific implementation of that stream client interface is generated in the GAPIC that utilizes the `gax.ProtoJSONStream` helper to decode the stream contents into the proto messages. The behavior of this generated stream client is the same as documented by the gRPC-defined interface.

When the stream is done/closed, `io.EOF` is returned.

If the `context.Context` used to initialize the stream is canceled, the invocations of `Recv` fail.

When the stream is finished or in the event of an error, the `gax.ProtoJSONStream` is closed by the generated stream client.

noahdietz: I regenerated the Firestore client with these changes and with REST transport generation enabled. I was then able to successfully consume a RunQuery REST stream.

Patch:
diff --git a/internal/gengapic/doc_file.go b/internal/gengapic/doc_file.go
index 3f94d20..179ac20 100644
--- a/internal/gengapic/doc_file.go
+++ b/internal/gengapic/doc_file.go
@@ -216,6 +216,7 @@ func (g *generator) genDocFile(year int, scopes []string, serv *descriptor.Servi
 	}
 
 	if hasREST {
+		// UnknownEnum error check helper.
 		p("// maybeUnknownEnum wraps the given proto-JSON parsing error if it is the result")
 		p("// of receiving an unknown enum value.")
 		p("func maybeUnknownEnum(err error) error {")
@@ -225,6 +226,8 @@ func (g *generator) genDocFile(year int, scopes []string, serv *descriptor.Servi
 		p("  return err")
 		p("}")
 		p("")
+
+		// buildHeaders from context and other metadata helper.
 		p("// buildHeaders extracts metadata from the outgoing context, joins it with any other")
 		p("// given metadata, and converts them into a http.Header. ")
 		p("func buildHeaders(ctx context.Context, mds ...metadata.MD) http.Header {")
@@ -234,6 +237,7 @@ func (g *generator) genDocFile(year int, scopes []string, serv *descriptor.Servi
 		p("  md := metadata.Join(mds...)")
 		p("  return http.Header(md)")
 		p("}")
+		p("")
 	}
 }
 
diff --git a/internal/gengapic/genrest.go b/internal/gengapic/genrest.go
index f9d53db..2b35c7b 100644
--- a/internal/gengapic/genrest.go
+++ b/internal/gengapic/genrest.go
@@ -489,9 +489,10 @@ func (g *generator) genRESTMethod(servName string, serv *descriptor.ServiceDescr
 }
 
 func (g *generator) serverStreamRESTCall(servName string, s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) error {
-	// Streaming calls are not currently supported for REST clients,
-	// but the interface signature must be preserved.
-	// Unimplemented REST methods will always error.
+	info := getHTTPInfo(m)
+	if info == nil {
+		return errors.E(nil, "method has no http info: %s", m.GetName())
+	}
 
 	inType := g.descInfo.Type[m.GetInputType()]
 
@@ -501,6 +502,14 @@ func (g *generator) serverStreamRESTCall(servName string, s *descriptor.ServiceD
 	}
 	g.imports[inSpec] = true
 
+	outType := g.descInfo.Type[m.GetOutputType()]
+
+	outSpec, err := g.descInfo.ImportSpec(outType)
+	if err != nil {
+		return err
+	}
+	g.imports[outSpec] = true
+
 	servSpec, err := g.descInfo.ImportSpec(s)
 	if err != nil {
 		return err
@@ -509,12 +518,127 @@ func (g *generator) serverStreamRESTCall(servName string, s *descriptor.ServiceD
 
 	p := g.printf
 	lowcaseServName := lowcaseRestClientName(servName)
+	streamClient := fmt.Sprintf("%sRESTClient", lowerFirst(m.GetName()))
+
+	// rest-client method
 	p("func (c *%s) %s(ctx context.Context, req *%s.%s, opts ...gax.CallOption) (%s.%s_%sClient, error) {",
 		lowcaseServName, m.GetName(), inSpec.Name, inType.GetName(), servSpec.Name, s.GetName(), m.GetName())
-	p(`  return nil, fmt.Errorf("%s not yet supported for REST clients")`, m.GetName())
+	body := "nil"
+	verb := strings.ToUpper(info.verb)
+
+	// Marshal body for HTTP methods that take a body.
+	if info.body != "" {
+		if verb == http.MethodGet || verb == http.MethodDelete {
+			return fmt.Errorf("invalid use of body parameter for a get/delete method %q", m.GetName())
+		}
+		p("m := protojson.MarshalOptions{AllowPartial: true}")
+		requestObject := "req"
+		if info.body != "*" {
+			requestObject = "body"
+			p("body := req%s", fieldGetter(info.body))
+		}
+		p("jsonReq, err := m.Marshal(%s)", requestObject)
+		p("if err != nil {")
+		p("  return nil, err")
+		p("}")
+		p("")
+
+		body = "bytes.NewReader(jsonReq)"
+		g.imports[pbinfo.ImportSpec{Path: "bytes"}] = true
+	}
+
+	g.generateURLString(m)
+	g.generateQueryString(m)
+	p("// Build HTTP headers from client and context metadata.")
+	p(`headers := buildHeaders(ctx, c.xGoogMetadata, metadata.Pairs("Content-Type", "application/json"))`)
+	p("var streamClient *%s", streamClient)
+	p("e := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {")
+	p(`  httpReq, err := http.NewRequest("%s", baseUrl.String(), %s)`, verb, body)
+	p("  if err != nil {")
+	p("      return err")
+	p("  }")
+	p("  httpReq = httpReq.WithContext(ctx)")
+	p("  httpReq.Header = headers")
+	p("")
+	p("  httpRsp, err := c.httpClient.Do(httpReq)")
+	p("  if err != nil{")
+	p("   return err")
+	p("  }")
+	p("")
+	p("  if err = googleapi.CheckResponse(httpRsp); err != nil {")
+	p("    return err")
+	p("  }")
+	p("")
+	p("  streamClient = &%s{", streamClient)
+	p("    ctx: ctx,")
+	p("    md: metadata.MD(httpRsp.Header),")
+	p("    stream: gax.NewProtoJSONStreamReader(httpRsp.Body, (&%s.%s{}).ProtoReflect().Type()),", outSpec.Name, outType.GetName())
+	p("  }")
+	p("  return nil")
+	p("}, opts...)")
+	p("")
+	p("return streamClient, e")
 	p("}")
 	p("")
 
+	g.imports[pbinfo.ImportSpec{Path: "google.golang.org/api/googleapi"}] = true
+	g.imports[pbinfo.ImportSpec{Path: "google.golang.org/protobuf/encoding/protojson"}] = true
+
+	// server-stream wrapper client
+	p("// %s is the stream client used to consume the server stream created by", streamClient)
+	p("// the REST implementation of %s.", m.GetName())
+	p("type %s struct {", streamClient)
+	p("  ctx context.Context")
+	p("  md metadata.MD")
+	p("  stream *gax.ProtoJSONStream")
+	p("}")
+	p("")
+	p("func (c *%s) Recv() (*%s.%s, error) {", streamClient, outSpec.Name, outType.GetName())
+	p("  if err := c.ctx.Err(); err != nil {")
+	p("    defer c.stream.Close()")
+	p("    return nil, err")
+	p("  }")
+	p("  msg, err := c.stream.Recv()")
+	p("  if err != nil {")
+	p("    defer c.stream.Close()")
+	p("    return nil, err")
+	p("  }")
+	p("  res := msg.(*%s.%s)", outSpec.Name, outType.GetName())
+	p("  return res, nil")
+	p("}")
+	p("")
+	p("func (c *%s) Header() (metadata.MD, error) {", streamClient)
+	p("  return c.md, nil")
+	p("}")
+	p("")
+	p("func (c *%s) Trailer() metadata.MD {", streamClient)
+	p("  return c.md")
+	p("}")
+	p("")
+	p("func (c *%s) CloseSend() error {", streamClient)
+	p("  // This is a no-op to fulfill the interface.")
+	p(`  return fmt.Errorf("this method is not implemented for a server-stream")`)
+	p("}")
+	p("")
+	p("func (c *%s) Context() context.Context {", streamClient)
+	p("  return c.ctx")
+	p("}")
+	p("")
+	p("func (c *%s) SendMsg(m interface{}) error {", streamClient)
+	p("  // This is a no-op to fulfill the interface.")
+	p(`  return fmt.Errorf("this method is not implemented for a server-stream")`)
+	p("}")
+	p("")
+	p("func (c *%s) RecvMsg(m interface{}) error {", streamClient)
+	p("  // This is a no-op to fulfill the interface.")
+	p(`  return fmt.Errorf("this method is not implemented, use Recv")`)
+	p("}")
+	p("")
+
+	g.imports[pbinfo.ImportSpec{Path: "context"}] = true
+	g.imports[pbinfo.ImportSpec{Path: "fmt"}] = true
+	g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc/metadata"}] = true
+
 	return nil
 }
 
diff --git a/showcase/go.sum b/showcase/go.sum
index 9a064e5..45c19da 100644
--- a/showcase/go.sum
+++ b/showcase/go.sum
@@ -126,6 +126,7 @@ github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20191227052852-215e87163ea7/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
+github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e h1:1r7pUrabqp18hOBcwBwiTsbnFeTZHV9eER/QT5JVZxY=
 github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=


Test patch:
diff --git a/internal/gengapic/genrest_test.go b/internal/gengapic/genrest_test.go
index 73aad26..12c2f74 100644
--- a/internal/gengapic/genrest_test.go
+++ b/internal/gengapic/genrest_test.go
@@ -533,6 +533,15 @@ func TestGenRestMethod(t *testing.T) {
 		Options:    pagingRPCOpt,
 	}
 
+	serverStreamRPC := &descriptor.MethodDescriptorProto{
+		Name:            proto.String("ServerStreamRPC"),
+		InputType:       proto.String(foofqn),
+		OutputType:      proto.String(foofqn),
+		ServerStreaming: proto.Bool(true),
+		// Reuse the unary RPC options because it's basically the same.
+		Options: unaryRPCOpt,
+	}
+
 	s := &descriptor.ServiceDescriptorProto{
 		Name: proto.String("FooService"),
 	}
@@ -570,13 +579,14 @@ func TestGenRestMethod(t *testing.T) {
 				pagedFooRes: f,
 			},
 			ParentElement: map[pbinfo.ProtoType]pbinfo.ProtoType{
-				opRPC:      s,
-				emptyRPC:   s,
-				unaryRPC:   s,
-				pagingRPC:  s,
-				nameField:  op,
-				sizeField:  foo,
-				otherField: foo,
+				opRPC:           s,
+				emptyRPC:        s,
+				unaryRPC:        s,
+				pagingRPC:       s,
+				serverStreamRPC: s,
+				nameField:       op,
+				sizeField:       foo,
+				otherField:      foo,
 			},
 			Type: map[string]pbinfo.ProtoType{
 				opfqn:          op,
@@ -637,6 +647,20 @@ func TestGenRestMethod(t *testing.T) {
 				{Name: "foopb", Path: "google.golang.org/genproto/cloud/foo/v1"}: true,
 			},
 		},
+		{
+			name:    "server_stream_rpc",
+			method:  serverStreamRPC,
+			options: &options{},
+			imports: map[pbinfo.ImportSpec]bool{
+				{Path: "bytes"}:   true,
+				{Path: "context"}: true,
+				{Path: "fmt"}:     true,
+				{Path: "google.golang.org/api/googleapi"}:                        true,
+				{Path: "google.golang.org/grpc/metadata"}:                        true,
+				{Path: "google.golang.org/protobuf/encoding/protojson"}:          true,
+				{Name: "foopb", Path: "google.golang.org/genproto/cloud/foo/v1"}: true,
+			},
+		},
 	} {
 		s.Method = []*descriptor.MethodDescriptorProto{tst.method}
 		g.opts = tst.options
diff --git a/internal/gengapic/testdata/doc_file.want b/internal/gengapic/testdata/doc_file.want
index 1272cbb..4911d05 100644
--- a/internal/gengapic/testdata/doc_file.want
+++ b/internal/gengapic/testdata/doc_file.want
@@ -181,3 +181,4 @@ func buildHeaders(ctx context.Context, mds ...metadata.MD) http.Header {
 	md := metadata.Join(mds...)
 	return http.Header(md)
 }
+
diff --git a/internal/gengapic/testdata/doc_file_alpha.want b/internal/gengapic/testdata/doc_file_alpha.want
index c6b4bba..ba3abe6 100644
--- a/internal/gengapic/testdata/doc_file_alpha.want
+++ b/internal/gengapic/testdata/doc_file_alpha.want
@@ -183,3 +183,4 @@ func buildHeaders(ctx context.Context, mds ...metadata.MD) http.Header {
 	md := metadata.Join(mds...)
 	return http.Header(md)
 }
+
diff --git a/internal/gengapic/testdata/doc_file_beta.want b/internal/gengapic/testdata/doc_file_beta.want
index 808723f..78dd79c 100644
--- a/internal/gengapic/testdata/doc_file_beta.want
+++ b/internal/gengapic/testdata/doc_file_beta.want
@@ -183,3 +183,4 @@ func buildHeaders(ctx context.Context, mds ...metadata.MD) http.Header {
 	md := metadata.Join(mds...)
 	return http.Header(md)
 }
+
diff --git a/internal/gengapic/testdata/rest_ServerStreamRPC.want b/internal/gengapic/testdata/rest_ServerStreamRPC.want
new file mode 100644
index 0000000..4c8fe87
--- /dev/null
+++ b/internal/gengapic/testdata/rest_ServerStreamRPC.want
@@ -0,0 +1,90 @@
+func (c *fooRESTClient) ServerStreamRPC(ctx context.Context, req *foopb.Foo, opts ...gax.CallOption) (foopb.FooService_ServerStreamRPCClient, error) {
+	m := protojson.MarshalOptions{AllowPartial: true}
+	jsonReq, err := m.Marshal(req)
+	if err != nil {
+		return nil, err
+	}
+
+	baseUrl, _ := url.Parse(c.endpoint)
+	baseUrl.Path += fmt.Sprintf("/v1/foo")
+
+	// Build HTTP headers from client and context metadata.
+	headers := buildHeaders(ctx, c.xGoogMetadata, metadata.Pairs("Content-Type", "application/json"))
+	var streamClient *serverStreamRPCRESTClient
+	e := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
+		httpReq, err := http.NewRequest("POST", baseUrl.String(), bytes.NewReader(jsonReq))
+		if err != nil {
+			return err
+		}
+		httpReq = httpReq.WithContext(ctx)
+		httpReq.Header = headers
+
+		httpRsp, err := c.httpClient.Do(httpReq)
+		if err != nil{
+			return err
+		}
+
+		if err = googleapi.CheckResponse(httpRsp); err != nil {
+			return err
+		}
+
+		streamClient = &serverStreamRPCRESTClient{
+			ctx: ctx,
+			md: metadata.MD(httpRsp.Header),
+			stream: gax.NewProtoJSONStreamReader(httpRsp.Body, (&foopb.Foo{}).ProtoReflect().Type()),
+		}
+		return nil
+	}, opts...)
+
+	return streamClient, e
+}
+
+// serverStreamRPCRESTClient is the stream client used to consume the server stream created by
+// the REST implementation of ServerStreamRPC.
+type serverStreamRPCRESTClient struct {
+	ctx context.Context
+	md metadata.MD
+	stream *gax.ProtoJSONStream
+}
+
+func (c *serverStreamRPCRESTClient) Recv() (*foopb.Foo, error) {
+	if err := c.ctx.Err(); err != nil {
+		defer c.stream.Close()
+		return nil, err
+	}
+	msg, err := c.stream.Recv()
+	if err != nil {
+		defer c.stream.Close()
+		return nil, err
+	}
+	res := msg.(*foopb.Foo)
+	return res, nil
+}
+
+func (c *serverStreamRPCRESTClient) Header() (metadata.MD, error) {
+	return c.md, nil
+}
+
+func (c *serverStreamRPCRESTClient) Trailer() metadata.MD {
+	return c.md
+}
+
+func (c *serverStreamRPCRESTClient) CloseSend() error {
+	// This is a no-op to fulfill the interface.
+	return fmt.Errorf("this method is not implemented for a server-stream")
+}
+
+func (c *serverStreamRPCRESTClient) Context() context.Context {
+	return c.ctx
+}
+
+func (c *serverStreamRPCRESTClient) SendMsg(m interface{}) error {
+	// This is a no-op to fulfill the interface.
+	return fmt.Errorf("this method is not implemented for a server-stream")
+}
+
+func (c *serverStreamRPCRESTClient) RecvMsg(m interface{}) error {
+	// This is a no-op to fulfill the interface.
+	return fmt.Errorf("this method is not implemented, use Recv")
+}
+
diff --git a/showcase/echo_test.go b/showcase/echo_test.go
index e4a252c..1a44ea2 100644
--- a/showcase/echo_test.go
+++ b/showcase/echo_test.go
@@ -148,24 +148,27 @@ func TestExpand(t *testing.T) {
 	defer check(t)
 	content := "The rain in Spain stays mainly on the plain!"
 	req := &showcasepb.ExpandRequest{Content: content}
-	s, err := echo.Expand(context.Background(), req)
-	if err != nil {
-		t.Fatal(err)
-	}
-	resps := []string{}
-	for {
-		resp, err := s.Recv()
-		if err == io.EOF {
-			break
-		}
+
+	for typ, client := range map[string]*showcase.EchoClient{"grpc": echo, "rest": echoREST} {
+		s, err := client.Expand(context.Background(), req)
 		if err != nil {
 			t.Fatal(err)
 		}
-		resps = append(resps, resp.GetContent())
-	}
-	got := strings.Join(resps, " ")
-	if content != got {
-		t.Errorf("Expand() = %q, want %q", got, content)
+		resps := []string{}
+		for {
+			resp, err := s.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			resps = append(resps, resp.GetContent())
+		}
+		got := strings.Join(resps, " ")
+		if content != got {
+			t.Errorf("%s Expand() = %q, want %q", typ, got, content)
+		}
 	}
 }