vitess.io/vitess@v0.16.2/go/vt/vtctl/vtctlclient/wrapper.go (about)

     1  /*
     2  Copyright 2019 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 vtctlclient
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"io"
    23  	"time"
    24  
    25  	"context"
    26  
    27  	logutilpb "vitess.io/vitess/go/vt/proto/logutil"
    28  )
    29  
    30  var (
    31  	defaultTimeout = time.Hour
    32  )
    33  
    34  // RunCommandAndWait executes a single command on a given vtctld and blocks until the command did return or timed out.
    35  // Output from vtctld is streamed as logutilpb.Event messages which
    36  // have to be consumed by the caller who has to specify a "recv" function.
    37  func RunCommandAndWait(ctx context.Context, server string, args []string, recv func(*logutilpb.Event)) error {
    38  	if recv == nil {
    39  		return errors.New("no function closure for Event stream specified")
    40  	}
    41  	// create the client
    42  	client, err := New(server)
    43  	if err != nil {
    44  		return fmt.Errorf("cannot dial to server %v: %v", server, err)
    45  	}
    46  	defer client.Close()
    47  
    48  	// run the command ( get the timeout from the context )
    49  	timeout := defaultTimeout
    50  	deadline, ok := ctx.Deadline()
    51  	if ok {
    52  		timeout = time.Until(deadline)
    53  	}
    54  	stream, err := client.ExecuteVtctlCommand(ctx, args, timeout)
    55  	if err != nil {
    56  		return fmt.Errorf("cannot execute remote command: %v", err)
    57  	}
    58  
    59  	// stream the result
    60  	for {
    61  		e, err := stream.Recv()
    62  		switch err {
    63  		case nil:
    64  			recv(e)
    65  		case io.EOF:
    66  			return nil
    67  		default:
    68  			return fmt.Errorf("remote error: %v", err)
    69  		}
    70  	}
    71  }