package logger import ( "bytes" "context" "database/sql" "fmt" "strings" "text/template" "time" "github.com/uptrace/bun" "go.uber.org/zap" "go.uber.org/zap/zapcore" ) // QueryHookOptions logging options type QueryHookOptions struct { LogSlow time.Duration Logger *zap.Logger QueryLevel zapcore.Level SlowLevel zapcore.Level ErrorLevel zapcore.Level MessageTemplate string ErrorTemplate string } // QueryHook wraps query hook type QueryHook struct { opts QueryHookOptions errorTemplate *template.Template messageTemplate *template.Template } // LogEntryVars variables made available t otemplate type LogEntryVars struct { Timestamp time.Time Query string Operation string Duration time.Duration Error error } // NewQueryHook returns new instance func NewQueryHook(opts QueryHookOptions) *QueryHook { h := new(QueryHook) if opts.ErrorTemplate == "" { opts.ErrorTemplate = "{{.Operation}}[{{.Duration}}]: {{.Query}}: {{.Error}}" } if opts.MessageTemplate == "" { opts.MessageTemplate = "{{.Operation}}[{{.Duration}}]: {{.Query}}" } h.opts = opts errorTemplate, err := template.New("ErrorTemplate").Parse(h.opts.ErrorTemplate) if err != nil { panic(err) } messageTemplate, err := template.New("MessageTemplate").Parse(h.opts.MessageTemplate) if err != nil { panic(err) } h.errorTemplate = errorTemplate h.messageTemplate = messageTemplate return h } // BeforeQuery does nothing tbh func (h *QueryHook) BeforeQuery(ctx context.Context, event *bun.QueryEvent) context.Context { return ctx } // AfterQuery convert a bun QueryEvent into a logrus message func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) { var level zapcore.Level var isError bool var msg bytes.Buffer now := time.Now() dur := now.Sub(event.StartTime) switch event.Err { case nil, sql.ErrNoRows: isError = false if h.opts.LogSlow > 0 && dur >= h.opts.LogSlow { level = h.opts.SlowLevel } else { level = h.opts.QueryLevel } default: isError = true level = h.opts.ErrorLevel } if level == 0 { return } args := &LogEntryVars{ Timestamp: now, Query: string(event.Query), Operation: eventOperation(event), Duration: dur, Error: event.Err, } if isError { if err := h.errorTemplate.Execute(&msg, args); err != nil { panic(err) } } else { if err := h.messageTemplate.Execute(&msg, args); err != nil { panic(err) } } switch level { case zapcore.DebugLevel: h.opts.Logger.Debug(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs)) case zapcore.InfoLevel: h.opts.Logger.Info(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs)) case zapcore.WarnLevel: h.opts.Logger.Warn(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs)) case zapcore.ErrorLevel: h.opts.Logger.Error(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs), zap.Error(event.Err)) case zapcore.FatalLevel: h.opts.Logger.Fatal(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs), zap.Error(event.Err)) case zapcore.PanicLevel: h.opts.Logger.Panic(msg.String(), zap.String("query", event.Query), zap.Any("args", event.QueryArgs), zap.Error(event.Err)) default: panic(fmt.Errorf("unsupported level: %v", level)) } } // taken from bun func eventOperation(event *bun.QueryEvent) string { switch event.QueryAppender.(type) { case *bun.SelectQuery: return "SELECT" case *bun.InsertQuery: return "INSERT" case *bun.UpdateQuery: return "UPDATE" case *bun.DeleteQuery: return "DELETE" case *bun.CreateTableQuery: return "CREATE TABLE" case *bun.DropTableQuery: return "DROP TABLE" } return queryOperation(event.Query) } // taken from bun func queryOperation(name string) string { if idx := strings.Index(name, " "); idx > 0 { name = name[:idx] } if len(name) > 16 { name = name[:16] } return string(name) }