bbgo_origin/pkg/cmd/root.go

package cmd

import (
	"net/http"
	"os"
	"path"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/heroku/rollrus"
	"github.com/joho/godotenv"
	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rifflock/lfshook"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	prefixed "github.com/x-cray/logrus-prefixed-formatter"

	"github.com/c9s/bbgo/pkg/bbgo"
	"github.com/c9s/bbgo/pkg/util"

	_ "github.com/go-sql-driver/mysql"
)

var cpuProfileFile *os.File

var userConfig *bbgo.Config

var RootCmd = &cobra.Command{
	Use:   "bbgo",
	Short: "bbgo is a crypto trading bot",

	// SilenceUsage is an option to silence usage when an error occurs.
	SilenceUsage: true,

	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		if err := cobraLoadDotenv(cmd, args); err != nil {
			return err
		}

		if viper.GetBool("debug") {
			log.Infof("debug mode is enabled")
			log.SetLevel(log.DebugLevel)
		}

		env := os.Getenv("BBGO_ENV")
		if env == "" {
			env = "development"
		}

		logFormatter, err := cmd.Flags().GetString("log-formatter")
		if err != nil {
			return err
		}

		if len(logFormatter) == 0 {
			switch env {
			case "production", "prod", "stag", "staging":
				// always use json formatter for production and staging
				log.SetFormatter(&log.JSONFormatter{})
			default:
				log.SetFormatter(&prefixed.TextFormatter{})
			}
		} else {
			switch logFormatter {
			case "prefixed":
				log.SetFormatter(&prefixed.TextFormatter{})
			case "text":
				log.SetFormatter(&log.TextFormatter{})
			case "json":
				log.SetFormatter(&log.JSONFormatter{})
			}
		}

		if token := viper.GetString("rollbar-token"); token != "" {
			log.Infof("found rollbar token %q, setting up rollbar hook...", util.MaskKey(token))

			log.AddHook(rollrus.NewHook(
				token,
				env,
			))
		}

		if viper.GetBool("metrics") {
			http.Handle("/metrics", promhttp.Handler())
			go func() {
				port := viper.GetString("metrics-port")
				log.Infof("starting metrics server at :%s", port)
				err := http.ListenAndServe(":"+port, nil)
				if err != nil {
					log.WithError(err).Errorf("metrics server error")
				}
			}()
		}

		cpuProfile, err := cmd.Flags().GetString("cpu-profile")
		if err != nil {
			return err
		}

		if cpuProfile != "" {
			log.Infof("starting cpu profiler, recording at %s", cpuProfile)

			cpuProfileFile, err = os.Create(cpuProfile)
			if err != nil {
				return errors.Wrap(err, "can not create file for CPU profile")
			}

			if err := pprof.StartCPUProfile(cpuProfileFile); err != nil {
				return errors.Wrap(err, "can not start CPU profile")
			}
		}

		return cobraLoadConfig(cmd, args)
	},
	PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
		pprof.StopCPUProfile()
		if cpuProfileFile != nil {
			return cpuProfileFile.Close() // error handling omitted for example
		}

		return nil
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		return nil
	},
}

func cobraLoadDotenv(cmd *cobra.Command, args []string) error {
	disableDotEnv, err := cmd.Flags().GetBool("no-dotenv")
	if err != nil {
		return err
	}

	if !disableDotEnv {
		dotenvFile, err := cmd.Flags().GetString("dotenv")
		if err != nil {
			return err
		}

		if _, err := os.Stat(dotenvFile); err == nil {
			if err := godotenv.Load(dotenvFile); err != nil {
				return errors.Wrap(err, "error loading dotenv file")
			}
		}
	}
	return nil
}

func cobraLoadConfig(cmd *cobra.Command, args []string) error {
	configFile, err := cmd.Flags().GetString("config")
	if err != nil {
		return errors.Wrapf(err, "failed to get the config flag")
	}

	// load config file nicely
	if len(configFile) > 0 {
		// if config file exists, use the config loaded from the config file.
		// otherwise, use an empty config object
		if _, err := os.Stat(configFile); err == nil {
			// load successfully
			userConfig, err = bbgo.Load(configFile, false)
			if err != nil {
				return errors.Wrapf(err, "can not load config file: %s", configFile)
			}

		} else if os.IsNotExist(err) {
			// config file doesn't exist, we should use the empty config
			userConfig = &bbgo.Config{}
		} else {
			// other error
			return errors.Wrapf(err, "config file load error: %s", configFile)
		}
	}

	return nil
}

