mirror of https://github.com/databricks/cli.git
136 lines
3.3 KiB
Go
136 lines
3.3 KiB
Go
package root
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
|
|
"github.com/databricks/cli/libs/cmdio"
|
|
"github.com/databricks/cli/libs/flags"
|
|
"github.com/databricks/cli/libs/log"
|
|
"github.com/fatih/color"
|
|
"golang.org/x/exp/slog"
|
|
)
|
|
|
|
const (
|
|
envLogFile = "DATABRICKS_LOG_FILE"
|
|
envLogLevel = "DATABRICKS_LOG_LEVEL"
|
|
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) {
|
|
attrs += fmt.Sprintf(" %s%s%s",
|
|
color.CyanString(a.Key),
|
|
color.CyanString("="),
|
|
color.YellowString(a.Value.String()))
|
|
})
|
|
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
|
|
}
|
|
|
|
func makeLogHandler(opts slog.HandlerOptions) (slog.Handler, error) {
|
|
switch logOutput {
|
|
case flags.OutputJSON:
|
|
return opts.NewJSONHandler(logFile.Writer()), nil
|
|
case flags.OutputText:
|
|
w := logFile.Writer()
|
|
if cmdio.IsTTY(w) {
|
|
return &friendlyHandler{
|
|
Handler: opts.NewTextHandler(w),
|
|
w: w,
|
|
}, nil
|
|
}
|
|
return opts.NewTextHandler(w), nil
|
|
|
|
default:
|
|
return nil, fmt.Errorf("invalid log output mode: %s", logOutput)
|
|
}
|
|
}
|
|
|
|
func initializeLogger(ctx context.Context) (context.Context, error) {
|
|
opts := slog.HandlerOptions{}
|
|
opts.Level = logLevel.Level()
|
|
opts.AddSource = true
|
|
opts.ReplaceAttr = log.ReplaceAttrFunctions{
|
|
log.ReplaceLevelAttr,
|
|
log.ReplaceSourceAttr,
|
|
}.ReplaceAttr
|
|
|
|
// Open the underlying log file if the user configured an actual file to log to.
|
|
err := logFile.Open()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
handler, err := makeLogHandler(opts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
slog.SetDefault(slog.New(handler))
|
|
return log.NewContext(ctx, slog.Default()), nil
|
|
}
|
|
|
|
var logFile = flags.NewLogFileFlag()
|
|
var logLevel = flags.NewLogLevelFlag()
|
|
var logOutput = flags.OutputText
|
|
|
|
func init() {
|
|
// Configure defaults from environment, if applicable.
|
|
// If the provided value is invalid it is ignored.
|
|
if v, ok := os.LookupEnv(envLogFile); ok {
|
|
logFile.Set(v)
|
|
}
|
|
if v, ok := os.LookupEnv(envLogLevel); ok {
|
|
logLevel.Set(v)
|
|
}
|
|
if v, ok := os.LookupEnv(envLogFormat); ok {
|
|
logOutput.Set(v)
|
|
}
|
|
|
|
RootCmd.PersistentFlags().Var(&logFile, "log-file", "file to write logs to")
|
|
RootCmd.PersistentFlags().Var(&logLevel, "log-level", "log level")
|
|
RootCmd.PersistentFlags().Var(&logOutput, "log-format", "log output format (text or json)")
|
|
RootCmd.RegisterFlagCompletionFunc("log-file", logFile.Complete)
|
|
RootCmd.RegisterFlagCompletionFunc("log-level", logLevel.Complete)
|
|
RootCmd.RegisterFlagCompletionFunc("log-format", logOutput.Complete)
|
|
}
|