vitess.io/vitess@v0.16.2/go/cmd/vtctldclient/command/root.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 command
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"fmt"
    23  	"io"
    24  	"strconv"
    25  	"time"
    26  
    27  	"github.com/spf13/cobra"
    28  
    29  	"vitess.io/vitess/go/trace"
    30  	"vitess.io/vitess/go/vt/logutil"
    31  	"vitess.io/vitess/go/vt/servenv"
    32  	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
    33  )
    34  
    35  var (
    36  	// VtctldClientProtocol is the protocol to use when creating the vtctldclient.VtctldClient.
    37  	VtctldClientProtocol = "grpc"
    38  
    39  	client        vtctldclient.VtctldClient
    40  	traceCloser   io.Closer
    41  	commandCtx    context.Context
    42  	commandCancel func()
    43  
    44  	server        string
    45  	actionTimeout time.Duration
    46  
    47  	// Root is the main entrypoint to the vtctldclient CLI.
    48  	Root = &cobra.Command{
    49  		Use:   "vtctldclient",
    50  		Short: "Executes a cluster management command on the remote vtctld server.",
    51  		// We use PersistentPreRun to set up the tracer, grpc client, and
    52  		// command context for every command.
    53  		PersistentPreRunE: func(cmd *cobra.Command, args []string) (err error) {
    54  			logutil.PurgeLogs()
    55  			traceCloser = trace.StartTracing("vtctldclient")
    56  			client, err = getClientForCommand(cmd)
    57  			ctx := cmd.Context()
    58  			if ctx == nil {
    59  				ctx = context.Background()
    60  			}
    61  			commandCtx, commandCancel = context.WithTimeout(ctx, actionTimeout)
    62  			return err
    63  		},
    64  		// Similarly, PersistentPostRun cleans up the resources spawned by
    65  		// PersistentPreRun.
    66  		PersistentPostRunE: func(cmd *cobra.Command, args []string) (err error) {
    67  			commandCancel()
    68  			if client != nil {
    69  				err = client.Close()
    70  			}
    71  			trace.LogErrorsWhenClosing(traceCloser)
    72  			return err
    73  		},
    74  		TraverseChildren: true,
    75  		// By default, cobra will print any error returned by a child command to
    76  		// stderr, and then return that error back up the call chain. Since we
    77  		// use vitess's log package to log any error we get back from
    78  		// Root.Execute() (in ../../main.go) this actually results in duplicate
    79  		// stderr lines. So, somewhat counterintuitively, we actually "silence"
    80  		// all errors in cobra (just from being output, they still get
    81  		// propagated).
    82  		SilenceErrors: true,
    83  		Version:       servenv.AppVersion.String(),
    84  		// If we've reached this function, it means that:
    85  		//
    86  		// (1) The user specified some positional arguments, which, for the way
    87  		// we've structured things can only be a subcommand name, **and**
    88  		//
    89  		// (2) Cobra was unable to find a subcommand with that name for which to
    90  		// call a Run or RunE function.
    91  		//
    92  		// From this we conclude that the user was trying to either run a
    93  		// command that doesn't exist (e.g. "vtctldclient delete-my-data") or
    94  		// has misspelled a legitimate command (e.g. "vtctldclient StapReplication").
    95  		// If we think this has happened, return an error, which will get
    96  		// displayed to the user in main.go along with the usage.
    97  		RunE: func(cmd *cobra.Command, args []string) error {
    98  			if cmd.Flags().NArg() > 0 {
    99  				return fmt.Errorf("unknown command: %s", cmd.Flags().Arg(0))
   100  			}
   101  
   102  			return nil
   103  		},
   104  	}
   105  )
   106  
   107  var errNoServer = errors.New("please specify --server <vtctld_host:vtctld_port> to specify the vtctld server to connect to")
   108  
   109  const skipClientCreationKey = "skip_client_creation"
   110  
   111  // getClientForCommand returns a vtctldclient.VtctldClient for a given command.
   112  // It validates that --server was passed to the CLI for commands that need it.
   113  func getClientForCommand(cmd *cobra.Command) (vtctldclient.VtctldClient, error) {
   114  	if skipStr, ok := cmd.Annotations[skipClientCreationKey]; ok {
   115  		skipClientCreation, err := strconv.ParseBool(skipStr)
   116  		if err != nil {
   117  			skipClientCreation = false
   118  		}
   119  
   120  		if skipClientCreation {
   121  			return nil, nil
   122  		}
   123  	}
   124  
   125  	if VtctldClientProtocol != "local" && server == "" {
   126  		return nil, errNoServer
   127  	}
   128  
   129  	return vtctldclient.New(VtctldClientProtocol, server)
   130  }
   131  
   132  func init() {
   133  	Root.PersistentFlags().StringVar(&server, "server", "", "server to use for connection (required)")
   134  	Root.PersistentFlags().DurationVar(&actionTimeout, "action_timeout", time.Hour, "timeout for the total command")
   135  }