本期的主要内容将手把手教会大家,编写probuf的go插件,以我自己编写的一个生成结构体的插件为例子。项目位于 https://github.com/hisheng/pr...

一、自定义ProtoBuf插件介绍

我们常用的go支持protobuf插件有

插件名称介绍
protoc-gen-go通过.proto文件生成.pb.go文件
protoc-gen-doc通过.proto文件生成文档
protoc-gen-go-errors通过.proto文件生成error
protoc-gen-go-errors通过.proto文件生成error
protoc-gen-go-grpc通过.proto文件生成grpc
protoc-gen-go-http通过.proto文件生成http
protoc-gen-openapi通过.proto文件生成openapi
protoc-gen-validate通过.proto文件生成validate验证
protoc-gen-go-enum通过.proto文件生成go自定义枚举

基本上该有的插件,都已经有写过了,所以我们可以查看别人的源码,来观察怎么来写一个插件。

二、手把手写自定义插件(以protoc-gen-go-struct为例)

2.1 新建protoc-gen-go-struct文件夹,并且进入到这个文件夹里。

mkdir protoc-gen-go-struct && cd protoc-gen-go-struct

2.2 新建go module项目

我们执行 go mod init modName 命名来生成项目如下:

go mod init github/hisheng/protoc-gen-go-struct

此时我们查看文件夹发现生产了一个 go.mod文件,查看一下代码如下:

module github/hisheng/protoc-gen-go-struct

go 1.19

2.3 写main函数

我们在项目根目录,写一个main.go如下

touch main.go

此时我们发现项目根目录下,生成了一个main.go文件。
我们写一个main函数,代码如下:

package main

import (
    "google.golang.org/protobuf/compiler/protogen"
)

func main() {
    protogen.Options{}.Run(func(gen *protogen.Plugin) error {
        for _, f := range gen.Files {
            if !f.Generate {
                continue
            }
            generateFile(gen, f)
        }
        return nil
    })
}

这个main()函数大家可以直接复制,基本所有的插件都是这样的格式,当然这个main也可以接受参数,这里我们简化,先不介绍,感兴趣的人,可以参考protoc-gen-go的main函数来写。
我们自己写的方法主要是 generateFile(gen, f) 这个函数。
这个函数用来读取.proto文件,并且生产go文件。

2.4 自定义generateFile(gen, f)函数

这个函数全称是 generateFile(gen protogen.Plugin, file protogen.File),接受的两个参数

gen *protogen.Plugin 为生成的插件,主要用来生成go文件
file *protogen.File 为.proto文件对他的file对象

我这里的主要代码是:

// 生成.struct.go文件,参数为 输出插件gen,以及读取的文件file
func generateFile(gen *protogen.Plugin, file *protogen.File) {
    filename := file.GeneratedFilenamePrefix + ".struct.go"
    g := gen.NewGeneratedFile(filename, file.GoImportPath)
    // 输出 package packageName
    g.P("package ", file.GoPackageName)
    g.P() // 换行

    for _, m := range file.Messages {
        // 输出 type m.GoIdent struct {
        g.P("type ", m.GoIdent, " struct {")
        for _, field := range m.Fields {
            leadingComment := field.Comments.Leading.String()
            trailingComment := field.Comments.Trailing.String()

            line := fmt.Sprintf("%s %s `json:\"%s\"` %s", field.GoName, field.Desc.Kind(), field.Desc.JSONName(), trailingComment)
            // 输出 行首注释
            g.P(leadingComment)
            // 输出 行内容
            g.P(line)
        }
        // 输出 }
        g.P("}")
    }
    g.P() // 换行
}

就是如此简单的10几行代码,就可以读取.proto文件,并生成.go文件了。接下来我们详细的介绍一下里面主要的变量以及对象。

2.4.1 第一步先生成 GeneratedFile 对象
filename := file.GeneratedFilenamePrefix + ".struct.go"
g := gen.NewGeneratedFile(filename, file.GoImportPath) 
2.4.2 第二步生成go代码的package
// 输出 package packageName
g.P("package ", file.GoPackageName)
g.P() // 换行
2.4.3 第三步便利proto文件的message

因为我们这里是找message然后生成struct,所以使用file.Messages来获取所有的message,然后遍历

    for _, m := range file.Messages {
        // 输出 type m.GoIdent struct {
        g.P("type ", m.GoIdent, " struct {")
        for _, field := range m.Fields {
            leadingComment := field.Comments.Leading.String()
            trailingComment := field.Comments.Trailing.String()

            line := fmt.Sprintf("%s %s `json:\"%s\"` %s", field.GoName, field.Desc.Kind(), field.Desc.JSONName(), trailingComment)
            // 输出 行首注释
            g.P(leadingComment)
            // 输出 行内容
            g.P(line)
        }
        // 输出 }
        g.P("}")
    }

这一步代码主要生产go的struct,生成后的样子如下:

type User struct {
    // xingming
    Name string `json:"name"` // 姓名
    // age
    Age int64 `json:"age"` // 年龄
}

三、编译插件

我们把这个代码,在本地编译安装,在项目根目录执行go install

go install

此时我们到自己的GOPATH目录查看是否生成
我们cd到$GOPATH的bin目录,一般go install安装的命令都在这里

cd  $GOPATH/bin && ls -al

image.png
我们看到了 protoc-gen-go-struct 二进制命令。

四、protoc使用protoc-gen-go-struct插件

我们在其他地方写测试方法,写一个 pt.proto文件

mkdir protoc_struct && cd protoc_struct && touch pt.proto

然后把pt.ptoto里面写入代码

syntax = "proto3";
package pt;
option go_package = "./protoc_struct";

message User {
  // xingming
  string name = 1;// 姓名
  // age
  int64 age = 2;// 年龄
}

最后我们在执行 protoc 命令

protoc --go-struct_out=./ pt.proto

此时我们发现生产了pt.struct.go文件
文件里面生成的go代码如下

package protoc_struct

type User struct {
    // xingming
    Name string `json:"name"` // 姓名
    // age
    Age int64 `json:"age"` // 年龄
}

和上面的.proto文件一一对应。
写到这里,我们自己写一个go对应的protobuf插件就完成。

本文参与了思否技术征文,欢迎正在阅读的你也加入。

海生
104 声望32 粉丝

与黑夜里,追求那一抹萤火。