zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

To generate grpc stream, fix issue #616 (#815)

Browse Source 
Co-authored-by: anqiansong <anqiansong@xiaoheiban.cn>
This commit is contained in:
anqiansong
2021-07-16 22:54:07 +08:00
committed by GitHub
parent 598fda0c97
commit db87fd3239
6 changed files with 123 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
tools/goctl/rpc/generator/gencall.go
View File

@@ -49,13 +49,13 @@ func New{{.serviceName}}(cli zrpc.Client) {{.serviceName}} {
`
callInterfaceFunctionTemplate = `{{if .hasComment}}{{.comment}}
{{end}}{{.method}}(ctx context.Context,in *{{.pbRequest}}) (*{{.pbResponse}},error)`
{{end}}{{.method}}(ctx context.Context{{if .hasReq}},in *{{.pbRequest}}{{end}}) ({{if .notStream}}*{{.pbResponse}}, {{else}}{{.streamBody}},{{end}} error)`
callFunctionTemplate = `
{{if .hasComment}}{{.comment}}{{end}}
func (m *default{{.serviceName}}) {{.method}}(ctx context.Context,in *{{.pbRequest}}) (*{{.pbResponse}}, error) {
func (m *default{{.serviceName}}) {{.method}}(ctx context.Context{{if .hasReq}},in *{{.pbRequest}}{{end}}) ({{if .notStream}}*{{.pbResponse}}, {{else}}{{.streamBody}},{{end}} error) {
client := {{.package}}.New{{.rpcServiceName}}Client(m.cli.Conn())
return client.{{.method}}(ctx, in)
return client.{{.method}}(ctx,{{if .hasReq}} in{{end}})
}
`
)
@@ -78,7 +78,7 @@ func (g *DefaultGenerator) GenCall(ctx DirContext, proto parser.Proto, cfg *conf
return err
}
iFunctions, err := g.getInterfaceFuncs(service)
iFunctions, err := g.getInterfaceFuncs(proto.PbPackage, service)
if err != nil {
return err
}
@@ -115,6 +115,14 @@ func (g *DefaultGenerator) genFunction(goPackage string, service parser.Service)
}
comment := parser.GetComment(rpc.Doc())
var streamServer string
if rpc.StreamsRequest && rpc.StreamsReturns {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_StreamClient")
} else if rpc.StreamsRequest {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_ClientStreamClient")
} else {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_ServerStreamClient")
}
buffer, err := util.With("sharedFn").Parse(text).Execute(map[string]interface{}{
"serviceName": stringx.From(service.Name).ToCamel(),
"rpcServiceName": parser.CamelCase(service.Name),
@@ -124,6 +132,9 @@ func (g *DefaultGenerator) genFunction(goPackage string, service parser.Service)
"pbResponse": parser.CamelCase(rpc.ReturnsType),
"hasComment": len(comment) > 0,
"comment": comment,
"hasReq": !rpc.StreamsRequest,
"notStream": !rpc.StreamsRequest && !rpc.StreamsReturns,
"streamBody": streamServer,
})
if err != nil {
return nil, err
@@ -134,7 +145,7 @@ func (g *DefaultGenerator) genFunction(goPackage string, service parser.Service)
return functions, nil
}
func (g *DefaultGenerator) getInterfaceFuncs(service parser.Service) ([]string, error) {
func (g *DefaultGenerator) getInterfaceFuncs(goPackage string, service parser.Service) ([]string, error) {
functions := make([]string, 0)
for _, rpc := range service.RPC {
@@ -144,13 +155,25 @@ func (g *DefaultGenerator) getInterfaceFuncs(service parser.Service) ([]string,
}
comment := parser.GetComment(rpc.Doc())
var streamServer string
if rpc.StreamsRequest && rpc.StreamsReturns {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_StreamClient")
} else if rpc.StreamsRequest {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_ClientStreamClient")
} else {
streamServer = fmt.Sprintf("%s.%s", goPackage, parser.CamelCase(service.Name)+"_ServerStreamClient")
}
buffer, err := util.With("interfaceFn").Parse(text).Execute(
map[string]interface{}{
"hasComment": len(comment) > 0,
"comment": comment,
"method": parser.CamelCase(rpc.Name),
"hasReq": !rpc.StreamsRequest,
"pbRequest": parser.CamelCase(rpc.RequestType),
"notStream": !rpc.StreamsRequest && !rpc.StreamsReturns,
"pbResponse": parser.CamelCase(rpc.ReturnsType),
"streamBody": streamServer,
})
if err != nil {
return nil, err