hook.go
2.0 KB
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
package pg
import (
"context"
"fmt"
"github.com/go-pg/pg/orm"
)
type dummyFormatter struct{}
func (dummyFormatter) FormatQuery(b []byte, query string, params ...interface{}) []byte {
return append(b, query...)
}
type QueryEvent struct {
Ctx context.Context
DB orm.DB
Query interface{}
Params []interface{}
Attempt int
Result Result
Error error
Data map[interface{}]interface{}
}
type QueryHook interface {
BeforeQuery(*QueryEvent)
AfterQuery(*QueryEvent)
}
func (ev *QueryEvent) UnformattedQuery() (string, error) {
b, err := queryString(ev.Query)
if err != nil {
return "", err
}
return string(b), nil
}
func (ev *QueryEvent) FormattedQuery() (string, error) {
b, err := appendQuery(nil, ev.DB, ev.Query, ev.Params...)
if err != nil {
return "", err
}
return string(b), nil
}
func queryString(query interface{}) ([]byte, error) {
switch query := query.(type) {
case orm.TemplateAppender:
return query.AppendTemplate(nil)
case string:
return dummyFormatter{}.FormatQuery(nil, query), nil
default:
return nil, fmt.Errorf("pg: can't append %T", query)
}
}
// AddQueryHook adds a hook into query processing.
func (db *baseDB) AddQueryHook(hook QueryHook) {
db.queryHooks = append(db.queryHooks, hook)
}
func (db *baseDB) queryStarted(
c context.Context,
ormDB orm.DB,
query interface{},
params []interface{},
attempt int,
) *QueryEvent {
if len(db.queryHooks) == 0 {
return nil
}
event := &QueryEvent{
Ctx: c,
DB: ormDB,
Query: query,
Params: params,
Attempt: attempt,
Data: make(map[interface{}]interface{}),
}
for _, hook := range db.queryHooks {
hook.BeforeQuery(event)
}
return event
}
func (db *baseDB) queryProcessed(
res Result,
err error,
event *QueryEvent,
) {
if event == nil {
return
}
event.Error = err
event.Result = res
for _, hook := range db.queryHooks {
hook.AfterQuery(event)
}
}
func copyQueryHooks(s []QueryHook) []QueryHook {
return s[:len(s):len(s)]
}