vitess.io/vitess@v0.16.2/go/vt/vtctl/localvtctldclient/client.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 localvtctldclient
    18  
    19  import (
    20  	"errors"
    21  	"sync"
    22  
    23  	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
    24  
    25  	vtctlservicepb "vitess.io/vitess/go/vt/proto/vtctlservice"
    26  )
    27  
    28  var (
    29  	m      sync.Mutex
    30  	server vtctlservicepb.VtctldServer
    31  )
    32  
    33  type localVtctldClient struct {
    34  	s vtctlservicepb.VtctldServer
    35  }
    36  
    37  // Close is part of the vtctldclient.VtctldClient interface.
    38  func (client *localVtctldClient) Close() error { return nil }
    39  
    40  //go:generate -command localvtctldclient go run ../vtctldclient/codegen
    41  //go:generate localvtctldclient --targetpkg localvtctldclient --impl localVtctldClient --out client_gen.go --local
    42  
    43  // New returns a local vtctldclient.VtctldClient that makes method calls on the
    44  // provided VtctldServer implementation. No network traffic takes place between
    45  // the client and server, and all CallOptions are ignored on RPCs.
    46  func New(s vtctlservicepb.VtctldServer) vtctldclient.VtctldClient {
    47  	return &localVtctldClient{s: s}
    48  }
    49  
    50  // SetServer sets the server implementation used when creating local clients
    51  // via the vtctldclient.Factory.
    52  //
    53  // This function must be called before calling vtctldclient.New.
    54  func SetServer(s vtctlservicepb.VtctldServer) {
    55  	m.Lock()
    56  	defer m.Unlock()
    57  
    58  	server = s
    59  }
    60  
    61  func localVtctldClientFactory(addr string) (vtctldclient.VtctldClient, error) {
    62  	m.Lock()
    63  	defer m.Unlock()
    64  
    65  	if server == nil {
    66  		return nil, errors.New("cannot create local vtctldclient without a server; call SetServer first")
    67  	}
    68  
    69  	return New(server), nil
    70  }
    71  
    72  func init() {
    73  	vtctldclient.Register("local", localVtctldClientFactory)
    74  }