github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/cmd/internal/svc/codegen/grpcproto.go

package codegen

import (
	"bufio"
	"fmt"
	"github.com/sirupsen/logrus"
	"github.com/unionj-cloud/go-doudou/cmd/internal/astutils"
	v3 "github.com/unionj-cloud/go-doudou/cmd/internal/protobuf/v3"
	"github.com/unionj-cloud/go-doudou/toolkit/templateutils"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"text/template"
)

var protoTmpl = `/**
* Generated by go-doudou {{.Version}}.
* Don't edit!
*
* Version No.: {{.ProtoVer}}
*/
syntax = "proto3";

package {{.Package}};
option go_package = "{{.GoPackage}}";
{{ range $i := .Imports }}
import "{{$i}}";
{{- end }}
{{ range $e := .Enums }}
enum {{$e.Name}} {
  {{- range $f := $e.Fields }}
  {{$f.Name}} = {{$f.Number}};
  {{- end }}
}
{{- end }}

{{- range $m := .Messages }}
{{ Eval "Message" $m }}
{{- end }}

{{- define "Message"}}
{{- if .Comments }}
{{ toComment .Comments }}
{{- end }}
message {{.Name}} {
  {{- range $f := .Fields }}
  {{- if $f.Type.Inner}}
  {{ Eval "Message" $f.Type }}
  {{- end }}
  {{- if $f.Comments }}
  {{ toComment $f.Comments }}
  {{- end }}
  {{$f.Type}} {{$f.Name}} = {{$f.Number}}{{if $f.JsonName}} [json_name="{{$f.JsonName}}"]{{end}};
  {{- end }}
}
{{- end}}

service {{.Name}} {
  {{- range $r := .Rpcs }}
  {{- if $r.Comments }}
  {{ toComment $r.Comments }}
  {{- end }}
  rpc {{$r.Name}}({{$r.Request}}) returns ({{$r.Response}});
  {{- end}}
}
`

func toComment(comments []string) string {
	if len(comments) == 0 {
		return ""
	}
	var b strings.Builder
	for i := range comments {
		b.WriteString(fmt.Sprintf("// %s\n", comments[i]))
	}
	return strings.TrimSuffix(b.String(), "\n")
}

func GenGrpcProto(dir string, ic astutils.InterfaceCollector) (service v3.Service, protoFile string) {
	var (
		err       error
		svcname   string
		fi        os.FileInfo
		tpl       *template.Template
		f         *os.File
		modfile   string
		modf      *os.File
		modName   string
		firstLine string
		grpcDir   string
	)
	grpcDir = filepath.Join(dir, "transport/grpc")
	if err = os.MkdirAll(grpcDir, os.ModePerm); err != nil {
		panic(err)
	}
	svcname = ic.Interfaces[0].Name
	protoFile = filepath.Join(grpcDir, strings.ToLower(svcname)+".proto")
	if fi, err = os.Stat(protoFile); err != nil && !os.IsNotExist(err) {
		panic(err)
	}
	if fi != nil {
		logrus.Warningln("file " + protoFile + " will be overwritten")
	}
	if f, err = os.Create(protoFile); err != nil {
		panic(err)
	}
	defer f.Close()
	modfile = filepath.Join(dir, "go.mod")
	if modf, err = Open(modfile); err != nil {
		panic(err)
	}
	reader := bufio.NewReader(modf)
	firstLine, _ = reader.ReadString('\n')
	modName = strings.TrimSpace(strings.TrimPrefix(firstLine, "module"))

	service = v3.NewService(svcname, modName+"/transport/grpc")
	service.Comments = ic.Interfaces[0].Comments
	for _, method := range ic.Interfaces[0].Methods {
		service.Rpcs = append(service.Rpcs, v3.NewRpc(method))
	}
	for k := range v3.ImportStore {
		service.Imports = append(service.Imports, k)
	}
	sort.Strings(service.Imports)
	for _, v := range v3.MessageStore {
		service.Messages = append(service.Messages, v)
	}
	sort.SliceStable(service.Messages, func(i, j int) bool {
		return service.Messages[i].Name < service.Messages[j].Name
	})
	for _, v := range v3.EnumStore {
		service.Enums = append(service.Enums, v)
	}
	sort.SliceStable(service.Enums, func(i, j int) bool {
		return service.Enums[i].Name < service.Enums[j].Name
	})
	tpl = template.New("proto.tmpl")
	funcMap := make(map[string]interface{})
	funcMap["toComment"] = toComment
	funcMap["Eval"] = templateutils.Eval(tpl)
	if tpl, err = tpl.Funcs(funcMap).Parse(protoTmpl); err != nil {
		panic(err)
	}
	if err = tpl.Execute(f, service); err != nil {
		panic(err)
	}
	return
}

func messagesOf(vofile string) []v3.Message {
	fset := token.NewFileSet()
	root, err := parser.ParseFile(fset, vofile, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}
	sc := astutils.NewStructCollector(ExprStringP)
	ast.Walk(sc, root)
	structs := sc.DocFlatEmbed()
	var ret []v3.Message
	for _, item := range structs {
		ret = append(ret, v3.NewMessage(item))
	}
	return ret
}

func ParseVoGrpc(dir string) {
	var (
		err        error
		messages   []v3.Message
		allMethods map[string][]astutils.MethodMeta
		allConsts  map[string][]string
	)
	vodir := filepath.Join(dir, "vo")
	var files []string
	err = filepath.Walk(vodir, astutils.Visit(&files))
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		v3.MessageNames = append(v3.MessageNames, getSchemaNames(file)...)
	}
	allMethods = make(map[string][]astutils.MethodMeta)
	allConsts = make(map[string][]string)
	for _, file := range files {
		methods, consts := enumsOf(file)
		for k, v := range methods {
			allMethods[k] = append(allMethods[k], v...)
		}
		for k, v := range consts {
			allConsts[k] = append(allConsts[k], v...)
		}
	}
	for k, v := range allMethods {
		if astutils.IsEnum(v) {
			v3.EnumStore[k] = v3.NewEnum(astutils.EnumMeta{
				Name:   k,
				Values: allConsts[k],
			})
		}
	}
	for _, file := range files {
		messages = append(messages, messagesOf(file)...)
	}
	for _, item := range messages {
		v3.MessageStore[item.Name] = item
	}
}