formatter.go

Documentation: github.com/cybertec-postgresql/pgwatch/v3/internal/log

     1  package log
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"runtime"
     7  	"sort"
     8  	"strings"
     9  	"time"
    10  
    11  	"github.com/sirupsen/logrus"
    12  )
    13  
    14  func newFormatter(noColors bool) *Formatter {
    15  	return &Formatter{
    16  		HideKeys: false,
    17  		FieldsOrder: []string{
    18  			"source", "metric",
    19  			"sink", "db", "filename", "address",
    20  			"rows", "sql", "params"},
    21  		TimestampFormat: "2006-01-02 15:04:05.000",
    22  		ShowFullLevel:   true,
    23  		NoColors:        noColors,
    24  	}
    25  }
    26  
    27  // Formatter - logrus formatter, implements logrus.Formatter
    28  // Forked from "github.com/antonfisher/nested-logrus-formatter"
    29  type Formatter struct {
    30  	// FieldsOrder - default: fields sorted alphabetically
    31  	FieldsOrder []string
    32  
    33  	// TimestampFormat - default: time.StampMilli = "Jan _2 15:04:05.000"
    34  	TimestampFormat string
    35  
    36  	// HideKeys - show [fieldValue] instead of [fieldKey:fieldValue]
    37  	HideKeys bool
    38  
    39  	// NoColors - disable colors
    40  	NoColors bool
    41  
    42  	// NoFieldsColors - apply colors only to the level, default is level + fields
    43  	NoFieldsColors bool
    44  
    45  	// NoFieldsSpace - no space between fields
    46  	NoFieldsSpace bool
    47  
    48  	// ShowFullLevel - show a full level [WARNING] instead of [WARN]
    49  	ShowFullLevel bool
    50  
    51  	// NoUppercaseLevel - no upper case for level value
    52  	NoUppercaseLevel bool
    53  
    54  	// TrimMessages - trim whitespaces on messages
    55  	TrimMessages bool
    56  
    57  	// CallerFirst - print caller info first
    58  	CallerFirst bool
    59  
    60  	// CustomCallerFormatter - set custom formatter for caller info
    61  	CustomCallerFormatter func(*runtime.Frame) string
    62  }
    63  
    64  // Format an log entry
    65  func (f *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
    66  	levelColor := getColorByLevel(entry.Level)
    67  
    68  	timestampFormat := f.TimestampFormat
    69  	if timestampFormat == "" {
    70  		timestampFormat = time.StampMilli
    71  	}
    72  
    73  	// output buffer
    74  	b := &bytes.Buffer{}
    75  
    76  	// write time
    77  	b.WriteString(entry.Time.Format(timestampFormat))
    78  
    79  	// write level
    80  	var level string
    81  	if f.NoUppercaseLevel {
    82  		level = entry.Level.String()
    83  	} else {
    84  		level = strings.ToUpper(entry.Level.String())
    85  	}
    86  
    87  	if f.CallerFirst {
    88  		f.writeCaller(b, entry)
    89  	}
    90  
    91  	if !f.NoColors {
    92  		fmt.Fprintf(b, "\x1b[%dm", levelColor)
    93  	}
    94  
    95  	b.WriteString(" [")
    96  	if f.ShowFullLevel {
    97  		b.WriteString(level)
    98  	} else {
    99  		b.WriteString(level[:4])
   100  	}
   101  	b.WriteString("]")
   102  
   103  	if !f.NoFieldsSpace {
   104  		b.WriteString(" ")
   105  	}
   106  
   107  	if !f.NoColors && f.NoFieldsColors {
   108  		b.WriteString("\x1b[0m")
   109  	}
   110  
   111  	// write fields
   112  	if f.FieldsOrder == nil {
   113  		f.writeFields(b, entry)
   114  	} else {
   115  		f.writeOrderedFields(b, entry)
   116  	}
   117  
   118  	if f.NoFieldsSpace {
   119  		b.WriteString(" ")
   120  	}
   121  
   122  	if !f.NoColors && !f.NoFieldsColors {
   123  		b.WriteString("\x1b[0m")
   124  	}
   125  
   126  	// write message
   127  	if f.TrimMessages {
   128  		b.WriteString(strings.TrimSpace(entry.Message))
   129  	} else {
   130  		b.WriteString(entry.Message)
   131  	}
   132  
   133  	if !f.CallerFirst {
   134  		f.writeCaller(b, entry)
   135  	}
   136  
   137  	b.WriteByte('\n')
   138  
   139  	return b.Bytes(), nil
   140  }
   141  
   142  func trimFilename(s string) string {
   143  	const sub = "internal/"
   144  	idx := strings.Index(s, sub)
   145  	if idx == -1 {
   146  		return s
   147  	}
   148  	return s[idx+len(sub):]
   149  }
   150  
   151  func (f *Formatter) writeCaller(b *bytes.Buffer, entry *logrus.Entry) {
   152  	if entry.HasCaller() {
   153  		if f.CustomCallerFormatter != nil {
   154  			fmt.Fprint(b, f.CustomCallerFormatter(entry.Caller))
   155  		} else {
   156  			if strings.Contains(entry.Caller.Function, "PgxLogger") {
   157  				return //skip internal logger function
   158  			}
   159  			fmt.Fprintf(
   160  				b,
   161  				" (%s:%d %s)",
   162  				trimFilename(entry.Caller.File),
   163  				entry.Caller.Line,
   164  				trimFilename(entry.Caller.Function),
   165  			)
   166  		}
   167  	}
   168  }
   169  
   170  func (f *Formatter) writeFields(b *bytes.Buffer, entry *logrus.Entry) {
   171  	if len(entry.Data) != 0 {
   172  		fields := make([]string, 0, len(entry.Data))
   173  		for field := range entry.Data {
   174  			fields = append(fields, field)
   175  		}
   176  
   177  		sort.Strings(fields)
   178  
   179  		for _, field := range fields {
   180  			f.writeField(b, entry, field)
   181  		}
   182  	}
   183  }
   184  
   185  func (f *Formatter) writeOrderedFields(b *bytes.Buffer, entry *logrus.Entry) {
   186  	length := len(entry.Data)
   187  	foundFieldsMap := map[string]bool{}
   188  	for _, field := range f.FieldsOrder {
   189  		if _, ok := entry.Data[field]; ok {
   190  			foundFieldsMap[field] = true
   191  			length--
   192  			f.writeField(b, entry, field)
   193  		}
   194  	}
   195  
   196  	if length > 0 {
   197  		notFoundFields := make([]string, 0, length)
   198  		for field := range entry.Data {
   199  			if !foundFieldsMap[field] {
   200  				notFoundFields = append(notFoundFields, field)
   201  			}
   202  		}
   203  
   204  		sort.Strings(notFoundFields)
   205  
   206  		for _, field := range notFoundFields {
   207  			f.writeField(b, entry, field)
   208  		}
   209  	}
   210  }
   211  
   212  func (f *Formatter) writeField(b *bytes.Buffer, entry *logrus.Entry, field string) {
   213  	if f.HideKeys {
   214  		fmt.Fprintf(b, "[%v]", entry.Data[field])
   215  	} else {
   216  		fmt.Fprintf(b, "[%s:%v]", field, entry.Data[field])
   217  	}
   218  
   219  	if !f.NoFieldsSpace {
   220  		b.WriteString(" ")
   221  	}
   222  }
   223  
   224  const (
   225  	colorRed     = 31
   226  	colorBlue    = 36
   227  	colorMagenta = 35
   228  	colorGreen   = 32
   229  	// colorYellow  = 33
   230  	// colorGray    = 37
   231  )
   232  
   233  func getColorByLevel(level logrus.Level) int {
   234  	switch level {
   235  	case logrus.DebugLevel, logrus.TraceLevel:
   236  		return colorBlue
   237  	case logrus.WarnLevel:
   238  		return colorMagenta
   239  	case logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel:
   240  		return colorRed
   241  	default:
   242  		return colorGreen
   243  	}
   244  }
   245