log.go

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

     1  package log
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"github.com/jackc/pgx/v5/tracelog"
     8  	"github.com/rifflock/lfshook"
     9  	"github.com/sirupsen/logrus"
    10  	"gopkg.in/natefinch/lumberjack.v2"
    11  )
    12  
    13  type (
    14  	// LoggerIface is the interface used by all components
    15  	LoggerIface logrus.FieldLogger
    16  	//LoggerHookerIface adds AddHook method to LoggerIface for database logging hook
    17  	LoggerHookerIface interface {
    18  		LoggerIface
    19  		AddHook(hook logrus.Hook)
    20  		AddSubscriber(msgCh MessageChanType)
    21  		RemoveSubscriber(msgCh MessageChanType)
    22  	}
    23  
    24  	loggerKey struct{}
    25  )
    26  
    27  type logger struct {
    28  	*logrus.Logger
    29  	*BrokerHook
    30  }
    31  
    32  func getLogFileWriter(opts CmdOpts) any {
    33  	if opts.LogFileRotate {
    34  		return &lumberjack.Logger{
    35  			Filename:   opts.LogFile,
    36  			MaxSize:    opts.LogFileSize,
    37  			MaxBackups: opts.LogFileNumber,
    38  			MaxAge:     opts.LogFileAge,
    39  		}
    40  	}
    41  	return opts.LogFile
    42  }
    43  
    44  const (
    45  	disableColors = true
    46  	enableColors  = false
    47  )
    48  
    49  func getLogFileFormatter(opts CmdOpts) logrus.Formatter {
    50  	if opts.LogFileFormat == "text" {
    51  		return newFormatter(disableColors)
    52  	}
    53  	return &logrus.JSONFormatter{}
    54  }
    55  
    56  // Init creates logging facilities for the application
    57  func Init(opts CmdOpts) LoggerHookerIface {
    58  	var err error
    59  	l := logger{logrus.New(), NewBrokerHook(context.Background(), opts.LogLevel)}
    60  	l.AddHook(l.BrokerHook)
    61  	l.Out = os.Stdout
    62  	if opts.LogFile > "" {
    63  		l.AddHook(lfshook.NewHook(getLogFileWriter(opts), getLogFileFormatter(opts)))
    64  	}
    65  	l.Level, err = logrus.ParseLevel(opts.LogLevel)
    66  	if err != nil {
    67  		l.Level = logrus.InfoLevel
    68  	}
    69  	l.SetFormatter(newFormatter(enableColors))
    70  	l.SetBrokerFormatter(newFormatter(disableColors))
    71  	l.SetReportCaller(l.Level > logrus.InfoLevel)
    72  	return l
    73  }
    74  
    75  // PgxLogger is the struct used to log using pgx postgres driver
    76  type PgxLogger struct {
    77  	l LoggerIface
    78  }
    79  
    80  // NewPgxLogger returns a new instance of PgxLogger
    81  func NewPgxLogger(l LoggerIface) *PgxLogger {
    82  	return &PgxLogger{l}
    83  }
    84  
    85  // Log transforms logging calls from pgx to logrus
    86  func (pgxlogger *PgxLogger) Log(ctx context.Context, level tracelog.LogLevel, msg string, data map[string]any) {
    87  	logger := GetLogger(ctx)
    88  	if logger == FallbackLogger { //switch from standard to specified
    89  		logger = pgxlogger.l
    90  	}
    91  	if data != nil {
    92  		logger = logger.WithFields(data)
    93  	}
    94  	switch level {
    95  	case tracelog.LogLevelTrace:
    96  		logger.WithField("PGX_LOG_LEVEL", level).Debug(msg)
    97  	case tracelog.LogLevelDebug, tracelog.LogLevelInfo: //pgx is way too chatty on INFO level
    98  		logger.Debug(msg)
    99  	case tracelog.LogLevelWarn:
   100  		logger.Warn(msg)
   101  	case tracelog.LogLevelError:
   102  		logger.Error(msg)
   103  	default:
   104  		logger.WithField("INVALID_PGX_LOG_LEVEL", level).Error(msg)
   105  	}
   106  }
   107  
   108  // WithLogger returns a new context with the provided logger. Use in
   109  // combination with logger.WithField(s) for great effect
   110  func WithLogger(ctx context.Context, logger LoggerIface) context.Context {
   111  	return context.WithValue(ctx, loggerKey{}, logger)
   112  }
   113  
   114  // FallbackLogger is an alias for the standard logger
   115  var FallbackLogger = Init(CmdOpts{})
   116  
   117  // GetLogger retrieves the current logger from the context. If no logger is
   118  // available, the default logger is returned
   119  func GetLogger(ctx context.Context) LoggerIface {
   120  	logger := ctx.Value(loggerKey{})
   121  	if logger == nil {
   122  		return FallbackLogger
   123  	}
   124  	return logger.(LoggerIface)
   125  }
   126