diff --git a/cmd/root/logger.go b/cmd/root/logger.go index be342a7a..494b28fc 100644 --- a/cmd/root/logger.go +++ b/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 diff --git a/libs/log/handler/colors.go b/libs/log/handler/colors.go new file mode 100644 index 00000000..a1b8e849 --- /dev/null +++ b/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 +} diff --git a/libs/log/handler/colors_test.go b/libs/log/handler/colors_test.go new file mode 100644 index 00000000..aa042fb0 --- /dev/null +++ b/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)) +} diff --git a/libs/log/handler/friendly.go b/libs/log/handler/friendly.go new file mode 100644 index 00000000..33b88a9e --- /dev/null +++ b/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}) +} diff --git a/libs/log/handler/friendly_test.go b/libs/log/handler/friendly_test.go new file mode 100644 index 00000000..ca6e823e --- /dev/null +++ b/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") + }) +} diff --git a/libs/log/handler/options.go b/libs/log/handler/options.go new file mode 100644 index 00000000..0b8cfbe2 --- /dev/null +++ b/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 +}