github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/cli/cli.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package cli
    12  
    13  import (
    14  	"os"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/workload"
    17  	"github.com/spf13/cobra"
    18  )
    19  
    20  // WorkloadCmd returns a new command that can serve as the root of the workload
    21  // command tree.
    22  func WorkloadCmd(userFacing bool) *cobra.Command {
    23  	rootCmd := SetCmdDefaults(&cobra.Command{
    24  		Use:   `workload`,
    25  		Short: `generators for data and query loads`,
    26  	})
    27  	for _, subCmdFn := range subCmdFns {
    28  		rootCmd.AddCommand(subCmdFn(userFacing))
    29  	}
    30  	if userFacing {
    31  		whitelist := map[string]struct{}{
    32  			`workload`: {},
    33  			`init`:     {},
    34  			`run`:      {},
    35  		}
    36  		for _, m := range workload.Registered() {
    37  			whitelist[m.Name] = struct{}{}
    38  		}
    39  		var addExperimental func(c *cobra.Command)
    40  		addExperimental = func(c *cobra.Command) {
    41  			c.Short = `[experimental] ` + c.Short
    42  			if _, ok := whitelist[c.Name()]; !ok {
    43  				c.Hidden = true
    44  			}
    45  			for _, sub := range c.Commands() {
    46  				addExperimental(sub)
    47  			}
    48  		}
    49  		addExperimental(rootCmd)
    50  	}
    51  	return rootCmd
    52  }
    53  
    54  var subCmdFns []func(userFacing bool) *cobra.Command
    55  
    56  // AddSubCmd adds a sub-command closure to the workload cli.
    57  //
    58  // Most commands are global singletons and hooked together in init functions but
    59  // the workload ones do fancy things with making each generator in
    60  // `workload.Registered` into a subcommand. This means the commands need to be
    61  // generated after `workload.Registered` is fully populated, but the latter is
    62  // now sometimes done in the outermost possible place (main.go) and so its init
    63  // runs last. Instead, we return a closure that is called in the main function,
    64  // which is guaranteed to run after all inits.
    65  func AddSubCmd(fn func(userFacing bool) *cobra.Command) {
    66  	subCmdFns = append(subCmdFns, fn)
    67  }
    68  
    69  // HandleErrs wraps a `RunE` cobra function to print an error but not the usage.
    70  func HandleErrs(
    71  	f func(cmd *cobra.Command, args []string) error,
    72  ) func(cmd *cobra.Command, args []string) {
    73  	return func(cmd *cobra.Command, args []string) {
    74  		err := f(cmd, args)
    75  		if err != nil {
    76  			cmd.Println("Error:", err.Error())
    77  			os.Exit(1)
    78  		}
    79  	}
    80  }