database/logger/gorm_v2.go

// Package logger ...
//
// Description : logger ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2025-10-12 17:23
package logger

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

	"git.zhangdeman.cn/zhangdeman/consts"

	"gorm.io/gorm"

	"go.uber.org/zap/zapcore"

	"go.uber.org/zap"

	loggerPkg "git.zhangdeman.cn/zhangdeman/logger"
	"gorm.io/gorm/logger"
)

// NewGormV2 获取日志实例
func NewGormV2(loggerLevel consts.LogLevel, consoleOutput bool, encoder zapcore.Encoder, splitConfig *loggerPkg.RotateLogConfig, traceIDField string, skip int) (logger.Interface, error) {
	logConfList := []loggerPkg.SetLoggerOptionFunc{loggerPkg.WithEncoder(encoder), loggerPkg.WithCallerSkip(skip), loggerPkg.WithCaller()}
	if consoleOutput {
		logConfList = append(logConfList, loggerPkg.WithConsoleOutput())
	}

	logInstance, err := loggerPkg.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日志实现
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 版本库日志实现
type Gorm struct {
	dbClient          *gorm.DB
	instance          *zap.Logger // 日志实例
	traceIDField      string      // 串联请求上下文的的ID
	extraCtxFieldList []string    // 从请求上线问提取的字段
	flag              string      // 数据库标识
	node              string      // 数据库节点 master / slave
	outCtx            context.Context
}

// LogMode ...
func (g *Gorm) LogMode(level logger.LogLevel) logger.Interface {
	return g
}

// Info 日志
func (g *Gorm) Info(ctx context.Context, s string, i ...any) {
	g.write(ctx, s, consts.LogLevelInfo, map[string]any{
		"log_data": i,
	})
}

// Warn ...
func (g *Gorm) Warn(ctx context.Context, s string, i ...any) {
	g.write(ctx, s, consts.LogLevelWarn, map[string]any{
		"log_data": i,
	})
}

// Error 日志
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 记录
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()
	}

	logData := 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":         "",
	}
	if nil != err {
		logData["err"] = err.Error()
	}
	g.write(ctx, "SQL执行记录", consts.LogLevelInfo, logData)
}

// write ...
func (g *Gorm) write(ctx context.Context, message string, level consts.LogLevel, 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 != data["sql"] {
		sqlStr := strings.TrimSpace(fmt.Sprintf("%v", data["sql"]))
		sqlArr := strings.Split(sqlStr, " ")
		if len(sqlArr) > 0 {
			data["sql_type"] = strings.ToUpper(sqlArr[0])
		}
	}
	if nil == g.outCtx {
		g.outCtx = context.Background()
	}

	dataList := loggerPkg.ZapLogDataList(loggerPkg.NewLogData(g.outCtx, consts.LogTypeDatabase, "", data))
	switch 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
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
	}
}