hook.go
2.9 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
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
package pg
import (
"context"
"fmt"
"time"
"github.com/go-pg/pg/v10/orm"
)
type BeforeScanHook = orm.BeforeScanHook
type AfterScanHook = orm.AfterScanHook
type AfterSelectHook = orm.AfterSelectHook
type BeforeInsertHook = orm.BeforeInsertHook
type AfterInsertHook = orm.AfterInsertHook
type BeforeUpdateHook = orm.BeforeUpdateHook
type AfterUpdateHook = orm.AfterUpdateHook
type BeforeDeleteHook = orm.BeforeDeleteHook
type AfterDeleteHook = orm.AfterDeleteHook
//------------------------------------------------------------------------------
type dummyFormatter struct{}
func (dummyFormatter) FormatQuery(b []byte, query string, params ...interface{}) []byte {
return append(b, query...)
}
// QueryEvent ...
type QueryEvent struct {
StartTime time.Time
DB orm.DB
Model interface{}
Query interface{}
Params []interface{}
fmtedQuery []byte
Result Result
Err error
Stash map[interface{}]interface{}
}
// QueryHook ...
type QueryHook interface {
BeforeQuery(context.Context, *QueryEvent) (context.Context, error)
AfterQuery(context.Context, *QueryEvent) error
}
// UnformattedQuery returns the unformatted query of a query event.
// The query is only valid until the query Result is returned to the user.
func (e *QueryEvent) UnformattedQuery() ([]byte, error) {
return queryString(e.Query)
}
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)
}
}
// FormattedQuery returns the formatted query of a query event.
// The query is only valid until the query Result is returned to the user.
func (e *QueryEvent) FormattedQuery() ([]byte, error) {
return e.fmtedQuery, nil
}
// AddQueryHook adds a hook into query processing.
func (db *baseDB) AddQueryHook(hook QueryHook) {
db.queryHooks = append(db.queryHooks, hook)
}
func (db *baseDB) beforeQuery(
ctx context.Context,
ormDB orm.DB,
model, query interface{},
params []interface{},
fmtedQuery []byte,
) (context.Context, *QueryEvent, error) {
if len(db.queryHooks) == 0 {
return ctx, nil, nil
}
event := &QueryEvent{
StartTime: time.Now(),
DB: ormDB,
Model: model,
Query: query,
Params: params,
fmtedQuery: fmtedQuery,
}
for _, hook := range db.queryHooks {
var err error
ctx, err = hook.BeforeQuery(ctx, event)
if err != nil {
return nil, nil, err
}
}
return ctx, event, nil
}
func (db *baseDB) afterQuery(
ctx context.Context,
event *QueryEvent,
res Result,
err error,
) error {
if event == nil {
return nil
}
event.Err = err
event.Result = res
for _, hook := range db.queryHooks {
err := hook.AfterQuery(ctx, event)
if err != nil {
return err
}
}
return nil
}
func copyQueryHooks(s []QueryHook) []QueryHook {
return s[:len(s):len(s)]
}