Add color logging capabilities for promlog

promlog/flag/flag.go

@@ -18,6 +18,12 @@ import (
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 )
 
+// ColorFlagName is the canonical flag name to enable color logging.
+const ColorFlagName = "log.color"
+
+// ColorFlagHelp is the help description for the log.color flag.
+const ColorFlagHelp = "Enable colors in the console log. Use --no-log.color to disable."
+
 // LevelFlagName is the canonical flag name to configure the allowed log level
 // within Prometheus projects.
 const LevelFlagName = "log.level"
@@ -35,6 +41,10 @@ const FormatFlagHelp = "Output format of log messages. One of: [logfmt, json]"
 // AddFlags adds the flags used by this package to the Kingpin application.
 // To use the default Kingpin application, call AddFlags(kingpin.CommandLine)
 func AddFlags(a *kingpin.Application, config *promlog.Config) {
+	config.Color = &promlog.Color{}
+	a.Flag(ColorFlagName, ColorFlagHelp).
+		Default("true").SetValue(config.Color)
+
 	config.Level = &promlog.AllowedLevel{}
 	a.Flag(LevelFlagName, LevelFlagHelp).
 		Default("info").SetValue(config.Level)

promlog/log.go

@@ -19,11 +19,13 @@ package promlog
 import (
 	"fmt"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/go-kit/log/term"
 )
 
 var (
@@ -36,6 +38,33 @@ var (
 	)
 )
 
+// Color is a settable boolean for controlling color output.
+type Color struct {
+	s       string
+	enabled bool
+}
+
+func (c *Color) Set(s string) error {
+	switch s {
+	case "true":
+		c.enabled = true
+	case "false":
+		c.enabled = false
+	default:
+		return fmt.Errorf("unrecognized boolean %q", s)
+	}
+	c.s = s
+	return nil
+}
+
+func (c *Color) Enabled() bool {
+	return c.enabled
+}
+
+func (c *Color) String() string {
+	return strconv.FormatBool(c.enabled)
+}
+
 // AllowedLevel is a settable identifier for the minimum level a log entry
 // must be have.
 type AllowedLevel struct {
@@ -60,6 +89,27 @@ func (l *AllowedLevel) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	return nil
 }
 
+func colorFn(keyvals ...interface{}) term.FgBgColor {
+	for i := 1; i < len(keyvals); i += 2 {
+		if keyvals[i] != "level" {
+			continue
+		}
+		switch keyvals[i+1] {
+		case "info":
+			return term.FgBgColor{Fg: term.Green}
+		case "debug":
+			return term.FgBgColor{Fg: term.Blue}
+		case "warn":
+			return term.FgBgColor{Fg: term.Yellow}
+		case "error":
+			return term.FgBgColor{Fg: term.Red}
+		default:
+			return term.FgBgColor{}
+		}
+	}
+	return term.FgBgColor{}
+}
+
 func (l *AllowedLevel) String() string {
 	return l.s
 }
@@ -104,18 +154,29 @@ func (f *AllowedFormat) Set(s string) error {
 
 // Config is a struct containing configurable settings for the logger
 type Config struct {
+	Color  *Color
 	Level  *AllowedLevel
 	Format *AllowedFormat
 }
 
 // New returns a new leveled oklog logger. Each logged line will be annotated
 // with a timestamp. The output always goes to stderr.
 func New(config *Config) log.Logger {
+	if config.Color == nil {
+		config.Color = &Color{s: "true", enabled: true}
+	}
 	var l log.Logger
+	syncWriter := log.NewSyncWriter(os.Stderr)
 	if config.Format != nil && config.Format.s == "json" {
-		l = log.NewJSONLogger(log.NewSyncWriter(os.Stderr))
+		l = log.NewJSONLogger(syncWriter)
 	} else {
-		l = log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
+		if config.Color.Enabled() {
+			// Returns a new logger with color logging capabilites if we're in a terminal, otherwise we
+			// just get a standard go-kit logger.
+			l = term.NewLogger(syncWriter, log.NewLogfmtLogger, colorFn)
+		} else {
+			l = log.NewJSONLogger(syncWriter)
+		}
 	}
 
 	if config.Level != nil {
@@ -131,11 +192,22 @@ func New(config *Config) log.Logger {
 // with a timestamp. The output always goes to stderr. Some properties can be
 // changed, like the level.
 func NewDynamic(config *Config) *logger {
+	if config.Color == nil {
+		config.Color = &Color{s: "true", enabled: true}
+	}
 	var l log.Logger
+	syncWriter := log.NewSyncWriter(os.Stderr)
+
 	if config.Format != nil && config.Format.s == "json" {
-		l = log.NewJSONLogger(log.NewSyncWriter(os.Stderr))
+		l = log.NewJSONLogger(syncWriter)
 	} else {
-		l = log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
+		if config.Color.Enabled() {
+			// Returns a new logger with color logging capabilites if we're in a terminal, otherwise we
+			// just get a standard go-kit logger.
+			l = term.NewLogger(syncWriter, log.NewLogfmtLogger, colorFn)
+		} else {
+			l = log.NewJSONLogger(syncWriter)
+		}
 	}
 
 	lo := &logger{