feat(gengapic): generate helpers for Go 1.23 iterators (#1542)

This changes starts generating build tag protected adapters for using our iterator types in Go range expressions.

Related: googleapis/gax-go#358

internal/gengapic/auxiliary.go

@@ -79,6 +79,10 @@ func (g *generator) genAuxFile() error {
 	g.commit(filepath.Join(g.opts.outDir, "auxiliary.go"), g.opts.pkgName)
 	g.reset()
 
+	g.genIteratorsGo123()
+	g.commitWithBuildTag(filepath.Join(g.opts.outDir, "auxiliary_go123.go"), g.opts.pkgName, "go1.23")
+	g.reset()
+
 	return nil
 }
 
@@ -123,6 +127,17 @@ func (g *generator) genIterators() error {
 	return nil
 }
 
+// genIteratorsGo123 generates adapters for Go iterators for Go versions 1.23+.
+func (g *generator) genIteratorsGo123() {
+	// Sort iterators to generate by type name to
+	// avoid spurious regenerations created by
+	// non-deterministic map traversal order.
+	iters := sortIteratorMap(g.aux.iters)
+	for _, iter := range iters {
+		g.pagingIterGo123(iter)
+	}
+}
+
 // sortIteratorMap sorts the map of iterator types by iterTypeName.
 func sortIteratorMap(m map[string]*iterType) []*iterType {
 	var iters []*iterType

internal/gengapic/auxiliary_test.go

@@ -444,6 +444,22 @@ func TestGenIterators(t *testing.T) {
 	}
 
 	txtdiff.Diff(t, g.pt.String(), filepath.Join("testdata", "gen_iterators.want"))
+
+	g.reset()
+
+	wantImports = map[pbinfo.ImportSpec]bool{
+		{Path: "iter"}: true,
+		{Path: "github.com/googleapis/gax-go/v2/iterator"}:                       true,
+		{Name: "examplepb", Path: "cloud.google.com/go/example/apiv1/examplepb"}: true,
+	}
+
+	g.genIteratorsGo123()
+
+	if diff := cmp.Diff(g.imports, wantImports); diff != "" {
+		t.Errorf("imports got(-),want(+):\n%s", diff)
+	}
+
+	txtdiff.Diff(t, g.pt.String(), filepath.Join("testdata", "gen_iterators_go123.want"))
 }
 
 func TestSortOperationWrapperMap(t *testing.T) {

internal/gengapic/example.go

@@ -40,6 +40,62 @@ func (g *generator) genExampleFile(serv *descriptorpb.ServiceDescriptorProto) er
 	return nil
 }
 
+func (g *generator) genExampleIteratorFile(serv *descriptorpb.ServiceDescriptorProto) error {
+	pkgName := g.opts.pkgName
+	servName := pbinfo.ReduceServName(serv.GetName(), pkgName)
+	methods := append(serv.GetMethod(), g.getMixinMethods()...)
+	for _, m := range methods {
+		// Don't need streaming RPCs
+		if m.GetClientStreaming() || m.GetServerStreaming() {
+			continue
+		}
+		pf, _, err := g.getPagingFields(m)
+		if err != nil {
+			return err
+		}
+		// Don't generate for non-list RPCs
+		if pf == nil {
+			continue
+		}
+
+		p := g.printf
+
+		inType := g.descInfo.Type[m.GetInputType()]
+		if inType == nil {
+			return fmt.Errorf("cannot find type %q, malformed descriptor?", m.GetInputType())
+		}
+
+		inSpec, err := g.descInfo.ImportSpec(inType)
+		if err != nil {
+			return err
+		}
+
+		g.imports[inSpec] = true
+		// Pick the first transport for simplicity. We don't need examples
+		// of each method for both transports when they have the same surface.
+		t := g.opts.transports[0]
+		s := servName
+		if t == rest {
+			s += "REST"
+		}
+		p("func Example%sClient_%s_all() {", servName, m.GetName())
+		g.exampleInitClient(pkgName, s)
+
+		p("")
+		p("req := &%s.%s{", inSpec.Name, inType.GetName())
+		p("  // TODO: Fill request struct fields.")
+		p("  // See https://pkg.go.dev/%s#%s.", inSpec.Path, inType.GetName())
+		p("}")
+
+		if err := g.examplePagingAllCall(m); err != nil {
+			return err
+		}
+		p("}")
+		p("")
+	}
+	return nil
+}
+
 func (g *generator) exampleClientFactory(pkgName, servName string) {
 	p := g.printf
 	for _, t := range g.opts.transports {
@@ -250,6 +306,30 @@ func (g *generator) examplePagingCall(m *descriptorpb.MethodDescriptorProto) err
 	return nil
 }
 
+func (g *generator) examplePagingAllCall(m *descriptorpb.MethodDescriptorProto) error {
+	outType := g.descInfo.Type[m.GetOutputType()]
+	if outType == nil {
+		return fmt.Errorf("cannot find type %q, malformed descriptor?", m.GetOutputType())
+	}
+
+	outSpec, err := g.descInfo.ImportSpec(outType)
+	if err != nil {
+		return err
+	}
+
+	p := g.printf
+
+	p("for resp, err := range c.%s(ctx, req).All() {", m.GetName())
+	p("  if err != nil {")
+	p("    // TODO: Handle error.")
+	p("  }")
+	p("  // TODO: Use resp.")
+	p("  _ = resp")
+	p("}")
+	g.imports[outSpec] = true
+	return nil
+}
+
 func (g *generator) exampleBidiCall(m *descriptorpb.MethodDescriptorProto, inType pbinfo.ProtoType, inSpec pbinfo.ImportSpec) {
 	p := g.printf
 

internal/gengapic/example_test.go

@@ -272,6 +272,24 @@ func TestExample(t *testing.T) {
 				t.Errorf("TestExample(%s): imports got(-),want(+):\n%s", tst.tstName, diff)
 			}
 			txtdiff.Diff(t, g.pt.String(), filepath.Join("testdata", tst.tstName+".want"))
+
+			g.reset()
+			// remove imports not used in iter example
+			delete(tst.imports, pbinfo.ImportSpec{Path: "google.golang.org/api/iterator"})
+			delete(tst.imports, pbinfo.ImportSpec{Path: "io"})
+			delete(tst.imports, pbinfo.ImportSpec{Name: "iampb", Path: "cloud.google.com/go/iam/apiv1/iampb"})
+
+			g.opts = &tst.options
+			g.mixins = mix
+			if tst.options.diregapic {
+				g.mixins = nil
+			}
+			g.aux.customOp = tst.op
+			g.genExampleIteratorFile(serv)
+			if diff := cmp.Diff(g.imports, tst.imports); diff != "" {
+				t.Errorf("TestExample(%s): imports got(-),want(+):\n%s", tst.tstName, diff)
+			}
+			txtdiff.Diff(t, g.pt.String(), filepath.Join("testdata", tst.tstName+"_all.want"))
 		})
 	}
 }

