generated from pkg/go-template
95 lines
2.9 KiB
Go
95 lines
2.9 KiB
Go
|
package orm
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
times "git.icechen.cn/monorepo/backend/pkg/time"
|
|||
|
|
"github.com/gofiber/fiber/v2"
|
|||
|
|
ctxLogger "github.com/luizsuper/ctxLoggers"
|
|||
|
|
"github.com/spf13/cast"
|
|||
|
|
"go.uber.org/zap"
|
|||
|
|
"gorm.io/gorm"
|
|||
|
|
"gorm.io/gorm/utils"
|
|||
|
|
"time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
func GetContextDB(ctx context.Context, db *gorm.DB) (*gorm.DB, error) {
|
|||
|
|
if err := db.Use(&TracePlugin{}); err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
return db.WithContext(ctx), nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type SQL struct {
|
|||
|
|
Timestamp string `json:"timestamp"` // 时间，格式：2006-01-02 15:04:05
|
|||
|
|
Stack string `json:"stack"` // 文件地址和行号
|
|||
|
|
SQL string `json:"sql"` // SQL 语句
|
|||
|
|
Rows int64 `json:"rows_affected"` // 影响行数
|
|||
|
|
CostSeconds float64 `json:"cost_seconds"` // 执行时长(单位秒)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
callBackBeforeName = "core:before"
|
|||
|
|
callBackAfterName = "core:after"
|
|||
|
|
startTime = "_start_time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type TracePlugin struct{}
|
|||
|
|
|
|||
|
|
func (op *TracePlugin) Name() string {
|
|||
|
|
return "tracePlugin"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
//记录db开始时间
|
|||
|
|
func before(db *gorm.DB) {
|
|||
|
|
db.InstanceSet(startTime, time.Now())
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func after(db *gorm.DB) {
|
|||
|
|
_ctx := db.Statement.Context
|
|||
|
|
ctx, ok := _ctx.(context.Context)
|
|||
|
|
if !ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
defer ctx.Done()
|
|||
|
|
_ts, isExist := db.InstanceGet(startTime)
|
|||
|
|
if !isExist {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
ts, ok := _ts.(time.Time)
|
|||
|
|
if !ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
sql := db.Dialector.Explain(db.Statement.SQL.String(), db.Statement.Vars...)
|
|||
|
|
sqlInfo := new(SQL)
|
|||
|
|
sqlInfo.Timestamp = times.CSTLayoutString()
|
|||
|
|
sqlInfo.SQL = sql
|
|||
|
|
sqlInfo.Stack = utils.FileWithLineNum()
|
|||
|
|
sqlInfo.Rows = db.Statement.RowsAffected
|
|||
|
|
sqlInfo.CostSeconds = time.Since(ts).Seconds()
|
|||
|
|
|
|||
|
|
ctxLogger.FInfo(nil, "sql", zap.Any("info", sqlInfo), zap.String("trace_id", cast.ToString(ctx.Value(fiber.HeaderXRequestID))))
|
|||
|
|
|
|||
|
|
return
|
|||
|
|
}
|