253 lines
6.5 KiB
Go
253 lines
6.5 KiB
Go
package logger
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"path"
|
||
"runtime"
|
||
"sort"
|
||
"strings"
|
||
|
||
"github.com/Tencent/WeKnora/internal/types"
|
||
"github.com/sirupsen/logrus"
|
||
)
|
||
|
||
// LogLevel 日志级别类型
|
||
type LogLevel string
|
||
|
||
// 日志级别常量
|
||
const (
|
||
LevelDebug LogLevel = "debug"
|
||
LevelInfo LogLevel = "info"
|
||
LevelWarn LogLevel = "warn"
|
||
LevelError LogLevel = "error"
|
||
LevelFatal LogLevel = "fatal"
|
||
)
|
||
|
||
// ANSI颜色代码
|
||
const (
|
||
colorRed = "\033[31m"
|
||
colorGreen = "\033[32m"
|
||
colorYellow = "\033[33m"
|
||
colorBlue = "\033[34m"
|
||
colorPurple = "\033[35m"
|
||
colorCyan = "\033[36m"
|
||
colorReset = "\033[0m"
|
||
)
|
||
|
||
type CustomFormatter struct {
|
||
ForceColor bool // 是否强制使用颜色，即使在非终端环境下
|
||
}
|
||
|
||
func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
|
||
timestamp := entry.Time.Format("2006-01-02 15:04:05.000")
|
||
level := strings.ToUpper(entry.Level.String())
|
||
|
||
// 根据日志级别设置颜色
|
||
var levelColor, resetColor string
|
||
if f.ForceColor {
|
||
switch entry.Level {
|
||
case logrus.DebugLevel:
|
||
levelColor = colorCyan
|
||
case logrus.InfoLevel:
|
||
levelColor = colorGreen
|
||
case logrus.WarnLevel:
|
||
levelColor = colorYellow
|
||
case logrus.ErrorLevel:
|
||
levelColor = colorRed
|
||
case logrus.FatalLevel:
|
||
levelColor = colorPurple
|
||
default:
|
||
levelColor = colorReset
|
||
}
|
||
resetColor = colorReset
|
||
}
|
||
|
||
// 取出 caller 字段
|
||
caller := ""
|
||
if val, ok := entry.Data["caller"]; ok {
|
||
caller = fmt.Sprintf("%v", val)
|
||
}
|
||
|
||
// 拼接字段部分：request_id 优先，其他排序后输出
|
||
fields := ""
|
||
|
||
// request_id 优先输出
|
||
if v, ok := entry.Data["request_id"]; ok {
|
||
fields += fmt.Sprintf("request_id=%v ", v)
|
||
}
|
||
|
||
// 其余字段排序后输出
|
||
keys := make([]string, 0, len(entry.Data))
|
||
for k := range entry.Data {
|
||
if k != "caller" && k != "request_id" {
|
||
keys = append(keys, k)
|
||
}
|
||
}
|
||
sort.Strings(keys)
|
||
for _, k := range keys {
|
||
fields += fmt.Sprintf("%s=%v ", k, entry.Data[k])
|
||
}
|
||
|
||
fields = strings.TrimSpace(fields)
|
||
|
||
// 拼接最终输出内容，添加颜色
|
||
return []byte(fmt.Sprintf("%s%-5s%s[%s] [%s] %-20s | %s\n",
|
||
levelColor, level, resetColor, timestamp, fields, caller, entry.Message)), nil
|
||
}
|
||
|
||
// 初始化全局日志设置
|
||
func init() {
|
||
// 设置日志格式而不修改全局时区
|
||
logrus.SetFormatter(&CustomFormatter{ForceColor: true})
|
||
logrus.SetReportCaller(false)
|
||
}
|
||
|
||
// GetLogger 获取日志实例
|
||
func GetLogger(c context.Context) *logrus.Entry {
|
||
if logger := c.Value(types.LoggerContextKey); logger != nil {
|
||
return logger.(*logrus.Entry)
|
||
}
|
||
newLogger := logrus.New()
|
||
newLogger.SetFormatter(&CustomFormatter{ForceColor: true})
|
||
// 设置默认日志级别
|
||
newLogger.SetLevel(logrus.DebugLevel)
|
||
// 启用调用者信息
|
||
return logrus.NewEntry(newLogger)
|
||
}
|
||
|
||
// SetLogLevel 设置日志级别
|
||
func SetLogLevel(level LogLevel) {
|
||
var logLevel logrus.Level
|
||
|
||
switch level {
|
||
case LevelDebug:
|
||
logLevel = logrus.DebugLevel
|
||
case LevelInfo:
|
||
logLevel = logrus.InfoLevel
|
||
case LevelWarn:
|
||
logLevel = logrus.WarnLevel
|
||
case LevelError:
|
||
logLevel = logrus.ErrorLevel
|
||
case LevelFatal:
|
||
logLevel = logrus.FatalLevel
|
||
default:
|
||
logLevel = logrus.InfoLevel
|
||
}
|
||
|
||
logrus.SetLevel(logLevel)
|
||
}
|
||
|
||
// 添加调用者字段
|
||
func addCaller(entry *logrus.Entry, skip int) *logrus.Entry {
|
||
pc, file, line, ok := runtime.Caller(skip)
|
||
if !ok {
|
||
return entry
|
||
}
|
||
shortFile := path.Base(file)
|
||
funcName := "unknown"
|
||
if fn := runtime.FuncForPC(pc); fn != nil {
|
||
// 只保留函数名，不带包路径（如 doSomething）
|
||
fullName := path.Base(fn.Name())
|
||
parts := strings.Split(fullName, ".")
|
||
funcName = parts[len(parts)-1]
|
||
}
|
||
return entry.WithField("caller", fmt.Sprintf("%s:%d[%s]", shortFile, line, funcName))
|
||
}
|
||
|
||
// WithRequestID 在日志中添加请求ID
|
||
func WithRequestID(c context.Context, requestID string) context.Context {
|
||
return WithField(c, "request_id", requestID)
|
||
}
|
||
|
||
// WithField 向日志中添加一个字段
|
||
func WithField(c context.Context, key string, value interface{}) context.Context {
|
||
logger := GetLogger(c).WithField(key, value)
|
||
return context.WithValue(c, types.LoggerContextKey, logger)
|
||
}
|
||
|
||
// WithFields 向日志中添加多个字段
|
||
func WithFields(c context.Context, fields logrus.Fields) context.Context {
|
||
logger := GetLogger(c).WithFields(fields)
|
||
return context.WithValue(c, types.LoggerContextKey, logger)
|
||
}
|
||
|
||
// Debug 输出调试级别的日志
|
||
func Debug(c context.Context, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Debug(args...)
|
||
}
|
||
|
||
// Debugf 使用格式化字符串输出调试级别的日志
|
||
func Debugf(c context.Context, format string, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Debugf(format, args...)
|
||
}
|
||
|
||
// Info 输出信息级别的日志
|
||
func Info(c context.Context, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Info(args...)
|
||
}
|
||
|
||
// Infof 使用格式化字符串输出信息级别的日志
|
||
func Infof(c context.Context, format string, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Infof(format, args...)
|
||
}
|
||
|
||
// Warn 输出警告级别的日志
|
||
func Warn(c context.Context, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Warn(args...)
|
||
}
|
||
|
||
// Warnf 使用格式化字符串输出警告级别的日志
|
||
func Warnf(c context.Context, format string, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Warnf(format, args...)
|
||
}
|
||
|
||
// Error 输出错误级别的日志
|
||
func Error(c context.Context, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Error(args...)
|
||
}
|
||
|
||
// Errorf 使用格式化字符串输出错误级别的日志
|
||
func Errorf(c context.Context, format string, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Errorf(format, args...)
|
||
}
|
||
|
||
// ErrorWithFields 输出带有额外字段的错误级别日志
|
||
func ErrorWithFields(c context.Context, err error, fields logrus.Fields) {
|
||
if fields == nil {
|
||
fields = logrus.Fields{}
|
||
}
|
||
if err != nil {
|
||
fields["error"] = err.Error()
|
||
}
|
||
addCaller(GetLogger(c), 2).WithFields(fields).Error("发生错误")
|
||
}
|
||
|
||
// Fatal 输出致命级别的日志并退出程序
|
||
func Fatal(c context.Context, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Fatal(args...)
|
||
}
|
||
|
||
// Fatalf 使用格式化字符串输出致命级别的日志并退出程序
|
||
func Fatalf(c context.Context, format string, args ...interface{}) {
|
||
addCaller(GetLogger(c), 2).Fatalf(format, args...)
|
||
}
|
||
|
||
// CloneContext 复制上下文中的关键信息到新上下文
|
||
func CloneContext(ctx context.Context) context.Context {
|
||
newCtx := context.Background()
|
||
|
||
for _, k := range []types.ContextKey{
|
||
types.LoggerContextKey,
|
||
types.TenantIDContextKey,
|
||
types.RequestIDContextKey,
|
||
types.TenantInfoContextKey,
|
||
} {
|
||
if v := ctx.Value(k); v != nil {
|
||
newCtx = context.WithValue(newCtx, k, v)
|
||
}
|
||
}
|
||
|
||
return newCtx
|
||
}
|