Repository: googleapis/gapic-generator-go
Pull request #1029 (2022-06-10 18:10:27)

feat(genpapic): support protobuf-go go_package mapping option
Per the [protobuf Go docs](https://developers.google.com/protocol-buffers/docs/reference/go-generated#package), there is an option to override the `go_package` option in specific proto files. If this is used during generation of the protobuf/grpc stubs, the GAPIC generator needs to know about it, so that the proper imports can be generated for the request/response/stub types.


Patch:
diff --git a/internal/gengapic/generator.go b/internal/gengapic/generator.go
index caf5a0f..a3548ba 100644
--- a/internal/gengapic/generator.go
+++ b/internal/gengapic/generator.go
@@ -141,6 +141,9 @@ func (g *generator) init(req *plugin.CodeGeneratorRequest) error {
 	g.opts = opts
 
 	g.descInfo = pbinfo.Of(files)
+	if len(g.opts.pkgOverrides) > 0 {
+		g.descInfo.PkgOverrides = g.opts.pkgOverrides
+	}
 
 	for _, f := range files {
 		for _, loc := range f.GetSourceCodeInfo().GetLocation() {
diff --git a/internal/gengapic/options.go b/internal/gengapic/options.go
index b4790a9..6c82eea 100644
--- a/internal/gengapic/options.go
+++ b/internal/gengapic/options.go
@@ -43,6 +43,7 @@ type options struct {
 	transports        []transport
 	metadata          bool
 	diregapic         bool
+	pkgOverrides      map[string]string
 }
 
 // parseOptions takes a string and parses it into a struct defining
@@ -53,6 +54,7 @@ type options struct {
 // * api-service-config (filepath)
 // * grpc-service-config (filepath)
 // * module (name)
+// * Mfile=import (e.g. Mgoogle/storage/v2/storage.proto=cloud.google.com/go/storage/internal/apiv2/stubs)
 // * release-level (one of 'alpha', 'beta', or empty)
 // * transport ('+' separated list of transport backends to generate)
 // * metadata (enable GAPIC metadata generation)
@@ -139,6 +141,16 @@ func parseOptions(parameter *string) (*options, error) {
 			sort.Slice(opts.transports, func(i, j int) bool {
 				return opts.transports[i] < opts.transports[j]
 			})
+		default:
+			// go_package override for the protobuf/grpc stubs.
+			// Mgoogle/storage/v2/storage.proto=cloud.google.com/go/storage/internal/apiv2/stubs
+			if key[0] == 'M' {
+				file := key[1:]
+				if opts.pkgOverrides == nil {
+					opts.pkgOverrides = make(map[string]string)
+				}
+				opts.pkgOverrides[file] = val
+			}
 		}
 	}
 
diff --git a/internal/pbinfo/pbinfo.go b/internal/pbinfo/pbinfo.go
index b4efdf5..212f363 100644
--- a/internal/pbinfo/pbinfo.go
+++ b/internal/pbinfo/pbinfo.go
@@ -54,6 +54,10 @@ type Info struct {
 
 	// Maps service names to their descriptors.
 	Serv map[string]*descriptor.ServiceDescriptorProto
+
+	// PkgOverrides is file-to-import mapping used to override the
+	// go_package option in the given proto file.
+	PkgOverrides map[string]string
 }
 
 // Of creates Info from given protobuf files.
@@ -63,6 +67,7 @@ func Of(files []*descriptor.FileDescriptorProto) Info {
 		ParentElement: map[ProtoType]ProtoType{},
 		Type:          map[string]ProtoType{},
 		Serv:          map[string]*descriptor.ServiceDescriptorProto{},
+		PkgOverrides:  map[string]string{},
 	}
 
 	for _, f := range files {
@@ -160,6 +165,9 @@ func (in *Info) NameSpec(e ProtoType) (string, ImportSpec, error) {
 	}
 
 	pkg := fdesc.GetOptions().GetGoPackage()
+	if pkgOverride, ok := in.PkgOverrides[fdesc.GetName()]; ok {
+		pkg = pkgOverride
+	}
 	if pkg == "" {
 		return "", ImportSpec{}, errors.E(nil, "can't determine import path for %v, file %q missing `option go_package`", eTxt, fdesc.GetName())
 	}


Test patch:
diff --git a/internal/gengapic/options_test.go b/internal/gengapic/options_test.go
index da4ed47..8dc7d40 100644
--- a/internal/gengapic/options_test.go
+++ b/internal/gengapic/options_test.go
@@ -75,6 +75,17 @@ func TestParseOptions(t *testing.T) {
 			},
 			expectErr: false,
 		},
+		{
+			param: "Mgoogle/example/library/v1/library.proto=new/import/path;pkg,go-gapic-package=path/to/out;pkg",
+			expectedOpts: &options{
+				pkgOverrides: map[string]string{"google/example/library/v1/library.proto": "new/import/path;pkg"},
+				transports:   []transport{grpc},
+				pkgPath:      "path/to/out",
+				pkgName:      "pkg",
+				outDir:       "path/to/out",
+			},
+			expectErr: false,
+		},
 		{
 			param:     "transport=tcp,go-gapic-package=path;pkg",
 			expectErr: true,
diff --git a/internal/pbinfo/pbinfo_test.go b/internal/pbinfo/pbinfo_test.go
index ea1afdf..00e5952 100644
--- a/internal/pbinfo/pbinfo_test.go
+++ b/internal/pbinfo/pbinfo_test.go
@@ -31,14 +31,27 @@ func TestNameSpec(t *testing.T) {
 		Name:       proto.String("Message"),
 		NestedType: []*descriptor.DescriptorProto{subMsg},
 	}
+	anotherMsg := &descriptor.DescriptorProto{
+		Name: proto.String("AnotherMessage"),
+	}
 	file := &descriptor.FileDescriptorProto{
 		Options: &descriptor.FileOptions{
 			GoPackage: proto.String("path.to/pb/foo;foo"),
 		},
 		MessageType: []*descriptor.DescriptorProto{msg},
 	}
+	anotherFile := &descriptor.FileDescriptorProto{
+		Name: proto.String("bar.proto"),
+		Options: &descriptor.FileOptions{
+			GoPackage: proto.String("path.to/pb/bar;bar"),
+		},
+		MessageType: []*descriptor.DescriptorProto{anotherMsg},
+	}
 
-	info := Of([]*descriptor.FileDescriptorProto{file})
+	info := Of([]*descriptor.FileDescriptorProto{file, anotherFile})
+	info.PkgOverrides = map[string]string{
+		anotherFile.GetName(): "path.to/pb/foo;foo",
+	}
 
 	for _, tst := range []struct {
 		e    ProtoType
@@ -46,6 +59,7 @@ func TestNameSpec(t *testing.T) {
 	}{
 		{msg, "Message"},
 		{subMsg, "Message_SubMessage"},
+		{anotherMsg, "AnotherMessage"},
 	} {
 		name, imp, err := info.NameSpec(tst.e)
 		if err != nil {