ClickHouse/programs/diagnostics/cmd/root.go

package cmd

import (
	"fmt"
	"github.com/ClickHouse/clickhouse-diagnostics/internal/platform/utils"
	"github.com/pkg/errors"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"net/http"
	_ "net/http/pprof"
	"os"
	"strings"
	"time"
)

func enableDebug() {
	if debug {
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
		go func() {
			err := http.ListenAndServe("localhost:8080", nil)
			if err != nil {
				log.Error().Err(err).Msg("unable to start debugger")
			} else {
				log.Debug().Msg("debugger has been started on port 8080")
			}
		}()
	}
}

var rootCmd = &cobra.Command{
	Use:   "clickhouse-diagnostics",
	Short: "Capture and convert ClickHouse diagnostic bundles.",
	Long:  `Captures ClickHouse diagnostic bundles to a number of supported formats, including file and ClickHouse itself. Converts bundles between formats.`,
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		enableDebug()
		err := initializeConfig()
		if err != nil {
			log.Error().Err(err)
			os.Exit(1)
		}
		return nil
	},
	Example: `clickhouse-diagnostics collect`,
}

const (
	colorRed = iota + 31
	colorGreen
	colorYellow
	colorMagenta = 35

	colorBold = 1
)

const TimeFormat = time.RFC3339

var debug bool
var configFiles []string

const (
	// The environment variable prefix of all environment variables bound to our command line flags.
	// For example, --output is bound to CLICKHOUSE_DIAGNOSTIC_OUTPUT.
	envPrefix = "CLICKHOUSE_DIAGNOSTIC"
)

func init() {
	rootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", false, "Enable debug mode")
	rootCmd.PersistentFlags().StringSliceVarP(&configFiles, "config", "f", []string{"clickhouse-diagnostics.yml", "/etc/clickhouse-diagnostics.yml"}, "Configuration file path")
	// set a usage template to ensure flags on root are listed as global
	rootCmd.SetUsageTemplate(`Usage:{{if .Runnable}}
  {{.UseLine}}{{end}}{{if .HasAvailableSubCommands}}
  {{.CommandPath}} [command]{{end}}{{if gt (len .Aliases) 0}}

Aliases:
  {{.NameAndAliases}}{{end}}{{if .HasExample}}

Examples:
{{.Example}}{{end}}{{if .HasAvailableSubCommands}}

Available Commands:{{range .Commands}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}

Global Flags:
{{.LocalFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .HasAvailableInheritedFlags}}

Additional help topics:{{range .Commands}}{{if .IsAdditionalHelpTopicCommand}}
  {{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableSubCommands}}

Use "{{.CommandPath}} [command] --help" for more information about a command.{{end}}
`)
	rootCmd.SetFlagErrorFunc(handleFlagErrors)

}

func Execute() {
	// logs go to stderr - stdout is exclusive for outputs e.g. tables
	output := zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: TimeFormat}
	// override the colors
	output.FormatLevel = func(i interface{}) string {
		var l string
		if ll, ok := i.(string); ok {
			switch ll {
			case zerolog.LevelTraceValue:
				l = colorize("TRC", colorMagenta)
			case zerolog.LevelDebugValue:
				l = colorize("DBG", colorMagenta)
			case zerolog.LevelInfoValue:
				l = colorize("INF", colorGreen)
			case zerolog.LevelWarnValue:
				l = colorize(colorize("WRN", colorYellow), colorBold)
			case zerolog.LevelErrorValue:
				l = colorize(colorize("ERR", colorRed), colorBold)
			case zerolog.LevelFatalValue:
				l = colorize(colorize("FTL", colorRed), colorBold)
			case zerolog.LevelPanicValue:
				l = colorize(colorize("PNC", colorRed), colorBold)
			default:
				l = colorize("???", colorBold)
			}
		} else {
			if i == nil {
				l = colorize("???", colorBold)
			} else {
				l = strings.ToUpper(fmt.Sprintf("%s", i))[0:3]
			}
		}
		return l
	}
	output.FormatTimestamp = func(i interface{}) string {
		tt := i.(string)
		return colorize(tt, colorGreen)
	}
	log.Logger = log.Output(output)
	zerolog.SetGlobalLevel(zerolog.InfoLevel)
	rootCmd.SetHelpCommand(helpCmd)
	if err := rootCmd.Execute(); err != nil {
		log.Fatal().Err(err)
	}
}

// colorize returns the string s wrapped in ANSI code c
func colorize(s interface{}, c int) string {
	return fmt.Sprintf("\x1b[%dm%v\x1b[0m", c, s)
}

func handleFlagErrors(cmd *cobra.Command, err error) error {
	fmt.Println(colorize(colorize(fmt.Sprintf("Error: %s\n", err), colorRed), colorBold))
	_ = cmd.Help()
	os.Exit(1)
	return nil
}

func initializeConfig() error {
	// we use the first config file we find
	var configFile string
	for _, confFile := range configFiles {
		if ok, _ := utils.FileExists(confFile); ok {
			configFile = confFile
			break
		}
	}
	if configFile == "" {
		log.Warn().Msgf("config file in %s not found - config file will be ignored", configFiles)
		return nil
	}
	viper.SetConfigFile(configFile)
	if err := viper.ReadInConfig(); err != nil {
		return errors.Wrapf(err, "Unable to read configuration file at %s", configFile)
	}
	return nil
}