-
-
Notifications
You must be signed in to change notification settings - Fork 201
/
debug.go
91 lines (75 loc) · 1.68 KB
/
debug.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
package bundebug
import (
"context"
"database/sql"
"fmt"
"reflect"
"time"
"github.com/fatih/color"
"github.com/uptrace/bun"
)
type ConfigOption func(*QueryHook)
func WithVerbose() ConfigOption {
return func(h *QueryHook) {
h.verbose = true
}
}
type QueryHook struct {
verbose bool
}
var _ bun.QueryHook = (*QueryHook)(nil)
func NewQueryHook(opts ...ConfigOption) *QueryHook {
h := new(QueryHook)
for _, opt := range opts {
opt(h)
}
return h
}
func (h *QueryHook) BeforeQuery(
ctx context.Context, event *bun.QueryEvent,
) context.Context {
return ctx
}
func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
if !h.verbose {
switch event.Err {
case nil, sql.ErrNoRows:
return
}
}
now := time.Now()
dur := now.Sub(event.StartTime)
args := []interface{}{
"[bun]",
now.Format(" 15:04:05.000 "),
formatOperation(event),
fmt.Sprintf(" %10s ", dur.Round(time.Microsecond)),
event.Query,
}
if event.Err != nil {
typ := reflect.TypeOf(event.Err).String()
args = append(args,
"\t",
color.New(color.BgRed).Sprintf(" %s ", typ+": "+event.Err.Error()),
)
}
fmt.Println(args...)
}
func formatOperation(event *bun.QueryEvent) string {
operation := event.Operation()
return operationColor(operation).Sprintf(" %-16s ", operation)
}
func operationColor(operation string) *color.Color {
switch operation {
case "SELECT":
return color.New(color.BgGreen, color.FgHiWhite)
case "INSERT":
return color.New(color.BgBlue, color.FgHiWhite)
case "UPDATE":
return color.New(color.BgYellow, color.FgHiBlack)
case "DELETE":
return color.New(color.BgMagenta, color.FgHiWhite)
default:
return color.New(color.BgWhite, color.FgHiBlack)
}
}