protoc-gen-go-hertz

Posted on Apr 4, 2024

背景:我有一个 API 项目 tzf-server,想利用 proto 文件生成 openapi.yaml 配合 Swagger 使用, 并且有与其一致的 HTTP API 供调用。很遗憾的是 Go 生态中尚未有能和 Python 生态中的 FastAPI 接近的 API 框架来简化业务代码编写和文档生成工作。 特别是 Hertz 框架,提供了很多功能甚至通过扩展支持了参数校验功能,但是这些并不能直接输出成 openapi.yaml 文件。 gRPC-Gateway 是一个不错的选择,但是尚不支持 OpenAPI V3。 还有一个隐藏的原因是在用 buf 管理 proto,但是 hertz 的生成工具怎么和 buf 一块使用没有相关的文档。所以我决定自己写一个。

目标:

  • proto 定义 HTTP API
  • 生成 openapi.yaml 供 Swagger 使用
  • 不引入其他框架,Kratos 内置了接近满足需求的功能,不过出于性能和场景考虑还是维持 Hertz

照着 Kratos 的 protoc-gen-go-http 写了 protoc-gen-go-hertz,请注意这个项目完全是为了 tzf-server 项目服务的,不保证通用性。

用起来也是简单:

go install github.com/ringsaturn/protoc-gen-go-hertz@latest

修改 buf.gen.yaml 文件:

version: v1
plugins:
  - plugin: buf.build/community/google-gnostic-openapi:v0.7.0
    out: tzf/v1
  - plugin: buf.build/protocolbuffers/go:v1.30.0
    out: .
    opt:
      - paths=source_relative
  - plugin: buf.build/grpc/go:v1.3.0
    out: .
    opt:
      - paths=source_relative
  - plugin: go-hertz
    out: .
    opt:
      - paths=source_relative
  - plugin: go-grpc-mock
    out: .
    opt: paths=source_relative

tzf-server 定义如下:

// 完整版本 https://github.com/ringsaturn/tzf-server/blob/v0.13.4/tzf/v1/api.proto

service TZFService {
  option (google.api.default_host) = "0.0.0.0:8080";

  rpc Ping(PingRequest) returns (PingResponse) {
    option (google.api.http) = {
      get: "/api/v1/ping"
    };
  }

  // GetTimezone returns timezone name by longitude and latitude.
  rpc GetTimezone(GetTimezoneRequest) returns (GetTimezoneResponse) {
    option (google.api.http) = {
      get: "/api/v1/tz"
    };
  }

  // GetTimezones returns timezone names by longitude and latitude.
  rpc GetTimezones(GetTimezonesRequest) returns (GetTimezonesResponse) {
    option (google.api.http) = {
      get: "/api/v1/tzs"
    };
  }

  // GetAllTimezones returns all timezone names.
  rpc GetAllTimezones(GetAllTimezonesRequest) returns (GetAllTimezonesResponse) {
    option (google.api.http) = {
      get: "/api/v1/tzs/all"
    };
  }
}

会生成对应的 interface

// 完整版本 https://github.com/ringsaturn/tzf-server/blob/v0.13.4/tzf/v1/api_hertz.pb.go

type TZFServiceHTTPServer interface {
	// GetAllTimezones GetAllTimezones returns all timezone names.
	GetAllTimezones(context.Context, *GetAllTimezonesRequest) (*GetAllTimezonesResponse, error)
	// GetTimezone GetTimezone returns timezone name by longitude and latitude.
	GetTimezone(context.Context, *GetTimezoneRequest) (*GetTimezoneResponse, error)
	// GetTimezones GetTimezones returns timezone names by longitude and latitude.
	GetTimezones(context.Context, *GetTimezonesRequest) (*GetTimezonesResponse, error)
	Ping(context.Context, *PingRequest) (*PingResponse, error)
}

只要实现这个 interface 就可以了。集成 Swagger 也可以很简单,用 Google gnostic 生成的 openapi.yaml 文件:

// 完整版本 https://github.com/ringsaturn/tzf-server/blob/v0.13.4/tzf/v1/api_hertz_swagger.go

//go:embed openapi.yaml
var openapiYAML []byte

func BindSwagger(h *server.Hertz, openAPIYAMLPath string, swaggerPath string) {
	h.GET(swaggerPath, swagger.WrapHandler(
		swaggerFiles.Handler,
		swagger.URL(openAPIYAMLPath),
	))

	h.GET(openAPIYAMLPath, func(c context.Context, ctx *app.RequestContext) {
		ctx.Header("Content-Type", "application/x-yaml")
		_, _ = ctx.Write(openapiYAML)
	})
}

func BindDefaultSwagger(h *server.Hertz) {
	BindSwagger(h, "/openapi.yaml", "/swagger/*any")
}