improve stub generator code

clubpay · Jul 14, 2022 · f1c22ac · f1c22ac
1 parent 936a9fd
commit f1c22ac
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 29 deletions.
diff --git a/desc/export_test.go b/desc/export_test.go
@@ -7,6 +7,6 @@ var (
 	NewStub = newStub
 )
 
-func (d *Stub) AddDTO(mTyp reflect.Type) error {
-	return d.addDTO(mTyp)
+func (d *Stub) AddDTO(mTyp reflect.Type, isErr bool) error {
+	return d.addDTO(mTyp, isErr)
 }
diff --git a/desc/service.go b/desc/service.go
@@ -166,26 +166,26 @@ func (s Service) Stub(pkgName string, tags ...string) (*Stub, error) {
 
 func (s Service) dtoStub(stub *Stub) error {
 	for _, c := range s.Contracts {
-		err := stub.addDTO(reflect.TypeOf(c.Input))
+		err := stub.addDTO(reflect.TypeOf(c.Input), false)
 		if err != nil {
 			return err
 		}
 		if c.Output != nil {
-			err = stub.addDTO(reflect.TypeOf(c.Output))
+			err = stub.addDTO(reflect.TypeOf(c.Output), false)
 			if err != nil {
 				return err
 			}
 		}
 
 		for _, pe := range s.PossibleErrors {
-			err = stub.addDTO(reflect.TypeOf(pe.Message))
+			err = stub.addDTO(reflect.TypeOf(pe.Message), true)
 			if err != nil {
 				return err
 			}
 		}
 
 		for _, pe := range c.PossibleErrors {
-			err = stub.addDTO(reflect.TypeOf(pe.Message))
+			err = stub.addDTO(reflect.TypeOf(pe.Message), true)
 			if err != nil {
 				return err
 			}
@@ -207,17 +207,18 @@ func (s Service) rpcStub(
 	if dto, ok := stub.getDTO(reflect.TypeOf(c.Input)); ok {
 		m.Request = dto
 	}
-	if dto, ok := stub.getDTO(reflect.TypeOf(c.Output)); ok {
-		m.Response = dto
+	if c.Output != nil {
+		if dto, ok := stub.getDTO(reflect.TypeOf(c.Output)); ok {
+			m.Response = dto
+		}
 	}
 
 	var possibleErrors []Error
 	possibleErrors = append(possibleErrors, s.PossibleErrors...)
 	possibleErrors = append(possibleErrors, c.PossibleErrors...)
 	for _, e := range possibleErrors {
 		if dto, ok := stub.getDTO(reflect.TypeOf(e.Message)); ok {
-			m.PossibleErrors = append(
-				m.PossibleErrors,
+			m.addPossibleError(
 				ErrorDTO{
 					Code: e.Code,
 					Item: e.Item,
@@ -241,17 +242,18 @@ func (s Service) restStub(
 	if dto, ok := stub.getDTO(reflect.TypeOf(c.Input)); ok {
 		m.Request = dto
 	}
-	if dto, ok := stub.getDTO(reflect.TypeOf(c.Output)); ok {
-		m.Response = dto
+	if c.Output != nil {
+		if dto, ok := stub.getDTO(reflect.TypeOf(c.Output)); ok {
+			m.Response = dto
+		}
 	}
 
 	var possibleErrors []Error
 	possibleErrors = append(possibleErrors, s.PossibleErrors...)
 	possibleErrors = append(possibleErrors, c.PossibleErrors...)
 	for _, e := range possibleErrors {
 		if dto, ok := stub.getDTO(reflect.TypeOf(e.Message)); ok {
-			m.PossibleErrors = append(
-				m.PossibleErrors,
+			m.addPossibleError(
 				ErrorDTO{
 					Code: e.Code,
 					Item: e.Item,

diff --git a/desc/stub.go b/desc/stub.go
@@ -3,6 +3,7 @@ package desc
 import (
 	"fmt"
 	"reflect"
+	"strings"
 
 	"github.com/clubpay/ronykit"
 )
@@ -12,9 +13,46 @@ type DTO struct {
 	Comments []string
 	Name     string
 	Type     string
+	IsErr    bool
 	Fields   []DTOField
 }
 
+func (dto DTO) CodeField() string {
+	var fn string
+	for _, f := range dto.Fields {
+		x := strings.ToLower(f.Name)
+		if f.Type != "int" {
+			continue
+		}
+		if x == "code" {
+			return f.Name
+		}
+		if strings.HasPrefix(f.Name, "code") {
+			fn = f.Name
+		}
+	}
+
+	return fn
+}
+
+func (dto DTO) ItemField() string {
+	var fn string
+	for _, f := range dto.Fields {
+		x := strings.ToLower(f.Name)
+		if f.Type != "string" {
+			continue
+		}
+		if x == "item" || x == "items" {
+			return f.Name
+		}
+		if strings.HasPrefix(f.Name, "item") {
+			fn = f.Name
+		}
+	}
+
+	return fn
+}
+
 // DTOField represents description of a field of the DTO
 type DTOField struct {
 	Name     string
@@ -49,6 +87,15 @@ type RESTMethod struct {
 	PossibleErrors []ErrorDTO
 }
 
+func (rm *RESTMethod) addPossibleError(dto ErrorDTO) {
+	for _, e := range rm.PossibleErrors {
+		if e.Code == dto.Code {
+			return
+		}
+	}
+	rm.PossibleErrors = append(rm.PossibleErrors, dto)
+}
+
 // RPCMethod represents description of a Contract with ronykit.RPCRouteSelector
 type RPCMethod struct {
 	Name           string
@@ -61,6 +108,15 @@ type RPCMethod struct {
 	ronykit.OutgoingRPCContainer
 }
 
+func (rm *RPCMethod) addPossibleError(dto ErrorDTO) {
+	for _, e := range rm.PossibleErrors {
+		if e.Code == dto.Code {
+			return
+		}
+	}
+	rm.PossibleErrors = append(rm.PossibleErrors, dto)
+}
+
 // Stub represents description of a stub of the service described by Service descriptor.
 type Stub struct {
 	tags  []string
@@ -78,8 +134,10 @@ func newStub(tags ...string) *Stub {
 	}
 }
 
-func (d *Stub) addDTO(mTyp reflect.Type) error {
-	dto := DTO{}
+func (d *Stub) addDTO(mTyp reflect.Type, isErr bool) error {
+	dto := DTO{
+		IsErr: isErr,
+	}
 	if mTyp.Kind() == reflect.Ptr {
 		mTyp = mTyp.Elem()
 	}
@@ -93,12 +151,12 @@ func (d *Stub) addDTO(mTyp reflect.Type) error {
 
 			switch {
 			case k == reflect.Struct:
-				err := d.addDTO(ft.Type)
+				err := d.addDTO(ft.Type, false)
 				if err != nil {
 					return err
 				}
 			case k == reflect.Ptr && ft.Type.Elem().Kind() == reflect.Struct:
-				err := d.addDTO(ft.Type.Elem())
+				err := d.addDTO(ft.Type.Elem(), false)
 				if err != nil {
 					return err
 				}

diff --git a/desc/stub_test.go b/desc/stub_test.go
@@ -12,7 +12,7 @@ var _ = Describe("Desc", func() {
 	It("should detect all DTOs", func() {
 		d := desc.NewStub("json")
 
-		Expect(d.AddDTO(reflect.TypeOf(&customStruct{}))).To(Succeed())
+		Expect(d.AddDTO(reflect.TypeOf(&customStruct{}), false)).To(Succeed())
 		Expect(d.DTOs).To(HaveLen(2))
 		Expect(d.DTOs["customSubStruct"].Fields).To(HaveLen(2))
 		Expect(d.DTOs["customSubStruct"].Fields[0].Name).To(Equal("SubParam1"))

diff --git a/exmples/simple-rest-server/stub/sampleservice.go b/exmples/simple-rest-server/stub/sampleservice.go
diff --git a/internal/tpl/go/stub.gotmpl b/internal/tpl/go/stub.gotmpl
@@ -16,6 +16,18 @@
 		{{- end }}
 		{{- end }}
 	}
+	{{ if .IsErr }}
+	{{- if ne .CodeField ""}}
+	func (x {{.Name}}) GetCode() int {
+		return x.{{.CodeField}}
+	}
+	{{- end }}
+	{{- if ne .ItemField ""}}
+	func (x {{.Name}}) GetItem() string {
+		return x.{{.ItemField}}
+	}
+	{{- end }}
+	{{- end }}
 {{ end }}
 // Code generated by RonyKIT Stub Generator (Golang); DO NOT EDIT.
 
@@ -75,7 +87,7 @@ func (s {{$serviceName}}Stub) {{$methodName}}(ctx context.Context, req *{{.Reque
 					return err
 				}
 
-				return stub.NewErrorWithMsg({{$errDto.Code}}, "{{$errDto.Item}}", res)
+				return stub.NewErrorWithMsg(res)
 			},
 		).
 		{{- end }}

diff --git a/stub/error.go b/stub/error.go
@@ -20,12 +20,18 @@ func NewError(code int, item string) *Error {
 	}
 }
 
-func NewErrorWithMsg(code int, item string, msg ronykit.Message) *Error {
-	return &Error{
-		code: code,
-		item: item,
-		msg:  msg,
+func NewErrorWithMsg(msg ronykit.Message) *Error {
+	wErr := &Error{
+		msg: msg,
 	}
+	if e, ok := msg.(interface{ GetCode() int }); ok {
+		wErr.code = e.GetCode()
+	}
+	if e, ok := msg.(interface{ GetItem() string }); ok {
+		wErr.item = e.GetItem()
+	}
+
+	return wErr
 }
 
 func WrapError(err error) *Error {