Migration to Cobra/Viper (#300)

* Cobra-preferred organization of the repo. Moved current cli functionality to cmd, made it to fit cobra ecosystem

* Added flag and env var support along with binding it to viper. Unmarshal from viper to config struct. Populate flags with defaults and deprecation handling. Started oath struct refactoring

* Added skip to nested oauth structures. Removed cli.go

* Added support for nested structures

* Resolve merge conflicts

* Fixed tests. Fixed env variables not getting loaded. Removed unnecessary comments

* Enabled both single and double dash flags. Fixed errors printed twice,usage always printed and logging.

* New command structure. Removed init functions

* Added access limiting for domains-google, groups-gitlab and organizations-github. Changes to config and oauth object structures

* Moved login days and JWT secret under Auth

* Added jwt secret auto-generation. Refactoring for flags.

* logging, help and usage changeS

* Separated viper instances and combined repeating parts into one function. Added support for bytes.ByteSize and time.Duration. Fixed tests.

* small fixes

* removes logrus dependency from profiler package

* fixes generate-sample-config script

Co-authored-by: Dmitry Filimonov <dmitry@pyroscope.io>

cmd/pyroscope/command/agent.go

@@ -0,0 +1,21 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/spf13/cobra"
+)
+
+func newAgentCmd(cfg *config.Agent) *cobra.Command {
+	vpr := newViper()
+	agentCmd := &cobra.Command{
+		Use:   "agent [flags]",
+		Short: "Start pyroscope agent.",
+		RunE: createCmdRunFn(cfg, vpr, false, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			return cli.StartAgent(cfg)
+		}),
+	}
+
+	cli.PopulateFlagSet(cfg, agentCmd.Flags(), vpr, cli.WithSkip("targets"))
+	return agentCmd
+}

pkg/cli/banner.go → cmd/pyroscope/command/banner.go

@@ -1,4 +1,4 @@
-package cli
+package command
 
 import (
 	"strings"

cmd/pyroscope/command/common.go

@@ -0,0 +1,123 @@
+package command
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+	"strings"
+
+	goexec "os/exec"
+
+	"github.com/mitchellh/mapstructure"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/util/bytesize"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
+)
+
+type cmdRunFn func(cmd *cobra.Command, args []string) error
+type cmdStartFn func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error
+
+func createCmdRunFn(cfg interface{}, vpr *viper.Viper, requiresArgs bool, fn cmdStartFn) cmdRunFn {
+	return func(cmd *cobra.Command, args []string) error {
+		var err error
+		if err = bindFlags(cfg, cmd, vpr); err != nil {
+			return fmt.Errorf("invalid configuration: %w", err)
+		}
+
+		var logger func(s string)
+		if l, ok := cfg.(config.LoggerConfiger); ok {
+			logger = l.InitializeLogging()
+		}
+
+		if c, ok := cfg.(config.FileConfiger); ok {
+			if err = loadConfigFile(c.ConfigFilePath(), cmd, vpr, cfg); err != nil {
+				return fmt.Errorf("loading configuration file: %w", err)
+			}
+		}
+
+		if (requiresArgs && len(args) == 0) || (len(args) > 0 && args[0] == "help") {
+			_ = cmd.Help()
+			return nil
+		}
+
+		if err = fn(cmd, args, logger); err != nil {
+			cmd.SilenceUsage = true
+		}
+
+		// Normally, if the program ran, the call should return ExitError and
+		// the exit code must be preserved. Otherwise, the error originates from
+		// pyroscope and will be printed.
+		if e, ok := err.(*goexec.ExitError); ok {
+			os.Exit(e.ExitCode())
+		}
+
+		return err
+	}
+}
+
+func bindFlags(cfg interface{}, cmd *cobra.Command, vpr *viper.Viper) error {
+	if err := vpr.BindPFlags(cmd.Flags()); err != nil {
+		return err
+	}
+	return viperUnmarshalWithBytesHook(vpr, cfg)
+}
+
+func newViper() *viper.Viper {
+	v := viper.New()
+	v.SetEnvPrefix("PYROSCOPE")
+	v.AutomaticEnv()
+	v.SetEnvKeyReplacer(strings.NewReplacer("-", "_", ".", "_"))
+	return v
+}
+
+func loadConfigFile(path string, cmd *cobra.Command, vpr *viper.Viper, v interface{}) error {
+	if path == "" {
+		return nil
+	}
+
+	vpr.SetConfigFile(path)
+	err := vpr.ReadInConfig()
+	switch {
+	case err == nil:
+		return viperUnmarshalWithBytesHook(vpr, v)
+	case isUserDefined(cmd.Flag("config"), vpr):
+		// User-defined configuration can not be read.
+		return err
+	case os.IsNotExist(err):
+		// Default configuration file not found.
+		return nil
+	default:
+		return err
+	}
+}
+
+func viperUnmarshalWithBytesHook(vpr *viper.Viper, cfg interface{}) error {
+	return vpr.Unmarshal(cfg, viper.DecodeHook(
+		mapstructure.ComposeDecodeHookFunc(
+			// Function to add a special type for «env. mode»
+			func(f reflect.Type, t reflect.Type, data interface{}) (interface{}, error) {
+				if t != reflect.TypeOf(bytesize.Byte) {
+					return data, nil
+				}
+
+				stringData, ok := data.(string)
+				if !ok {
+					return data, nil
+				}
+
+				return bytesize.Parse(stringData)
+			},
+			// Function to support net.IP
+			mapstructure.StringToIPHookFunc(),
+			// Appended by the two default functions
+			mapstructure.StringToTimeDurationHookFunc(),
+			mapstructure.StringToSliceHookFunc(","),
+		),
+	))
+}
+
+func isUserDefined(f *pflag.Flag, v *viper.Viper) bool {
+	return f.Changed || (f.DefValue != "" && f.DefValue != v.GetString(f.Name))
+}

