feat(genpapic): support protobuf-go go_package mapping option (#1029)

googleapis · Jun 10, 2022 · f40c830 · f40c830
1 parent c9e3ce7
commit f40c830
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 1 deletion.
diff --git 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
@@ -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/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.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())
 	}

diff --git a/internal/pbinfo/pbinfo_test.go b/internal/pbinfo/pbinfo_test.go
@@ -31,21 +31,35 @@ 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
 		name string
 	}{
 		{msg, "Message"},
 		{subMsg, "Message_SubMessage"},
+		{anotherMsg, "AnotherMessage"},
 	} {
 		name, imp, err := info.NameSpec(tst.e)
 		if err != nil {