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"
	"strings"
	"time"

	"git.zhangdeman.cn/zhangdeman/consts"

	"gorm.io/gorm"

	"go.uber.org/zap/zapcore"

	"go.uber.org/zap"

	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 string, 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(outCtx context.Context, dbClient *gorm.DB, instance *zap.Logger, node string, extraCtxFieldList []string) logger.Interface {
	nodeArr := strings.Split(node, "|")
	i := &Gorm{
		dbClient:          dbClient,
		instance:          instance,
		traceIDField:      consts.GinTraceIDField,
		extraCtxFieldList: extraCtxFieldList,
		flag:              "",
		node:              node,
		outCtx:            outCtx,
	}
	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
	outCtx            context.Context
}

// 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 ...any) {
	g.write(ctx, s, consts.LogLevelInfo, map[string]any{
		"log_data": i,
	})
}

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

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

// 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()
	}

	g.write(ctx, "SQL执行记录", consts.LogLevelInfo, map[string]any{
		"db_flag":     g.flag,
		"db_node":     g.node,
		"begin_time":  start,
		"finish_time": end,
		"used_time":   (end - start) / 1e6,
		"sql":         sql,
		"affect_rows": affectRows,
		"err":         err.Error(),
	})
}

// write ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 4:11 PM 2021/12/24
func (g *Gorm) write(ctx context.Context, message string, level string, data map[string]any) {
	if len(message) == 0 {
		message = "SQL执行记录"
	}
	if nil == g.instance {
		// 未设置日志实例
		return
	}
	if nil == data {
		data = make(map[string]any)
	}
	if nil == g.outCtx {
		g.outCtx = context.Background()
	}
	for _, extraField := range g.extraCtxFieldList {
		if len(extraField) == 0 {
			continue
		}
		val := g.outCtx.Value(extraField)
		if nil == val {
			val = ""
		}
		if val, exist := data[extraField]; !exist || nil == val {
			data[extraField] = val
		}
	}

	dataList := logger2.ZapLogDataList(logger2.NewLogData(consts.LogTypeDatabase, "", data))
	switch strings.ToUpper(level) {
	case consts.LogLevelDebug:
		g.instance.Debug(message, dataList...)
	case consts.LogLevelInfo:
		g.instance.Info(message, dataList...)
	case consts.LogLevelWarn:
		g.instance.Warn(message, dataList...)
	case consts.LogLevelError:
		g.instance.Error(message, dataList...)
	case consts.LogLevelPanic:
		g.instance.Panic(message, dataList...)
	default:
		g.instance.Info(message, dataList...)
	}
}

// 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
	}
}