...
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
15 Logger logrus.FieldLogger
16
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
27
28
29
30 suppressPgxErrorsKey struct{}
31 )
32
33
34
35
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
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
94 type PgxLogger struct {
95 l Logger
96 }
97
98
99 func NewPgxLogger(l Logger) *PgxLogger {
100 return &PgxLogger{l}
101 }
102
103
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 {
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:
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
131
132 func WithLogger(ctx context.Context, logger Logger) context.Context {
133 return context.WithValue(ctx, loggerKey{}, logger)
134 }
135
136
137 var FallbackLogger = Init(CmdOpts{})
138
139
140
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)
152 return l
153 }
154