logger/wrapper/gorm_v2.go

// Package wrapper...
//
// Description : gorm v2 版本接口实现
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021/03/01 9:52 下午
package wrapper

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"gorm.io/gorm"

	"go.uber.org/zap/zapcore"

	"go.uber.org/zap"

	pkgGinDefine "git.zhangdeman.cn/zhangdeman/gin/define"
	logger2 "git.zhangdeman.cn/zhangdeman/logger"
	"gorm.io/gorm/logger"
)

// NewGormV2 获取日志实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 9:56 下午 2021/3/1
func NewGormV2(loggerLevel zapcore.Level, consoleOutput bool, encoder zapcore.Encoder, splitConfig *logger2.RotateLogConfig, traceIDField string, skip int) (logger.Interface, error) {
	logConfList := []logger2.SetLoggerOptionFunc{logger2.WithEncoder(encoder), logger2.WithCallerSkip(skip), logger2.WithCaller()}
	if consoleOutput {
		logConfList = append(logConfList, logger2.WithConsoleOutput())
	}

	logInstance, err := logger2.NewLogger(loggerLevel, splitConfig, logConfList...)
	if nil != err {
		return nil, err
	}
	if len(traceIDField) == 0 {
		traceIDField = "trace_id"
	}

	return &Gorm{
		instance:     logInstance,
		traceIDField: traceIDField,
	}, nil
}

// NewGormLoggerWithInstance 获取gorm日志实现
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 3:36 PM 2021/12/24
func NewGormLoggerWithInstance(ctx *gin.Context, dbClient *gorm.DB, instance *zap.Logger, node string, extraCtxFieldList []string) logger.Interface {
	nodeArr := strings.Split(node, "|")
	i := &Gorm{
		dbClient:          dbClient,
		instance:          instance,
		traceIDField:      pkgGinDefine.TraceIDField,
		extraCtxFieldList: extraCtxFieldList,
		flag:              "",
		node:              node,
		ctx:               ctx,
	}
	if len(nodeArr) >= 2 {
		i.node = nodeArr[1]
		i.flag = nodeArr[0]
	}
	return i
}

// Gorm v2 版本库日志实现
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 9:55 下午 2021/3/1
type Gorm struct {
	dbClient          *gorm.DB
	instance          *zap.Logger  // 日志实例
	traceIDField      string       // 串联请求上下文的的ID
	extraCtxFieldList []string     // 从请求上线问提取的字段
	flag              string       // 数据库标识
	node              string       // 数据库节点 master / slave
	ctx               *gin.Context // gin上下文
}

// LogMode ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:08 下午 2021/3/1
func (g *Gorm) LogMode(level logger.LogLevel) logger.Interface {
	return g
}

// Info 日志
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:18 下午 2021/3/1
func (g *Gorm) Info(ctx context.Context, s string, i ...interface{}) {
	g.write(nil, "info")
}

// Warn ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:16 下午 2021/3/1
func (g *Gorm) Warn(ctx context.Context, s string, i ...interface{}) {
	g.write(nil, "warn")
}

// Error 日志
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:18 下午 2021/3/1
func (g *Gorm) Error(ctx context.Context, s string, i ...interface{}) {
	g.write(nil, "error")
}

// Trace Trace 记录
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:19 下午 2021/3/1
func (g *Gorm) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
	start := begin.UnixNano()
	end := time.Now().UnixNano()
	sql := ""
	affectRows := int64(0)
	if nil != fc {
		sql, affectRows = fc()
	}

	dataList := []zap.Field{
		zap.String("db_flag", g.flag),
		zap.String("db_node", g.node),
		zap.Int64("begin_time", start),
		zap.Int64("finish_time", end),
		zap.String("used_time", fmt.Sprintf("%fms", float64(end-start)/1e6)),
		zap.String("sql", sql),
		zap.Int64("affect_rows", affectRows),
		zap.Error(err),
	}
	if len(g.traceIDField) > 0 {
		dataList = append(dataList, zap.String(g.traceIDField, g.getTraceID(ctx)))
	}
	g.write(dataList, "info")
}

// write ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 4:11 PM 2021/12/24
func (g *Gorm) write(dataList []zap.Field, level string) {
	if nil == g.instance {
		// 未设置日志实例
		return
	}
	if nil == dataList {
		dataList = make([]zap.Field, 0)
	}
	if nil != g.ctx {
		for _, extraField := range g.extraCtxFieldList {
			if len(extraField) == 0 {
				continue
			}
			dataList = append(dataList, zap.Any(extraField, g.ctx.Value(extraField)))
		}
	}

	message := "SQL执行记录"
	switch strings.ToLower(level) {
	case "info":
		g.instance.Info(message, dataList...)
	case "warn":
		g.instance.Warn(message, dataList...)
	case "error":
		g.instance.Error(message, dataList...)
	default:
		g.instance.Info(message, dataList...)
	}
}

// getTraceID 获取traceID
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:11 下午 2021/3/1
func (g *Gorm) getTraceID(ctx context.Context) string {
	if nil == g.ctx {
		return ""
	}
	return g.ctx.GetString(g.traceIDField)
}

// GetGormSQL 获取trace fn
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:38 下午 2021/3/1
func GetGormSQL(dbClient *gorm.DB) func() (string, int64) {
	return func() (string, int64) {
		return dbClient.Dialector.Explain(dbClient.Statement.SQL.String(), dbClient.Statement.Vars...), dbClient.RowsAffected
	}
}