From 53b72409d8e7fd53dcac17c6e8d8149bb3c8bc43 Mon Sep 17 00:00:00 2001
From: Sean Liao <seankhliao@gmail.com>
Date: Thu, 9 Apr 2020 13:26:02 +0200
Subject: [PATCH] use same system as level for format

---
 logging/format.go | 52 +++++++++++++++++++++++++++++++++++++++++++++++
 logging/gokit.go  | 16 +++++++++++----
 logging/logrus.go | 12 +++++++----
 server/server.go  |  5 +++--
 4 files changed, 75 insertions(+), 10 deletions(-)
 create mode 100644 logging/format.go

diff --git a/logging/format.go b/logging/format.go
new file mode 100644
index 00000000..701b32f5
--- /dev/null
+++ b/logging/format.go
@@ -0,0 +1,52 @@
+package logging
+
+import (
+	"flag"
+
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+)
+
+// Format is a settable identifier for the output format of logs
+type Format struct {
+	s      string
+	Logrus logrus.Formatter
+}
+
+// RegisterFlags adds the log format flag to the provided flagset.
+func (f *Format) RegisterFlags(fs *flag.FlagSet) {
+	f.Set("logfmt")
+	fs.Var(f, "log.format", "Output log messages in the given format. Valid formats: [logfmt, json]")
+}
+
+func (f Format) String() string {
+	return f.s
+}
+
+// UnmarshalYAML implements yaml.Unmarshaler.
+func (f *Format) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var format string
+	if err := unmarshal(&format); err != nil {
+		return err
+	}
+	return f.Set(format)
+}
+
+// MarshalYAML implements yaml.Marshaler.
+func (f Format) MarshalYAML() (interface{}, error) {
+	return f.String(), nil
+}
+
+// Set updates the value of the output format.  Implements flag.Value
+func (f *Format) Set(s string) error {
+	switch s {
+	case "logfmt":
+		f.Logrus = &logrus.JSONFormatter{}
+	case "json":
+		f.Logrus = &logrus.JSONFormatter{}
+	default:
+		return errors.Errorf("unrecognized log format %q", s)
+	}
+	f.s = s
+	return nil
+}
diff --git a/logging/gokit.go b/logging/gokit.go
index dac4ce8b..508fc385 100644
--- a/logging/gokit.go
+++ b/logging/gokit.go
@@ -8,17 +8,25 @@ import (
 	"github.com/go-kit/kit/log/level"
 )
 
-// NewGoKit creates a new Interface backed by a GoKit logger
-func NewGoKit(l Level, format string) Interface {
-	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
-	if format == "json" {
+// NewGoKitFormat creates a new Interface backed by a GoKit logger
+// format can be "json" or defaults to logfmt
+func NewGoKitFormat(l Level, f Format) Interface {
+	var logger log.Logger
+	if f.s == "json" {
 		logger = log.NewJSONLogger(log.NewSyncWriter(os.Stderr))
+	} else {
+		logger = log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
 	}
 	logger = level.NewFilter(logger, l.Gokit)
 	logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", log.DefaultCaller)
 	return gokit{logger}
 }
 
+// NewGoKit creates a new Interface backed by a GoKit logger
+func NewGoKit(l Level) Interface {
+	return NewGoKitFormat(l, Format{s: "logfmt"})
+}
+
 // GoKit wraps an existing gokit Logger.
 func GoKit(logger log.Logger) Interface {
 	return gokit{logger}
diff --git a/logging/logrus.go b/logging/logrus.go
index b4f9d8e9..2c8c3410 100644
--- a/logging/logrus.go
+++ b/logging/logrus.go
@@ -7,16 +7,20 @@ import (
 )
 
 // NewLogrus makes a new Interface backed by a logrus logger
-func NewLogrus(level Level, format string) Interface {
+// format can be "json" or defaults to logfmt
+func NewLogrusFormat(level Level, f Format) Interface {
 	log := logrus.New()
 	log.Out = os.Stderr
 	log.Level = level.Logrus
-	if format == "json" {
-		log.Formatter = &logrus.JSONFormatter{}
-	}
+	log.Formatter = f.Logrus
 	return logrusLogger{log}
 }
 
+// NewLogrus makes a new Interface backed by a logrus logger
+func NewLogrus(level Level) Interface {
+	return NewLogrusFormat(level, Format{Logrus: &logrus.TextFormatter{}})
+}
+
 // Logrus wraps an existing Logrus logger.
 func Logrus(l *logrus.Logger) Interface {
 	return logrusLogger{l}
diff --git a/server/server.go b/server/server.go
index 1a4428f9..ba01cf49 100644
--- a/server/server.go
+++ b/server/server.go
@@ -60,7 +60,7 @@ type Config struct {
 	GRPCServerTime                  time.Duration `yaml:"grpc_server_keepalive_time"`
 	GRPCServerTimeout               time.Duration `yaml:"grpc_server_keepalive_timeout"`
 
-	LogFormat string            `yaml:"log_format"`
+	LogFormat logging.Format    `yaml:"log_format"`
 	LogLevel  logging.Level     `yaml:"log_level"`
 	Log       logging.Interface `yaml:"-"`
 
@@ -91,6 +91,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.GRPCServerTime, "server.grpc.keepalive.time", time.Hour*2, "Duration after which a keepalive probe is sent in case of no activity over the connection., Default: 2h")
 	f.DurationVar(&cfg.GRPCServerTimeout, "server.grpc.keepalive.timeout", time.Second*20, "After having pinged for keepalive check, the duration after which an idle connection should be closed, Default: 20s")
 	f.StringVar(&cfg.PathPrefix, "server.path-prefix", "", "Base path to serve all API routes from (e.g. /v1/)")
+	cfg.LogFormat.RegisterFlags(f)
 	cfg.LogLevel.RegisterFlags(f)
 }
 
@@ -142,7 +143,7 @@ func New(cfg Config) (*Server, error) {
 	// logrus.
 	log := cfg.Log
 	if log == nil {
-		log = logging.NewLogrus(cfg.LogLevel, cfg.LogFormat)
+		log = logging.NewLogrus(cfg.LogLevel)
 	}
 
 	log.WithField("http", httpListener.Addr()).WithField("grpc", grpcListener.Addr()).Infof("server listening on addresses")