Repository: googleapis/gapic-generator-go
Pull request #866 (2022-01-11 23:34:04)

feat(gengapic): implement diregapic lro foundation + polling
This is the foundation of Full LRO support for DIREGAPIC/Compute:
* operation service initialization
* operation service-specific operation wrappers
* Poll, Done, Name methods

Missing APIs compared to Standard LRO:
* Cancel (all Compute operation services expose a Delete rpc, but nothing exactly like Cancel)
* Wait (only 3/4 Compute operation services expose a Wait rpc, need to follow up here)
* Metadata (No plans to support)
* Response (No plans to support)

Specifically, the Operation wrapper(s) generated for DIREGAPIC APIs will create a single Operation wrapper type, that embeds an (internal) operationHandle interface consisting of `Proto` and `Poll` methods. Every operation service will get an operationHandle implementation, which will be embedded in Operation wrappers created by RPCs that create Operations. Each method knows which operation service handle they need to create. The operation service handle is instantiated with the consumer API client and used for every Operation wrapper created by that client.

Companion PR for compute is: https://github.com/googleapis/google-cloud-go/pull/5346
Specifically, the [`operations.go`](https://github.com/googleapis/google-cloud-go/pull/5346/files#diff-f4a0c12e2a5a0ec101424679218e0c5270a8d0b61e697b4cb3edaedaf835d790) file contains the new operation handles.

noahdietz: @frodothetrue can you please give this new generated surface (and the companion PR) a look when you get a chance?

Patch:
diff --git a/WORKSPACE b/WORKSPACE
index f410cc7..1c7fe83 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -22,12 +22,6 @@ http_archive(
     ],
 )
 
-load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_dependencies")
-
-go_rules_dependencies()
-
-go_register_toolchains(version = "1.15.8")
-
 http_archive(
     name = "bazel_gazelle",
     sha256 = "de69a09dc70417580aabf20a28619bb3ef60d038470c7cf8442fafcf627c21cb",
@@ -37,8 +31,26 @@ http_archive(
     ],
 )
 
+load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_dependencies")
+
 # gazelle:repo bazel_gazelle
-load("@bazel_gazelle//:deps.bzl", "gazelle_dependencies")
+load("@bazel_gazelle//:deps.bzl", "gazelle_dependencies", "go_repository")
+
+# TODO(noahdietz): We should remove this eventually, it will complicate bazel dep updates.
+# This is necesary to include extendedops and routing annotations because rules_go doesn't
+# provde a recent enough version.
+go_repository(
+    name = "org_golang_google_genproto",
+    importpath = "google.golang.org/genproto",
+    sum = "h1:I0YcKz0I7OAhddo7ya8kMnvprhcWM045PmkBdMO9zN0=",
+    version = "v0.0.0-20211208223120-3a66f561d7aa",
+)
+
+go_rules_dependencies()
+
+go_register_toolchains(
+    version = "1.15.8",
+)
 
 gazelle_dependencies()
 
