github.com/mloves0824/enron/cmd/enron@v0.0.0-20230830012320-113bbf6be3c8/internal/proto/server/template.go (about)

     1  package server
     2  
     3  import (
     4  	"bytes"
     5  	"html/template"
     6  )
     7  
     8  //nolint:lll
     9  var serviceTemplate = `
    10  {{- /* delete empty line */ -}}
    11  package service
    12  
    13  import (
    14  	{{- if .UseContext }}
    15  	"context"
    16  	{{- end }}
    17  	{{- if .UseIO }}
    18  	"io"
    19  	{{- end }}
    20  
    21  	pb "{{ .Package }}"
    22  	{{- if .GoogleEmpty }}
    23  	"google.golang.org/protobuf/types/known/emptypb"
    24  	{{- end }}
    25  )
    26  
    27  type {{ .Service }}Service struct {
    28  	pb.Unimplemented{{ .Service }}Server
    29  }
    30  
    31  func New{{ .Service }}Service() *{{ .Service }}Service {
    32  	return &{{ .Service }}Service{}
    33  }
    34  
    35  {{- $s1 := "google.protobuf.Empty" }}
    36  {{ range .Methods }}
    37  {{- if eq .Type 1 }}
    38  func (s *{{ .Service }}Service) {{ .Name }}(ctx context.Context, req {{ if eq .Request $s1 }}*emptypb.Empty{{ else }}*pb.{{ .Request }}{{ end }}) ({{ if eq .Reply $s1 }}*emptypb.Empty{{ else }}*pb.{{ .Reply }}{{ end }}, error) {
    39  	return {{ if eq .Reply $s1 }}&emptypb.Empty{}{{ else }}&pb.{{ .Reply }}{}{{ end }}, nil
    40  }
    41  
    42  {{- else if eq .Type 2 }}
    43  func (s *{{ .Service }}Service) {{ .Name }}(conn pb.{{ .Service }}_{{ .Name }}Server) error {
    44  	for {
    45  		req, err := conn.Recv()
    46  		if err == io.EOF {
    47  			return nil
    48  		}
    49  		if err != nil {
    50  			return err
    51  		}
    52  		
    53  		err = conn.Send(&pb.{{ .Reply }}{})
    54  		if err != nil {
    55  			return err
    56  		}
    57  	}
    58  }
    59  
    60  {{- else if eq .Type 3 }}
    61  func (s *{{ .Service }}Service) {{ .Name }}(conn pb.{{ .Service }}_{{ .Name }}Server) error {
    62  	for {
    63  		req, err := conn.Recv()
    64  		if err == io.EOF {
    65  			return conn.SendAndClose(&pb.{{ .Reply }}{})
    66  		}
    67  		if err != nil {
    68  			return err
    69  		}
    70  	}
    71  }
    72  
    73  {{- else if eq .Type 4 }}
    74  func (s *{{ .Service }}Service) {{ .Name }}(req {{ if eq .Request $s1 }}*emptypb.Empty
    75  {{ else }}*pb.{{ .Request }}{{ end }}, conn pb.{{ .Service }}_{{ .Name }}Server) error {
    76  	for {
    77  		err := conn.Send(&pb.{{ .Reply }}{})
    78  		if err != nil {
    79  			return err
    80  		}
    81  	}
    82  }
    83  
    84  {{- end }}
    85  {{- end }}
    86  `
    87  
    88  type MethodType uint8
    89  
    90  const (
    91  	unaryType          MethodType = 1
    92  	twoWayStreamsType  MethodType = 2
    93  	requestStreamsType MethodType = 3
    94  	returnsStreamsType MethodType = 4
    95  )
    96  
    97  // Service is a proto service.
    98  type Service struct {
    99  	Package     string
   100  	Service     string
   101  	Methods     []*Method
   102  	GoogleEmpty bool
   103  
   104  	UseIO      bool
   105  	UseContext bool
   106  }
   107  
   108  // Method is a proto method.
   109  type Method struct {
   110  	Service string
   111  	Name    string
   112  	Request string
   113  	Reply   string
   114  
   115  	// type: unary or stream
   116  	Type MethodType
   117  }
   118  
   119  func (s *Service) execute() ([]byte, error) {
   120  	const empty = "google.protobuf.Empty"
   121  	buf := new(bytes.Buffer)
   122  	for _, method := range s.Methods {
   123  		if (method.Type == unaryType && (method.Request == empty || method.Reply == empty)) ||
   124  			(method.Type == returnsStreamsType && method.Request == empty) {
   125  			s.GoogleEmpty = true
   126  		}
   127  		if method.Type == twoWayStreamsType || method.Type == requestStreamsType {
   128  			s.UseIO = true
   129  		}
   130  		if method.Type == unaryType {
   131  			s.UseContext = true
   132  		}
   133  	}
   134  	tmpl, err := template.New("service").Parse(serviceTemplate)
   135  	if err != nil {
   136  		return nil, err
   137  	}
   138  	if err := tmpl.Execute(buf, s); err != nil {
   139  		return nil, err
   140  	}
   141  	return buf.Bytes(), nil
   142  }