github.com/unionj-cloud/go-doudou/v2@v2.3.5/toolkit/protobuf/v3/service.go (about)

     1  package v3
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/iancoleman/strcase"
     6  	"github.com/unionj-cloud/go-doudou/v2/toolkit/astutils"
     7  	"github.com/unionj-cloud/go-doudou/v2/toolkit/constants"
     8  	"github.com/unionj-cloud/go-doudou/v2/version"
     9  	"reflect"
    10  	"strings"
    11  	"time"
    12  )
    13  
    14  type ProtoGenerator struct {
    15  	fieldNamingFunc func(string) string
    16  }
    17  
    18  type ProtoGeneratorOption func(*ProtoGenerator)
    19  
    20  func WithFieldNamingFunc(fn func(string) string) ProtoGeneratorOption {
    21  	return func(p *ProtoGenerator) {
    22  		p.fieldNamingFunc = fn
    23  	}
    24  }
    25  
    26  func NewProtoGenerator(options ...ProtoGeneratorOption) ProtoGenerator {
    27  	var p ProtoGenerator
    28  	for _, opt := range options {
    29  		opt(&p)
    30  	}
    31  	return p
    32  }
    33  
    34  const Syntax = "proto3"
    35  
    36  type Service struct {
    37  	Name      string
    38  	Package   string
    39  	GoPackage string
    40  	Syntax    string
    41  	// go-doudou version
    42  	Version  string
    43  	ProtoVer string
    44  	Rpcs     []Rpc
    45  	Messages []Message
    46  	Enums    []Enum
    47  	Comments []string
    48  	Imports  []string
    49  }
    50  
    51  func (receiver ProtoGenerator) NewService(name, goPackage string) Service {
    52  	return Service{
    53  		Name:      strcase.ToCamel(name) + "Service",
    54  		Package:   strcase.ToSnake(name),
    55  		GoPackage: goPackage,
    56  		Syntax:    Syntax,
    57  		Version:   version.Release,
    58  		ProtoVer:  fmt.Sprintf("v%s", time.Now().Local().Format(constants.FORMAT10)),
    59  	}
    60  }
    61  
    62  type StreamType int
    63  
    64  const (
    65  	biStream = iota + 1
    66  	clientStream
    67  	serverStream
    68  )
    69  
    70  type Rpc struct {
    71  	Name       string
    72  	Request    Message
    73  	Response   Message
    74  	Comments   []string
    75  	StreamType StreamType
    76  }
    77  
    78  func (receiver ProtoGenerator) NewRpc(method astutils.MethodMeta) Rpc {
    79  	rpcName := strcase.ToCamel(method.Name) + "Rpc"
    80  	rpcRequest := receiver.newRequest(rpcName, method.Params)
    81  	if reflect.DeepEqual(rpcRequest, Empty) {
    82  		ImportStore["google/protobuf/empty.proto"] = struct{}{}
    83  	}
    84  	if !strings.HasPrefix(rpcRequest.Name, "stream ") && !rpcRequest.IsImported {
    85  		if _, ok := MessageStore[rpcRequest.Name]; !ok {
    86  			MessageStore[rpcRequest.Name] = rpcRequest
    87  		}
    88  	}
    89  	rpcResponse := receiver.newResponse(rpcName, method.Results)
    90  	if reflect.DeepEqual(rpcResponse, Empty) {
    91  		ImportStore["google/protobuf/empty.proto"] = struct{}{}
    92  	}
    93  	if !strings.HasPrefix(rpcResponse.Name, "stream ") && !rpcResponse.IsImported {
    94  		if _, ok := MessageStore[rpcResponse.Name]; !ok {
    95  			MessageStore[rpcResponse.Name] = rpcResponse
    96  		}
    97  	}
    98  	var st StreamType
    99  	if strings.HasPrefix(rpcRequest.Name, "stream ") && strings.HasPrefix(rpcResponse.Name, "stream ") {
   100  		st = biStream
   101  	} else if strings.HasPrefix(rpcRequest.Name, "stream ") {
   102  		st = clientStream
   103  	} else if strings.HasPrefix(rpcResponse.Name, "stream ") {
   104  		st = serverStream
   105  	}
   106  	return Rpc{
   107  		Name:       rpcName,
   108  		Request:    rpcRequest,
   109  		Response:   rpcResponse,
   110  		Comments:   method.Comments,
   111  		StreamType: st,
   112  	}
   113  }
   114  
   115  func (receiver ProtoGenerator) newRequest(rpcName string, params []astutils.FieldMeta) Message {
   116  	if len(params) == 0 {
   117  		return Empty
   118  	}
   119  	if len(params) == 1 && params[0].Type == "context.Context" {
   120  		return Empty
   121  	}
   122  	if params[0].Type == "context.Context" {
   123  		params = params[1:]
   124  	}
   125  	if len(params) == 1 {
   126  		if m, ok := receiver.MessageOf(params[0].Type).(Message); ok && m.IsTopLevel {
   127  			if strings.HasPrefix(params[0].Name, "stream") {
   128  				m.Name = "stream " + m.Name
   129  			}
   130  			return m
   131  		}
   132  	}
   133  	var fields []Field
   134  	for i, field := range params {
   135  		fields = append(fields, receiver.newField(field, i+1))
   136  	}
   137  	return Message{
   138  		Name:       strcase.ToCamel(rpcName + "Request"),
   139  		Fields:     fields,
   140  		IsTopLevel: true,
   141  	}
   142  }
   143  
   144  func (receiver ProtoGenerator) newResponse(rpcName string, params []astutils.FieldMeta) Message {
   145  	if len(params) == 0 {
   146  		return Empty
   147  	}
   148  	if len(params) == 1 && params[0].Type == "error" {
   149  		return Empty
   150  	}
   151  	if params[len(params)-1].Type == "error" {
   152  		params = params[:len(params)-1]
   153  	}
   154  	if len(params) == 1 {
   155  		if m, ok := receiver.MessageOf(params[0].Type).(Message); ok && m.IsTopLevel {
   156  			if strings.HasPrefix(params[0].Name, "stream") {
   157  				m.Name = "stream " + m.Name
   158  			}
   159  			return m
   160  		}
   161  	}
   162  	var fields []Field
   163  	for i, field := range params {
   164  		fields = append(fields, receiver.newField(field, i+1))
   165  	}
   166  	return Message{
   167  		Name:       strcase.ToCamel(rpcName + "Response"),
   168  		Fields:     fields,
   169  		IsTopLevel: true,
   170  	}
   171  }