diff --git a/internal/gengapic/BUILD.bazel b/internal/gengapic/BUILD.bazel
index 8cf2841..3005f76 100644
--- a/internal/gengapic/BUILD.bazel
+++ b/internal/gengapic/BUILD.bazel
@@ -40,6 +40,7 @@ go_library(
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@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",
diff --git a/internal/gengapic/custom_operation.go b/internal/gengapic/custom_operation.go
index fbee539..36bddbc 100644
--- a/internal/gengapic/custom_operation.go
+++ b/internal/gengapic/custom_operation.go
@@ -18,12 +18,15 @@ import (
 	"fmt"
 
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
+	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
+	"google.golang.org/genproto/googleapis/cloud/extendedops"
+	"google.golang.org/protobuf/proto"
 )
 
 // customOp represents a custom operation type for long running operations.
 type customOp struct {
-	message   *descriptor.DescriptorProto
-	generated bool
+	message *descriptor.DescriptorProto
+	handles []*descriptor.ServiceDescriptorProto
 }
 
 // isCustomOp determines if the given method should return a custom operation wrapper.
@@ -63,39 +66,219 @@ 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 string) string {
+func (g *generator) customOpInit(h, p string) string {
 	opName := g.aux.customOp.message.GetName()
 
-	s := fmt.Sprintf("&%s{proto: %s}", opName, p)
+	s := fmt.Sprintf("&%s{&%s{c: c.operationClient, proto: %s}}", opName, h, p)
 
 	return s
 }
 
-// customOperationType generates the custom operation wrapper type using the
-// generators current printer. This should only be called once per package.
+// customOperationType generates the custom operation wrapper type and operation
+// service handle implementations using the generators current printer. This
+// should only be called once per package.
 func (g *generator) customOperationType() error {
 	op := g.aux.customOp
 	if op == nil {
 		return nil
 	}
 	opName := op.message.GetName()
+	handleInt := lowerFirst(opName + "Handle")
 
 	ptyp, err := g.customOpPointerType()
 	if err != nil {
 		return err
 	}
+	_, opImp, err := g.descInfo.NameSpec(op.message)
+	if err != nil {
+		return err
+	}
+	g.imports[opImp] = true
+
+	statusField := operationField(op.message, extendedops.OperationResponseMapping_STATUS)
+	if statusField == nil {
+		return fmt.Errorf("operation message %s is missing an annotated status field", op.message.GetName())
+	}
+
+	opNameField := operationField(op.message, extendedops.OperationResponseMapping_NAME)
+	if opNameField == nil {
+		return fmt.Errorf("operation message %s is missing an annotated name field", op.message.GetName())
+	}
+
+	opNameGetter := fieldGetter(opNameField.GetName())
 
 	p := g.printf
 
-	p("// %s represents a long running operation for this API.", opName)
+	p("// %s represents a long-running operation for this API.", opName)
 	p("type %s struct {", opName)
-	p("  proto %s", ptyp)
+	p("  %s", handleInt)
 	p("}")
 	p("")
-	p("// Proto returns the raw type this wraps.")
-	p("func (o *%s) Proto() %s {", opName, ptyp)
-	p("  return o.proto")
+
+	// Done
+	p("// Done reports whether the long-running operation has completed.")
+	p("func (o *%s) Done() bool {", opName)
+	p(g.customOpStatusCheck(statusField))
 	p("}")
+	p("")
+
+	// Name
+	p("// Name returns the name of the long-running operation.")
+	p("// The name is assigned by the server and is unique within the service from which the operation is created.")
+	p("func (o *%s) Name() string {", opName)
+	p("  return o.Proto()%s", opNameGetter)
+	p("}")
+	p("")
+
+	p("type %s interface {", handleInt)
+	p("  // Poll retrieves the operation.")
+	p("  Poll(ctx context.Context, opts ...gax.CallOption) error")
+	p("")
+	p("  // Proto returns the long-running operation message.")
+	p("  Proto() %s", ptyp)
+	p("}")
+	p("")
+	g.imports[pbinfo.ImportSpec{Path: "context"}] = true
+	g.imports[pbinfo.ImportSpec{Name: "gax", Path: "github.com/googleapis/gax-go/v2"}] = true
+
+	for _, handle := range op.handles {
+		s := pbinfo.ReduceServName(handle.GetName(), opImp.Name)
+		n := lowerFirst(s + "Handle")
+
+		// 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())
+
+		// type
+		p("// Implements the %s interface for %s.", handleInt, handle.GetName())
+		p("type %s struct {", n)
+		p("  c *%sClient", s)
+		p("  proto %s", ptyp)
+		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("  if err != nil {")
+		p("    return err")
+		p("  }")
+		p("  h.proto = resp")
+		p("  return nil")
+		p("}")
+		p("")
+
+		// Proto
+		p("// Proto returns the raw type this wraps.")
+		p("func (h *%s) Proto() %s {", n, ptyp)
+		p("  return h.proto")
+		p("}")
+		p("")
+	}
+
+	return nil
+}
+
+// loadCustomOpServices maps the service declared as a google.cloud.operation_service
+// to the service that owns the method(s) declaring it.
+func (g *generator) loadCustomOpServices(servs []*descriptor.ServiceDescriptorProto) {
+	handles := g.aux.customOp.handles
+	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)
+				}
+				break
+			}
+		}
+	}
+	g.aux.customOp.handles = handles
+}
+
+// customOpService loads the ServiceDescriptorProto for the google.cloud.operation_service
+// named on the given method.
+func (g *generator) customOpService(m *descriptor.MethodDescriptorProto) *descriptor.ServiceDescriptorProto {
+	opServName := proto.GetExtension(m.GetOptions(), extendedops.E_OperationService).(string)
+	if opServName == "" {
+		return nil
+	}
+
+	file := g.descInfo.ParentFile[m]
+	fqn := fmt.Sprintf(".%s.%s", file.GetPackage(), opServName)
+
+	return g.descInfo.Serv[fqn]
+}
+
+// customOpStatusCheck constructs a return statement that checks if the operation's Status
+// field indicates it is done.
+func (g *generator) customOpStatusCheck(st *descriptor.FieldDescriptorProto) string {
+	ret := fmt.Sprintf("return o.Proto()%s", fieldGetter(st.GetName()))
+	if st.GetType() == descriptor.FieldDescriptorProto_TYPE_ENUM {
+		done := g.customOpStatusEnumDone()
+		ret = fmt.Sprintf("%s == %s", ret, done)
+	}
+
+	return ret
+}
+
+// customOpStatusEnumDone constructs the Go name of the operation's status enum
+// DONE value.
+func (g *generator) customOpStatusEnumDone() string {
+	op := g.aux.customOp.message
+
+	// Ignore the error here, it would failed much earlier if the
+	// operation type was unresolvable.
+	_, imp, _ := g.descInfo.NameSpec(op)
 
+	// Ignore the nil case here, it would failed earlier if the
+	// status field was not present.
+	statusField := operationField(op, extendedops.OperationResponseMapping_STATUS)
+	statusEnum := g.descInfo.Type[statusField.GetTypeName()]
+
+	enum := fmt.Sprintf("%s_DONE", g.nestedName(g.descInfo.ParentElement[statusEnum]))
+
+	s := fmt.Sprintf("%s.%s", imp.Name, enum)
+
+	return s
+}
+
+// operationField is a helper for loading the target google.cloud.operation_field annotation value
+// if present on a field in the given message.
+func operationField(m *descriptor.DescriptorProto, target extendedops.OperationResponseMapping) *descriptor.FieldDescriptorProto {
+	for _, f := range m.GetField() {
+		mapping := proto.GetExtension(f.GetOptions(), extendedops.E_OperationField).(extendedops.OperationResponseMapping)
+		if mapping == target {
+			return f
+		}
+	}
 	return nil
 }
