databricks-cli/cmd/root/logger.go

113 lines
2.8 KiB
Go

package root
import (
"context"
"fmt"
"log/slog"
"os"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/env"
"github.com/databricks/cli/libs/flags"
"github.com/databricks/cli/libs/log"
"github.com/databricks/cli/libs/log/handler"
"github.com/spf13/cobra"
)
const (
envLogFile = "DATABRICKS_LOG_FILE"
envLogLevel = "DATABRICKS_LOG_LEVEL"
envLogFormat = "DATABRICKS_LOG_FORMAT"
)
type logFlags struct {
file flags.LogFileFlag
level flags.LogLevelFlag
output flags.Output
debug bool
}
func (f *logFlags) makeLogHandler(opts slog.HandlerOptions) (slog.Handler, error) {
switch f.output {
case flags.OutputJSON:
return slog.NewJSONHandler(f.file.Writer(), &opts), nil
case flags.OutputText:
w := f.file.Writer()
if cmdio.IsTTY(w) {
return handler.NewFriendlyHandler(w, &handler.Options{
Color: true,
Level: opts.Level,
ReplaceAttr: opts.ReplaceAttr,
}), nil
}
return slog.NewTextHandler(w, &opts), nil
default:
return nil, fmt.Errorf("invalid log output mode: %s", f.output)
}
}
func (f *logFlags) initializeContext(ctx context.Context) (context.Context, error) {
if f.debug {
f.level.Set("debug")
}
opts := slog.HandlerOptions{}
opts.Level = f.level.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 := f.file.Open()
if err != nil {
return nil, err
}
handler, err := f.makeLogHandler(opts)
if err != nil {
return nil, err
}
slog.SetDefault(slog.New(handler).With(slog.Int("pid", os.Getpid())))
return log.NewContext(ctx, slog.Default()), nil
}
func initLogFlags(cmd *cobra.Command) *logFlags {
f := logFlags{
file: flags.NewLogFileFlag(),
level: flags.NewLogLevelFlag(),
output: flags.OutputText,
}
// Configure defaults from environment, if applicable.
// If the provided value is invalid it is ignored.
if v, ok := env.Lookup(cmd.Context(), envLogFile); ok {
f.file.Set(v)
}
if v, ok := env.Lookup(cmd.Context(), envLogLevel); ok {
f.level.Set(v)
}
if v, ok := env.Lookup(cmd.Context(), envLogFormat); ok {
f.output.Set(v)
}
flags := cmd.PersistentFlags()
flags.BoolVar(&f.debug, "debug", false, "enable debug logging")
flags.Var(&f.file, "log-file", "file to write logs to")
flags.Var(&f.level, "log-level", "log level")
flags.Var(&f.output, "log-format", "log output format (text or json)")
// mark fine-grained flags hidden from global --help
flags.MarkHidden("log-file")
flags.MarkHidden("log-level")
flags.MarkHidden("log-format")
cmd.RegisterFlagCompletionFunc("log-file", f.file.Complete)
cmd.RegisterFlagCompletionFunc("log-level", f.level.Complete)
cmd.RegisterFlagCompletionFunc("log-format", f.output.Complete)
return &f
}