...

Source file src/github.com/cybertec-postgresql/pgwatch/v5/internal/log/log.go

Documentation: github.com/cybertec-postgresql/pgwatch/v5/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  	// Logger is the interface used by all components
    15  	Logger logrus.FieldLogger
    16  	//LoggerHooker adds AddHook method to LoggerIface for database logging hook
    17  	LoggerHooker interface {
    18  		Logger
    19  		AddHook(hook logrus.Hook)
    20  		AddSubscriber(msgCh MessageChanType)
    21  		RemoveSubscriber(msgCh MessageChanType)
    22  	}
    23  
    24  	loggerKey struct{}
    25  
    26  	// suppressPgxErrorsKey marks a context in which pgx tracer errors are
    27  	// expected (e.g. a batch that deliberately probes for cascade failures and
    28  	// then retries). Such errors are downgraded to debug so they don't flood
    29  	// the log; the caller is responsible for logging the meaningful outcome.
    30  	suppressPgxErrorsKey struct{}
    31  )
    32  
    33  // WithSuppressedPgxErrors returns a context in which pgx tracer log entries at
    34  // error level are downgraded to debug. Use it around operations where query or
    35  // batch failures are an expected, handled part of the control flow.
    36  func WithSuppressedPgxErrors(ctx context.Context) context.Context {
    37  	return context.WithValue(ctx, suppressPgxErrorsKey{}, true)
    38  }
    39  
    40  func pgxErrorsSuppressed(ctx context.Context) bool {
    41  	v, _ := ctx.Value(suppressPgxErrorsKey{}).(bool)
    42  	return v
    43  }
    44  
    45  type logger struct {
    46  	*logrus.Logger
    47  	*BrokerHook
    48  }
    49  
    50  func getLogFileWriter(opts CmdOpts) any {
    51  	if opts.LogFileRotate {
    52  		return &lumberjack.Logger{
    53  			Filename:   opts.LogFile,
    54  			MaxSize:    opts.LogFileSize,
    55  			MaxBackups: opts.LogFileNumber,
    56  			MaxAge:     opts.LogFileAge,
    57  		}
    58  	}
    59  	return opts.LogFile
    60  }
    61  
    62  const (
    63  	disableColors = true
    64  	enableColors  = false
    65  )
    66  
    67  func getLogFileFormatter(opts CmdOpts) logrus.Formatter {
    68  	if opts.LogFileFormat == "text" {
    69  		return newFormatter(disableColors)
    70  	}
    71  	return &logrus.JSONFormatter{}
    72  }
    73  
    74  // Init creates logging facilities for the application
    75  func Init(opts CmdOpts) LoggerHooker {
    76  	var err error
    77  	l := logger{logrus.New(), NewBrokerHook(context.Background(), opts.LogLevel)}
    78  	l.AddHook(l.BrokerHook)
    79  	l.Out = os.Stdout
    80  	if opts.LogFile > "" {
    81  		l.AddHook(lfshook.NewHook(getLogFileWriter(opts), getLogFileFormatter(opts)))
    82  	}
    83  	l.Level, err = logrus.ParseLevel(opts.LogLevel)
    84  	if err != nil {
    85  		l.Level = logrus.InfoLevel
    86  	}
    87  	l.SetFormatter(newFormatter(enableColors))
    88  	l.SetBrokerFormatter(newFormatter(disableColors))
    89  	l.SetReportCaller(l.Level > logrus.InfoLevel)
    90  	return l
    91  }
    92  
    93  // PgxLogger is the struct used to log using pgx postgres driver
    94  type PgxLogger struct {
    95  	l Logger
    96  }
    97  
    98  // NewPgxLogger returns a new instance of PgxLogger
    99  func NewPgxLogger(l Logger) *PgxLogger {
   100  	return &PgxLogger{l}
   101  }
   102  
   103  // Log transforms logging calls from pgx to logrus
   104  func (pgxlogger *PgxLogger) Log(ctx context.Context, level tracelog.LogLevel, msg string, data map[string]any) {
   105  	logger := GetLogger(ctx)
   106  	if logger == FallbackLogger { //switch from standard to specified
   107  		logger = pgxlogger.l
   108  	}
   109  	if data != nil {
   110  		logger = logger.WithFields(data)
   111  	}
   112  	switch level {
   113  	case tracelog.LogLevelTrace:
   114  		logger.WithField("PGX_LOG_LEVEL", level).Debug(msg)
   115  	case tracelog.LogLevelDebug, tracelog.LogLevelInfo: //pgx is way too chatty on INFO level
   116  		logger.Debug(msg)
   117  	case tracelog.LogLevelWarn:
   118  		logger.Warn(msg)
   119  	case tracelog.LogLevelError:
   120  		if pgxErrorsSuppressed(ctx) {
   121  			logger.Debug(msg)
   122  		} else {
   123  			logger.Error(msg)
   124  		}
   125  	default:
   126  		logger.WithField("INVALID_PGX_LOG_LEVEL", level).Error(msg)
   127  	}
   128  }
   129  
   130  // WithLogger returns a new context with the provided logger. Use in
   131  // combination with logger.WithField(s) for great effect
   132  func WithLogger(ctx context.Context, logger Logger) context.Context {
   133  	return context.WithValue(ctx, loggerKey{}, logger)
   134  }
   135  
   136  // FallbackLogger is an alias for the standard logger
   137  var FallbackLogger = Init(CmdOpts{})
   138  
   139  // GetLogger retrieves the current logger from the context. If no logger is
   140  // available, the default logger is returned
   141  func GetLogger(ctx context.Context) Logger {
   142  	logger := ctx.Value(loggerKey{})
   143  	if logger == nil {
   144  		return FallbackLogger
   145  	}
   146  	return logger.(Logger)
   147  }
   148  
   149  func NewNoopLogger() Logger {
   150  	l := logrus.New()
   151  	l.SetLevel(logrus.PanicLevel) // Noop logger should not output anything
   152  	return l
   153  }
   154