cmd/pyroscope/command/connect.go

@@ -0,0 +1,22 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/exec"
+	"github.com/spf13/cobra"
+)
+
+func newConnectCmd(cfg *config.Exec) *cobra.Command {
+	vpr := newViper()
+	connectCmd := &cobra.Command{
+		Use:   "connect [flags]",
+		Short: "Connect to an existing process and profile it",
+		RunE: createCmdRunFn(cfg, vpr, true, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			return exec.Cli(cfg, args)
+		}),
+	}
+
+	cli.PopulateFlagSet(cfg, connectCmd.Flags(), vpr, cli.WithSkip("group-name", "user-name", "no-root-drop"))
+	return connectCmd
+}

cmd/pyroscope/command/convert.go

@@ -0,0 +1,58 @@
+package command
+
+import (
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/convert"
+	"github.com/spf13/cobra"
+
+	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
+	"github.com/pyroscope-io/pyroscope/pkg/structs/transporttrie"
+)
+
+func newConvertCmd(cfg *config.Convert) *cobra.Command {
+	vpr := newViper()
+	convertCmd := &cobra.Command{
+		Use:   "convert [flags] <input-file>",
+		Short: "Convert between different profiling formats",
+		RunE: createCmdRunFn(cfg, vpr, false, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			var input io.Reader
+			if len(args) == 0 {
+				input = os.Stdin
+			} else {
+				logger("not implemented yet")
+				return nil
+			}
+
+			parser := convert.ParseGroups
+			switch cfg.Format {
+			case "tree":
+				t := tree.New()
+				parser(input, func(name []byte, val int) {
+					t.Insert(name, uint64(val))
+				})
+
+				t.SerializeNoDict(4096, os.Stdout)
+			case "trie":
+				t := transporttrie.New()
+				parser(input, func(name []byte, val int) {
+					t.Insert(name, uint64(val), true)
+				})
+
+				t.Serialize(os.Stdout)
+			default:
+				logger(fmt.Sprintf("unknown format: %s", cfg.Format))
+			}
+
+			return nil
+		}),
+		Hidden: true,
+	}
+
+	cli.PopulateFlagSet(cfg, convertCmd.Flags(), vpr)
+	return convertCmd
+}

cmd/pyroscope/command/dbmanager.go

@@ -0,0 +1,25 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/dbmanager"
+	"github.com/spf13/cobra"
+)
+
+func newDbManagerCmd(cfg *config.CombinedDbManager) *cobra.Command {
+	vpr := newViper()
+	dbmanagerCmd := &cobra.Command{
+		Use:   "dbmanager [flags] <args>",
+		Short: "tools for managing database",
+		RunE: createCmdRunFn(cfg, vpr, false, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			return dbmanager.Cli(cfg.DbManager, cfg.Server, args)
+
+		}),
+		Hidden: true,
+	}
+
+	cli.PopulateFlagSet(cfg.DbManager, dbmanagerCmd.Flags(), vpr)
+	cli.PopulateFlagSet(cfg.Server, dbmanagerCmd.Flags(), vpr, cli.WithSkip("log-level", "storage-path", "metric-export-rules"))
+	return dbmanagerCmd
+}