+
+// operationResponseField is a helper for finding the message field that declares the target field name
+// in the google.cloud.operation_response_field annotation.
+func operationResponseField(m *descriptor.DescriptorProto, target string) *descriptor.FieldDescriptorProto {
+	for _, f := range m.GetField() {
+		mapping := proto.GetExtension(f.GetOptions(), extendedops.E_OperationResponseField).(string)
+		if mapping == target {
+			return f
+		}
+	}
+	return nil
+}
+
+// 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 {
+	s = pbinfo.ReduceServName(s, pkg)
+	return lowerFirst(s + "Handle")
+}
diff --git a/internal/gengapic/generator.go b/internal/gengapic/generator.go
index ba74d8e..38a072b 100644
--- a/internal/gengapic/generator.go
+++ b/internal/gengapic/generator.go
@@ -69,6 +69,10 @@ type generator struct {
 	mixins mixins
 
 	hasIAMPolicyOverrides bool
+
+	// customOpServices is a map of service descriptors with methods that create custom operations
+	// to the service descriptors of the custom operation services that manage those custom operation instances.
+	customOpServices map[*descriptor.ServiceDescriptorProto]*descriptor.ServiceDescriptorProto
 }
 
 func (g *generator) init(req *plugin.CodeGeneratorRequest) error {
@@ -82,6 +86,7 @@ func (g *generator) init(req *plugin.CodeGeneratorRequest) error {
 	g.mixins = make(mixins)
 	g.comments = map[protoiface.MessageV1]string{}
 	g.imports = map[pbinfo.ImportSpec]bool{}
+	g.customOpServices = map[*descriptor.ServiceDescriptorProto]*descriptor.ServiceDescriptorProto{}
 	g.aux = &auxTypes{
 		iters: map[string]*iterType{},
 		lros:  map[*descriptor.MethodDescriptorProto]bool{},
@@ -247,3 +252,15 @@ func (g *generator) fqn(p pbinfo.ProtoType) string {
 	}
 	return fmt.Sprintf("%s.%s", g.fqn(parent), p.GetName())
 }
+
+func (g *generator) nestedName(nested pbinfo.ProtoType) string {
+	name := nested.GetName()
+
+	parent, hasParent := g.descInfo.ParentElement[nested]
+	for hasParent {
+		name = fmt.Sprintf("%s_%s", parent.GetName(), name)
+		parent, hasParent = g.descInfo.ParentElement[parent]
+	}
+
+	return name
+}
diff --git a/internal/gengapic/gengapic.go b/internal/gengapic/gengapic.go
index 9f9f922..f4de0bc 100644
--- a/internal/gengapic/gengapic.go
+++ b/internal/gengapic/gengapic.go
@@ -75,7 +75,8 @@ 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), false}
+		g.aux.customOp = &customOp{op.(*descriptor.DescriptorProto), []*descriptor.ServiceDescriptorProto{}}
+		g.loadCustomOpServices(genServs)
 	}
 
 	// Use the proto package from the parent file of the first Service seen.
@@ -130,6 +131,14 @@ func Gen(genReq *plugin.CodeGeneratorRequest) (*plugin.CodeGeneratorResponse, er
 		})
 	}
 
+	if g.aux.customOp != nil {
+		g.reset()
+		if err := g.customOperationType(); err != nil {
+			return &g.resp, err
+		}
+		g.commit(filepath.Join(g.opts.outDir, "operations.go"), g.opts.pkgName)
+	}
+
 	return &g.resp, nil
 }
 
@@ -197,12 +206,6 @@ func (g *generator) gen(serv *descriptor.ServiceDescriptorProto) error {
 	for _, iter := range iters {
 		g.pagingIter(iter)
 	}
-	if g.aux.customOp != nil && !g.aux.customOp.generated {
-		if err := g.customOperationType(); err != nil {
-			return err
-		}
-		g.aux.customOp.generated = true
-	}
 
 	return nil
 }
