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.
This commit is contained in:
Pieter Noordhuis 2023-12-01 13:17:04 +01:00 committed by GitHub
parent 83d50001fc
commit 60a8abdcd7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 466 additions and 54 deletions

View File

@ -3,7 +3,6 @@ package root
import ( import (
"context" "context"
"fmt" "fmt"
"io"
"log/slog" "log/slog"
"os" "os"
@ -11,7 +10,7 @@ import (
"github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/env"
"github.com/databricks/cli/libs/flags" "github.com/databricks/cli/libs/flags"
"github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/log"
"github.com/fatih/color" "github.com/databricks/cli/libs/log/handler"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
@ -21,54 +20,6 @@ const (
envLogFormat = "DATABRICKS_LOG_FORMAT" 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 { type logFlags struct {
file flags.LogFileFlag file flags.LogFileFlag
level flags.LogLevelFlag level flags.LogLevelFlag
@ -83,10 +34,11 @@ func (f *logFlags) makeLogHandler(opts slog.HandlerOptions) (slog.Handler, error
case flags.OutputText: case flags.OutputText:
w := f.file.Writer() w := f.file.Writer()
if cmdio.IsTTY(w) { if cmdio.IsTTY(w) {
return &friendlyHandler{ return handler.NewFriendlyHandler(w, &handler.Options{
Handler: slog.NewTextHandler(w, &opts), Color: true,
w: w, Level: opts.Level,
}, nil ReplaceAttr: opts.ReplaceAttr,
}), nil
} }
return slog.NewTextHandler(w, &opts), nil return slog.NewTextHandler(w, &opts), nil

View File

@ -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
}

View File

@ -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))
}

View File

@ -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})
}

View File

@ -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")
})
}

View File

@ -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
}