Skip to content

Commit 0e99bc3

Browse files
Marina-SakaiMarina-Sakai
authored
feat(thrift): unwrap struct for streaming type descriptor (#81)
* feat: unwrap struct for streaming type descriptor * fix: only add IsWithoutWrapping method * fix: not need streaming mode
1 parent 800ccab commit 0e99bc3

File tree

4 files changed

+157
-83
lines changed

4 files changed

+157
-83
lines changed

‎go.sum‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,8 +186,6 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
186186
google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
187187
google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
188188
google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
189-
google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
190-
google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
191189
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
192190
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
193191
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

‎thrift/descriptor.go‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -331,13 +331,14 @@ func (f FieldDescriptor) DefaultValue() *DefaultValue {
331331

332332
// FunctionDescriptor idl function descriptor
333333
type FunctionDescriptor struct {
334-
oneway bool
335-
hasRequestBase bool
336-
request *TypeDescriptor
337-
response *TypeDescriptor
338-
name string
339-
endpoints []http.Endpoint
340-
annotations []parser.Annotation
334+
oneway bool
335+
hasRequestBase bool
336+
request *TypeDescriptor
337+
response *TypeDescriptor
338+
name string
339+
endpoints []http.Endpoint
340+
annotations []parser.Annotation
341+
isWithoutWrapping bool
341342
}
342343

343344
// Name returns the name of the function
@@ -377,6 +378,11 @@ func (f FunctionDescriptor) Annotations() []parser.Annotation {
377378
return f.annotations
378379
}
379380

381+
// IsWithoutWrapping returns if the request and response are not wrapped in struct
382+
func (f FunctionDescriptor) IsWithoutWrapping() bool {
383+
return f.isWithoutWrapping
384+
}
385+
380386
// ServiceDescriptor is the runtime descriptor of a service
381387
type ServiceDescriptor struct {
382388
name string

‎thrift/idl.go‎

Lines changed: 109 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import (
2727
"time"
2828
"unsafe"
2929

30+
"github.com/cloudwego/thriftgo/generator/golang/streaming"
31+
3032
"github.com/cloudwego/dynamicgo/http"
3133
"github.com/cloudwego/dynamicgo/internal/json"
3234
"github.com/cloudwego/dynamicgo/internal/rt"
@@ -371,104 +373,137 @@ func addFunction(ctx context.Context, fn *parser.Function, tree *parser.Thrift,
371373

372374
}
373375

376+
st, err := streaming.ParseStreaming(fn)
377+
if err != nil {
378+
return err
379+
}
380+
isStreaming := st.ClientStreaming || st.ServerStreaming
381+
374382
var hasRequestBase bool
375383
var req *TypeDescriptor
376384
var resp *TypeDescriptor
377385

378386
// parse request field
379387
if opts.ParseFunctionMode != meta.ParseResponseOnly {
380-
// WARN: only support single argument
381-
reqAst := fn.Arguments[0]
382-
req = &TypeDescriptor{
383-
typ: STRUCT,
384-
struc: &StructDescriptor{
385-
baseID: FieldID(math.MaxUint16),
386-
ids: util.FieldIDMap{},
387-
names: util.FieldNameMap{},
388-
requires: make(RequiresBitmap, 1),
389-
},
390-
}
391-
392-
reqType, err := parseType(ctx, reqAst.Type, tree, structsCache, 0, opts, nextAnns, Request)
388+
req, hasRequestBase, err = parseRequest(ctx, isStreaming, fn, tree, structsCache, nextAnns, opts)
393389
if err != nil {
394390
return err
395391
}
396-
if reqType.Type() == STRUCT {
397-
for _, f := range reqType.Struct().names.All() {
398-
x := (*FieldDescriptor)(f.Val)
399-
if x.isRequestBase {
400-
hasRequestBase = true
401-
break
402-
}
403-
}
404-
}
405-
reqField := &FieldDescriptor{
406-
name: reqAst.Name,
407-
id: FieldID(reqAst.ID),
408-
typ: reqType,
409-
}
410-
req.Struct().ids.Set(int32(reqAst.ID), unsafe.Pointer(reqField))
411-
req.Struct().names.Set(reqAst.Name, unsafe.Pointer(reqField))
412-
req.Struct().names.Build()
413392
}
414393

415394
// parse response filed
416395
if opts.ParseFunctionMode != meta.ParseRequestOnly {
417-
respAst := fn.FunctionType
418-
resp = &TypeDescriptor{
419-
typ: STRUCT,
420-
struc: &StructDescriptor{
421-
baseID: FieldID(math.MaxUint16),
422-
ids: util.FieldIDMap{},
423-
names: util.FieldNameMap{},
424-
requires: make(RequiresBitmap, 1),
425-
},
426-
}
427-
respType, err := parseType(ctx, respAst, tree, structsCache, 0, opts, nextAnns, Response)
396+
resp, err = parseResponse(ctx, isStreaming, fn, tree, structsCache, nextAnns, opts)
428397
if err != nil {
429398
return err
430399
}
431-
respField := &FieldDescriptor{
432-
typ: respType,
433-
}
434-
resp.Struct().ids.Set(0, unsafe.Pointer(respField))
435-
// response has no name or id
436-
resp.Struct().names.Set("", unsafe.Pointer(respField))
437-
438-
// parse exceptions
439-
if len(fn.Throws) > 0 {
440-
// only support single exception
441-
exp := fn.Throws[0]
442-
exceptionType, err := parseType(ctx, exp.Type, tree, structsCache, 0, opts, nextAnns, Exception)
443-
if err != nil {
444-
return err
445-
}
446-
exceptionField := &FieldDescriptor{
447-
name: exp.Name,
448-
alias: exp.Name,
449-
id: FieldID(exp.ID),
450-
// isException: true,
451-
typ: exceptionType,
452-
}
453-
resp.Struct().ids.Set(int32(exp.ID), unsafe.Pointer(exceptionField))
454-
resp.Struct().names.Set(exp.Name, unsafe.Pointer(exceptionField))
455-
}
456-
resp.Struct().names.Build()
457400
}
458401

459402
fnDsc := &FunctionDescriptor{
460-
name: fn.Name,
461-
oneway: fn.Oneway,
462-
request: req,
463-
response: resp,
464-
hasRequestBase: hasRequestBase,
465-
endpoints: enpdoints,
466-
annotations: annos,
403+
name: fn.Name,
404+
oneway: fn.Oneway,
405+
request: req,
406+
response: resp,
407+
hasRequestBase: hasRequestBase,
408+
endpoints: enpdoints,
409+
annotations: annos,
410+
isWithoutWrapping: isStreaming,
467411
}
468412
sDsc.functions[fn.Name] = fnDsc
469413
return nil
470414
}
471415

416+
func parseRequest(ctx context.Context, isStreaming bool, fn *parser.Function, tree *parser.Thrift, structsCache compilingCache, nextAnns []parser.Annotation, opts Options) (req *TypeDescriptor, hasRequestBase bool, err error) {
417+
// WARN: only support single argument
418+
reqAst := fn.Arguments[0]
419+
reqType, err := parseType(ctx, reqAst.Type, tree, structsCache, 0, opts, nextAnns, Request)
420+
if err != nil {
421+
return nil, hasRequestBase, err
422+
}
423+
if reqType.Type() == STRUCT {
424+
for _, f := range reqType.Struct().names.All() {
425+
x := (*FieldDescriptor)(f.Val)
426+
if x.isRequestBase {
427+
hasRequestBase = true
428+
break
429+
}
430+
}
431+
}
432+
433+
if isStreaming {
434+
return reqType, hasRequestBase, nil
435+
}
436+
437+
// wrap with a struct
438+
wrappedTyDsc := &TypeDescriptor{
439+
typ: STRUCT,
440+
struc: &StructDescriptor{
441+
baseID: FieldID(math.MaxUint16),
442+
ids: util.FieldIDMap{},
443+
names: util.FieldNameMap{},
444+
requires: make(RequiresBitmap, 1),
445+
},
446+
}
447+
reqField := &FieldDescriptor{
448+
name: reqAst.Name,
449+
id: FieldID(reqAst.ID),
450+
typ: reqType,
451+
}
452+
wrappedTyDsc.Struct().ids.Set(int32(reqAst.ID), unsafe.Pointer(reqField))
453+
wrappedTyDsc.Struct().names.Set(reqAst.Name, unsafe.Pointer(reqField))
454+
wrappedTyDsc.Struct().names.Build()
455+
return wrappedTyDsc, hasRequestBase, nil
456+
}
457+
458+
func parseResponse(ctx context.Context, isStreaming bool, fn *parser.Function, tree *parser.Thrift, structsCache compilingCache, nextAnns []parser.Annotation, opts Options) (resp *TypeDescriptor, err error) {
459+
respAst := fn.FunctionType
460+
respType, err := parseType(ctx, respAst, tree, structsCache, 0, opts, nextAnns, Response)
461+
if err != nil {
462+
return nil, err
463+
}
464+
465+
if isStreaming {
466+
return respType, nil
467+
}
468+
469+
wrappedResp := &TypeDescriptor{
470+
typ: STRUCT,
471+
struc: &StructDescriptor{
472+
baseID: FieldID(math.MaxUint16),
473+
ids: util.FieldIDMap{},
474+
names: util.FieldNameMap{},
475+
requires: make(RequiresBitmap, 1),
476+
},
477+
}
478+
respField := &FieldDescriptor{
479+
typ: respType,
480+
}
481+
wrappedResp.Struct().ids.Set(0, unsafe.Pointer(respField))
482+
// response has no name or id
483+
wrappedResp.Struct().names.Set("", unsafe.Pointer(respField))
484+
485+
// parse exceptions
486+
if len(fn.Throws) > 0 {
487+
// only support single exception
488+
exp := fn.Throws[0]
489+
exceptionType, err := parseType(ctx, exp.Type, tree, structsCache, 0, opts, nextAnns, Exception)
490+
if err != nil {
491+
return nil, err
492+
}
493+
exceptionField := &FieldDescriptor{
494+
name: exp.Name,
495+
alias: exp.Name,
496+
id: FieldID(exp.ID),
497+
// isException: true,
498+
typ: exceptionType,
499+
}
500+
wrappedResp.Struct().ids.Set(int32(exp.ID), unsafe.Pointer(exceptionField))
501+
wrappedResp.Struct().names.Set(exp.Name, unsafe.Pointer(exceptionField))
502+
}
503+
wrappedResp.Struct().names.Build()
504+
return wrappedResp, nil
505+
}
506+
472507
// reuse builtin types
473508
var builtinTypes = map[string]*TypeDescriptor{
474509
"void": {name: "void", typ: VOID, struc: new(StructDescriptor)},

‎thrift/idl_test.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,3 +402,38 @@ func TestNewFunctionDescriptorFromPath(t *testing.T) {
402402
require.NotNil(t, p.Functions()["ExampleMethod"])
403403
require.Nil(t, p.Functions()["Ping"])
404404
}
405+
406+
func TestStreamingFunctionDescriptorFromContent(t *testing.T) {
407+
path := "a/b/main.thrift"
408+
content := `
409+
namespace go thrift
410+
411+
struct Request {
412+
1: required string message,
413+
}
414+
415+
struct Response {
416+
1: required string message,
417+
}
418+
419+
service TestService {
420+
Response Echo (1: Request req) (streaming.mode="bidirectional"),
421+
Response EchoClient (1: Request req) (streaming.mode="client"),
422+
Response EchoServer (1: Request req) (streaming.mode="server"),
423+
Response EchoUnary (1: Request req) (streaming.mode="unary"), // not recommended
424+
Response EchoBizException (1: Request req) (streaming.mode="client"),
425+
426+
Response EchoPingPong (1: Request req), // KitexThrift, non-streaming
427+
}
428+
`
429+
dsc, err := NewDescritorFromContent(context.Background(), path, content, nil, false)
430+
require.Nil(t, err)
431+
require.Equal(t, true, dsc.Functions()["Echo"].IsWithoutWrapping())
432+
require.Equal(t, true, dsc.Functions()["EchoClient"].IsWithoutWrapping())
433+
require.Equal(t, true, dsc.Functions()["EchoServer"].IsWithoutWrapping())
434+
require.Equal(t, false, dsc.Functions()["EchoUnary"].IsWithoutWrapping())
435+
require.Equal(t, true, dsc.Functions()["EchoBizException"].IsWithoutWrapping())
436+
require.Equal(t, false, dsc.Functions()["EchoPingPong"].IsWithoutWrapping())
437+
require.Equal(t, "Request", dsc.Functions()["EchoClient"].Request().Struct().Name())
438+
require.Equal(t, "", dsc.Functions()["EchoUnary"].Request().Struct().Name())
439+
}

0 commit comments

Comments
 (0)