Repository: googleapis/gapic-generator-go
Pull request #873 (2022-01-26 23:53:21)

feat(gengapic): diregapic add operation Wait helper
Adds the Operation helper `Wait` to DIREGAPIC operation wrappers. This uses `Poll` to check the status of the operation until it is done or the context is canceled. It uses exponential backoff until that backoff reaches 1 min, when it will then poll regularly every minute. This is the same for standard GAPICs using the LRO helper. 

cc: @vchudnov-g @vam-google FYI


Patch:
diff --git a/internal/gengapic/custom_operation.go b/internal/gengapic/custom_operation.go
index 36bddbc..c7b62eb 100644
--- a/internal/gengapic/custom_operation.go
+++ b/internal/gengapic/custom_operation.go
@@ -130,6 +130,28 @@ func (g *generator) customOperationType() error {
 	p("}")
 	p("")
 
+	// Wait
+	p("// Wait blocks until the operation is complete, polling regularly")
+	p("// after an intial period of backing off between attempts.")
+	p("func (o *%s) Wait(ctx context.Context, opts ...gax.CallOption) error {", opName)
+	p("  bo := gax.Backoff{")
+	p("    Initial: %s,", defaultPollInitialDelay)
+	p("    Max:     %s,", defaultPollMaxDelay)
+	p("  }")
+	p("  for {")
+	p("    if err := o.Poll(ctx, opts...); err != nil {")
+	p("      return err")
+	p("    }")
+	p("    if o.Done() {")
+	p("      return nil")
+	p("    }")
+	p("    if err := gax.Sleep(ctx, bo.Pause()); err != nil {")
+	p("      return err")
+	p("    }")
+	p("  }")
+	p("}")
+	p("")
+
 	p("type %s interface {", handleInt)
 	p("  // Poll retrieves the operation.")
 	p("  Poll(ctx context.Context, opts ...gax.CallOption) error")
diff --git a/internal/gengapic/gengapic.go b/internal/gengapic/gengapic.go
index f4de0bc..e198af1 100644
--- a/internal/gengapic/gengapic.go
+++ b/internal/gengapic/gengapic.go
@@ -34,16 +34,18 @@ import (
 const (
 	emptyValue = "google.protobuf.Empty"
 	// protoc puts a dot in front of name, signaling that the name is fully qualified.
-	emptyType           = "." + emptyValue
-	lroType             = ".google.longrunning.Operation"
-	alpha               = "alpha"
-	beta                = "beta"
-	disableDeadlinesVar = "GOOGLE_API_GO_EXPERIMENTAL_DISABLE_DEFAULT_DEADLINE"
-	fieldTypeBool       = descriptor.FieldDescriptorProto_TYPE_BOOL
-	fieldTypeString     = descriptor.FieldDescriptorProto_TYPE_STRING
-	fieldTypeBytes      = descriptor.FieldDescriptorProto_TYPE_BYTES
-	fieldTypeMessage    = descriptor.FieldDescriptorProto_TYPE_MESSAGE
-	fieldLabelRepeated  = descriptor.FieldDescriptorProto_LABEL_REPEATED
+	emptyType               = "." + emptyValue
+	lroType                 = ".google.longrunning.Operation"
+	alpha                   = "alpha"
+	beta                    = "beta"
+	disableDeadlinesVar     = "GOOGLE_API_GO_EXPERIMENTAL_DISABLE_DEFAULT_DEADLINE"
+	fieldTypeBool           = descriptor.FieldDescriptorProto_TYPE_BOOL
+	fieldTypeString         = descriptor.FieldDescriptorProto_TYPE_STRING
+	fieldTypeBytes          = descriptor.FieldDescriptorProto_TYPE_BYTES
+	fieldTypeMessage        = descriptor.FieldDescriptorProto_TYPE_MESSAGE
+	fieldLabelRepeated      = descriptor.FieldDescriptorProto_LABEL_REPEATED
+	defaultPollInitialDelay = "time.Second" // 1 second
+	defaultPollMaxDelay     = "time.Minute" // 1 minute
 )
 
 var headerParamRegexp = regexp.MustCompile(`{([_.a-z0-9]+)`)
diff --git a/internal/gengapic/lro.go b/internal/gengapic/lro.go
index a897dde..7cd393d 100644
--- a/internal/gengapic/lro.go
+++ b/internal/gengapic/lro.go
@@ -179,11 +179,11 @@ func (g *generator) lroType(servName string, serv *descriptor.ServiceDescriptorP
 		p("// See documentation of Poll for error-handling information.")
 		if opInfo.GetResponseType() == emptyValue {
 			p("func (op *%s) Wait(ctx context.Context, opts ...gax.CallOption) error {", lroType)
-			p("  return op.lro.WaitWithInterval(ctx, nil, time.Minute, opts...)")
+			p("  return op.lro.WaitWithInterval(ctx, nil, %s, opts...)", defaultPollMaxDelay)
 		} else {
 			p("func (op *%s) Wait(ctx context.Context, opts ...gax.CallOption) (*%s, error) {", lroType, respType)
 			p("  var resp %s", respType)
-			p("  if err := op.lro.WaitWithInterval(ctx, &resp, time.Minute, opts...); err != nil {")
+			p("  if err := op.lro.WaitWithInterval(ctx, &resp, %s, opts...); err != nil {", defaultPollMaxDelay)
 			p("    return nil, err")
 			p("  }")
 			p("  return &resp, nil")


Test patch:
diff --git a/internal/gengapic/testdata/custom_op_type_bool.want b/internal/gengapic/testdata/custom_op_type_bool.want
index 11171c0..ff0995e 100644
--- a/internal/gengapic/testdata/custom_op_type_bool.want
+++ b/internal/gengapic/testdata/custom_op_type_bool.want
@@ -14,6 +14,26 @@ func (o *Operation) Name() string {
 	return o.Proto().GetName()
 }
 
+// Wait blocks until the operation is complete, polling regularly
+// after an intial period of backing off between attempts.
+func (o *Operation) Wait(ctx context.Context, opts ...gax.CallOption) error {
+	bo := gax.Backoff{
+		Initial: time.Second,
+		Max:     time.Minute,
+	}
+	for {
+		if err := o.Poll(ctx, opts...); err != nil {
+			return err
+		}
+		if o.Done() {
+			return nil
+		}
+		if err := gax.Sleep(ctx, bo.Pause()); err != nil {
+			return err
+		}
+	}
+}
+
 type operationHandle interface {
 	// Poll retrieves the operation.
 	Poll(ctx context.Context, opts ...gax.CallOption) error
diff --git a/internal/gengapic/testdata/custom_op_type_enum.want b/internal/gengapic/testdata/custom_op_type_enum.want
index 8085027..c134faf 100644
--- a/internal/gengapic/testdata/custom_op_type_enum.want
+++ b/internal/gengapic/testdata/custom_op_type_enum.want
@@ -14,6 +14,26 @@ func (o *Operation) Name() string {
 	return o.Proto().GetName()
 }
 
+// Wait blocks until the operation is complete, polling regularly
+// after an intial period of backing off between attempts.
+func (o *Operation) Wait(ctx context.Context, opts ...gax.CallOption) error {
+	bo := gax.Backoff{
+		Initial: time.Second,
+		Max:     time.Minute,
+	}
+	for {
+		if err := o.Poll(ctx, opts...); err != nil {
+			return err
+		}
+		if o.Done() {
+			return nil
+		}
+		if err := gax.Sleep(ctx, bo.Pause()); err != nil {
+			return err
+		}
+	}
+}
+
 type operationHandle interface {
 	// Poll retrieves the operation.
 	Poll(ctx context.Context, opts ...gax.CallOption) error