Go基于GORM 获取当前请求所执行的 SQL 信息(思路详解)

脚本专栏 2024/11/15 佚名

3 2 1

前言

为了便于精准排查问题，需要将当前的请求信息与当前执行的 SQL 信息设置对应关系记录下来，记录的 SQL 信息包括：

执行 SQL 的当前时间；
执行 SQL 的文件地址和行号；
执行 SQL 的花费时长；
执行 SQL 的影响行数；
执行的 SQL 语句；

数据库组件使用的是GORM。

思路

1、在执行 SQL 前，设置开始执行时间（计算执行时长会用到）；

2、在执行 SQL 后，第一，获取当前请求的上下文，为什么获取上下文，因为需要从上下文中获取本次请求信息，第二，获取 SQL 执行前的时间，用来计算执行时长，第三，获取执行的 SQL 信息，然后将数据设置到Trace中，Trace是项目中链路包，后面文章会对其介绍；

上面需要用到GORM两个知识点 Callbacks和Context，这两个是在GORM V2才有的，需要 import 的包为gorm.io/gorm。

演示代码

Context的传递需要使用GORM V2提供的WithContext()方法。

func (u *userRepo) getUserByID(ctx core.Context, id uint) (*user_model.UserDemo, error) {
 data := new(user_model.UserDemo)
 err := u.db.GetDbR().WithContext(ctx).First(data, id).Error
 if err != nil {
 return nil, errors.Wrap(err, "[user_demo] get user data err")
 }
 return data, nil
}

编写CallBacks插件代码，GORM 的 Plugin 接口的编写非常简单，只需要实现两个方法即可。

// Plugin GORM plugin interface
type Plugin interface {
 Name() string
 Initialize(*DB) error
}

下面是我写的插件代码：

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
}

var _ gorm.Plugin = &TracePlugin{}

func before(db *gorm.DB) {
 db.InstanceSet(startTime, time.Now())
 return
}

func after(db *gorm.DB) {
 _ctx := db.Statement.Context
 ctx, ok := _ctx.(core.Context)
 if !ok {
 return
 }

 _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(trace.SQL)
 sqlInfo.Timestamp = time_parse.CSTLayoutString()
 sqlInfo.SQL = sql
 sqlInfo.Stack = utils.FileWithLineNum()
 sqlInfo.Rows = db.Statement.RowsAffected
 sqlInfo.CostSeconds = time.Since(ts).Seconds()
 ctx.Trace().AppendSQL(sqlInfo)

 return
}

最后，在 db 连接的时候使用这个插件：