diff --git a/internal/gengapic/genrest.go b/internal/gengapic/genrest.go
index 6c1940d..0c4ca7a 100644
--- a/internal/gengapic/genrest.go
+++ b/internal/gengapic/genrest.go
@@ -48,6 +48,12 @@ func (g *generator) restClientInit(serv *descriptor.ServiceDescriptorProto, serv
 	p("  // The http client.")
 	p("  httpClient *http.Client")
 	p("")
+	if opServ, ok := g.customOpServices[serv]; ok {
+		opServName := pbinfo.ReduceServName(opServ.GetName(), g.opts.pkgName)
+		p("// operationClient is used to call the operation-specific management service.")
+		p("operationClient *%sClient", opServName)
+		p("")
+	}
 	p("	 // The x-goog-* metadata to be sent with each request.")
 	p("	 xGoogMetadata metadata.MD")
 	p("}")
@@ -109,6 +115,7 @@ func (g *generator) restClientUtilities(serv *descriptor.ServiceDescriptorProto,
 	lowcaseServName := lowcaseRestClientName(servName)
 	clientName := camelToSnake(serv.GetName())
 	clientName = strings.Replace(clientName, "_", " ", -1)
+	opServ, hasCustomOp := g.customOpServices[serv]
 
 	p("// New%sRESTClient creates a new %s rest client.", servName, clientName)
 	g.serviceDoc(serv)
@@ -125,6 +132,20 @@ func (g *generator) restClientUtilities(serv *descriptor.ServiceDescriptorProto,
 	p("    }")
 	p("    c.setGoogleClientInfo()")
 	p("")
+	if hasCustomOp {
+		opServName := pbinfo.ReduceServName(opServ.GetName(), g.opts.pkgName)
+		p("o := []option.ClientOption{")
+		p("  option.WithHTTPClient(httpClient),")
+		p("  option.WithEndpoint(endpoint),")
+		p("}")
+		p("opC, err := New%sRESTClient(ctx, o...)", opServName)
+		p("if err != nil {")
+		p("  return nil, err")
+		p("}")
+		p("c.operationClient = opC")
+		p("")
+		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/api/option"}] = true
+	}
 	// TODO(dovs): make rest default call options
 	// TODO(dovs): set the LRO client
 	p("    return &%[1]sClient{internalClient: c, CallOptions: &%[1]sCallOptions{}}, nil", servName)
@@ -150,6 +171,11 @@ func (g *generator) restClientUtilities(serv *descriptor.ServiceDescriptorProto,
 	p("func (c *%s) Close() error {", lowcaseServName)
 	p("    // Replace httpClient with nil to force cleanup.")
 	p("    c.httpClient = nil")
+	if hasCustomOp {
+		p("if err := c.operationClient.Close(); err != nil {")
+		p("  return err")
+		p("}")
+	}
 	p("    return nil")
 	p("}")
 	p("")
@@ -771,6 +797,8 @@ func (g *generator) unaryRESTCall(servName string, m *descriptor.MethodDescripto
 	// the ImportSpec for the OutputType, which has already happened above.
 	retTyp, _ := g.returnType(m)
 
+	isCustomOp := g.isCustomOp(m, info)
+
 	p := g.printf
 	lowcaseServName := lowcaseRestClientName(servName)
 	p("func (c *%s) %s(ctx context.Context, req *%s.%s, opts ...gax.CallOption) (%s, error) {",
@@ -855,8 +883,10 @@ func (g *generator) unaryRESTCall(servName string, m *descriptor.MethodDescripto
 	p("  return nil, e")
 	p("}")
 	ret := "return resp, nil"
-	if g.isCustomOp(m, info) {
-		p("op := %s", g.customOpInit("resp"))
+	if isCustomOp {
+		s := g.customOpService(m)
+		handleName := handleName(s.GetName(), g.opts.pkgName)
+		p("op := %s", g.customOpInit(handleName, "resp"))
 		ret = "return op, nil"
 	}
 	p(ret)
diff --git a/internal/gengapic/helpers.go b/internal/gengapic/helpers.go
index 4148d2f..4e55f91 100644
--- a/internal/gengapic/helpers.go
+++ b/internal/gengapic/helpers.go
@@ -133,6 +133,18 @@ func containsTransport(t []transport, tr transport) bool {
 	return false
 }
 
+// containsService determines if a set of services contains a specific service,
+// by simple name.
+func containsService(s []*descriptor.ServiceDescriptorProto, srv *descriptor.ServiceDescriptorProto) bool {
+	for _, x := range s {
+		if x.GetName() == srv.GetName() {
+			return true
+		}
+	}
+
+	return false
+}
+
 // isRequired returns if a field is annotated as REQUIRED or not.
 func isRequired(field *descriptor.FieldDescriptorProto) bool {
 	if field.GetOptions() == nil {
diff --git a/internal/gengapic/paging.go b/internal/gengapic/paging.go
index 5a350a3..ee0a46e 100644
--- a/internal/gengapic/paging.go
+++ b/internal/gengapic/paging.go
@@ -50,12 +50,7 @@ func (g *generator) iterTypeOf(elemField *descriptor.FieldDescriptorProto) (*ite
 
 		// Prepend parent Message name for nested Messages
 		// to match the generated Go type name.
-		typ := eType
-		typeName := typ.GetName()
-		for parent, ok := g.descInfo.ParentElement[typ]; ok; parent, ok = g.descInfo.ParentElement[typ] {
-			typeName = fmt.Sprintf("%s_%s", parent.GetName(), typeName)
-			typ = parent
-		}
+		typeName := g.nestedName(eType)
 
 		eMsg, ok := eType.(*descriptor.DescriptorProto)
 		if !ok {
diff --git a/repositories.bzl b/repositories.bzl
index ae364a7..d928c2d 100644
--- a/repositories.bzl
+++ b/repositories.bzl
@@ -987,12 +987,6 @@ def com_googleapis_gapic_generator_go_repositories():
         sum = "h1:FZR1q0exgwxzPzp/aF+VccGrSfxfPpkBqjIIEq3ru6c=",
         version = "v1.6.7",
     )
-    go_repository(
-        name = "org_golang_google_genproto",
-        importpath = "google.golang.org/genproto",
-        sum = "h1:I0YcKz0I7OAhddo7ya8kMnvprhcWM045PmkBdMO9zN0=",
-        version = "v0.0.0-20211208223120-3a66f561d7aa",
-    )
     go_repository(
         name = "org_golang_google_grpc",
         importpath = "google.golang.org/grpc",


Test patch:
diff --git a/internal/gengapic/client_init_test.go b/internal/gengapic/client_init_test.go
index 9ef8354..7914686 100644
--- a/internal/gengapic/client_init_test.go
+++ b/internal/gengapic/client_init_test.go
@@ -27,6 +27,7 @@ import (
 	"github.com/googleapis/gapic-generator-go/internal/txtdiff"
 	"google.golang.org/genproto/googleapis/api/annotations"
 	"google.golang.org/genproto/googleapis/api/serviceconfig"
+	"google.golang.org/genproto/googleapis/cloud/extendedops"
 	code "google.golang.org/genproto/googleapis/rpc/code"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
@@ -348,6 +349,14 @@ func TestClientInit(t *testing.T) {
 			Deprecated: proto.Bool(true),
 		},
 	}
+
+	opS := &descriptor.ServiceDescriptorProto{
+		Name:   proto.String("FooOperationService"),
+		Method: []*descriptor.MethodDescriptorProto{},
+	}
+
+	customOpOpts := &descriptor.MethodOptions{}
+	proto.SetExtension(customOpOpts, extendedops.E_OperationService, opS.GetName())
 	servCustomOp := &descriptor.ServiceDescriptorProto{
 		Name: proto.String("Foo"),
 		Method: []*descriptor.MethodDescriptorProto{
@@ -355,7 +364,7 @@ func TestClientInit(t *testing.T) {
 				Name:       proto.String("Zip"),
 				InputType:  proto.String(".mypackage.Bar"),
 				OutputType: proto.String(".mypackage.Operation"),
-				Options:    &descriptor.MethodOptions{},
+				Options:    customOpOpts,
 			},
 		},
 	}
@@ -368,12 +377,13 @@ func TestClientInit(t *testing.T) {
 	}
 
 	for _, tst := range []struct {
-		tstName   string
-		servName  string
-		mixins    mixins
-		serv      *descriptor.ServiceDescriptorProto
-		parameter *string
-		imports   map[pbinfo.ImportSpec]bool
+		tstName      string
+		servName     string
+		mixins       mixins
+		serv         *descriptor.ServiceDescriptorProto
+		customOpServ *descriptor.ServiceDescriptorProto
+		parameter    *string
+		imports      map[pbinfo.ImportSpec]bool
 	}{
 		{
 			tstName: "foo_client_init",
@@ -477,10 +487,11 @@ func TestClientInit(t *testing.T) {
 			},
 		},
 		{
-			tstName:   "custom_op_init",
-			servName:  "",
-			serv:      servCustomOp,
-			parameter: proto.String("go-gapic-package=path;mypackage,transport=rest,diregapic"),
+			tstName:      "custom_op_init",
+			servName:     "",
+			serv:         servCustomOp,
+			customOpServ: opS,
+			parameter:    proto.String("go-gapic-package=path;mypackage,transport=rest,diregapic"),
 			imports: map[pbinfo.ImportSpec]bool{
 				{Path: "context"}:                      true,
 				{Path: "fmt"}:                          true,
@@ -500,7 +511,7 @@ func TestClientInit(t *testing.T) {
 			Options: &descriptor.FileOptions{
 				GoPackage: proto.String("github.com/googleapis/mypackage/v1"),
 			},
-			Service: []*descriptor.ServiceDescriptorProto{tst.serv},
+			Service: []*descriptor.ServiceDescriptorProto{tst.serv, tst.customOpServ},
 			MessageType: []*descriptor.DescriptorProto{
 				{
 					Name: proto.String("Bar"),
@@ -532,6 +543,11 @@ func TestClientInit(t *testing.T) {
 		g.aux.customOp = &customOp{
 			message: cop,
 		}
+		if tst.customOpServ != nil {
+			g.customOpServices = map[*descriptor.ServiceDescriptorProto]*descriptor.ServiceDescriptorProto{
+				tst.serv: tst.customOpServ,
+			}
+		}
 
 		g.reset()
 		g.makeClients(tst.serv, tst.servName)
diff --git a/internal/gengapic/custom_operation_test.go b/internal/gengapic/custom_operation_test.go
index 76d9b40..182c91e 100644
--- a/internal/gengapic/custom_operation_test.go
+++ b/internal/gengapic/custom_operation_test.go
@@ -23,6 +23,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
 	"github.com/googleapis/gapic-generator-go/internal/txtdiff"
+	"google.golang.org/genproto/googleapis/cloud/extendedops"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/runtime/protoiface"
 )
@@ -95,37 +96,125 @@ func TestCustomOpInit(t *testing.T) {
 			},
 		},
 	}
-	got := g.customOpInit("foo")
-	want := "&Operation{proto: foo}"
+	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)
 	}
 }
 
 func TestCustomOperationType(t *testing.T) {
+	nameOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(nameOpts, extendedops.E_OperationField, extendedops.OperationResponseMapping_NAME)
+	nameField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("name"),
+		Type:    descriptor.FieldDescriptorProto_TYPE_STRING.Enum(),
+		Options: nameOpts,
+	}
+
+	statusEnum := &descriptor.EnumDescriptorProto{
+		Name: proto.String("Status"),
+		Value: []*descriptor.EnumValueDescriptorProto{
+			{
+				Name:   proto.String("DONE"),
+				Number: proto.Int32(0),
+			},
+		},
+	}
+
+	statusOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(statusOpts, extendedops.E_OperationField, extendedops.OperationResponseMapping_STATUS)
+	statusEnumField := &descriptor.FieldDescriptorProto{
+		Name:     proto.String("status"),
+		Type:     descriptor.FieldDescriptorProto_TYPE_ENUM.Enum(),
+		TypeName: proto.String(".google.cloud.foo.v1.Operation.Status"),
+		Options:  statusOpts,
+	}
+
+	statusBoolField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("status"),
+		Type:    descriptor.FieldDescriptorProto_TYPE_BOOL.Enum(),
+		Options: statusOpts,
+	}
+
 	op := &descriptor.DescriptorProto{
-		Name: proto.String("Operation"),
+		Name:     proto.String("Operation"),
+		EnumType: []*descriptor.EnumDescriptorProto{statusEnum},
+	}
+
+	inNameOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(inNameOpts, extendedops.E_OperationResponseField, nameField.GetName())
+	inNameField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("operation"),
+		Type:    descriptor.FieldDescriptorProto_TYPE_STRING.Enum(),
+		Options: inNameOpts,
+	}
+
+	getInput := &descriptor.DescriptorProto{
+		Name:  proto.String("GetFooOperationRequest"),
+		Field: []*descriptor.FieldDescriptorProto{inNameField},
+	}
+
+	fooOpServ := &descriptor.ServiceDescriptorProto{
+		Name: proto.String("FooOperationsService"),
+		Method: []*descriptor.MethodDescriptorProto{
+			{
+				Name:      proto.String("Get"),
+				InputType: proto.String(".google.cloud.foo.v1.GetFooOperationRequest"),
+			},
+		},
+	}
+
+	f := &descriptor.FileDescriptorProto{
+		Package: proto.String("google.cloud.foo.v1"),
+		Options: &descriptor.FileOptions{
+			GoPackage: proto.String("google.golang.org/genproto/cloud/foo/v1;foo"),
+		},
+		Service: []*descriptor.ServiceDescriptorProto{fooOpServ},
 	}
 	g := &generator{
 		aux: &auxTypes{
 			customOp: &customOp{
 				message: op,
+				handles: []*descriptor.ServiceDescriptorProto{fooOpServ},
 			},
 		},
 		descInfo: pbinfo.Info{
 			ParentFile: map[protoiface.MessageV1]*descriptor.FileDescriptorProto{
-				op: {
-					Package: proto.String("google.cloud.foo.v1"),
-					Options: &descriptor.FileOptions{
-						GoPackage: proto.String("google.golang.org/genproto/cloud/foo/v1;foo"),
-					},
-				},
+				op:        f,
+				fooOpServ: f,
+				getInput:  f,
+			},
+			ParentElement: map[pbinfo.ProtoType]pbinfo.ProtoType{
+				statusEnum: op,
+			},
+			Type: map[string]pbinfo.ProtoType{
+				statusEnumField.GetTypeName():                 statusEnum,
+				".google.cloud.foo.v1.GetFooOperationRequest": getInput,
 			},
 		},
+		imports: map[pbinfo.ImportSpec]bool{},
 	}
-	err := g.customOperationType()
-	if err != nil {
-		t.Fatal(err)
+	for _, tst := range []struct {
+		name string
+		st   *descriptor.FieldDescriptorProto
+	}{
+		{
+			name: "enum",
+			st:   statusEnumField,
+		},
+		{
+			name: "bool",
+			st:   statusBoolField,
+		},
+	} {
+		op.Field = []*descriptor.FieldDescriptorProto{nameField, tst.st}
+		err := g.customOperationType()
+		if err != nil {
+			t.Fatal(err)
+		}
+		tn := "custom_op_type_" + tst.name
+		txtdiff.Diff(t, tn, g.pt.String(), filepath.Join("testdata", tn+".want"))
+		g.reset()
 	}
-	txtdiff.Diff(t, "custom_op_type", g.pt.String(), filepath.Join("testdata", "custom_op_type.want"))
 }
diff --git a/internal/gengapic/genrest_test.go b/internal/gengapic/genrest_test.go
index 4876b70..73aad26 100644
--- a/internal/gengapic/genrest_test.go
+++ b/internal/gengapic/genrest_test.go
@@ -25,6 +25,7 @@ import (
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
 	"github.com/googleapis/gapic-generator-go/internal/txtdiff"
 	"google.golang.org/genproto/googleapis/api/annotations"
+	"google.golang.org/genproto/googleapis/cloud/extendedops"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protodesc"
 	"google.golang.org/protobuf/runtime/protoiface"
@@ -461,8 +462,16 @@ func TestGenRestMethod(t *testing.T) {
 	}
 	pagedFooResFQN := fmt.Sprintf(".%s.PagedFooResponse", pkg)
 
+	nameOpts := &descriptor.FieldOptions{}
+	proto.SetExtension(nameOpts, extendedops.E_OperationField, extendedops.OperationResponseMapping_NAME)
+	nameField := &descriptor.FieldDescriptorProto{
+		Name:    proto.String("name"),
+		Type:    descriptor.FieldDescriptorProto_TYPE_STRING.Enum(),
+		Options: nameOpts,
+	}
 	op := &descriptor.DescriptorProto{
-		Name: proto.String("Operation"),
+		Name:  proto.String("Operation"),
+		Field: []*descriptor.FieldDescriptorProto{nameField},
 	}
 	opfqn := fmt.Sprintf(".%s.Operation", pkg)
 
@@ -472,6 +481,7 @@ func TestGenRestMethod(t *testing.T) {
 			Post: "/v1/foo",
 		},
 	})
+	proto.SetExtension(opRPCOpt, extendedops.E_OperationService, "FooOperationService")
 
 	opRPC := &descriptor.MethodDescriptorProto{
 		Name:       proto.String("CustomOp"),
@@ -526,13 +536,16 @@ func TestGenRestMethod(t *testing.T) {
 	s := &descriptor.ServiceDescriptorProto{
 		Name: proto.String("FooService"),
 	}
+	opS := &descriptor.ServiceDescriptorProto{
+		Name: proto.String("FooOperationService"),
+	}
 
 	f := &descriptor.FileDescriptorProto{
 		Package: proto.String(pkg),
 		Options: &descriptor.FileOptions{
 			GoPackage: proto.String("google.golang.org/genproto/cloud/foo/v1;foo"),
 		},
-		Service: []*descriptor.ServiceDescriptorProto{s},
+		Service: []*descriptor.ServiceDescriptorProto{s, opS},
 	}
 
 	g := &generator{
@@ -543,9 +556,14 @@ func TestGenRestMethod(t *testing.T) {
 			iters: map[string]*iterType{},
 		},
 		opts: &options{},
+		customOpServices: map[*descriptor.ServiceDescriptorProto]*descriptor.ServiceDescriptorProto{
+			s: opS,
+		},
 		descInfo: pbinfo.Info{
 			ParentFile: map[protoiface.MessageV1]*descriptor.FileDescriptorProto{
 				op:          f,
+				opS:         f,
+				opRPC:       f,
 				foo:         f,
 				s:           f,
 				pagedFooReq: f,
@@ -556,6 +574,7 @@ func TestGenRestMethod(t *testing.T) {
 				emptyRPC:   s,
 				unaryRPC:   s,
 				pagingRPC:  s,
+				nameField:  op,
 				sizeField:  foo,
 				otherField: foo,
 			},
diff --git a/internal/gengapic/testdata/custom_op_init.want b/internal/gengapic/testdata/custom_op_init.want
index 3a86ccd..a889309 100644
--- a/internal/gengapic/testdata/custom_op_init.want
+++ b/internal/gengapic/testdata/custom_op_init.want
@@ -54,6 +54,9 @@ type restClient struct {
 	// The http client.
 	httpClient *http.Client
 
+	// operationClient is used to call the operation-specific management service.
+	operationClient *FooOperationClient
+
 	// The x-goog-* metadata to be sent with each request.
 	xGoogMetadata metadata.MD
 }
@@ -74,6 +77,16 @@ func NewRESTClient(ctx context.Context, opts ...option.ClientOption) (*Client, e
 	}
 	c.setGoogleClientInfo()
 
+	o := []option.ClientOption{
+		option.WithHTTPClient(httpClient),
+		option.WithEndpoint(endpoint),
+	}
+	opC, err := NewFooOperationRESTClient(ctx, o...)
+	if err != nil {
+		return nil, err
+	}
+	c.operationClient = opC
+
 	return &Client{internalClient: c, CallOptions: &CallOptions{}}, nil
 }
 
@@ -91,6 +104,9 @@ func (c *restClient) setGoogleClientInfo(keyval ...string) {
 func (c *restClient) Close() error {
 	// Replace httpClient with nil to force cleanup.
 	c.httpClient = nil
+	if err := c.operationClient.Close(); err != nil {
+		return err
+	}
 	return nil
 }
 
diff --git a/internal/gengapic/testdata/custom_op_type.want b/internal/gengapic/testdata/custom_op_type.want
deleted file mode 100644
index f33765d..0000000
--- a/internal/gengapic/testdata/custom_op_type.want
+++ /dev/null
@@ -1,9 +0,0 @@
-// Operation represents a long running operation for this API.
-type Operation struct {
-	proto *foopb.Operation
-}
-
-// Proto returns the raw type this wraps.
-func (o *Operation) Proto() *foopb.Operation {
-	return o.proto
-}
diff --git a/internal/gengapic/testdata/custom_op_type_bool.want b/internal/gengapic/testdata/custom_op_type_bool.want
new file mode 100644
index 0000000..11171c0
--- /dev/null
+++ b/internal/gengapic/testdata/custom_op_type_bool.want
@@ -0,0 +1,45 @@
+// Operation represents a long-running operation for this API.
+type Operation struct {
+	operationHandle
+}
+
+// Done reports whether the long-running operation has completed.
+func (o *Operation) Done() bool {
+	return o.Proto().GetStatus()
+}
+
+// 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 (o *Operation) Name() string {
+	return o.Proto().GetName()
+}
+
+type operationHandle interface {
+	// Poll retrieves the operation.
+	Poll(ctx context.Context, opts ...gax.CallOption) error
+
+	// Proto returns the long-running operation message.
+	Proto() *foopb.Operation
+}
+
+// Implements the operationHandle interface for FooOperationsService.
+type fooOperationsHandle struct {
+	c *FooOperationsClient
+	proto *foopb.Operation
+}
+
+// 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...)
+	if err != nil {
+		return err
+	}
+	h.proto = resp
+	return nil
+}
+
+// Proto returns the raw type this wraps.
+func (h *fooOperationsHandle) Proto() *foopb.Operation {
+	return h.proto
+}
+
diff --git a/internal/gengapic/testdata/custom_op_type_enum.want b/internal/gengapic/testdata/custom_op_type_enum.want
new file mode 100644
index 0000000..8085027
--- /dev/null
+++ b/internal/gengapic/testdata/custom_op_type_enum.want
@@ -0,0 +1,45 @@
+// Operation represents a long-running operation for this API.
+type Operation struct {
+	operationHandle
+}
+
+// Done reports whether the long-running operation has completed.
+func (o *Operation) Done() bool {
+	return o.Proto().GetStatus() == foopb.Operation_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 (o *Operation) Name() string {
+	return o.Proto().GetName()
+}
+
+type operationHandle interface {
+	// Poll retrieves the operation.
+	Poll(ctx context.Context, opts ...gax.CallOption) error
+
+	// Proto returns the long-running operation message.
+	Proto() *foopb.Operation
+}
+
+// Implements the operationHandle interface for FooOperationsService.
+type fooOperationsHandle struct {
+	c *FooOperationsClient
+	proto *foopb.Operation
+}
+
+// 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...)
+	if err != nil {
+		return err
+	}
+	h.proto = resp
+	return nil
+}
+
+// Proto returns the raw type this wraps.
+func (h *fooOperationsHandle) Proto() *foopb.Operation {
+	return h.proto
+}
+
diff --git a/internal/gengapic/testdata/rest_CustomOp.want b/internal/gengapic/testdata/rest_CustomOp.want
index 2193ccb..e71bf37 100644
--- a/internal/gengapic/testdata/rest_CustomOp.want
+++ b/internal/gengapic/testdata/rest_CustomOp.want
@@ -46,6 +46,6 @@ func (c *fooRESTClient) CustomOp(ctx context.Context, req *foopb.Foo, opts ...ga
 	if e != nil {
 		return nil, e
 	}
-	op := &Operation{proto: resp}
+	op := &Operation{&handle{c: c.operationClient, proto: resp}}
 	return op, nil
 }