From 6cbb17fb763c15d4e4dbf3239e4aa29bfc2f6f18 Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Wed, 13 Dec 2023 09:47:18 +0800
Subject: [PATCH] all: refactor log formatter functions

---
 pkg/bbgo/envvar.go              | 48 +++++++++++++++++++++++++++++++++
 pkg/cmd/root.go                 | 25 ++++-------------
 pkg/exchange/bitget/exchange.go | 12 ---------
 3 files changed, 53 insertions(+), 32 deletions(-)
 create mode 100644 pkg/bbgo/envvar.go

diff --git a/pkg/bbgo/envvar.go b/pkg/bbgo/envvar.go
new file mode 100644
index 000000000..10ac022b3
--- /dev/null
+++ b/pkg/bbgo/envvar.go
@@ -0,0 +1,48 @@
+package bbgo
+
+import (
+	"os"
+
+	log "github.com/sirupsen/logrus"
+	prefixed "github.com/x-cray/logrus-prefixed-formatter"
+)
+
+func GetCurrentEnv() string {
+	env := os.Getenv("BBGO_ENV")
+	if env == "" {
+		env = "development"
+	}
+
+	return env
+}
+
+func NewLogFormatterWithEnv(env string) log.Formatter {
+	switch env {
+	case "production", "prod", "stag", "staging":
+		// always use json formatter for production and staging
+		return &log.JSONFormatter{}
+	}
+
+	return &prefixed.TextFormatter{}
+}
+
+type LogFormatterType string
+
+const (
+	LogFormatterTypePrefixed LogFormatterType = "prefixed"
+	LogFormatterTypeText     LogFormatterType = "text"
+	LogFormatterTypeJson     LogFormatterType = "json"
+)
+
+func NewLogFormatter(logFormatter LogFormatterType) log.Formatter {
+	switch logFormatter {
+	case LogFormatterTypePrefixed:
+		return &prefixed.TextFormatter{}
+	case LogFormatterTypeText:
+		return &log.TextFormatter{}
+	case LogFormatterTypeJson:
+		return &log.JSONFormatter{}
+	}
+
+	return &prefixed.TextFormatter{}
+}
diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
index 67ccd6bba..417e6196c 100644
--- a/pkg/cmd/root.go
+++ b/pkg/cmd/root.go
@@ -17,7 +17,6 @@ import (
 	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"
@@ -46,10 +45,7 @@ var RootCmd = &cobra.Command{
 			log.SetLevel(log.DebugLevel)
 		}
 
-		env := os.Getenv("BBGO_ENV")
-		if env == "" {
-			env = "development"
-		}
+		env := bbgo.GetCurrentEnv()
 
 		logFormatter, err := cmd.Flags().GetString("log-formatter")
 		if err != nil {
@@ -57,22 +53,11 @@ var RootCmd = &cobra.Command{
 		}
 
 		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{})
-			}
+			formatter := bbgo.NewLogFormatterWithEnv(env)
+			log.SetFormatter(formatter)
 		} else {
-			switch logFormatter {
-			case "prefixed":
-				log.SetFormatter(&prefixed.TextFormatter{})
-			case "text":
-				log.SetFormatter(&log.TextFormatter{})
-			case "json":
-				log.SetFormatter(&log.JSONFormatter{})
-			}
+			formatter := bbgo.NewLogFormatter(bbgo.LogFormatterType(logFormatter))
+			log.SetFormatter(formatter)
 		}
 
 		if token := viper.GetString("rollbar-token"); token != "" {
diff --git a/pkg/exchange/bitget/exchange.go b/pkg/exchange/bitget/exchange.go
index 217406326..dc0ca8e7a 100644
--- a/pkg/exchange/bitget/exchange.go
+++ b/pkg/exchange/bitget/exchange.go
@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/sirupsen/logrus"
-	prefixed "github.com/x-cray/logrus-prefixed-formatter"
 	"go.uber.org/multierr"
 	"golang.org/x/time/rate"
 
@@ -69,19 +68,8 @@ type LogFunction func(msg string, args ...interface{})
 
 var debugf LogFunction
 
-func isPrefixFormatterConfigured() bool {
-	_, isPrefixFormatter := logrus.StandardLogger().Formatter.(*prefixed.TextFormatter)
-	return isPrefixFormatter
-}
-
 func getDebugFunction() LogFunction {
 	if v, ok := util.GetEnvVarBool("DEBUG_BITGET"); ok && v {
-		if isPrefixFormatterConfigured() {
-			return func(msg string, args ...interface{}) {
-				log.Infof("[BITGET] "+msg, args...)
-			}
-		}
-
 		return log.Infof
 	}