使用 protobuf 自定义选项生成自定义代码

前言

在《protoc-gen-go 介绍与源代码分析》中介绍了基于 protoc-gen-go 制作自定义 protoc go 插件

想要更好的编写自定义插件，必然会遇到如何在 proto 定义文件中，做些自定义内容

本文介绍，如何通过 protobuf 自定义选项生成自定义代码

protobuf 自定义选项

protobuf 自定义选项可以用在 proto 文件全局域、消息、字段、服务、方法等等上面

并能通过官方提供的 API ： proto.GetExtension 来获取这些选项；

也可以自己写 protoc go 插件时，根据选项值不同，自动生成不同的代码；

除此以外，自定义选项还有不少高级应用，本文不会对此展开。

本文，主要介绍下面这个实际用法的制作过程：

比如， proto 文件定义：

syntax = "proto3";

import "broadcast.proto";

package proto;

service Say {
    rpc Hello(Request) returns (mypack.NoReply) { option (mypack.Broadcast) = true; }
    rpc Ping(Request) returns (mypack.NoReply) { }
}

message Request {
    string name = 1;
}

option (mypack.Broadcast) = true; 就是一个自定义选项

自定义 proto go 插件，可以根据这些自定义选项，生成一些自定义代码

比如，生成如下代码：

// For example

type SayService interface {
	BroadcastHello(ctx context.Context) error
	Ping(ctx context.Context) error
}

有 option (mypack.Broadcast) = true; 标记的，让方法名加上前缀 Broadcast

下面就来介绍下，如何定义 option (mypack.Broadcast) 与如何自定义插件让方法名加上前缀 Broadcast

制作自定义选项

一、定义选项

直接举例说明，类似如下 broadcast.proto 定义：

syntax = "proto3";
package mypack;
option go_package = "github.com/fananchong/test_protobuf_options;mypack";
import "protoc-gen-go/descriptor/descriptor.proto";
extend google.protobuf.MethodOptions { bool Broadcast = 50000; }
message NoReply {}

syntax = "proto3";
- 使用 proto3 语法。也可以使用 proto2 看自己喜好
package mypack;
- 包名
- 必须要有
option go_package = "github.com/fananchong/test_protobuf_options;mypack";
- 其他 proto 文件根据这个定义 import 这个 proto
- 必须要有
import "protoc-gen-go/descriptor/descriptor.proto";
- 导入 proto 元数据定义
- 全路径为 github.com/golang/protobuf/protoc-gen-go/descriptor/descriptor.proto
- 项目名 github.com/golang/protobuf 部分需要省略
- 必须要有
- import 第3方开源库的 proto ，都遵循这个规则
extend google.protobuf.MethodOptions { bool Broadcast = 50000; }
- 自定义 method 选项： Broadcast
- 核心语句，本 proto 为的就是定义这个
- 必须要有
message NoReply {}
- 非必须，定义一个消息，其他 proto 可以拿去用

二、生成 pb.go 文件

直接举例说明，类似如下脚本：

set GOPROXY=https://goproxy.io
go get github.com/golang/protobuf
go list -m -f "{{.Dir}}" github.com/golang/protobuf > temp
set /P DEP=<temp
echo %DEP%
protoc -I. -I%DEP% --go_out=. broadcast.proto
copy /y github.com\fananchong\test_protobuf_options\broadcast.pb.go .
rd /s /q github.com

前提，项目是 go mod 项目。即项目根目录下有 go.mod 文件
set GOPROXY=https://goproxy.io && go get github.com/golang/protobuf
- 使 go.mod 文件 require github.com/golang/protobuf 库
go list -m -f "{{.Dir}}" github.com/golang/protobuf > temp && set /P DEP=<temp && echo %DEP%
- 获取 github.com/golang/protobuf 库绝对路径
protoc -I. -I%DEP% --go_out=. broadcast.proto
- 生成 broadcast.pb.go
- 生成出来的在本项目的 github.com\fananchong\test_protobuf_options\ 下
copy /y github.com\fananchong\test_protobuf_options\broadcast.pb.go . && rd /s /q github.com
- 拷贝至本目录

三、具体参考例子

https://github.com/fananchong/test_protobuf_options

基于 proto_gen_go 自定义插件，根据选项来生成自定义代码

直接举例说明，类似如下代码：

func (g *test) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
	origMethName := method.GetName()
	methName := generator.CamelCase(origMethName)

	isBroadcast := false
	if v, err := proto.GetExtension(method.GetOptions(), mypack.E_Broadcast); err == nil {
		isBroadcast = *(v.(*bool))
	}
	if isBroadcast {
		return fmt.Sprintf("Broadcast%s(ctx %s.Context) error", methName, contextPkg)
	}
	return fmt.Sprintf("%s(ctx %s.Context) error", methName, contextPkg)
}

v, err := proto.GetExtension(method.GetOptions(), mypack.E_Broadcast)
- 该行语句可以获取自定义选项的内容

完整插件例子，请参考： https://github.com/fananchong/test_protobuf_options/tree/master/protoc-gen-test

如何在 proto 文件中使用自定义选项

如，定义 test.proto 文件

syntax = "proto3";

import "broadcast.proto";

package proto;

service Say {
    rpc Hello(Request) returns (mypack.NoReply) { option (mypack.Broadcast) = true; }
    rpc Ping(Request) returns (mypack.NoReply) { }
}

message Request {
    string name = 1;
}

并通过 proto_gen_test 生成 test.pb.go ：

// 无关代码略

// For example

type SayService interface {
	BroadcastHello(ctx context.Context) error
	Ping(ctx context.Context) error
}

具体例子，请参见： https://github.com/fananchong/test_protobuf_options_example

其他说明

本文中，都是举例说明（东西很少，跑通流程又要多次试错，故直接上代码）

需要了解的一些背景知识：

protobuf 自定义选项规范细节
proto-gen-go 如何更好的支持 go mod
proto go 代码如何获取 Service 、 Method 的自定义选项内容

请参见下面的【参考资料】

参考资料

由于 protobuf 官方文档需要翻墙才能查阅，因此，参考资源均来自

百度
github.com/golang/protobuf 及其 issue
github 其他开源代码

主要有：

介绍 protobuf 自定义选项语法规范
- https://blog.csdn.net/tycoon1988/article/details/45268449
- https://github.com/golang/protobuf
protoc-gen-go 如何支持 go mod
- https://github.com/golang/protobuf/issues/748
- https://github.com/golang/protobuf/issues/897
代码中如何获取 Service 、 Method 的自定义选项
- https://github.com/golang/protobuf/issues/489
- https://github.com/prasek/go-grpc-info

以上

fananchong2

发布了129 篇原创文章 · 获赞 73 · 访问量 16万+

私信关注