feat(gengapic): diregapic lro polling request params

internal/gengapic/BUILD.bazel

@@ -78,6 +78,7 @@ go_test(
         "@go_googleapis//google/rpc:code_go_proto",
         "@io_bazel_rules_go//proto/wkt:compiler_plugin_go_proto",
         "@io_bazel_rules_go//proto/wkt:descriptor_go_proto",
+        "@org_golang_google_genproto//googleapis/cloud/extendedops",
         "@org_golang_google_genproto//googleapis/gapic/metadata",
         "@org_golang_google_protobuf//encoding/protojson",
         "@org_golang_google_protobuf//proto",

internal/gengapic/custom_operation.go

@@ -16,6 +16,7 @@ package gengapic
 
 import (
 	"fmt"
+	"sort"
 
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
@@ -25,8 +26,9 @@ import (
 
 // customOp represents a custom operation type for long running operations.
 type customOp struct {
-	message *descriptor.DescriptorProto
-	handles []*descriptor.ServiceDescriptorProto
+	message       *descriptor.DescriptorProto
+	handles       []*descriptor.ServiceDescriptorProto
+	pollingParams map[*descriptor.ServiceDescriptorProto][]string
 }
 
 // isCustomOp determines if the given method should return a custom operation wrapper.
@@ -66,12 +68,40 @@ func (g *generator) customOpPointerType() (string, error) {
 // customOpInit builds a string containing the Go code for initializing the
 // operation wrapper type with the Go identifier for a variable that is the
 // proto-defined operation type.
-func (g *generator) customOpInit(h, p string) string {
+func (g *generator) customOpInit(rspVar, reqVar, opVar string, req *descriptor.DescriptorProto, s *descriptor.ServiceDescriptorProto) {
+	h := handleName(s.GetName(), g.opts.pkgName)
 	opName := g.aux.customOp.message.GetName()
-
-	s := fmt.Sprintf("&%s{&%s{c: c.operationClient, proto: %s}}", opName, h, p)
-
-	return s
+	pt := g.pt.Printf
+
+	// Collect all of the fields marked with google.cloud.operation_request_field
+	// and map the getter method to the polling request's field name, while also
+	// collecting these field name keys for sorted access.
+	keys := []string{}
+	paramToGetter := map[string]string{}
+	for _, f := range req.GetField() {
+		param, ok := operationRequestField(f)
+		if !ok {
+			continue
+		}
+		// Only include those operation_request_fields that are also tracked as
+		// polling params for the operation service.
+		param = lowerFirst(snakeToCamel(param))
+		if params := g.aux.customOp.pollingParams[s]; strContains(params, param) {
+			keys = append(keys, param)
+			paramToGetter[param] = fmt.Sprintf("%s%s", reqVar, fieldGetter(f.GetName()))
+		}
+	}
+	sort.Strings(keys)
+
+	pt("%s := &%s{", opVar, opName)
+	pt("  &%s{", h)
+	pt("    c: c.operationClient,")
+	pt("    proto: %s,", rspVar)
+	for _, param := range keys {
+		pt("    %s: %s,", param, paramToGetter[param])
+	}
+	pt("  },")
+	pt("}")
 }
 
 // customOperationType generates the custom operation wrapper type and operation
@@ -164,32 +194,35 @@ func (g *generator) customOperationType() error {
 	g.imports[pbinfo.ImportSpec{Name: "gax", Path: "github.com/googleapis/gax-go/v2"}] = true
 
 	for _, handle := range op.handles {
+		pollingParams := op.pollingParams[handle]
 		s := pbinfo.ReduceServName(handle.GetName(), opImp.Name)
-		n := lowerFirst(s + "Handle")
+		n := handleName(handle.GetName(), opImp.Name)
 
 		// Look up polling method and its input.
-		var get *descriptor.MethodDescriptorProto
-		for _, m := range handle.GetMethod() {
-			if m.GetName() == "Get" {
-				get = m
-				break
-			}
-		}
-		getInput := g.descInfo.Type[get.GetInputType()]
-		inNameField := operationResponseField(getInput.(*descriptor.DescriptorProto), opNameField.GetName())
+		poll := operationPollingMethod(handle)
+		pollReq := g.descInfo.Type[poll.GetInputType()].(*descriptor.DescriptorProto)
+		pollNameField := operationResponseField(pollReq, opNameField.GetName())
 
 		// type
 		p("// Implements the %s interface for %s.", handleInt, handle.GetName())
 		p("type %s struct {", n)
 		p("  c *%sClient", s)
 		p("  proto %s", ptyp)
+		for _, param := range pollingParams {
+			p("  %s string", param)
+		}
 		p("}")
 		p("")
 
 		// Poll
 		p("// Poll retrieves the latest data for the long-running operation.")
 		p("func (h *%s) Poll(ctx context.Context, opts ...gax.CallOption) error {", n)
-		p("  resp, err := h.c.Get(ctx, &%s.%s{%s: h.proto%s}, opts...)", opImp.Name, upperFirst(getInput.GetName()), upperFirst(inNameField.GetName()), opNameGetter)
+		p("  resp, err := h.c.Get(ctx, &%s.%s{", opImp.Name, upperFirst(pollReq.GetName()))
+		p("    %s: h.proto%s,", snakeToCamel(pollNameField.GetName()), opNameGetter)
+		for _, f := range pollingParams {
+			p("    %s: h.%s,", upperFirst(f), f)
+		}
+		p("  }, opts...)")
 		p("  if err != nil {")
 		p("    return err")
 		p("  }")
@@ -210,21 +243,26 @@ func (g *generator) customOperationType() error {
 }
 
 // loadCustomOpServices maps the service declared as a google.cloud.operation_service
-// to the service that owns the method(s) declaring it.
+// to the service that owns the method(s) declaring it, as well as collects the set of
+// operation services for handle generation, and maps the polling request parameters to
+// that same operation service descriptor.
 func (g *generator) loadCustomOpServices(servs []*descriptor.ServiceDescriptorProto) {
 	handles := g.aux.customOp.handles
+	pollingParams := map[*descriptor.ServiceDescriptorProto][]string{}
 	for _, serv := range servs {
 		for _, meth := range serv.GetMethod() {
 			if opServ := g.customOpService(meth); opServ != nil {
 				g.customOpServices[serv] = opServ
 				if !containsService(handles, opServ) {
 					handles = append(handles, opServ)
+					pollingParams[opServ] = g.pollingRequestParameters(meth, opServ)
 				}
 				break
 			}
 		}
 	}
 	g.aux.customOp.handles = handles
+	g.aux.customOp.pollingParams = pollingParams
 }
 
 // customOpService loads the ServiceDescriptorProto for the google.cloud.operation_service
@@ -298,6 +336,57 @@ func operationResponseField(m *descriptor.DescriptorProto, target string) *descr
 	return nil
 }
 
+// operationRequestField is a helper for extracting the operation_request_field annotation from a field.
+func operationRequestField(f *descriptor.FieldDescriptorProto) (string, bool) {
+	mapping := proto.GetExtension(f.GetOptions(), extendedops.E_OperationRequestField).(string)
+	if mapping != "" {
+		return mapping, true
+	}
+
+	return "", false
+}
+
+// operationPollingMethod is a helper for finding the operation service RPC annotated with operation_polling_method.
+func operationPollingMethod(s *descriptor.ServiceDescriptorProto) *descriptor.MethodDescriptorProto {
+	for _, m := range s.GetMethod() {
+		if proto.GetExtension(m.GetOptions(), extendedops.E_OperationPollingMethod).(bool) {
+			return m
+		}
+	}
+
+	return nil
+}
+
+// pollingRequestParamters collects the polling request parameters for an operation service's polling method
+// based on which are annotated in an initiating RPC's request message with operation_request_field and that
+// are also marked as required on the polling request message. Specifically, this weeds out the parent_id field
+// of the GlobalOrganizationOperations polling params.
+func (g *generator) pollingRequestParameters(m *descriptor.MethodDescriptorProto, opServ *descriptor.ServiceDescriptorProto) []string {
+	var params []string
+	poll := operationPollingMethod(opServ)
+	if poll == nil {
+		return params
+	}
+	pollReqName := poll.GetInputType()
+
+	inType := g.descInfo.Type[m.GetInputType()].(*descriptor.DescriptorProto)
+
+	for _, f := range inType.GetField() {
+
+		mapping, ok := operationRequestField(f)
+		if !ok {
+			continue
+		}
+		pollField := g.lookupField(pollReqName, mapping)
+		if pollField != nil && isRequired(pollField) {
+			params = append(params, lowerFirst(snakeToCamel(mapping)))
+		}
+	}
+	sort.Strings(params)
+
+	return params
+}
+
 // handleName is a helper for constructing a operation handle name from the
 // operation service name and Go package name.
 func handleName(s, pkg string) string {

internal/gengapic/custom_operation_test.go

@@ -89,18 +89,35 @@ func TestCustomOpInit(t *testing.T) {
 	op := &descriptor.DescriptorProto{
 		Name: proto.String("Operation"),
 	}
+	projFieldOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(projFieldOpts, extendedops.E_OperationRequestField, "project")
+	projField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("request_project"),
+		Options: projFieldOpts,
+	}
+	zoneFieldOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(zoneFieldOpts, extendedops.E_OperationRequestField, "zone")
+	zoneField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("request_zone"),
+		Options: zoneFieldOpts,
+	}
+	req := &descriptor.DescriptorProto{
+		Field: []*descriptor.FieldDescriptorProto{projField, zoneField},
+	}
+	opServ := &descriptor.ServiceDescriptorProto{Name: proto.String("FooOperationService")}
 	g := &generator{
 		aux: &auxTypes{
 			customOp: &customOp{
 				message: op,
+				pollingParams: map[*descriptor.ServiceDescriptorProto][]string{
+					opServ: {"project", "zone"},
+				},
 			},
 		},
+		opts: &options{pkgName: "bar"},
 	}
-	got := g.customOpInit("fooOperationHandle", "foo")
-	want := "&Operation{&fooOperationHandle{c: c.operationClient, proto: foo}}"
-	if diff := cmp.Diff(got, want); diff != "" {
-		t.Errorf("got(-),want(+):\n%s", diff)
-	}
+	g.customOpInit("foo", "req", "op", req, opServ)
+	txtdiff.Diff(t, "custom_op_init_helper", g.pt.String(), filepath.Join("testdata", "custom_op_init_helper.want"))
 }
 
 func TestCustomOperationType(t *testing.T) {
@@ -150,6 +167,8 @@ func TestCustomOperationType(t *testing.T) {
 		Options: inNameOpts,
 	}
 
+	getOpts := &descriptor.MethodOptions{}
+	proto.SetExtension(getOpts, extendedops.E_OperationPollingMethod, true)
 	getInput := &descriptor.DescriptorProto{
 		Name:  proto.String("GetFooOperationRequest"),
 		Field: []*descriptor.FieldDescriptorProto{inNameField},
@@ -161,6 +180,7 @@ func TestCustomOperationType(t *testing.T) {
 			{
 				Name:      proto.String("Get"),
 				InputType: proto.String(".google.cloud.foo.v1.GetFooOperationRequest"),
+				Options:   getOpts,
 			},
 		},
 	}
@@ -177,6 +197,9 @@ func TestCustomOperationType(t *testing.T) {
 			customOp: &customOp{
 				message: op,
 				handles: []*descriptor.ServiceDescriptorProto{fooOpServ},
+				pollingParams: map[*descriptor.ServiceDescriptorProto][]string{
+					fooOpServ: {"project", "zone"},
+				},
 			},
 		},
 		descInfo: pbinfo.Info{

internal/gengapic/gengapic.go

@@ -77,7 +77,10 @@ func Gen(genReq *plugin.CodeGeneratorRequest) (*plugin.CodeGeneratorResponse, er
 	protoPkg := g.descInfo.ParentFile[genServs[0]].GetPackage()
 
 	if op, ok := g.descInfo.Type[fmt.Sprintf(".%s.Operation", protoPkg)]; g.opts.diregapic && ok {
-		g.aux.customOp = &customOp{op.(*descriptor.DescriptorProto), []*descriptor.ServiceDescriptorProto{}}
+		g.aux.customOp = &customOp{
+			message:       op.(*descriptor.DescriptorProto),
+			handles:       []*descriptor.ServiceDescriptorProto{},
+			pollingParams: map[*descriptor.ServiceDescriptorProto][]string{}}
 		g.loadCustomOpServices(genServs)
 	}
 

internal/gengapic/genrest.go

@@ -884,10 +884,9 @@ func (g *generator) unaryRESTCall(servName string, m *descriptor.MethodDescripto
 	p("}")
 	ret := "return resp, nil"
 	if isCustomOp {
-		s := g.customOpService(m)
-		handleName := handleName(s.GetName(), g.opts.pkgName)
-		p("op := %s", g.customOpInit(handleName, "resp"))
-		ret = "return op, nil"
+		opVar := "op"
+		g.customOpInit("resp", "req", opVar, inType.(*descriptor.DescriptorProto), g.customOpService(m))
+		ret = fmt.Sprintf("return %s, nil", opVar)
 	}
 	p(ret)
 	p("}")

internal/gengapic/testdata/custom_op_init_helper.want

@@ -0,0 +1,8 @@
+op := &Operation{
+	&fooOperationHandle{
+		c: c.operationClient,
+		proto: foo,
+		project: req.GetRequestProject(),
+		zone: req.GetRequestZone(),
+	},
+}

internal/gengapic/testdata/custom_op_type_bool.want

@@ -46,11 +46,17 @@ type operationHandle interface {
 type fooOperationsHandle struct {
 	c *FooOperationsClient
 	proto *foopb.Operation
+	project string
+	zone string
 }
 
 // Poll retrieves the latest data for the long-running operation.
 func (h *fooOperationsHandle) Poll(ctx context.Context, opts ...gax.CallOption) error {
-	resp, err := h.c.Get(ctx, &foopb.GetFooOperationRequest{Operation: h.proto.GetName()}, opts...)
+	resp, err := h.c.Get(ctx, &foopb.GetFooOperationRequest{
+		Operation: h.proto.GetName(),
+		Project: h.project,
+		Zone: h.zone,
+	}, opts...)
 	if err != nil {
 		return err
 	}

internal/gengapic/testdata/custom_op_type_enum.want

@@ -46,11 +46,17 @@ type operationHandle interface {
 type fooOperationsHandle struct {
 	c *FooOperationsClient
 	proto *foopb.Operation
+	project string
+	zone string
 }
 
 // Poll retrieves the latest data for the long-running operation.
 func (h *fooOperationsHandle) Poll(ctx context.Context, opts ...gax.CallOption) error {
-	resp, err := h.c.Get(ctx, &foopb.GetFooOperationRequest{Operation: h.proto.GetName()}, opts...)
+	resp, err := h.c.Get(ctx, &foopb.GetFooOperationRequest{
+		Operation: h.proto.GetName(),
+		Project: h.project,
+		Zone: h.zone,
+	}, opts...)
 	if err != nil {
 		return err
 	}

internal/gengapic/testdata/rest_CustomOp.want

@@ -46,6 +46,11 @@ func (c *fooRESTClient) CustomOp(ctx context.Context, req *foopb.Foo, opts ...ga
 	if e != nil {
 		return nil, e
 	}
-	op := &Operation{&handle{c: c.operationClient, proto: resp}}
+	op := &Operation{
+		&handle{
+			c: c.operationClient,
+			proto: resp,
+		},
+	}
 	return op, nil
 }