-
-
Notifications
You must be signed in to change notification settings - Fork 202
/
debug.go
139 lines (119 loc) · 2.71 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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package bundebug
import (
"context"
"database/sql"
"fmt"
"io"
"os"
"reflect"
"time"
"github.com/fatih/color"
"github.com/uptrace/bun"
)
type Option func(*QueryHook)
// WithEnabled enables/disables the hook.
func WithEnabled(on bool) Option {
return func(h *QueryHook) {
h.enabled = on
}
}
// WithVerbose configures the hook to log all queries
// (by default, only failed queries are logged).
func WithVerbose(on bool) Option {
return func(h *QueryHook) {
h.verbose = on
}
}
// WithWriter sets the log output to an io.Writer
// the default is os.Stderr
func WithWriter(w io.Writer) Option {
return func(h *QueryHook) {
h.writer = w
}
}
// FromEnv configures the hook using the environment variable value.
// For example, WithEnv("BUNDEBUG"):
// - BUNDEBUG=0 - disables the hook.
// - BUNDEBUG=1 - enables the hook.
// - BUNDEBUG=2 - enables the hook and verbose mode.
func FromEnv(keys ...string) Option {
if len(keys) == 0 {
keys = []string{"BUNDEBUG"}
}
return func(h *QueryHook) {
for _, key := range keys {
if env, ok := os.LookupEnv(key); ok {
h.enabled = env != "" && env != "0"
h.verbose = env == "2"
break
}
}
}
}
type QueryHook struct {
enabled bool
verbose bool
writer io.Writer
}
var _ bun.QueryHook = (*QueryHook)(nil)
func NewQueryHook(opts ...Option) *QueryHook {
h := &QueryHook{
enabled: true,
writer: os.Stderr,
}
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.enabled {
return
}
if !h.verbose {
switch event.Err {
case nil, sql.ErrNoRows, sql.ErrTxDone:
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.Fprintln(h.writer, 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)
}
}