Add --go_opt flag to adapt to the version after 1.4.0 of protoc-gen-go (#767)

Co-authored-by: anqiansong <anqiansong@xiaoheiban.cn>

tools/goctl/goctl.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"runtime"
 
+	"github.com/logrusorgru/aurora"
 	"github.com/tal-tech/go-zero/core/load"
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/stat"
@@ -31,7 +32,7 @@ import (
 )
 
 var (
-	buildVersion = "1.1.8"
+	buildVersion = "1.1.9-pre"
 	commands     = []cli.Command{
 		{
 			Name:   "upgrade",
@@ -365,6 +366,10 @@ var (
 							Name:  "proto_path, I",
 							Usage: `native command of protoc, specify the directory in which to search for imports. [optional]`,
 						},
+						cli.StringSliceFlag{
+							Name:  "go_opt",
+							Usage: `native command of protoc-gen-go, specify the mapping from proto to go, eg --go_opt=proto_import=go_package_import. [optional]`,
+						},
 						cli.StringFlag{
 							Name:  "dir, d",
 							Usage: `the target path of the code`,
@@ -542,6 +547,6 @@ func main() {
 	app.Commands = commands
 	// cli already print error messages
 	if err := app.Run(os.Args); err != nil {
-		fmt.Println("error:", err)
+		fmt.Println(aurora.Red("error: " + err.Error()))
 	}
 }

tools/goctl/rpc/README.md

@@ -134,15 +134,21 @@ USAGE:
 OPTIONS:
    --src value, -s value         the file path of the proto source file
    --proto_path value, -I value  native command of protoc, specify the directory in which to search for imports. [optional]
+   --go_opt value                native command of protoc-gen-go, specify the mapping from proto to go, eg --go_opt=proto_import=go_package_import. [optional]
    --dir value, -d value         the target path of the code
+   --style value                 the file naming format, see [https://github.com/tal-tech/go-zero/tree/master/tools/goctl/config/readme.md]
    --idea                        whether the command execution environment is from idea plugin. [optional]
+
 ```
 
 ### 参数说明
 
 * --src 必填，proto数据源，目前暂时支持单个proto文件生成
-* --proto_path 可选，protoc原生子命令，用于指定proto import从何处查找，可指定多个路径,如`goctl rpc -I={path1} -I={path2} ...`,在没有import时可不填。当前proto路径不用指定，已经内置，`-I`的详细用法请参考`protoc -h`
+* --proto_path 可选，protoc原生子命令，用于指定proto import从何处查找，可指定多个路径,如`goctl rpc -I={path1} -I={path2} ...`
+  ,在没有import时可不填。当前proto路径不用指定，已经内置，`-I`的详细用法请参考`protoc -h`
+* --go_opt 可选，protoc-gen-go插件原生flag，用于指定go_package
 * --dir 可选，默认为proto文件所在目录，生成代码的目标目录
+* --style 可选，指定生成文件名的命名风格
 * --idea 可选，是否为idea插件中执行，终端执行可以忽略
 
 

tools/goctl/rpc/cli/cli.go

@@ -17,6 +17,7 @@ func RPC(c *cli.Context) error {
 	out := c.String("dir")
 	style := c.String("style")
 	protoImportPath := c.StringSlice("proto_path")
+	goOptions := c.StringSlice("go_opt")
 	if len(src) == 0 {
 		return errors.New("missing -src")
 	}
@@ -29,7 +30,7 @@ func RPC(c *cli.Context) error {
 		return err
 	}
 
-	return g.Generate(src, out, protoImportPath)
+	return g.Generate(src, out, protoImportPath, goOptions...)
 }
 
 // RPCNew is to generate rpc greet service, this greet service can speed

tools/goctl/rpc/generator/defaultgenerator.go

@@ -11,8 +11,11 @@ type DefaultGenerator struct {
 	log console.Console
 }
 
+// just test interface implement
+var _ Generator = (*DefaultGenerator)(nil)
+
 // NewDefaultGenerator returns an instance of DefaultGenerator
-func NewDefaultGenerator() *DefaultGenerator {
+func NewDefaultGenerator() Generator {
 	log := console.NewColorConsole()
 	return &DefaultGenerator{
 		log: log,
@@ -33,5 +36,6 @@ func (g *DefaultGenerator) Prepare() error {
 	}
 
 	_, err = exec.LookPath("protoc-gen-go")
+
 	return err
 }

tools/goctl/rpc/generator/gen.go

@@ -36,7 +36,7 @@ func NewRPCGenerator(g Generator, cfg *conf.Config) *RPCGenerator {
 // Generate generates an rpc service, through the proto file,
 // code storage directory, and proto import parameters to control
 // the source file and target location of the rpc service that needs to be generated
-func (g *RPCGenerator) Generate(src, target string, protoImportPath []string) error {
+func (g *RPCGenerator) Generate(src, target string, protoImportPath []string, goOptions ...string) error {
 	abs, err := filepath.Abs(target)
 	if err != nil {
 		return err
@@ -73,7 +73,7 @@ func (g *RPCGenerator) Generate(src, target string, protoImportPath []string) er
 		return err
 	}
 
-	err = g.g.GenPb(dirCtx, protoImportPath, proto, g.cfg)
+	err = g.g.GenPb(dirCtx, protoImportPath, proto, g.cfg, goOptions...)
 	if err != nil {
 		return err
 	}

tools/goctl/rpc/generator/gen_test.go

@@ -41,7 +41,11 @@ func TestRpcGenerate(t *testing.T) {
 	defer func() {
 		_ = os.RemoveAll(srcDir)
 	}()
-	err = g.Generate("./test.proto", projectDir, []string{src})
+
+	common, err := filepath.Abs(".")
+	assert.Nil(t, err)
+
+	err = g.Generate("./test.proto", projectDir, []string{common, src}, "Mbase/common.proto=./base")
 	assert.Nil(t, err)
 	_, err = execx.Run("go test "+projectName, projectDir)
 	if err != nil {
@@ -60,7 +64,7 @@ func TestRpcGenerate(t *testing.T) {
 	}
 
 	projectDir = filepath.Join(workDir, projectName)
-	err = g.Generate("./test.proto", projectDir, []string{src})
+	err = g.Generate("./test.proto", projectDir, []string{common, src}, "Mbase/common.proto=./base")
 	assert.Nil(t, err)
 	_, err = execx.Run("go test "+projectName, projectDir)
 	if err != nil {
@@ -70,7 +74,7 @@ func TestRpcGenerate(t *testing.T) {
 	}
 
 	// case not in go mod and go path
-	err = g.Generate("./test.proto", projectDir, []string{src})
+	err = g.Generate("./test.proto", projectDir, []string{common, src}, "Mbase/common.proto=./base")
 	assert.Nil(t, err)
 	_, err = execx.Run("go test "+projectName, projectDir)
 	if err != nil {

tools/goctl/rpc/generator/generator.go

@@ -15,5 +15,5 @@ type Generator interface {
 	GenLogic(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
 	GenServer(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
 	GenSvc(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
-	GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, cfg *conf.Config) error
+	GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, cfg *conf.Config, goOptions ...string) error
 }

tools/goctl/rpc/generator/genpb.go

@@ -2,33 +2,86 @@ package generator
 
 import (
 	"bytes"
+	"errors"
 	"path/filepath"
 	"strings"
 
+	"github.com/tal-tech/go-zero/core/collection"
 	conf "github.com/tal-tech/go-zero/tools/goctl/config"
 	"github.com/tal-tech/go-zero/tools/goctl/rpc/execx"
 	"github.com/tal-tech/go-zero/tools/goctl/rpc/parser"
 )
 
+const googleProtocGenGoErr = `--go_out: protoc-gen-go: plugins are not supported; use 'protoc --go-grpc_out=...' to generate gRPC`
+
 // GenPb generates the pb.go file, which is a layer of packaging for protoc to generate gprc,
 // but the commands and flags in protoc are not completely joined in goctl. At present, proto_path(-I) is introduced
-func (g *DefaultGenerator) GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, _ *conf.Config) error {
+func (g *DefaultGenerator) GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, _ *conf.Config, goOptions ...string) error {
 	dir := ctx.GetPb()
 	cw := new(bytes.Buffer)
-	base := filepath.Dir(proto.Src)
+	directory, base := filepath.Split(proto.Src)
+	directory = filepath.Clean(directory)
 	cw.WriteString("protoc ")
+	protoImportPathSet := collection.NewSet()
 	for _, ip := range protoImportPath {
-		cw.WriteString(" -I=" + ip)
+		pip := " --proto_path=" + ip
+		if protoImportPathSet.Contains(pip) {
+			continue
+		}
+
+		protoImportPathSet.AddStr(pip)
+		cw.WriteString(pip)
+	}
+	currentPath := " --proto_path=" + directory
+	if !protoImportPathSet.Contains(currentPath) {
+		cw.WriteString(currentPath)
 	}
-	cw.WriteString(" -I=" + base)
 	cw.WriteString(" " + proto.Name)
 	if strings.Contains(proto.GoPackage, "/") {
 		cw.WriteString(" --go_out=plugins=grpc:" + ctx.GetMain().Filename)
 	} else {
 		cw.WriteString(" --go_out=plugins=grpc:" + dir.Filename)
 	}
+
+	// Compatible with version 1.4.0，github.com/golang/protobuf/protoc-gen-go@v1.4.0
+	// --go_opt usage please see https://developers.google.com/protocol-buffers/docs/reference/go-generated#package
+	optSet := collection.NewSet()
+	for _, op := range goOptions {
+		opt := " --go_opt=" + op
+		if optSet.Contains(opt) {
+			continue
+		}
+
+		optSet.AddStr(op)
+		cw.WriteString(" --go_opt=" + op)
+	}
+
+	var currentFileOpt string
+	if filepath.IsAbs(proto.GoPackage) {
+		currentFileOpt = " --go_opt=M" + base + "=" + proto.GoPackage
+	} else if strings.Contains(proto.GoPackage, string(filepath.Separator)) {
+		currentFileOpt = " --go_opt=M" + base + "=./" + proto.GoPackage
+	} else {
+		currentFileOpt = " --go_opt=M" + base + "=../" + proto.GoPackage
+	}
+	if !optSet.Contains(currentFileOpt) {
+		cw.WriteString(currentFileOpt)
+	}
+
 	command := cw.String()
 	g.log.Debug(command)
 	_, err := execx.Run(command, "")
+	if err != nil {
+		if strings.Contains(err.Error(), googleProtocGenGoErr) {
+			return errors.New(`Unsupported plugin protoc-gen-go which installed from the following source:
+google.golang.org/protobuf/cmd/protoc-gen-go, 
+github.com/protocolbuffers/protobuf-go/cmd/protoc-gen-go;
+
+Please replace it by the following command, we recommend to use version before v1.3.5:
+go get -u github.com/golang/protobuf/protoc-gen-go`)
+		}
+
 		return err
+	}
+	return nil
 }