Skip to content

Commit

Permalink
refactor(tool): rm thriftgo dep from pkg generator (#1680)
Browse files Browse the repository at this point in the history
  • Loading branch information
@xiaost
xiaost authored Jan 22, 2025
1 parent 4f3c3fc commit b2a9665
Show file tree
Hide file tree
Showing 9 changed files with 67 additions and 90 deletions.
12 changes: 6 additions & 6 deletions tool/internal_pkg/generator/generator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -569,7 +569,7 @@ func (g *generator) setImports(name string, pkg *PackageInfo) {
if g.StreamX || (!m.ServerStreaming && !m.ClientStreaming) {
pkg.AddImports("context")
}
if g.StreamX && m.Streaming.IsStreaming {
if g.StreamX && m.IsStreaming {
pkg.AddImports("github.com/cloudwego/kitex/pkg/streamx")
}
for _, a := range m.Args {
Expand All @@ -587,7 +587,7 @@ func (g *generator) setImports(name string, pkg *PackageInfo) {
// for StreamX, if there is streaming method, generate Server Interface in server.go
if g.StreamX {
for _, method := range pkg.AllMethods() {
if method.Streaming.IsStreaming {
if method.IsStreaming {
pkg.AddImports("context")
pkg.AddImports("github.com/cloudwego/kitex/pkg/streamx")
}
Expand Down Expand Up @@ -619,15 +619,15 @@ func (g *generator) setImports(name string, pkg *PackageInfo) {
}
// streaming imports
if !g.StreamX {
if m.Streaming.IsStreaming || pkg.Codec == "protobuf" {
if m.IsStreaming || pkg.Codec == "protobuf" {
// protobuf handler support both PingPong and Unary (streaming) requests
pkg.AddImport("streaming", "github.com/cloudwego/kitex/pkg/streaming")
}
if m.ClientStreaming || m.ServerStreaming {
if m.IsStreaming {
pkg.AddImports("fmt")
}
} else {
if m.Streaming.IsStreaming {
if m.IsStreaming {
pkg.AddImports("github.com/cloudwego/kitex/client/streamxclient")
pkg.AddImports("github.com/cloudwego/kitex/client/streamxclient/streamxcallopt")
pkg.AddImports("github.com/cloudwego/kitex/pkg/streamx")
Expand Down Expand Up @@ -671,7 +671,7 @@ func needCallOpt(pkg *PackageInfo) bool {
switch pkg.Codec {
case "thrift":
for _, m := range pkg.ServiceInfo.AllMethods() {
if !m.Streaming.IsStreaming {
if !m.IsStreaming {
needCallOpt = true
break
}
Expand Down
19 changes: 4 additions & 15 deletions tool/internal_pkg/generator/generator_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@ import (
"testing"
"time"

"github.com/cloudwego/thriftgo/generator/golang/streaming"

"github.com/cloudwego/kitex/internal/test"
"github.com/cloudwego/kitex/tool/internal_pkg/util"
)
Expand Down Expand Up @@ -202,9 +200,8 @@ func Test_needCallOpt(t *testing.T) {
Codec: "thrift",
ServiceInfo: &ServiceInfo{
Methods: []*MethodInfo{{
Streaming: &streaming.Streaming{
IsStreaming: true,
},
IsStreaming: true,
ServerStreaming: true,
}},
},
}
Expand All @@ -215,16 +212,8 @@ func Test_needCallOpt(t *testing.T) {
Codec: "thrift",
ServiceInfo: &ServiceInfo{
Methods: []*MethodInfo{
{
Streaming: &streaming.Streaming{
IsStreaming: true,
},
},
{
Streaming: &streaming.Streaming{
IsStreaming: false,
},
},
{IsStreaming: true, ServerStreaming: true},
{},
},
},
}
Expand Down
20 changes: 17 additions & 3 deletions tool/internal_pkg/generator/type.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,6 @@ import (
"strings"
"text/template"

"github.com/cloudwego/thriftgo/generator/golang/streaming"

"github.com/cloudwego/kitex/tool/internal_pkg/util"
"github.com/cloudwego/kitex/transport"
)
Expand Down Expand Up @@ -211,12 +209,28 @@ type MethodInfo struct {
ResStructName string
IsResponseNeedRedirect bool // int -> int*
GenArgResultStruct bool
IsStreaming bool
ClientStreaming bool
ServerStreaming bool
Streaming *streaming.Streaming
StreamX bool
}

func (m *MethodInfo) StreamingMode() string {
if !m.IsStreaming {
return ""
}
if m.ClientStreaming {
if m.ServerStreaming {
return "bidirectional"
}
return "client"
}
if m.ServerStreaming {
return "server"
}
return "unary"
}

// Parameter .
type Parameter struct {
Deps []PkgInfo
Expand Down
30 changes: 2 additions & 28 deletions tool/internal_pkg/pluginmode/protoc/plugin.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import (
"text/template"

genfastpb "github.com/cloudwego/fastpb/protoc-gen-fastpb/generator"
"github.com/cloudwego/thriftgo/generator/golang/streaming"
gengo "google.golang.org/protobuf/cmd/protoc-gen-go/internal_gengo"
"google.golang.org/protobuf/compiler/protogen"

Expand Down Expand Up @@ -297,13 +296,11 @@ func (pp *protocPlugin) convertTypes(file *protogen.File) (ss []*generator.Servi
ServerStreaming: m.Desc.IsStreamingServer(),
}
si.Methods = append(si.Methods, mi)
if !si.HasStreaming && (mi.ClientStreaming || mi.ServerStreaming) {
if mi.ClientStreaming || mi.ServerStreaming {
mi.IsStreaming = true
si.HasStreaming = true
}
}
for _, m := range si.Methods {
BuildStreaming(m, si.HasStreaming)
}
if file.Generate {
si.GenerateHandler = true
}
Expand Down Expand Up @@ -348,29 +345,6 @@ func (pp *protocPlugin) convertTypes(file *protogen.File) (ss []*generator.Servi
return
}

// BuildStreaming builds protobuf MethodInfo.Streaming as for Thrift, to simplify codegen
func BuildStreaming(mi *generator.MethodInfo, serviceHasStreaming bool) {
s := &streaming.Streaming{
// pb: if one method is streaming, then the service is streaming, making all methods streaming
IsStreaming: serviceHasStreaming,
}
if mi.ClientStreaming && mi.ServerStreaming {
s.Mode = streaming.StreamingBidirectional
s.BidirectionalStreaming = true
s.ClientStreaming = true
s.ServerStreaming = true
} else if mi.ClientStreaming && !mi.ServerStreaming {
s.Mode = streaming.StreamingClientSide
s.ClientStreaming = true
} else if !mi.ClientStreaming && mi.ServerStreaming {
s.Mode = streaming.StreamingServerSide
s.ServerStreaming = true
} else if serviceHasStreaming {
s.Mode = streaming.StreamingUnary // Unary APIs over HTTP2
}
mi.Streaming = s
}

func (pp *protocPlugin) getCombineServiceName(name string, svcs []*generator.ServiceInfo) string {
for _, svc := range svcs {
if svc.ServiceName == name {
Expand Down
2 changes: 1 addition & 1 deletion tool/internal_pkg/pluginmode/thriftgo/convertor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -455,7 +455,7 @@ func (c *converter) makeMethod(si *generator.ServiceInfo, f *golang.Function) (*
Void: f.Void,
ArgStructName: f.ArgType().GoName().String(),
GenArgResultStruct: false,
Streaming: st,
IsStreaming: st.IsStreaming,
ClientStreaming: st.ClientStreaming,
ServerStreaming: st.ServerStreaming,
ArgsLength: len(f.Arguments()),
Expand Down
32 changes: 16 additions & 16 deletions tool/internal_pkg/tpl/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,23 +39,23 @@ type Client interface {
{{- if and (eq $.Codec "protobuf") (or .ClientStreaming .ServerStreaming)}}{{/* protobuf: generate streaming calls in Client, to keep compatibility */}}
{{.Name}}(ctx context.Context {{if not .ClientStreaming}}{{range .Args}}, {{.RawName}} {{.Type}}{{end}}{{end}}, callOptions ...callopt.Option ) (stream {{.ServiceName}}_{{.RawName}}Client, err error)
{{- else}}
{{- if or (eq $.Codec "protobuf") (eq .Streaming.Mode "")}}
{{- if or (eq $.Codec "protobuf") (eq .StreamingMode "")}}
{{.Name}}(ctx context.Context {{range .Args}}, {{.RawName}} {{.Type}}{{end}}, callOptions ...callopt.Option ) ({{if not .Void}}r {{.Resp.Type}}, {{end}}err error)
{{- end}}
{{- end}}
{{- /* Streamx interface for streaming method */}}
{{- if and .StreamX (eq $.Codec "thrift") .Streaming.IsStreaming}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $clientSide := (eq .Streaming.Mode "client")}}
{{- $serverSide := (eq .Streaming.Mode "server")}}
{{- $bidiSide := (eq .Streaming.Mode "bidirectional")}}
{{- if and .StreamX (eq $.Codec "thrift") .IsStreaming}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $clientSide := (eq .StreamingMode "client")}}
{{- $serverSide := (eq .StreamingMode "server")}}
{{- $bidiSide := (eq .StreamingMode "bidirectional")}}
{{- $arg := index .Args 0}}
{{.Name}}{{- if $streamingUnary}}(ctx context.Context, req {{$arg.Type}}, callOptions ...streamxcallopt.CallOption) (r {{.Resp.Type}}, err error)
{{- else if $clientSide}}(ctx context.Context, callOptions ...streamxcallopt.CallOption) (context.Context, streamx.ClientStreamingClient[{{NotPtr $arg.Type}}, {{NotPtr .Resp.Type}}], error)
{{- else if $serverSide}}(ctx context.Context, req {{$arg.Type}}, callOptions ...streamxcallopt.CallOption) (context.Context, streamx.ServerStreamingClient[{{NotPtr .Resp.Type}}], error)
{{- else if $bidiSide}}(ctx context.Context, callOptions ...streamxcallopt.CallOption) (context.Context, streamx.BidiStreamingClient[{{NotPtr $arg.Type}}, {{NotPtr .Resp.Type}}], error)
{{- end}}
{{- end}}{{- /* if and .StreamX (eq $.Codec "thrift") .Streaming.IsStreaming end */}}
{{- end}}{{- /* if and .StreamX (eq $.Codec "thrift") .IsStreaming end */}}
{{- end}}
}
Expand All @@ -66,7 +66,7 @@ type StreamClient interface {
{{- range .AllMethods}}
{{- if or .ClientStreaming .ServerStreaming}}
{{.Name}}(ctx context.Context {{if not .ClientStreaming}}{{range .Args}}, {{.RawName}} {{.Type}}{{end}}{{end}}, callOptions ...streamcall.Option ) (stream {{.ServiceName}}_{{.RawName}}Client, err error)
{{- else if .Streaming.Unary}}
{{- else if eq .StreamingMode "unary"}}
{{.Name}}(ctx context.Context {{range .Args}}, {{.RawName}} {{.Type}}{{end}}, callOptions ...streamcall.Option ) ({{if not .Void}}r {{.Resp.Type}}, {{end}}err error)
{{- end}}
{{- end}}
Expand Down Expand Up @@ -139,19 +139,19 @@ func (p *k{{$.ServiceName}}Client) {{.Name}}(ctx context.Context {{if not .Clien
}
{{- end}}
{{- else}}
{{- if or (eq $.Codec "protobuf") (eq .Streaming.Mode "")}}
{{- if or (eq $.Codec "protobuf") (eq .StreamingMode "")}}
func (p *k{{$.ServiceName}}Client) {{.Name}}(ctx context.Context {{range .Args}}, {{.RawName}} {{.Type}}{{end}}, callOptions ...callopt.Option ) ({{if not .Void}}r {{.Resp.Type}}, {{end}}err error) {
ctx = client.NewCtxWithCallOptions(ctx, callOptions)
return p.kClient.{{.Name}}(ctx{{range .Args}}, {{.RawName}}{{end}})
}
{{- end}}
{{- end}}
{{- /* Streamx interface for streaming method */}}
{{- if and .StreamX (eq $.Codec "thrift") .Streaming.IsStreaming}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $clientSide := (eq .Streaming.Mode "client")}}
{{- $serverSide := (eq .Streaming.Mode "server")}}
{{- $bidiSide := (eq .Streaming.Mode "bidirectional")}}
{{- if and .StreamX (eq $.Codec "thrift") .IsStreaming}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $clientSide := (eq .StreamingMode "client")}}
{{- $serverSide := (eq .StreamingMode "server")}}
{{- $bidiSide := (eq .StreamingMode "bidirectional")}}
{{- $arg := index .Args 0}}
func (p *k{{$.ServiceName}}Client) {{.Name}}{{- if $streamingUnary}}(ctx context.Context, req {{$arg.Type}}, callOptions ...streamxcallopt.CallOption) (r {{.Resp.Type}}, err error) {
return p.kClient.{{.Name}}(ctx, req, callOptions...)
Expand All @@ -163,7 +163,7 @@ func (p *k{{$.ServiceName}}Client) {{.Name}}{{- if $streamingUnary}}(ctx context
return p.kClient.{{.Name}}(ctx, callOptions...)
{{- end}}
}
{{- end}}{{- /* if and .StreamX (eq $.Codec "thrift") .Streaming.IsStreaming end */}}
{{- end}}{{- /* if and .StreamX (eq $.Codec "thrift") .IsStreaming end */}}
{{end}}
{{- if not .StreamX}}
Expand Down Expand Up @@ -207,7 +207,7 @@ func (p *k{{$.ServiceName}}StreamClient) {{.Name}}(ctx context.Context {{if not
ctx = client.NewCtxWithCallOptions(ctx, streamcall.GetCallOptions(callOptions))
return p.kClient.{{.Name}}(ctx{{if not .ClientStreaming}}{{range .Args}}, {{.RawName}}{{end}}{{end}})
}
{{else if .Streaming.Unary}}
{{else if eq .StreamingMode "unary"}}
func (p *k{{$.ServiceName}}StreamClient) {{.Name}}(ctx context.Context {{range .Args}}, {{.RawName}} {{.Type}}{{end}}, callOptions ...streamcall.Option ) ({{if not .Void}}r {{.Resp.Type}}, {{end}}err error) {
ctx = client.NewCtxWithCallOptions(ctx, streamcall.GetCallOptions(callOptions))
return p.kClient.{{.Name}}(ctx{{range .Args}}, {{.RawName}}{{end}})
Expand Down
12 changes: 6 additions & 6 deletions tool/internal_pkg/tpl/handler.method.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@ package tpl
// HandlerMethodsTpl is the template for generating methods in handler.go.
var HandlerMethodsTpl string = `{{define "HandlerMethod"}}
{{range .AllMethods}}
{{- if and .StreamX .Streaming.IsStreaming}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $clientSide := (eq .Streaming.Mode "client")}}
{{- $serverSide := (eq .Streaming.Mode "server")}}
{{- $bidiSide := (eq .Streaming.Mode "bidirectional")}}
{{- if and .StreamX .IsStreaming}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $clientSide := (eq .StreamingMode "client")}}
{{- $serverSide := (eq .StreamingMode "server")}}
{{- $bidiSide := (eq .StreamingMode "bidirectional")}}
{{- $arg := index .Args 0}}
func (s *{{.ServiceName}}Impl) {{.Name}}{{- if $streamingUnary}}(ctx context.Context, req {{$arg.Type}}) (resp {{.Resp.Type}}, err error) {
{{- else if $clientSide}}(ctx context.Context, stream streamx.ClientStreamingServer[{{NotPtr $arg.Type}}, {{NotPtr .Resp.Type}}]) (resp {{.Resp.Type}}, err error) {
Expand Down Expand Up @@ -57,7 +57,7 @@ func (s *{{$.ServiceName}}Impl) {{.Name}}(ctx context.Context {{range .Args}}, {
}
{{end}}{{/* if .Void end */}}
{{end}}{{/* if or .ClientStreaming .ServerStreaming end */}}
{{end}}{{/* if and .StreamX .Streaming.IsStreaming end */}}
{{end}}{{/* if and .StreamX .IsStreaming end */}}
{{end}}{{/* range .AllMethods end */}}
{{end}}{{/* define "HandlerMethod" */}}
`
8 changes: 4 additions & 4 deletions tool/internal_pkg/tpl/server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,10 @@ import (
{{- $serverInterfaceName = .ServiceName }}
type {{.ServiceName}} interface {
{{- range .AllMethods}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $clientSide := (eq .Streaming.Mode "client")}}
{{- $serverSide := (eq .Streaming.Mode "server")}}
{{- $bidiSide := (eq .Streaming.Mode "bidirectional")}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $clientSide := (eq .StreamingMode "client")}}
{{- $serverSide := (eq .StreamingMode "server")}}
{{- $bidiSide := (eq .StreamingMode "bidirectional")}}
{{- $arg := index .Args 0}}
{{.Name}}{{- if $streamingUnary}}(ctx context.Context, req {{$arg.Type}}) ({{.Resp.Type}}, error)
{{- else if $clientSide}}(ctx context.Context, stream streamx.ClientStreamingServer[{{NotPtr $arg.Type}}, {{NotPtr .Resp.Type}}]) ({{.Resp.Type}}, error)
Expand Down
22 changes: 11 additions & 11 deletions tool/internal_pkg/tpl/service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,11 @@ var serviceMethods = map[string]kitex.MethodInfo{
{{- else if .ServerStreaming -}} kitex.StreamingServer
{{- else if .ClientStreaming -}} kitex.StreamingClient
{{- else -}}
{{- if or (eq $.Codec "protobuf") (eq .Streaming.Mode "unary") -}} kitex.StreamingUnary
{{- if or (eq $.Codec "protobuf") (eq .StreamingMode "unary") -}} kitex.StreamingUnary
{{- else -}} kitex.StreamingNone
{{- end -}}
{{- end}}),
{{- if and .StreamX .Streaming.IsStreaming}}
{{- if and .StreamX .IsStreaming}}
kitex.WithMethodExtra("streamx", "true"),
{{- end}}
),
Expand Down Expand Up @@ -146,7 +146,7 @@ func newServiceInfo(hasStreaming bool, keepStreamingMethods bool, keepNonStreami
{{range .AllMethods}}
{{- $isStreaming := or .ClientStreaming .ServerStreaming}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $unary := and (not .ServerStreaming) (not .ClientStreaming)}}
{{- $clientSide := and .ClientStreaming (not .ServerStreaming)}}
{{- $serverSide := and (not .ClientStreaming) .ServerStreaming}}
Expand All @@ -159,7 +159,7 @@ func newServiceInfo(hasStreaming bool, keepStreamingMethods bool, keepNonStreami
{{- else if $clientSide -}} {{- $mode = "serviceinfo.StreamingClient" }} {{- $handlerFunc = "InvokeClientStreamHandler" }}
{{- else if $bidiSide -}} {{- $mode = "serviceinfo.StreamingBidirectional" }} {{- $handlerFunc = "InvokeBidiStreamHandler" }}
{{- end}}
{{- if or (eq $.Codec "protobuf") ($isStreaming) (.Streaming.IsStreaming) }}
{{- if or (eq $.Codec "protobuf") ($isStreaming) (.IsStreaming) }}
{{- $arg = index .Args 0}}{{/* streaming api only supports exactly one argument */}}
{{- end}}
Expand Down Expand Up @@ -205,7 +205,7 @@ func {{LowerFirst .Name}}Handler(ctx context.Context, handler interface{}, arg,
return handler.({{.PkgRefName}}.{{.ServiceName}}).{{.Name}}({{if $serverSide}}req, {{end}}stream)
{{- end}} {{/* $unary end */}}
{{- else}} {{/* thrift logic */}}
{{- if and .StreamX .Streaming.IsStreaming}}
{{- if and .StreamX .IsStreaming}}
return streamxserver.{{$handlerFunc}}[{{NotPtr $arg.Type}}, {{NotPtr .Resp.Type}}](
ctx, arg.(streamx.StreamReqArgs), result.(streamx.StreamResArgs),
{{- if $streamingUnary }}func(ctx context.Context, req {{$arg.Type}}) ({{.Resp.Type}}, error) {
Expand Down Expand Up @@ -488,11 +488,11 @@ func newServiceClient(c client.Client) *kClient {
{{range .AllMethods}}
{{- /* streaming logic */}}
{{- if and .StreamX .Streaming.IsStreaming}}
{{- $streamingUnary := (eq .Streaming.Mode "unary")}}
{{- $clientSide := (eq .Streaming.Mode "client")}}
{{- $serverSide := (eq .Streaming.Mode "server")}}
{{- $bidiSide := (eq .Streaming.Mode "bidirectional")}}
{{- if and .StreamX .IsStreaming}}
{{- $streamingUnary := (eq .StreamingMode "unary")}}
{{- $clientSide := (eq .StreamingMode "client")}}
{{- $serverSide := (eq .StreamingMode "server")}}
{{- $bidiSide := (eq .StreamingMode "bidirectional")}}
{{- $mode := ""}}
{{- if $bidiSide -}} {{- $mode = "kitex.StreamingBidirectional" }}
{{- else if $serverSide -}} {{- $mode = "kitex.StreamingServer" }}
Expand Down Expand Up @@ -595,7 +595,7 @@ func (p *kClient) {{.Name}}(ctx context.Context {{range .Args}}, {{.RawName}} {{
{{end -}}
}
{{- end}}
{{- end}}{{/* if and .StreamX .Streaming.IsStreaming end */}}
{{- end}}{{/* if and .StreamX .IsStreaming end */}}
{{end}}
{{- if .FrugalPretouch}}
Expand Down

0 comments on commit b2a9665

Please sign in to comment.