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))