cmd/pyroscope/command/exec.go

@@ -0,0 +1,22 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/exec"
+	"github.com/spf13/cobra"
+)
+
+func newExecCmd(cfg *config.Exec) *cobra.Command {
+	vpr := newViper()
+	execCmd := &cobra.Command{
+		Use:   "exec [flags] <args>",
+		Short: "Start a new process from arguments and profile it",
+		RunE: createCmdRunFn(cfg, vpr, true, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			return exec.Cli(cfg, args)
+		}),
+	}
+
+	cli.PopulateFlagSet(cfg, execCmd.Flags(), vpr, cli.WithSkip("pid"))
+	return execCmd
+}

cmd/pyroscope/command/root.go

@@ -0,0 +1,83 @@
+package command
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"strings"
+
+	"github.com/pyroscope-io/pyroscope/pkg/build"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+func newRootCmd(cfg *config.Config) *cobra.Command {
+	rootCmd := &cobra.Command{
+		Use: "pyroscope [flags] <subcommand>",
+		Run: func(cmd *cobra.Command, args []string) {
+			if cfg.Version {
+				fmt.Println(gradientBanner())
+				fmt.Println(build.Summary())
+				fmt.Println("")
+			} else {
+				fmt.Println(gradientBanner())
+				fmt.Println(DefaultUsageFunc(cmd.Flags(), cmd))
+			}
+		},
+	}
+
+	rootCmd.SetUsageFunc(func(cmd *cobra.Command) error {
+		fmt.Println(gradientBanner())
+		fmt.Println(DefaultUsageFunc(cmd.Flags(), cmd))
+		return nil
+	})
+
+	rootCmd.SetHelpFunc(func(cmd *cobra.Command, a []string) {
+		fmt.Println(gradientBanner())
+		fmt.Println(DefaultUsageFunc(cmd.Flags(), cmd))
+	})
+
+	rootCmd.PersistentFlags().BoolVarP(&cfg.Version, "version", "v", false, "print pyroscope version details")
+	return rootCmd
+}
+
+// Initialize adds all child commands to the root command and sets flags appropriately
+func Initialize() error {
+	var cfg config.Config
+
+	rootCmd := newRootCmd(&cfg)
+	rootCmd.SilenceErrors = true
+	rootCmd.AddCommand(
+		newAgentCmd(&cfg.Agent),
+		newConnectCmd(&cfg.Exec),
+		newConvertCmd(&cfg.Convert),
+		newDbManagerCmd(&config.CombinedDbManager{DbManager: &cfg.DbManager, Server: &cfg.Server}),
+		newExecCmd(&cfg.Exec),
+		newServerCmd(&cfg.Server),
+		newVersionCmd(),
+	)
+
+	logrus.SetReportCaller(true)
+	logrus.SetFormatter(&logrus.TextFormatter{
+		TimestampFormat: "2006-01-02T15:04:05.000000",
+		FullTimestamp:   true,
+		CallerPrettyfier: func(f *runtime.Frame) (string, string) {
+			filename := f.File
+			if len(filename) > 38 {
+				filename = filename[38:]
+			}
+			return "", fmt.Sprintf(" %s:%d", filename, f.Line)
+		},
+	})
+
+	args := os.Args[1:]
+	for i, arg := range args {
+		if len(arg) > 2 && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") {
+			args[i] = fmt.Sprintf("-%s", arg)
+		}
+	}
+
+	rootCmd.SetArgs(args)
+	return rootCmd.Execute()
+}

cmd/pyroscope/command/server.go

@@ -0,0 +1,21 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/spf13/cobra"
+)
+
+func newServerCmd(cfg *config.Server) *cobra.Command {
+	vpr := newViper()
+	serverCmd := &cobra.Command{
+		Use:   "server [flags]",
+		Short: "Start pyroscope server. This is the database + web-based user interface",
+		RunE: createCmdRunFn(cfg, vpr, false, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
+			return cli.StartServer(cfg)
+		}),
+	}
+
+	cli.PopulateFlagSet(cfg, serverCmd.Flags(), vpr, cli.WithSkip("metric-export-rules"))
+	return serverCmd
+}