/
plugin.go
68 lines (55 loc) · 2.09 KB
/
plugin.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
// Copyright 2022 Lingfei Kong <colin404@foxmail.com>. All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file. The original repo for
// this file is https://github.com/superproj/onex.
//
package db
import (
"time"
"gorm.io/gorm"
"k8s.io/klog/v2"
)
const (
callBackBeforeName = "core:before"
callBackAfterName = "core:after"
startTime = "_start_time"
)
// TracePlugin defines gorm plugin used to trace sql.
type TracePlugin struct{}
// Name returns the name of trace plugin.
func (op *TracePlugin) Name() string {
return "tracePlugin"
}
// Initialize initialize the trace plugin.
func (op *TracePlugin) Initialize(db *gorm.DB) (err error) {
// 开始前
_ = db.Callback().Create().Before("gorm:before_create").Register(callBackBeforeName, before)
_ = db.Callback().Query().Before("gorm:query").Register(callBackBeforeName, before)
_ = db.Callback().Delete().Before("gorm:before_delete").Register(callBackBeforeName, before)
_ = db.Callback().Update().Before("gorm:setup_reflect_value").Register(callBackBeforeName, before)
_ = db.Callback().Row().Before("gorm:row").Register(callBackBeforeName, before)
_ = db.Callback().Raw().Before("gorm:raw").Register(callBackBeforeName, before)
// 结束后
_ = db.Callback().Create().After("gorm:after_create").Register(callBackAfterName, after)
_ = db.Callback().Query().After("gorm:after_query").Register(callBackAfterName, after)
_ = db.Callback().Delete().After("gorm:after_delete").Register(callBackAfterName, after)
_ = db.Callback().Update().After("gorm:after_update").Register(callBackAfterName, after)
_ = db.Callback().Row().After("gorm:row").Register(callBackAfterName, after)
_ = db.Callback().Raw().After("gorm:raw").Register(callBackAfterName, after)
return
}
var _ gorm.Plugin = &TracePlugin{}
func before(db *gorm.DB) {
db.InstanceSet(startTime, time.Now())
}
func after(db *gorm.DB) {
_ts, isExist := db.InstanceGet(startTime)
if !isExist {
return
}
ts, ok := _ts.(time.Time)
if !ok {
return
}
klog.V(4).Infof("sql cost time: %fs", time.Since(ts).Seconds())
}