126 lines
3.2 KiB
Go
126 lines
3.2 KiB
Go
package plugin
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
|
|
"github.com/gogo/protobuf/protoc-gen-gogo/generator"
|
|
)
|
|
|
|
type ttrpcGenerator struct {
|
|
*generator.Generator
|
|
generator.PluginImports
|
|
|
|
typeurlPkg generator.Single
|
|
ttrpcPkg generator.Single
|
|
contextPkg generator.Single
|
|
}
|
|
|
|
func init() {
|
|
generator.RegisterPlugin(new(ttrpcGenerator))
|
|
}
|
|
|
|
func (p *ttrpcGenerator) Name() string {
|
|
return "ttrpc"
|
|
}
|
|
|
|
func (p *ttrpcGenerator) Init(g *generator.Generator) {
|
|
p.Generator = g
|
|
}
|
|
|
|
func (p *ttrpcGenerator) Generate(file *generator.FileDescriptor) {
|
|
p.PluginImports = generator.NewPluginImports(p.Generator)
|
|
p.contextPkg = p.NewImport("context")
|
|
p.typeurlPkg = p.NewImport("github.com/containerd/typeurl")
|
|
p.ttrpcPkg = p.NewImport("github.com/stevvooe/ttrpc")
|
|
|
|
for _, service := range file.GetService() {
|
|
serviceName := service.GetName()
|
|
if pkg := file.GetPackage(); pkg != "" {
|
|
serviceName = pkg + "." + serviceName
|
|
}
|
|
|
|
p.genService(serviceName, service)
|
|
}
|
|
}
|
|
|
|
func (p *ttrpcGenerator) genService(fullName string, service *descriptor.ServiceDescriptorProto) {
|
|
serviceName := service.GetName() + "Service"
|
|
p.P()
|
|
p.P("type ", serviceName, " interface{")
|
|
p.In()
|
|
for _, method := range service.Method {
|
|
p.P(method.GetName(),
|
|
"(ctx ", p.contextPkg.Use(), ".Context, ",
|
|
"req *", p.typeName(method.GetInputType()), ") ",
|
|
"(*", p.typeName(method.GetOutputType()), ", error)")
|
|
|
|
}
|
|
p.Out()
|
|
p.P("}")
|
|
|
|
p.P()
|
|
// registration method
|
|
p.P("func Register", serviceName, "(srv *", p.ttrpcPkg.Use(), ".Server, svc ", serviceName, ") error {")
|
|
p.In()
|
|
p.P(`return srv.Register("`, fullName, `", map[string]`, p.ttrpcPkg.Use(), ".Handler{")
|
|
p.In()
|
|
for _, method := range service.Method {
|
|
p.P(`"`, method.GetName(), `": `, p.ttrpcPkg.Use(), `.HandlerFunc(func(ctx context.Context, req interface{}) (interface{}, error) {`)
|
|
p.In()
|
|
p.P("return svc.", method.GetName(), "(ctx, req.(*", p.typeName(method.GetInputType()), "))")
|
|
p.Out()
|
|
p.P("}),")
|
|
}
|
|
p.Out()
|
|
p.P("})")
|
|
p.Out()
|
|
p.P("}")
|
|
|
|
clientType := service.GetName() + "Client"
|
|
clientStructType := strings.ToLower(clientType[:1]) + clientType[1:]
|
|
p.P()
|
|
p.P("type ", clientStructType, " struct{")
|
|
p.In()
|
|
p.P("client *", p.ttrpcPkg.Use(), ".Client")
|
|
p.Out()
|
|
p.P("}")
|
|
p.P()
|
|
p.P("func New", clientType, "(client *", p.ttrpcPkg.Use(), ".Client)", serviceName, "{")
|
|
p.In()
|
|
p.P("return &", clientStructType, "{")
|
|
p.In()
|
|
p.P("client: client,")
|
|
p.Out()
|
|
p.P("}")
|
|
p.Out()
|
|
p.P("}")
|
|
p.P()
|
|
for _, method := range service.Method {
|
|
p.P()
|
|
p.P("func (c *", clientStructType, ") ", method.GetName(),
|
|
"(ctx ", p.contextPkg.Use(), ".Context, ",
|
|
"req *", p.typeName(method.GetInputType()), ") ",
|
|
"(*", p.typeName(method.GetOutputType()), ", error) {")
|
|
p.In()
|
|
p.P("resp, err := c.client.Call(ctx, ", `"`+fullName+`", `, `"`+method.GetName()+`"`, ", req)")
|
|
p.P("if err != nil {")
|
|
p.In()
|
|
p.P("return nil, err")
|
|
p.Out()
|
|
p.P("}")
|
|
p.P("return resp.(*", p.typeName(method.GetOutputType()), "), nil")
|
|
p.Out()
|
|
p.P("}")
|
|
}
|
|
}
|
|
|
|
func (p *ttrpcGenerator) objectNamed(name string) generator.Object {
|
|
p.Generator.RecordTypeUse(name)
|
|
return p.Generator.ObjectNamed(name)
|
|
}
|
|
|
|
func (p *ttrpcGenerator) typeName(str string) string {
|
|
return p.Generator.TypeName(p.objectNamed(str))
|
|
}
|