Rewrite the friendly log handler (#1038)

## Changes

It wasn't working because it deferred to the regular `slog.TextHandler`
for the `WithAttr` and `WithGroup` functions. Both of these functions
don't mutate the handler but return a new one. When the top-level logger
called one of these, log records in that context used the standard
handler instead of ours.

To implement tracking of attributes and groups, I followed the guide at
https://github.com/golang/example/blob/master/slog-handler-guide/README.md
for writing custom handlers.

## Tests

The new tests demonstrate formatting through `t.Log` and look good.

cmd/root/logger.go

@@ -3,7 +3,6 @@ package root
 import (
 	"context"
 	"fmt"
-	"io"
 	"log/slog"
 	"os"
 
@@ -11,7 +10,7 @@ import (
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/log"
-	"github.com/fatih/color"
+	"github.com/databricks/cli/libs/log/handler"
 	"github.com/spf13/cobra"
 )
 
@@ -21,54 +20,6 @@ const (
 	envLogFormat = "DATABRICKS_LOG_FORMAT"
 )
 
-type friendlyHandler struct {
-	slog.Handler
-	w io.Writer
-}
-
-var (
-	levelTrace = color.New(color.FgYellow).Sprint("TRACE")
-	levelDebug = color.New(color.FgYellow).Sprint("DEBUG")
-	levelInfo  = color.New(color.FgGreen).Sprintf("%5s", "INFO")
-	levelWarn  = color.New(color.FgMagenta).Sprintf("%5s", "WARN")
-	levelError = color.New(color.FgRed).Sprint("ERROR")
-)
-
-func (l *friendlyHandler) coloredLevel(rec slog.Record) string {
-	switch rec.Level {
-	case log.LevelTrace:
-		return levelTrace
-	case slog.LevelDebug:
-		return levelDebug
-	case slog.LevelInfo:
-		return levelInfo
-	case slog.LevelWarn:
-		return levelWarn
-	case log.LevelError:
-		return levelError
-	}
-	return ""
-}
-
-func (l *friendlyHandler) Handle(ctx context.Context, rec slog.Record) error {
-	t := fmt.Sprintf("%02d:%02d", rec.Time.Hour(), rec.Time.Minute())
-	attrs := ""
-	rec.Attrs(func(a slog.Attr) bool {
-		attrs += fmt.Sprintf(" %s%s%s",
-			color.CyanString(a.Key),
-			color.CyanString("="),
-			color.YellowString(a.Value.String()))
-		return true
-	})
-	msg := fmt.Sprintf("%s %s %s%s\n",
-		color.MagentaString(t),
-		l.coloredLevel(rec),
-		rec.Message,
-		attrs)
-	_, err := l.w.Write([]byte(msg))
-	return err
-}
-
 type logFlags struct {
 	file   flags.LogFileFlag
 	level  flags.LogLevelFlag
@@ -83,10 +34,11 @@ func (f *logFlags) makeLogHandler(opts slog.HandlerOptions) (slog.Handler, error
 	case flags.OutputText:
 		w := f.file.Writer()
 		if cmdio.IsTTY(w) {
-			return &friendlyHandler{
-				Handler: slog.NewTextHandler(w, &opts),
-				w:       w,
-			}, nil
+			return handler.NewFriendlyHandler(w, &handler.Options{
+				Color:       true,
+				Level:       opts.Level,
+				ReplaceAttr: opts.ReplaceAttr,
+			}), nil
 		}
 		return slog.NewTextHandler(w, &opts), nil
 

libs/log/handler/colors.go

@@ -0,0 +1,56 @@
+package handler
+
+import "github.com/fatih/color"
+
+// ttyColors is a slice of colors that can be enabled or disabled.
+// This adds one level of indirection to the colors such that they
+// can be easily be enabled or disabled together, regardless of
+// global settings in the color package.
+type ttyColors []*color.Color
+
+// ttyColor is an enum for the colors in ttyColors.
+type ttyColor int
+
+const (
+	ttyColorInvalid ttyColor = iota
+	ttyColorTime
+	ttyColorMessage
+	ttyColorAttrKey
+	ttyColorAttrSeparator
+	ttyColorAttrValue
+	ttyColorLevelTrace
+	ttyColorLevelDebug
+	ttyColorLevelInfo
+	ttyColorLevelWarn
+	ttyColorLevelError
+
+	// Marker for the last value to know how many colors there are.
+	ttyColorLevelLast
+)
+
+func newColors(enable bool) ttyColors {
+	ttyColors := make(ttyColors, ttyColorLevelLast)
+	ttyColors[ttyColorInvalid] = color.New(color.FgWhite)
+	ttyColors[ttyColorTime] = color.New(color.FgBlack, color.Bold)
+	ttyColors[ttyColorMessage] = color.New(color.Reset)
+	ttyColors[ttyColorAttrKey] = color.New(color.Faint)
+	ttyColors[ttyColorAttrSeparator] = color.New(color.Faint)
+	ttyColors[ttyColorAttrValue] = color.New(color.Reset)
+	ttyColors[ttyColorLevelTrace] = color.New(color.FgMagenta)
+	ttyColors[ttyColorLevelDebug] = color.New(color.FgCyan)
+	ttyColors[ttyColorLevelInfo] = color.New(color.FgBlue)
+	ttyColors[ttyColorLevelWarn] = color.New(color.FgYellow)
+	ttyColors[ttyColorLevelError] = color.New(color.FgRed)
+
+	if enable {
+		for _, color := range ttyColors {
+			color.EnableColor()
+		}
+	} else {
+		for _, color := range ttyColors {
+			color.DisableColor()
+		}
+	}
+
+	return ttyColors
+}

libs/log/handler/colors_test.go

@@ -0,0 +1,31 @@
+package handler
+
+import (
+	"fmt"
+	"testing"
+)
+
+func showColors(t *testing.T, colors ttyColors) {
+	t.Log(colors[ttyColorInvalid].Sprint("invalid"))
+	t.Log(colors[ttyColorTime].Sprint("time"))
+	t.Log(
+		fmt.Sprint(
+			colors[ttyColorAttrKey].Sprint("key"),
+			colors[ttyColorAttrSeparator].Sprint("="),
+			colors[ttyColorAttrValue].Sprint("value"),
+		),
+	)
+	t.Log(colors[ttyColorLevelTrace].Sprint("trace"))
+	t.Log(colors[ttyColorLevelDebug].Sprint("debug"))
+	t.Log(colors[ttyColorLevelInfo].Sprint("info"))
+	t.Log(colors[ttyColorLevelWarn].Sprint("warn"))
+	t.Log(colors[ttyColorLevelError].Sprint("error"))
+}
+
+func TestTTYColorsEnabled(t *testing.T) {
+	showColors(t, newColors(true))
+}
+
+func TestTTYColorsDisabled(t *testing.T) {
+	showColors(t, newColors(false))
+}

libs/log/handler/friendly.go

@@ -0,0 +1,248 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+// friendlyHandler implements a custom [slog.Handler] that writes
+// human readable (and colorized) log lines to a terminal.
+//
+// The implementation is based on the guide at:
+// https://github.com/golang/example/blob/master/slog-handler-guide/README.md
+type friendlyHandler struct {
+	opts Options
+	goas []groupOrAttrs
+	mu   *sync.Mutex
+	out  io.Writer
+
+	// List of colors to use for formatting.
+	ttyColors
+
+	// Cache (colorized) level strings.
+	levelTrace string
+	levelDebug string
+	levelInfo  string
+	levelWarn  string
+	levelError string
+}
+
+// groupOrAttrs holds either a group name or a list of slog.Attrs.
+type groupOrAttrs struct {
+	group string      // group name if non-empty
+	attrs []slog.Attr // attrs if non-empty
+}
+
+func NewFriendlyHandler(out io.Writer, opts *Options) slog.Handler {
+	h := &friendlyHandler{out: out, mu: &sync.Mutex{}}
+	if opts != nil {
+		h.opts = *opts
+	}
+	if h.opts.Level == nil {
+		h.opts.Level = slog.LevelInfo
+	}
+
+	h.ttyColors = newColors(opts.Color)
+
+	// Cache (colorized) level strings.
+	// The colors to use for each level are configured in `colors.go`.
+	h.levelTrace = h.sprintf(ttyColorLevelTrace, "%5s", "TRACE")
+	h.levelDebug = h.sprintf(ttyColorLevelDebug, "%5s", "DEBUG")
+	h.levelInfo = h.sprintf(ttyColorLevelInfo, "%5s", "INFO")
+	h.levelWarn = h.sprintf(ttyColorLevelWarn, "%5s", "WARN")
+	h.levelError = h.sprintf(ttyColorLevelError, "%5s", "ERROR")
+	return h
+}
+
+func (h *friendlyHandler) sprint(color ttyColor, args ...any) string {
+	return h.ttyColors[color].Sprint(args...)
+}
+
+func (h *friendlyHandler) sprintf(color ttyColor, format string, args ...any) string {
+	return h.ttyColors[color].Sprintf(format, args...)
+}
+
+func (h *friendlyHandler) coloredLevel(r slog.Record) string {
+	switch r.Level {
+	case log.LevelTrace:
+		return h.levelTrace
+	case log.LevelDebug:
+		return h.levelDebug
+	case log.LevelInfo:
+		return h.levelInfo
+	case log.LevelWarn:
+		return h.levelWarn
+	case log.LevelError:
+		return h.levelError
+	}
+	return ""
+}
+
+// Enabled implements slog.Handler.
+func (h *friendlyHandler) Enabled(ctx context.Context, level slog.Level) bool {
+	return level >= h.opts.Level.Level()
+}
+
+type handleState struct {
+	h *friendlyHandler
+
+	buf    []byte
+	prefix string
+
+	// Keep stack of groups to pass to [slog.ReplaceAttr] function.
+	groups []string
+}
+
+func (h *friendlyHandler) handleState() *handleState {
+	return &handleState{
+		h: h,
+
+		buf:    make([]byte, 0, 1024),
+		prefix: "",
+	}
+}
+
+func (s *handleState) openGroup(name string) {
+	s.groups = append(s.groups, name)
+	s.prefix += name + "."
+}
+
+func (s *handleState) closeGroup(name string) {
+	s.prefix = s.prefix[:len(s.prefix)-len(name)-1]
+	s.groups = s.groups[:len(s.groups)-1]
+}
+
+func (s *handleState) append(args ...any) {
+	s.buf = fmt.Append(s.buf, args...)
+}
+
+func (s *handleState) appendf(format string, args ...any) {
+	s.buf = fmt.Appendf(s.buf, format, args...)
+}
+
+func (s *handleState) appendAttr(a slog.Attr) {
+	if rep := s.h.opts.ReplaceAttr; rep != nil && a.Value.Kind() != slog.KindGroup {
+		// Resolve before calling ReplaceAttr, so the user doesn't have to.
+		a.Value = a.Value.Resolve()
+		a = rep(s.groups, a)
+	}
+
+	// Resolve the Attr's value before doing anything else.
+	a.Value = a.Value.Resolve()
+
+	// Ignore empty Attrs.
+	if a.Equal(slog.Attr{}) {
+		return
+	}
+
+	switch a.Value.Kind() {
+	case slog.KindGroup:
+		attrs := a.Value.Group()
+		// Output only non-empty groups.
+		if len(attrs) > 0 {
+			if a.Key != "" {
+				s.openGroup(a.Key)
+			}
+			for _, aa := range attrs {
+				s.appendAttr(aa)
+			}
+			if a.Key != "" {
+				s.closeGroup(a.Key)
+			}
+		}
+	case slog.KindTime:
+		s.append(
+			" ",
+			s.h.sprint(ttyColorAttrKey, s.prefix, a.Key),
+			s.h.sprint(ttyColorAttrSeparator, "="),
+			s.h.sprint(ttyColorAttrValue, a.Value.Time().Format(time.RFC3339Nano)),
+		)
+	default:
+		str := a.Value.String()
+		format := "%s"
+
+		// Quote values wih spaces, to make them easy to parse.
+		if strings.ContainsAny(str, " \t\n") {
+			format = "%q"
+		}
+
+		s.append(
+			" ",
+			s.h.sprint(ttyColorAttrKey, s.prefix, a.Key),
+			s.h.sprint(ttyColorAttrSeparator, "="),
+			s.h.sprint(ttyColorAttrValue, fmt.Sprintf(format, str)),
+		)
+	}
+}
+
+// Handle implements slog.Handler.
+func (h *friendlyHandler) Handle(ctx context.Context, r slog.Record) error {
+	state := h.handleState()
+	state.append(h.sprintf(ttyColorTime, "%02d:%02d:%02d ", r.Time.Hour(), r.Time.Minute(), r.Time.Second()))
+	state.appendf("%s ", h.coloredLevel(r))
+	state.append(h.sprint(ttyColorMessage, r.Message))
+
+	// Handle state from WithGroup and WithAttrs.
+	goas := h.goas
+	if r.NumAttrs() == 0 {
+		// If the record has no Attrs, remove groups at the end of the list; they are empty.
+		for len(goas) > 0 && goas[len(goas)-1].group != "" {
+			goas = goas[:len(goas)-1]
+		}
+	}
+	for _, goa := range goas {
+		if goa.group != "" {
+			state.openGroup(goa.group)
+		} else {
+			for _, a := range goa.attrs {
+				state.appendAttr(a)
+			}
+		}
+	}
+
+	// Add attributes from the record.
+	r.Attrs(func(a slog.Attr) bool {
+		state.appendAttr(a)
+		return true
+	})
+
+	// Add newline.
+	state.append("\n")
+
+	// Write the log line.
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	_, err := h.out.Write(state.buf)
+	return err
+}
+
+func (h *friendlyHandler) withGroupOrAttrs(goa groupOrAttrs) *friendlyHandler {
+	h2 := *h
+	h2.goas = make([]groupOrAttrs, len(h.goas)+1)
+	copy(h2.goas, h.goas)
+	h2.goas[len(h2.goas)-1] = goa
+	return &h2
+}
+
+// WithGroup implements slog.Handler.
+func (h *friendlyHandler) WithGroup(name string) slog.Handler {
+	if name == "" {
+		return h
+	}
+	return h.withGroupOrAttrs(groupOrAttrs{group: name})
+}
+
+// WithAttrs implements slog.Handler.
+func (h *friendlyHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	if len(attrs) == 0 {
+		return h
+	}
+	return h.withGroupOrAttrs(groupOrAttrs{attrs: attrs})
+}

libs/log/handler/friendly_test.go

@@ -0,0 +1,110 @@
+package handler
+
+import (
+	"bytes"
+	"context"
+	"log/slog"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+func TestFriendlyHandler(t *testing.T) {
+	var out bytes.Buffer
+
+	handler := NewFriendlyHandler(&out, &Options{
+		Color: true,
+		Level: log.LevelTrace,
+	})
+
+	logger := slog.New(handler)
+
+	// Helper function to run a test case and print the output.
+	run := func(fn func()) {
+		out.Reset()
+		fn()
+		t.Log(strings.TrimSpace(out.String()))
+	}
+
+	// One line per level.
+	for _, level := range []slog.Level{
+		log.LevelTrace,
+		log.LevelDebug,
+		log.LevelInfo,
+		log.LevelWarn,
+		log.LevelError,
+	} {
+		run(func() {
+			logger.Log(context.Background(), level, "simple message")
+		})
+	}
+
+	// Single key/value pair.
+	run(func() {
+		logger.Info("simple message", "key", "value")
+	})
+
+	// Multiple key/value pairs.
+	run(func() {
+		logger.Info("simple message", "key1", "value", "key2", "value")
+	})
+
+	// Multiple key/value pairs with duplicate keys.
+	run(func() {
+		logger.Info("simple message", "key", "value", "key", "value")
+	})
+
+	// Log message with time.
+	run(func() {
+		logger.Info("simple message", "time", time.Now())
+	})
+
+	// Log message with grouped key/value pairs.
+	run(func() {
+		logger.Info("simple message", slog.Group("group", slog.String("key", "value")))
+	})
+
+	// Add key/value pairs to logger.
+	run(func() {
+		logger.With("logger_key", "value").Info("simple message")
+	})
+
+	// Add group to logger.
+	run(func() {
+		logger.WithGroup("logger_group").Info("simple message", "key", "value")
+	})
+
+	// Add group and key/value pairs to logger.
+	run(func() {
+		logger.WithGroup("logger_group").With("logger_key", "value").Info("simple message")
+	})
+}
+
+func TestFriendlyHandlerReplaceAttr(t *testing.T) {
+	var out bytes.Buffer
+
+	handler := NewFriendlyHandler(&out, &Options{
+		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
+			if a.Key == "key" {
+				a.Key = "replaced"
+			}
+			return a
+		},
+	})
+
+	logger := slog.New(handler)
+
+	// Helper function to run a test case and print the output.
+	run := func(fn func()) {
+		out.Reset()
+		fn()
+		t.Log(strings.TrimSpace(out.String()))
+	}
+
+	// ReplaceAttr replaces attributes.
+	run(func() {
+		logger.Info("simple message", "key", "value")
+	})
+}

libs/log/handler/options.go

@@ -0,0 +1,15 @@
+package handler
+
+import "log/slog"
+
+type Options struct {
+	// Color enables colorized output.
+	Color bool
+
+	// Level is the minimum enabled logging level.
+	Level slog.Leveler
+
+	// ReplaceAttr is a function that can be used to replace attributes.
+	// Semantics are identical to [slog.ReplaceAttr].
+	ReplaceAttr func(groups []string, a slog.Attr) slog.Attr
+}