handler.go raw

   1  package log
   2  
   3  import (
   4  	"bytes"
   5  	"context"
   6  	"encoding/json"
   7  	"fmt"
   8  	"io"
   9  	"log/slog"
  10  	"strings"
  11  )
  12  
  13  var timeFormat = "2006/01/02 15:04:05"
  14  
  15  var levels = map[slog.Level]string{
  16  	LevelFatal:      "FATAL",
  17  	slog.LevelError: "ERROR",
  18  	slog.LevelWarn:  "WARN",
  19  	slog.LevelInfo:  "INFO",
  20  	slog.LevelDebug: "DEBUG",
  21  	LevelTrace:      "TRACE",
  22  }
  23  
  24  // SlogHandler is custom slog handler
  25  type SlogHandler struct {
  26  	h      slog.Handler
  27  	writer io.Writer
  28  	goas   []groupOrAttrs
  29  }
  30  type groupOrAttrs struct {
  31  	group string
  32  	attrs []slog.Attr
  33  }
  34  
  35  // Enabled is a wrapper over slog.Handler.Enabled method
  36  func (h *SlogHandler) Enabled(ctx context.Context, level slog.Level) bool {
  37  	return h.h.Enabled(ctx, level)
  38  }
  39  
  40  // Handle handles the Record
  41  // It will only be called when Enabled returns true
  42  func (h *SlogHandler) Handle(_ context.Context, r slog.Record) error {
  43  
  44  	buf := make([]byte, 0, 1024)
  45  
  46  	// If the record has no Attrs, remove groups at the end of the list; they are empty.
  47  	goas := normalizeGroupsAndAttributes(h.goas, r.NumAttrs())
  48  
  49  	for _, goa := range goas {
  50  		if goa.group != "" {
  51  			buf = fmt.Appendf(buf, "%s:\n", goa.group)
  52  		} else {
  53  			for _, a := range goa.attrs {
  54  				buf = h.appendAttr(buf, a)
  55  			}
  56  		}
  57  	}
  58  
  59  	r.Attrs(func(a slog.Attr) bool {
  60  		buf = h.appendAttr(buf, a)
  61  		return true
  62  	})
  63  
  64  	levelString := levels[r.Level]
  65  
  66  	if len(buf) > 0 {
  67  		_, err := fmt.Fprintf(h.writer, "%s [%s] %s: %s\n", r.Time.Format(timeFormat), levelString, r.Message, buf)
  68  		return err
  69  	}
  70  	_, err := fmt.Fprintf(h.writer, "%s [%s] %s\n", r.Time.Format(timeFormat), levelString, r.Message)
  71  	return err
  72  }
  73  
  74  // WithAttrs returns a new Handler whose attributes consist of
  75  // both the receiver's attributes and the arguments.
  76  // The Handler owns the slice: it may retain, modify or discard it.
  77  func (h *SlogHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
  78  	if len(attrs) == 0 {
  79  		return h
  80  	}
  81  	return h.withGroupOrAttrs(groupOrAttrs{attrs: attrs})
  82  }
  83  
  84  // WithGroup returns a new Handler with the given group appended to
  85  // the receiver's existing groups.
  86  func (h *SlogHandler) WithGroup(name string) slog.Handler {
  87  	if name == "" {
  88  		return h
  89  	}
  90  	return h.withGroupOrAttrs(groupOrAttrs{group: name})
  91  }
  92  
  93  func (h *SlogHandler) withGroupOrAttrs(goa groupOrAttrs) *SlogHandler {
  94  	h2 := *h
  95  	h2.goas = make([]groupOrAttrs, len(h.goas)+1)
  96  	copy(h2.goas, h.goas)
  97  	h2.goas[len(h2.goas)-1] = goa
  98  	return &h2
  99  }
 100  
 101  func normalizeGroupsAndAttributes(groupOfAttrs []groupOrAttrs, numAttrs int) []groupOrAttrs {
 102  	goas := groupOfAttrs
 103  	if numAttrs == 0 {
 104  		for len(goas) > 0 && goas[len(goas)-1].group != "" {
 105  			goas = goas[:len(goas)-1]
 106  		}
 107  	}
 108  	return goas
 109  }
 110  
 111  func (h *SlogHandler) appendAttr(buf []byte, a slog.Attr) []byte {
 112  	a.Value = a.Value.Resolve()
 113  	if a.Equal(slog.Attr{}) {
 114  		return buf
 115  	}
 116  
 117  	if len(buf) != 0 {
 118  		buf = fmt.Append(buf, " ")
 119  	}
 120  	switch a.Value.Kind() {
 121  	case slog.KindTime:
 122  		buf = fmt.Appendf(buf, "%s=%s", a.Key, a.Value.Time().Format(timeFormat))
 123  
 124  	default:
 125  		buf = parseDefault(buf, a.Key, a.Value.Any())
 126  	}
 127  
 128  	return buf
 129  
 130  }
 131  
 132  func parseDefault(buf []byte, key string, value any) []byte {
 133  
 134  	if f, ok := value.(Fields); ok {
 135  		// Inline Fields data if key is empty
 136  		if key == "" {
 137  			inlined := []string{}
 138  			args := f.Get()
 139  			for i := 0; i < len(args); i += 2 {
 140  				inlined = append(inlined, fmt.Sprintf("%s=%v", args[i], args[i+1]))
 141  			}
 142  			buf = fmt.Appendf(buf, "%v", strings.Join(inlined, " "))
 143  			return buf
 144  		}
 145  		t, _ := json.Marshal(f)
 146  		value = string(t)
 147  	}
 148  
 149  	buf = fmt.Appendf(buf, "%v=%v", key, value)
 150  	return buf
 151  }
 152  
 153  func suppressDefaults(
 154  	next func([]string, slog.Attr) slog.Attr,
 155  ) func([]string, slog.Attr) slog.Attr {
 156  	return func(groups []string, a slog.Attr) slog.Attr {
 157  		if a.Key == slog.TimeKey ||
 158  			a.Key == slog.LevelKey ||
 159  			a.Key == slog.MessageKey {
 160  			return slog.Attr{}
 161  		}
 162  		if next == nil {
 163  			return a
 164  		}
 165  		return next(groups, a)
 166  	}
 167  }
 168  
 169  // NewSlogHandler returns configured SlogHandler
 170  func NewSlogHandler(w io.Writer, opts *slog.HandlerOptions) *SlogHandler {
 171  
 172  	if opts == nil {
 173  		opts = &slog.HandlerOptions{}
 174  	}
 175  	b := &bytes.Buffer{}
 176  
 177  	return &SlogHandler{
 178  		h: slog.NewJSONHandler(b, &slog.HandlerOptions{
 179  			Level:       opts.Level,
 180  			AddSource:   opts.AddSource,
 181  			ReplaceAttr: suppressDefaults(opts.ReplaceAttr),
 182  		}),
 183  		writer: w,
 184  	}
 185  }
 186