Go 自动生成代码工具 一( go-zero 中 goctl rpc 命令代码生成原理)

发布时间 2023-11-27 15:44:33作者: 杨阳的技术博客

总共分为三篇:

1. 分析`go-zero`中 ` coctl rpc`  通过一个` proto`文件生成一系列文件。
2. 模仿这个原理,结合`protoc` 生成代码的特性,把gin的接口定义,也放入proto文件中,自动生成gin的接口代码。
3. 自动生成项目中error错误定义文档。(通过go源码自动生成文档)

go-zero 中 goctl rpc 命令代码生成原理

一、 使用效果对比

go-zero 与 Kratos 是国内两个主流的go微服务框架,都对微服务开发中常见的 服务发现、认证、监控、日志、链路追踪等功能进行了封装。

分析下,当使用go-zero时,当我们定义了一个 .proto文件后,可以通过命令生成一个 go-zero的项目

goctl rpc protoc greet.proto --go_out=.  --go-grpc_out=.  --zrpc_out=.
官方例子

其实protoc命令的使用,很相似,特别是里面的一些参数:

 protoc --go_out=.  --go-grpc_out=.  ./*.proto

不过, protoc 只会生成一个 pd.go (rpc) _grpc.pd.go ,但是 goctl rpc 能生成一系列文件。

   demo
  ├── etc
  │   └── greet.yaml
  ├── go.mod
  ├── greet
  │   ├── greet.pb.go
  │   └── greet_grpc.pb.go
  ├── greet.go
  ├── greet.proto
  ├── greetclient
  │   └── greet.go
  └── internal
      ├── config
      │   └── config.go
      ├── logic
      │   └── pinglogic.go
      ├── server
      │   └── greetserver.go
      └── svc
          └── servicecontext.go

  8 directories, 11 files

二、造成不同的原因

这里需要去看,go-zero的源码,地址:https://github.com/zeromicro/go-zero

tools -> goctl下有很多很多目录,对应的都是goctl丰富的命令。

看下 goctl.go,跟进能看到 是采用 cobra 实现的对命令的接收。

func main() {
	logx.Disable()
	load.Disable()
	cmd.Execute()
}

var (
	//go:embed usage.tpl   绑定模板,后面还会看到
	usageTpl string
	rootCmd  = cobrax.NewCommand("goctl") // 采用 cobra 库实现对命令的接收,这个库使用非常简洁方便。
)

// Execute executes the given command
func Execute() {
	os.Args = supportGoStdFlag(os.Args)
	if err := rootCmd.Execute(); err != nil {
		fmt.Println(color.Red.Render(err.Error()))
		os.Exit(codeFailure)
	}
}

接下来,关注rpc目录,查看 cmd.go

截取了,关注的代码:
var (
	// Cmd describes a rpc command.
	Cmd = cobrax.NewCommand("rpc", cobrax.WithRunE(func(command *cobra.Command, strings []string) error {
		return cli.RPCTemplate(true)
	}))
	templateCmd = cobrax.NewCommand("template", cobrax.WithRunE(func(command *cobra.Command, strings []string) error {
		return cli.RPCTemplate(false)
	}))

	newCmd    = cobrax.NewCommand("new", cobrax.WithRunE(cli.RPCNew), cobrax.WithArgs(cobra.MatchAll(cobra.ExactArgs(1), cobra.OnlyValidArgs)))
	protocCmd = cobrax.NewCommand("protoc", cobrax.WithRunE(cli.ZRPC), cobrax.WithArgs(cobra.MatchAll(cobra.ExactArgs(1), cobra.OnlyValidArgs)))
)

 func init() {
        protocCmdFlags.BoolVarP(&cli.VarBoolMultiple, "multiple", "m")
  	    protocCmdFlags.StringSliceVar(&cli.VarStringSliceGoOut, "go_out")
    	protocCmdFlags.StringSliceVar(&cli.VarStringSliceGoGRPCOut, "go-grpc_out")
    	protocCmdFlags.StringSliceVar(&cli.VarStringSliceGoOpt, "go_opt")
    	protocCmdFlags.StringSliceVar(&cli.VarStringSliceGoGRPCOpt, "go-grpc_opt")
    	protocCmdFlags.StringSliceVar(&cli.VarStringSlicePlugin, "plugin")
    	protocCmdFlags.StringSliceVarP(&cli.VarStringSliceProtoPath, "proto_path", "I")
    	protocCmdFlags.StringVar(&cli.VarStringStyle, "style")
    	protocCmdFlags.StringVar(&cli.VarStringZRPCOut, "zrpc_out")
}

当收到 proto后, 触发这个函数 cli.ZRPC,只截取关键代码:

  // ZRPC generates grpc code directly by protoc and generates
  // zrpc code by goctl.
  func ZRPC(_ *cobra.Command, args []string) error {
  	
    // 1. 先获取参数,判断那些参数有传入
    // 2. 拼凑出 自动生成代码需要的 参数
  	var ctx generator.ZRpcContext
  	ctx.Multiple = VarBoolMultiple
  	ctx.Src = source
  	ctx.GoOutput = goOut
  	ctx.GrpcOutput = grpcOut
  	ctx.IsGooglePlugin = isGooglePlugin
  	ctx.Output = zrpcOut
  	ctx.ProtocCmd = strings.Join(protocArgs, " ")
  	ctx.IsGenClient = VarBoolClient
     // 核心部分
  	g := generator.NewGenerator(style, verbose)  
  	return g.Generate(&ctx)
  }

小结:

  命令不同的原因,因为 go-zero使用 cobra ,实现了一套命令。接下来看,具体是如何实现的,生成不同的代码。

三、入口函数

  代码中去除了一些 err的判断
  // Generate generates a 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 *Generator) Generate(zctx *ZRpcContext) error {
  	abs, err := filepath.Abs(zctx.Output)
  	err = pathx.MkdirIfNotExist(abs)
    // 创建目录      	

  	err = g.Prepare()
  	
  	projectCtx, err := ctx.Prepare(abs)
  	
  	p := parser.NewDefaultProtoParser()
  	proto, err := p.Parse(zctx.Src, zctx.Multiple) // 拿到了proto文件的内容
 
  	dirCtx, err := mkdir(projectCtx, proto, g.cfg, zctx) // 创建各个子模块的目录,后面可以跟进看下
  
  	err = g.GenEtc(dirCtx, proto, g.cfg) // 生成etc文件

  	err = g.GenPb(dirCtx, zctx) // 生成pd文件
  	
  	err = g.GenConfig(dirCtx, proto, g.cfg) // 生成config.go
  	
  	err = g.GenSvc(dirCtx, proto, g.cfg) // 生成 ServiceContext.go
  	
  	err = g.GenLogic(dirCtx, proto, g.cfg, zctx) // 生成 logic.go

  	err = g.GenServer(dirCtx, proto, g.cfg, zctx) // 生成server.go
  
  	err = g.GenMain(dirCtx, proto, g.cfg, zctx) // 生成main.go
  	
   if zctx.IsGenClient {
	  err = g.GenCall(dirCtx, proto, g.cfg, zctx)  // 生成 pb
   }
  	return err
  }

3.1 先看生成各个文件目录的代码:

func mkdir(ctx *ctx.ProjectContext, proto parser.Proto, conf *conf.Config, c *ZRpcContext) (DirContext,
error) {
  inner := make(map[string]Dir)
	etcDir := filepath.Join(ctx.WorkDir, "etc")
	clientDir := filepath.Join(ctx.WorkDir, "client")
	internalDir := filepath.Join(ctx.WorkDir, "internal")
	configDir := filepath.Join(internalDir, "config")
	logicDir := filepath.Join(internalDir, "logic")
	serverDir := filepath.Join(internalDir, "server")
	svcDir := filepath.Join(internalDir, "svc")
	pbDir := filepath.Join(ctx.WorkDir, proto.GoPackage)

  inner[etc] = Dir{
	Filename: etcDir,
	Package:  filepath.ToSlash(filepath.Join(ctx.Path, strings.TrimPrefix(etcDir, ctx.Dir))),
	Base:     filepath.Base(etcDir),
	GetChildPackage: func(childPath string) (string, error) {
		return getChildPackage(etcDir, childPath)
	},
}
 	return &defaultDirContext{
	ctx:         ctx,
	inner:       inner,
	serviceName: stringx.From(strings.ReplaceAll(serviceName, "-", "")),
}, nil

}

能看到这么把创建目录的逻辑已经生成好,放入了inner这map中,最后返回给 dirCtx 这个变量,
这个变量在生成各种文件内容时候,都有传递。

3.2 如何拿到proto文件的内容的

入口方法哪里能看到:

  p := parser.NewDefaultProtoParser()
  proto, err := p.Parse(zctx.Src, zctx.Multiple) // 拿到了proto文件的内容

跟进看下:

import (
	"github.com/emicklei/proto" // 关键第三方库
)

// Parse provides to parse the proto file into a golang structure,
// which is convenient for subsequent rpc generation and use
func (p *DefaultProtoParser) Parse(src string, multiple ...bool) (Proto, error) {
	var ret Proto

	abs, err := filepath.Abs(src)
	
	r, err := os.Open(abs)    	
	defer r.Close()

	parser := proto.NewParser(r)
	set, err := parser.Parse()
	
	var serviceList Services
	proto.Walk(
		set,
		proto.WithImport(func(i *proto.Import) {  
			ret.Import = append(ret.Import, Import{Import: i})
		}),
		proto.WithMessage(func(message *proto.Message) {
			ret.Message = append(ret.Message, Message{Message: message})
		}),
		proto.WithPackage(func(p *proto.Package) {
			ret.Package = Package{Package: p}
		}),
		proto.WithService(func(service *proto.Service) {
			serv := Service{Service: service}
			elements := service.Elements
			for _, el := range elements {
				v, _ := el.(*proto.RPC)
				if v == nil {
					continue
				}
				serv.RPC = append(serv.RPC, &RPC{RPC: v})
			}
			serviceList = append(serviceList, serv)
		}),
		proto.WithOption(func(option *proto.Option) {
			if option.Name == "go_package" {
				ret.GoPackage = option.Constant.Source
			}
		}),
	)
	ret.PbPackage = GoSanitized(filepath.Base(ret.GoPackage))
	ret.Src = abs
	ret.Name = filepath.Base(abs)
	ret.Service = serviceList

	return ret, nil
}

go-zero使用第三方库,对proto文件进行了解析, github.com/emicklei/proto,能从代码中看到,分别对message、import、service、goPackage进行解析。这个库star数并不是很高,但是 go-zero能采用它,说明还是很不错的。

四、生成go文件的原理

拿生成service.go文件举例:

const logicFunctionTemplate = `{{if .hasComment}}{{.comment}}{{end}}
func (l *{{.logicName}}) {{.method}} ({{if .hasReq}}in {{.request}}{{if .stream}},stream {{.streamBody}}{{end}}{{else}}stream {{.streamBody}}{{end}}) ({{if .hasReply}}{{.response}},{{end}} error) {
	// todo: add your logic here and delete this line
	
	return {{if .hasReply}}&{{.responseType}}{},{{end}} nil
}
`
//  这里大量用到了模板,第一个 函数模板;
// 第二个采用   go:embed logic.tpl 定义了一个模板文件,和 logicTemplate 进行绑定

//go:embed logic.tpl
var logicTemplate string

logic.tpl 模板的内容

  package {{.packageName}}

  import (
  	"context"

  	{{.imports}}

  	"github.com/zeromicro/go-zero/core/logx"
  )

  type {{.logicName}} struct {
  	ctx    context.Context
  	svcCtx *svc.ServiceContext
  	logx.Logger
  }

func New{{.logicName}}(ctx context.Context,svcCtx *svc.ServiceContext) *{{.logicName}} {
	return &{{.logicName}}{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}
{{.functions}}

拿一个生成好了的 logic进行比对:

package logic

import (
	"context"
	"go_zero_micro/api/user"
	"go_zero_micro/app/user/svc"

	"github.com/zeromicro/go-zero/core/logx"
)

type GetUserNameLogic struct {
	ctx    context.Context
	svcCtx *svc.ServiceContext
	logx.Logger
}

func NewGetUserNameLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetUserNameLogic {
	return &GetUserNameLogic{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}

func (l *GetUserNameLogic) GetUserName(in *user.Request) (*user.Response, error) {

	logx.Info("received GetUserName", in.UserId)
	if len(in.UserId) > 0 {
		return &user.Response{UserName: "frank"}, nil
	}
	return &user.Response{}, nil
}

基本都能和模板对应上,那些 {{.imports}} 这些变量,其实都是从go的代码中定义的变量传递过来,go模板在前后端不分离的项目中,会使用更多,常用来给 html或js中传递变量。

  // logic 核心代码
   func (g *Generator) genLogicGroup(ctx DirContext, proto parser.Proto, cfg *conf.Config) error {
	dir := ctx.GetLogic() // 获取logic的目录
	for _, item := range proto.Service {
		serviceName := item.Name
		for _, rpc := range item.RPC {
          // 声明了 模板中,需要变动的变量,可以和模板对应看,都是一一对应的
			var (
				err           error
				filename      string
				logicName     string
				logicFilename string
				packageName   string
			)

			logicName = fmt.Sprintf("%sLogic", stringx.From(rpc.Name).ToCamel())
			childPkg, err := dir.GetChildPackage(serviceName)
		
			serviceDir := filepath.Base(childPkg)
			nameJoin := fmt.Sprintf("%s_logic", serviceName)
			packageName = strings.ToLower(stringx.From(nameJoin).ToCamel())
			logicFilename, err = format.FileNamingFormat(cfg.NamingFormat, rpc.Name+"_logic")
			
            // 确定文件名
			filename = filepath.Join(dir.Filename, serviceDir, logicFilename+".go")
            // 生成函数,也是采用上面的函数模板 
			functions, err := g.genLogicFunction(serviceName, proto.PbPackage, logicName, rpc)
			imports := collection.NewSet()
			imports.AddStr(fmt.Sprintf(`"%v"`, ctx.GetSvc().Package))
			imports.AddStr(fmt.Sprintf(`"%v"`, ctx.GetPb().Package))
			text, err := pathx.LoadTemplate(category, logicTemplateFileFile, logicTemplate)
			
			if err = util.With("logic").GoFmt(true).Parse(text).SaveTo(map[string]any{ // 写入文件
				"logicName":   logicName,
				"functions":   functions,
				"packageName": packageName,
				"imports":     strings.Join(imports.KeysStr(), pathx.NL),
			}, filename, false); err != nil {
				return err
			}
		}
	}
	return nil
}

  // SaveTo writes the codes to the target path
  func (t *DefaultTemplate) SaveTo(data any, path string, forceUpdate bool) error {
  	if pathx.FileExists(path) && !forceUpdate {
  		return nil
  	}
  	output, err := t.Execute(data)
  	// 最后使用 将要生成的内容写入文件
  	return os.WriteFile(path, output.Bytes(), regularPerm)
  }

其它go文件的生成就不一个个看了,基本都是采用同样的方式。

小结:

  1. 先定义好了模板。
  2. 通过前面获取到的 proto文件内容,对具体的函数名,变量名进行替换。
  3. 写入文件

五、总结

  1. 采用 cobra 定义了命令goctl rpc protoc
  2. 生成已经定义好的 目录文件,这个在代码中是固定的。
  3. 使用第三方库解析proto文件。github.com/emicklei/proto
  4. 定义好各个文件的模板,根据解析后proto参数,重新生成go代码。
  5. 写入文件。