feat(cmd/gen): support mapping _grpc.pb.go files to different Go packages from .pb.go

grpc#8233

Author: qianlongzt

cmd/protoc-gen-go-grpc/grpc.go

@@ -132,12 +132,16 @@ const fileDescriptorProtoPackageFieldNumber = 2
 const fileDescriptorProtoSyntaxFieldNumber = 12
 
 // generateFile generates a _grpc.pb.go file containing gRPC service definitions.
-func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
+func generateFile(gen *protogen.Plugin, mapping map[string]string, file *protogen.File) *protogen.GeneratedFile {
 	if len(file.Services) == 0 {
 		return nil
 	}
 	filename := file.GeneratedFilenamePrefix + "_grpc.pb.go"
-	g := gen.NewGeneratedFile(filename, file.GoImportPath)
+	path := file.GoImportPath
+	if n, ok := mapping[file.Desc.Path()]; ok {
+		path = protogen.GoImportPath(n)
+	}
+	g := gen.NewGeneratedFile(filename, path)
 	// Attach all comments associated with the syntax field.
 	genLeadingComments(g, file.Desc.SourceLocations().ByPath(protoreflect.SourcePath{fileDescriptorProtoSyntaxFieldNumber}))
 	g.P("// Code generated by protoc-gen-go-grpc. DO NOT EDIT.")

cmd/protoc-gen-go-grpc/main.go

@@ -36,6 +36,7 @@ package main
 import (
 	"flag"
 	"fmt"
+	"strings"
 
 	"google.golang.org/protobuf/compiler/protogen"
 	"google.golang.org/protobuf/types/descriptorpb"
@@ -59,6 +60,9 @@ func main() {
 	requireUnimplemented = flags.Bool("require_unimplemented_servers", true, "set to false to match legacy behavior")
 	useGenericStreams = flags.Bool("use_generic_streams_experimental", true, "set to true to use generic types for streaming client and server objects; this flag is EXPERIMENTAL and may be changed or removed in a future release")
 
+	var grpcMapping mappingFlag
+	flags.Var(&grpcMapping, "grpc_mapping", "change proto file go_package option, should be like xxx.proto=go_package. Multiple OK")
+
 	protogen.Options{
 		ParamFunc: flags.Set,
 	}.Run(func(gen *protogen.Plugin) error {
@@ -69,8 +73,29 @@ func main() {
 			if !f.Generate {
 				continue
 			}
-			generateFile(gen, f)
+			generateFile(gen, grpcMapping, f)
 		}
 		return nil
 	})
 }
+
+type mappingFlag map[string]string
+
+func (a *mappingFlag) String() string {
+	return fmt.Sprintf("%v", *a)
+}
+
+func (a *mappingFlag) Set(value string) error {
+	k, v, ok := strings.Cut(value, "=")
+	if !ok {
+		return fmt.Errorf("invalid grpc_mapping: %v", value)
+	}
+	if k == "" || v == "" {
+		return fmt.Errorf("invalid grpc_mapping: %v", value)
+	}
+	if *a == nil {
+		*a = make(mappingFlag)
+	}
+	(*a)[k] = v
+	return nil
+}