internal/gengapic/generator.go

@@ -219,12 +219,19 @@ func (g *generator) printf(s string, a ...interface{}) {
 
 // TODO(chrisdsmith): Add generator_test.go with TestCommit
 
+func (g *generator) commit(fileName, pkgName string) int {
+	return g.commitWithBuildTag(fileName, pkgName, "")
+}
+
 // commit adds header, etc to current pt and returns the line length of the
 // final file output.
-func (g *generator) commit(fileName, pkgName string) int {
+func (g *generator) commitWithBuildTag(fileName, pkgName, buildTag string) int {
 	var header strings.Builder
 	fmt.Fprintf(&header, license.Apache, time.Now().Year())
 	header.WriteString(g.headerComments.String() + "\n")
+	if buildTag != "" {
+		fmt.Fprintf(&header, "//go:build %s\n\n", buildTag)
+	}
 	fmt.Fprintf(&header, "package %s\n\n", pkgName)
 
 	var imps []pbinfo.ImportSpec

internal/gengapic/gengapic.go

@@ -133,6 +133,13 @@ func gen(genReq *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse
 		g.imports[pbinfo.ImportSpec{Name: g.opts.pkgName, Path: g.opts.pkgPath}] = true
 		g.commit(outFile+"_client_example_test.go", g.opts.pkgName+"_test")
 
+		g.reset()
+		if err := g.genExampleIteratorFile(s); err != nil {
+			return &g.resp, fmt.Errorf("error generating iter example for %q; %v", s.GetName(), err)
+		}
+		g.imports[pbinfo.ImportSpec{Name: g.opts.pkgName, Path: g.opts.pkgPath}] = true
+		g.commitWithBuildTag(outFile+"_client_example_go123_test.go", g.opts.pkgName+"_test", "go1.23")
+
 		// Replace original set of transports for the next service that may have
 		// REST-able RPCs.
 		if !hasREST {

internal/gengapic/paging.go

@@ -484,3 +484,20 @@ func (g *generator) pagingIter(pt *iterType) {
 		g.imports[spec] = true
 	}
 }
+
+func (g *generator) pagingIterGo123(pt *iterType) {
+	p := g.printf
+
+	p("// All returns an iterator. If an error is returned by the iterator, the")
+	p("// iterator will stop after that iteration.")
+	p("func (it *%s) All() iter.Seq2[%s, error] {", pt.iterTypeName, pt.elemTypeName)
+	p("  return iterator.RangeAdapter(it.Next)")
+	p("}")
+	p("")
+
+	g.imports[pbinfo.ImportSpec{Path: "iter"}] = true
+	g.imports[pbinfo.ImportSpec{Path: "github.com/googleapis/gax-go/v2/iterator"}] = true
+	for _, spec := range pt.elemImports {
+		g.imports[spec] = true
+	}
+}

internal/gengapic/testdata/custom_op_example_all.want

@@ -0,0 +1,26 @@
+func ExampleFooClient_GetManyThings_all() {
+	ctx := context.Background()
+	// This snippet has been automatically generated and should be regarded as a code template only.
+	// It will require modifications to work:
+	// - It may require correct/in-range values for request initialization.
+	// - It may require specifying regional endpoints when creating the service client as shown in:
+	//   https://pkg.go.dev/cloud.google.com/go#hdr-Client_Options
+	c, err := Bar.NewFooRESTClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.PageInputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#PageInputType.
+	}
+	for resp, err := range c.GetManyThings(ctx, req).All() {
+		if err != nil {
+			// TODO: Handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+

internal/gengapic/testdata/empty_example_all.want

@@ -0,0 +1,78 @@
+func ExampleClient_GetManyThings_all() {
+	ctx := context.Background()
+	// This snippet has been automatically generated and should be regarded as a code template only.
+	// It will require modifications to work:
+	// - It may require correct/in-range values for request initialization.
+	// - It may require specifying regional endpoints when creating the service client as shown in:
+	//   https://pkg.go.dev/cloud.google.com/go#hdr-Client_Options
+	c, err := Foo.NewClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &mypackagepb.PageInputType{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/mypackage#PageInputType.
+	}
+	for resp, err := range c.GetManyThings(ctx, req).All() {
+		if err != nil {
+			// TODO: Handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+
+func ExampleClient_ListLocations_all() {
+	ctx := context.Background()
+	// This snippet has been automatically generated and should be regarded as a code template only.
+	// It will require modifications to work:
+	// - It may require correct/in-range values for request initialization.
+	// - It may require specifying regional endpoints when creating the service client as shown in:
+	//   https://pkg.go.dev/cloud.google.com/go#hdr-Client_Options
+	c, err := Foo.NewClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &locationpb.ListLocationsRequest{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/google.golang.org/genproto/googleapis/cloud/location#ListLocationsRequest.
+	}
+	for resp, err := range c.ListLocations(ctx, req).All() {
+		if err != nil {
+			// TODO: Handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+
+func ExampleClient_ListOperations_all() {
+	ctx := context.Background()
+	// This snippet has been automatically generated and should be regarded as a code template only.
+	// It will require modifications to work:
+	// - It may require correct/in-range values for request initialization.
+	// - It may require specifying regional endpoints when creating the service client as shown in:
+	//   https://pkg.go.dev/cloud.google.com/go#hdr-Client_Options
+	c, err := Foo.NewClient(ctx)
+	if err != nil {
+		// TODO: Handle error.
+	}
+	defer c.Close()
+
+	req := &longrunningpb.ListOperationsRequest{
+		// TODO: Fill request struct fields.
+		// See https://pkg.go.dev/cloud.google.com/go/longrunning/autogen/longrunningpb#ListOperationsRequest.
+	}
+	for resp, err := range c.ListOperations(ctx, req).All() {
+		if err != nil {
+			// TODO: Handle error.
+		}
+		// TODO: Use resp.
+		_ = resp
+	}
+}
+