func init() {
	RootCmd.PersistentFlags().Bool("debug", false, "debug mode")
	RootCmd.PersistentFlags().Bool("metrics", false, "enable prometheus metrics")
	RootCmd.PersistentFlags().String("metrics-port", "9090", "prometheus http server port")

	RootCmd.PersistentFlags().Bool("no-dotenv", false, "disable built-in dotenv")
	RootCmd.PersistentFlags().String("dotenv", ".env.local", "the dotenv file you want to load")

	RootCmd.PersistentFlags().String("config", "bbgo.yaml", "config file")

	RootCmd.PersistentFlags().String("log-formatter", "", "configure log formatter")

	RootCmd.PersistentFlags().String("rollbar-token", "", "rollbar token")

	// A flag can be 'persistent' meaning that this flag will be available to
	// the command it's assigned to as well as every command under that command.
	// For global flags, assign a flag as a persistent flag on the root.
	RootCmd.PersistentFlags().String("slack-token", "", "slack token")
	RootCmd.PersistentFlags().String("slack-channel", "dev-bbgo", "slack trading channel")
	RootCmd.PersistentFlags().String("slack-error-channel", "bbgo-error", "slack error channel")

	RootCmd.PersistentFlags().String("telegram-bot-token", "", "telegram bot token from bot father")
	RootCmd.PersistentFlags().String("telegram-bot-auth-token", "", "telegram auth token")

	RootCmd.PersistentFlags().String("binance-api-key", "", "binance api key")
	RootCmd.PersistentFlags().String("binance-api-secret", "", "binance api secret")

	RootCmd.PersistentFlags().String("max-api-key", "", "max api key")
	RootCmd.PersistentFlags().String("max-api-secret", "", "max api secret")

	RootCmd.PersistentFlags().String("cpu-profile", "", "cpu profile")

	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))

	// Enable environment variable binding, the env vars are not overloaded yet.
	viper.AutomaticEnv()

	// setup the config paths for looking up the config file
	/*
		viper.AddConfigPath("config")
		viper.AddConfigPath("$HOME/.bbgo")
		viper.AddConfigPath("/etc/bbgo")

		// set the config file name and format for loading the config file.
		viper.SetConfigName("bbgo")
		viper.SetConfigType("yaml")

		err := viper.ReadInConfig()
		if err != nil {
			log.WithError(err).Fatal("failed to load config file")
		}
	*/
	// Once the flags are defined, we can bind config keys with flags.
	if err := viper.BindPFlags(RootCmd.PersistentFlags()); err != nil {
		log.WithError(err).Errorf("failed to bind persistent flags. please check the flag settings.")
		return
	}

	environment := os.Getenv("BBGO_ENV")
	logDir := "log"
	switch environment {
	case "production", "prod":
		if err := os.MkdirAll(logDir, 0777); err != nil {
			log.Panic(err)
		}
		writer, err := rotatelogs.New(
			path.Join(logDir, "access_log.%Y%m%d"),
			rotatelogs.WithLinkName("access_log"),
			// rotatelogs.WithMaxAge(24 * time.Hour),
			rotatelogs.WithRotationTime(time.Duration(24)*time.Hour),
		)
		if err != nil {
			log.Panic(err)
		}

		log.AddHook(
			lfshook.NewHook(
				lfshook.WriterMap{
					log.DebugLevel: writer,
					log.InfoLevel:  writer,
					log.WarnLevel:  writer,
					log.ErrorLevel: writer,
					log.FatalLevel: writer,
				}, &log.JSONFormatter{}),
		)

	}
}

func Execute() {
	if err := RootCmd.Execute(); err != nil {
		log.WithError(err).Fatalf("cannot execute command")
	}
}