From cb8916c697c959fe3100531de59d9556960b1266 Mon Sep 17 00:00:00 2001 From: noahdietz Date: Fri, 28 Jan 2022 11:57:54 -0800 Subject: [PATCH 1/3] feat(gengapic): diregapic lro polling request params --- internal/gengapic/custom_operation.go | 127 +++++++++++++++--- internal/gengapic/custom_operation_test.go | 33 ++++- internal/gengapic/gengapic.go | 5 +- internal/gengapic/genrest.go | 7 +- .../testdata/custom_op_init_helper.want | 8 ++ .../testdata/custom_op_type_bool.want | 8 +- .../testdata/custom_op_type_enum.want | 8 +- internal/gengapic/testdata/rest_CustomOp.want | 7 +- 8 files changed, 172 insertions(+), 31 deletions(-) create mode 100644 internal/gengapic/testdata/custom_op_init_helper.want diff --git a/internal/gengapic/custom_operation.go b/internal/gengapic/custom_operation.go index c7b62ebc2..1b1736791 100644 --- a/internal/gengapic/custom_operation.go +++ b/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,42 @@ 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(p, r, o string, req *descriptor.DescriptorProto, s *descriptor.ServiceDescriptorProto) { + h := handleName(s.GetName(), g.opts.pkgName) opName := g.aux.customOp.message.GetName() + pt := func(f string, v ...interface{}) { + g.pt.Printf(f, v...) + } - s := fmt.Sprintf("&%s{&%s{c: c.operationClient, proto: %s}}", opName, h, p) - - return s + // 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", r, fieldGetter(f.GetName())) + } + } + sort.Strings(keys) + + pt("%s := &%s{", o, opName) + pt(" &%s{", h) + pt(" c: c.operationClient,") + pt(" proto: %s,", p) + for _, param := range keys { + pt(" %s: %s,", param, paramToGetter[param]) + } + pt(" },") + pt("}") } // customOperationType generates the custom operation wrapper type and operation @@ -164,32 +196,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 +245,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 +338,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 { + poll := operationPollingMethod(opServ) + if poll == nil { + return nil + } + pollReqName := poll.GetInputType() + + inType := g.descInfo.Type[m.GetInputType()].(*descriptor.DescriptorProto) + + params := []string{} + 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 { diff --git a/internal/gengapic/custom_operation_test.go b/internal/gengapic/custom_operation_test.go index 182c91e2f..a030e5feb 100644 --- a/internal/gengapic/custom_operation_test.go +++ b/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{ diff --git a/internal/gengapic/gengapic.go b/internal/gengapic/gengapic.go index e198af10a..4a7c2d271 100644 --- a/internal/gengapic/gengapic.go +++ b/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) } diff --git a/internal/gengapic/genrest.go b/internal/gengapic/genrest.go index 0c4ca7a1e..44b45ff1e 100644 --- a/internal/gengapic/genrest.go +++ b/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("}") diff --git a/internal/gengapic/testdata/custom_op_init_helper.want b/internal/gengapic/testdata/custom_op_init_helper.want new file mode 100644 index 000000000..b1f198c1a --- /dev/null +++ b/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(), + }, +} diff --git a/internal/gengapic/testdata/custom_op_type_bool.want b/internal/gengapic/testdata/custom_op_type_bool.want index ff0995ea7..14c90252c 100644 --- a/internal/gengapic/testdata/custom_op_type_bool.want +++ b/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 } diff --git a/internal/gengapic/testdata/custom_op_type_enum.want b/internal/gengapic/testdata/custom_op_type_enum.want index c134fafc9..b747113a6 100644 --- a/internal/gengapic/testdata/custom_op_type_enum.want +++ b/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 } diff --git a/internal/gengapic/testdata/rest_CustomOp.want b/internal/gengapic/testdata/rest_CustomOp.want index e71bf3795..b95d20408 100644 --- a/internal/gengapic/testdata/rest_CustomOp.want +++ b/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 } From b092c1a92567ff173b5cfb6669a25ed7b380614f Mon Sep 17 00:00:00 2001 From: noahdietz Date: Fri, 28 Jan 2022 12:00:43 -0800 Subject: [PATCH 2/3] updat bazel deps --- internal/gengapic/BUILD.bazel | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/gengapic/BUILD.bazel b/internal/gengapic/BUILD.bazel index 3005f768a..8a88bb13a 100644 --- a/internal/gengapic/BUILD.bazel +++ b/internal/gengapic/BUILD.bazel @@ -78,7 +78,8 @@ 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/gapic/metadata", + "@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", "@org_golang_google_protobuf//reflect/protodesc", From 5eb7a29a8d3f9633791d85eb2f30d83ba51a755c Mon Sep 17 00:00:00 2001 From: noahdietz Date: Fri, 28 Jan 2022 12:50:00 -0800 Subject: [PATCH 3/3] address feedback --- internal/gengapic/BUILD.bazel | 2 +- internal/gengapic/custom_operation.go | 16 +++++++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/internal/gengapic/BUILD.bazel b/internal/gengapic/BUILD.bazel index 8a88bb13a..2dfb8f66b 100644 --- a/internal/gengapic/BUILD.bazel +++ b/internal/gengapic/BUILD.bazel @@ -79,7 +79,7 @@ go_test( "@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_genproto//googleapis/gapic/metadata", "@org_golang_google_protobuf//encoding/protojson", "@org_golang_google_protobuf//proto", "@org_golang_google_protobuf//reflect/protodesc", diff --git a/internal/gengapic/custom_operation.go b/internal/gengapic/custom_operation.go index 1b1736791..011b297d6 100644 --- a/internal/gengapic/custom_operation.go +++ b/internal/gengapic/custom_operation.go @@ -68,12 +68,10 @@ 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(p, r, o string, req *descriptor.DescriptorProto, s *descriptor.ServiceDescriptorProto) { +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() - pt := func(f string, v ...interface{}) { - g.pt.Printf(f, v...) - } + 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 @@ -90,15 +88,15 @@ func (g *generator) customOpInit(p, r, o string, req *descriptor.DescriptorProto param = lowerFirst(snakeToCamel(param)) if params := g.aux.customOp.pollingParams[s]; strContains(params, param) { keys = append(keys, param) - paramToGetter[param] = fmt.Sprintf("%s%s", r, fieldGetter(f.GetName())) + paramToGetter[param] = fmt.Sprintf("%s%s", reqVar, fieldGetter(f.GetName())) } } sort.Strings(keys) - pt("%s := &%s{", o, opName) + pt("%s := &%s{", opVar, opName) pt(" &%s{", h) pt(" c: c.operationClient,") - pt(" proto: %s,", p) + pt(" proto: %s,", rspVar) for _, param := range keys { pt(" %s: %s,", param, paramToGetter[param]) } @@ -364,15 +362,15 @@ func operationPollingMethod(s *descriptor.ServiceDescriptorProto) *descriptor.Me // 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 nil + return params } pollReqName := poll.GetInputType() inType := g.descInfo.Type[m.GetInputType()].(*descriptor.DescriptorProto) - params := []string{} for _, f := range inType.GetField() { mapping, ok := operationRequestField(f)