[系列] Go - 基於 GORM 獲取當前請求所執行的 SQL 資訊

新亮筆記發表於2021-01-23
GoORMSQL

前言

為了便於精準排查問題，需要將當前的請求資訊與當前執行的 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 連線的時候使用這個外掛：

// 使用外掛
db.Use(&TracePlugin{})

效果

小結

這是編寫的 trace 包的一部分，這個包可以記錄這些資訊：

  • 支援設定 trace_id
  • 支援設定 request 資訊
  • 支援設定 response 資訊
  • 支援設定 third_party_requests 三方請求資訊
  • 支援設定 debugs 列印除錯資訊
  • 支援設定 sqls 執行 SQL 資訊
  • 可記錄 cost_seconds 執行時長

以上程式碼在 go-gin-api 專案中，地址：https://github.com/xinliangnote/go-gin-api

相關文章