vitess.io/vitess@v0.16.2/go/vt/vtctl/vtctldclient/codegen/template.go (about) 1 /* 2 Copyright 2021 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package main 18 19 import ( 20 "strings" 21 "text/template" 22 ) 23 24 const tmplStr = `// Code generated by {{ .ClientName }}-generator. DO NOT EDIT. 25 26 /* 27 Copyright 2021 The Vitess Authors. 28 29 Licensed under the Apache License, Version 2.0 (the "License"); 30 you may not use this file except in compliance with the License. 31 You may obtain a copy of the License at 32 33 http://www.apache.org/licenses/LICENSE-2.0 34 35 Unless required by applicable law or agreed to in writing, software 36 distributed under the License is distributed on an "AS IS" BASIS, 37 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 38 See the License for the specific language governing permissions and 39 limitations under the License. 40 */ 41 42 package {{ .PackageName }} 43 44 import ( 45 "context" 46 47 "google.golang.org/grpc" 48 {{ if not .Local -}} 49 "google.golang.org/grpc/codes" 50 "google.golang.org/grpc/status" 51 {{- end }} 52 {{ if .NeedsGRPCShim -}} 53 "vitess.io/vitess/go/vt/vtctl/internal/grpcshim" 54 {{- end }} 55 56 {{ range .Imports -}} 57 {{ if ne .Alias "" }}{{ .Alias }} {{ end }}"{{ .Path }}" 58 {{ end -}} 59 ) 60 {{ range .Methods }} 61 {{ if and $.Local .IsStreaming -}} 62 type {{ streamAdapterName .Name }} struct { 63 *grpcshim.BidiStream 64 ch chan {{ .StreamMessage.Type }} 65 } 66 67 func (stream *{{ streamAdapterName .Name }}) Recv() ({{ .StreamMessage.Type }}, error) { 68 select { 69 case <-stream.Context().Done(): 70 return nil, stream.Context().Err() 71 case <-stream.Closed(): 72 // Stream has been closed for future sends. If there are messages that 73 // have already been sent, receive them until there are no more. After 74 // all sent messages have been received, Recv will return the CloseErr. 75 select { 76 case msg := <-stream.ch: 77 return msg, nil 78 default: 79 return nil, stream.CloseErr() 80 } 81 case err := <-stream.ErrCh: 82 return nil, err 83 case msg := <-stream.ch: 84 return msg, nil 85 } 86 } 87 88 func (stream *{{ streamAdapterName .Name }}) Send(msg {{ .StreamMessage.Type }}) error { 89 select { 90 case <-stream.Context().Done(): 91 return stream.Context().Err() 92 case <-stream.Closed(): 93 return grpcshim.ErrStreamClosed 94 case stream.ch <- msg: 95 return nil 96 } 97 } 98 {{ end -}} 99 // {{ .Name }} is part of the vtctlservicepb.VtctldClient interface. 100 func (client *{{ $.Type }}) {{ .Name }}(ctx context.Context, {{ .Param.Name }} {{ .Param.Type }}, opts ...grpc.CallOption) ({{ .Result.Type }}, error) { 101 {{ if not $.Local -}} 102 if client.c == nil { 103 return nil, status.Error(codes.Unavailable, connClosedMsg) 104 } 105 106 return client.c.{{ .Name }}(ctx, in, opts...) 107 {{- else -}} 108 {{- if .IsStreaming -}} 109 stream := &{{ streamAdapterName .Name }}{ 110 BidiStream: grpcshim.NewBidiStream(ctx), 111 ch: make(chan {{ .StreamMessage.Type }}, 1), 112 } 113 go func() { 114 err := client.s.{{ .Name }}(in, stream) 115 stream.CloseWithError(err) 116 }() 117 118 return stream, nil 119 {{- else -}} 120 return client.s.{{ .Name }}(ctx, in) 121 {{- end -}} 122 {{- end }} 123 } 124 {{ end }}` 125 126 var tmpl = template.Must(template.New("vtctldclient-generator").Funcs(map[string]any{ 127 "streamAdapterName": func(s string) string { 128 if len(s) == 0 { 129 return s 130 } 131 132 head := s[:1] 133 tail := s[1:] 134 return strings.ToLower(head) + tail + "StreamAdapter" 135 }, 136 }).Parse(tmplStr))