properly read LRO type from proto annotation

internal/gengapic/gengapic.go

@@ -263,7 +263,9 @@ func (g *generator) gen(serv *descriptor.ServiceDescriptorProto, pkgName string)
 		return aux.lros[i].GetName() < aux.lros[j].GetName()
 	})
 	for _, m := range aux.lros {
-		g.lroType(servName, m)
+		if err := g.lroType(servName, serv, m); err != nil {
+			return errors.E(err, "while generating LRO type for %q", m.GetName())
+		}
 	}
 
 	var iters []iterType

internal/gengapic/gengapic_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
+	"google.golang.org/genproto/googleapis/api/annotations"
 )
 
 func TestComment(t *testing.T) {
@@ -162,6 +163,7 @@ func TestGenMethod(t *testing.T) {
 	}
 
 	file := &descriptor.FileDescriptorProto{
+		Package: proto.String("my.pkg"),
 		Options: &descriptor.FileOptions{
 			GoPackage: proto.String("mypackage"),
 		},
@@ -195,6 +197,7 @@ func TestGenMethod(t *testing.T) {
 			Name:       proto.String("GetBigThing"),
 			InputType:  proto.String(".my.pkg.InputType"),
 			OutputType: proto.String(".google.longrunning.Operation"),
+			Options:    &descriptor.MethodOptions{},
 		},
 		{
 			Name:       proto.String("GetManyThings"),
@@ -210,14 +213,37 @@ func TestGenMethod(t *testing.T) {
 		},
 	}
 
+methods:
 	for _, m := range meths {
+		g.pt.Reset()
+
+		// Just add this everywhere. Only LRO method will pick it up.
+		if m.Options != nil {
+			lroType := &annotations.LongrunningOperationTypes{
+				Response: "OutputType",
+			}
+			proto.SetExtension(m.Options, annotations.E_LongrunningOperationTypes, lroType)
+		}
+
 		aux := auxTypes{
 			iters: map[string]iterType{},
 		}
 		if err := g.genMethod("Foo", serv, m, &aux); err != nil {
 			t.Error(err)
-		} else {
-			diff(t, m.GetName(), g.pt.String(), filepath.Join("testdata", "method_"+m.GetName()+".want"))
+			continue
 		}
+
+		for _, m := range aux.lros {
+			if err := g.lroType("MyService", serv, m); err != nil {
+				t.Error(err)
+				continue methods
+			}
+		}
+
+		for _, iter := range aux.iters {
+			g.pagingIter(iter)
+		}
+
+		diff(t, m.GetName(), g.pt.String(), filepath.Join("testdata", "method_"+m.GetName()+".want"))
 	}
 }

internal/gengapic/lro.go

@@ -15,8 +15,14 @@
 package gengapic
 
 import (
+	"fmt"
+	"strings"
+
+	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
+	"github.com/googleapis/gapic-generator-go/internal/errors"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
+	"google.golang.org/genproto/googleapis/api/annotations"
 )
 
 func (g *generator) lroCall(servName string, m *descriptor.MethodDescriptorProto) error {
@@ -63,14 +69,55 @@ func (g *generator) lroCall(servName string, m *descriptor.MethodDescriptorProto
 	return nil
 }
 
-func (g *generator) lroType(servName string, m *descriptor.MethodDescriptorProto) {
+func (g *generator) lroType(servName string, serv *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) error {
 	lroType := lroTypeName(*m.Name)
 	p := g.printf
 
-	// TODO(pongad): programmatically fill these.
-	respType := "Foo.Bar"
-	metaType := "Foo.MetaBar"
-	hasMeta := true
+	eLRO, err := proto.GetExtension(m.Options, annotations.E_LongrunningOperationTypes)
+	if err != nil {
+		return errors.E(err, "cannot read LRO types")
+	}
+	eLROType := eLRO.(*annotations.LongrunningOperationTypes)
+
+	var respType string
+	{
+		fullName := eLROType.Response
+
+		// eLRO.ResponseType is either fully-qualified or in the same package as the method.
+		if strings.IndexByte(fullName, '.') < 0 {
+			fullName = g.descInfo.ParentFile[serv].GetPackage() + "." + fullName
+		}
+
+		// When we build a map[name]Type in pbinfo, we prefix names with '.' to signify that they are fully qualified.
+		// The string in ResponseType does not have the prefix, so we add it.
+		fullName = "." + fullName
+
+		typ := g.descInfo.Type[fullName]
+		respSpec, err := g.descInfo.ImportSpec(typ)
+		if err != nil {
+			return errors.E(err, "cannot find LRO type %q; type not linked?", fullName)
+		}
+		g.imports[respSpec] = true
+		respType = fmt.Sprintf("%s.%s", respSpec.Name, typ.GetName())
+	}
+
+	hasMeta := eLROType.Metadata != ""
+	var metaType string
+	if hasMeta {
+		fullName := eLROType.Metadata
+		if strings.IndexByte(fullName, '.') < 0 {
+			fullName = g.descInfo.ParentFile[serv].GetPackage() + "." + fullName
+		}
+		fullName = "." + fullName
+
+		typ := g.descInfo.Type[fullName]
+		meta, err := g.descInfo.ImportSpec(typ)
+		if err != nil {
+			return errors.E(err, "cannot find LRO metadata type %q; type not linked?", fullName)
+		}
+		g.imports[meta] = true
+		metaType = fmt.Sprintf("%s.%s", meta.Name, typ.GetName())
+	}
 
 	// Type definition
 	{
@@ -144,7 +191,7 @@ func (g *generator) lroType(servName string, m *descriptor.MethodDescriptorProto
 		p("// Metadata itself does not contact the server, but Poll does.")
 		p("// To get the latest metadata, call this method after a successful call to Poll.")
 		p("// If the metadata is not available, the returned metadata and error are both nil.")
-		p("func (op *%s) Metadata() (*%s, error) {", lroType, respType)
+		p("func (op *%s) Metadata() (*%s, error) {", lroType, metaType)
 		p("  var meta %s", metaType)
 		p("  if err := op.lro.Metadata(&meta); err == longrunning.ErrNoMetadata {")
 		p("    return nil, nil")
@@ -174,7 +221,7 @@ func (g *generator) lroType(servName string, m *descriptor.MethodDescriptorProto
 		p("}")
 		p("")
 	}
-
+	return nil
 }
 
 func lroTypeName(methodName string) string {

internal/gengapic/testdata/method_BidiThings.want

@@ -1,82 +1,3 @@
-func (c *FooClient) GetEmptyThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) error {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetEmptyThing[0:len(c.CallOptions.GetEmptyThing):len(c.CallOptions.GetEmptyThing)], opts...)
-	err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-		var err error
-		_, err = c.fooClient.GetEmptyThing(ctx, req, settings.GRPC...)
-		return err
-	}, opts...)
-	return err
-}
-
-func (c *FooClient) GetOneThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) (*mypackagepb.OutputType, error) {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetOneThing[0:len(c.CallOptions.GetOneThing):len(c.CallOptions.GetOneThing)], opts...)
-	var resp *mypackagepb.OutputType
-	err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-		var err error
-		resp, err = c.fooClient.GetOneThing(ctx, req, settings.GRPC...)
-		return err
-	}, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return resp, nil
-}
-
-func (c *FooClient) GetBigThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) (*GetBigThingOperation, error) {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetBigThing[0:len(c.CallOptions.GetBigThing):len(c.CallOptions.GetBigThing)], opts...)
-	var resp *longrunningpb.Operation
-	err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-		var err error
-		resp, err = c.fooClient.GetBigThing(ctx, req, settings.GRPC...)
-		return err
-	}, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return &GetBigThingOperation{
-		lro: longrunning.InternalNewOperation(c.LROClient, resp),
-	}, nil
-}
-
-func (c *FooClient) GetManyThings(ctx context.Context, req *mypackagepb.PageInputType, opts ...gax.CallOption) *StringIterator {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetManyThings[0:len(c.CallOptions.GetManyThings):len(c.CallOptions.GetManyThings)], opts...)
-	it := &StringIterator{}
-	req = proto.Clone(req).(*mypackagepb.PageInputType)
-	it.InternalFetch = func(pageSize int, pageToken string) ([]string, string, error) {
-		var resp *mypackagepb.PageOutputType
-		req.PageToken = pageToken
-		if pageSize > math.MaxInt32 {
-			req.PageSize = math.MaxInt32
-		} else {
-			req.PageSize = int32(pageSize)
-		}
-		err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-			var err error
-			resp, err = c.fooClient.GetManyThings(ctx, req, settings.GRPC...)
-			return err
-		}, opts...)
-		if err != nil {
-			return nil, "", err
-		}
-		return resp.Items, resp.NextPageToken, nil
-	}
-	fetch := func(pageSize int, pageToken string) (string, error) {
-		items, nextPageToken, err := it.InternalFetch(pageSize, pageToken)
-		if err != nil {
-			return "", err
-		}
-		it.items = append(it.items, items...)
-		return nextPageToken, nil
-	}
-	it.pageInfo, it.nextFunc = iterator.NewPageInfo(fetch, it.bufLen, it.takeBuf)
-	it.pageInfo.MaxSize = int(req.PageSize)
-	return it
-}
-
 func (c *FooClient) BidiThings(ctx context.Context, opts ...gax.CallOption) (mypackagepb._BidiThingsClient, error) {
 	ctx = insertMetadata(ctx, c.xGoogMetadata)
 	opts = append(c.CallOptions.BidiThings[0:len(c.CallOptions.BidiThings):len(c.CallOptions.BidiThings)], opts...)

internal/gengapic/testdata/method_GetBigThing.want

@@ -1,29 +1,3 @@
-func (c *FooClient) GetEmptyThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) error {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetEmptyThing[0:len(c.CallOptions.GetEmptyThing):len(c.CallOptions.GetEmptyThing)], opts...)
-	err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-		var err error
-		_, err = c.fooClient.GetEmptyThing(ctx, req, settings.GRPC...)
-		return err
-	}, opts...)
-	return err
-}
-
-func (c *FooClient) GetOneThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) (*mypackagepb.OutputType, error) {
-	ctx = insertMetadata(ctx, c.xGoogMetadata)
-	opts = append(c.CallOptions.GetOneThing[0:len(c.CallOptions.GetOneThing):len(c.CallOptions.GetOneThing)], opts...)
-	var resp *mypackagepb.OutputType
-	err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
-		var err error
-		resp, err = c.fooClient.GetOneThing(ctx, req, settings.GRPC...)
-		return err
-	}, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return resp, nil
-}
-
 func (c *FooClient) GetBigThing(ctx context.Context, req *mypackagepb.InputType, opts ...gax.CallOption) (*GetBigThingOperation, error) {
 	ctx = insertMetadata(ctx, c.xGoogMetadata)
 	opts = append(c.CallOptions.GetBigThing[0:len(c.CallOptions.GetBigThing):len(c.CallOptions.GetBigThing)], opts...)
@@ -41,3 +15,56 @@ func (c *FooClient) GetBigThing(ctx context.Context, req *mypackagepb.InputType,
 	}, nil
 }
 
+// GetBigThingOperation manages a long-running operation from GetBigThing.
+type GetBigThingOperation struct {
+	lro *longrunning.Operation
+}
+
+// GetBigThingOperation returns a new GetBigThingOperation from a given name.
+// The name must be that of a previously created GetBigThingOperation, possibly from a different process.
+func (c *MyServiceClient) GetBigThingOperation(name string) *GetBigThingOperation {
+	return &GetBigThingOperation{
+		lro: longrunning.InternalNewOperation(c.LROClient, &longrunningpb.Operation{Name: name}),
+	}
+}
+
+// Wait blocks until the long-running operation is completed, returning the response and any errors encountered.
+//
+// See documentation of Poll for error-handling information.
+func (op *GetBigThingOperation) Wait(ctx context.Context, opts ...gax.CallOption) (*mypackagepb.OutputType, error) {
+	var resp mypackagepb.OutputType
+	if err := op.lro.WaitWithInterval(ctx, &resp, time.Minute, opts...); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+// Poll fetches the latest state of the long-running operation.
+//
+// If Poll fails, the error is returned and op is unmodified. If Poll succeeds and
+// the operation has completed with failure, the error is returned and op.Done will return true.
+// If Poll succeeds and the operation has completed successfully,
+// op.Done will return true, and the response of the operation is returned.
+// If Poll succeeds and the operation has not completed, the returned response and error are both nil.
+func (op *GetBigThingOperation) Poll(ctx context.Context, opts ...gax.CallOption) (*mypackagepb.OutputType, error) {
+	var resp mypackagepb.OutputType
+	if err := op.lro.Poll(ctx, &resp, opts...); err != nil {
+		return nil, err
+	}
+	if !op.Done() {
+		return nil, nil
+	}
+	return &resp, nil
+}
+
+// Done reports whether the long-running operation has completed.
+func (op *GetBigThingOperation) Done() bool {
+	return op.lro.Done()
+}
+
+// Name returns the name of the long-running operation.
+// The name is assigned by the server and is unique within the service from which the operation is created.
+func (op *GetBigThingOperation) Name() string {
+	return op.lro.Name()
+}
+