hpc/internal/logger/gorm.go

package logger

import (
	"context"
	"errors"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gorm.io/gorm"
	gormlogger "gorm.io/gorm/logger"
)

const slowQueryThreshold = 200 * time.Millisecond

// GormLogger implements gorm's logger.Interface backed by zap.
type GormLogger struct {
	logger *zap.Logger
	level  zapcore.Level
	silent bool
}

// Compile-time interface check.
var _ gormlogger.Interface = (*GormLogger)(nil)

// NewGormLogger creates a new GormLogger wrapping the given zap logger.
// The level string maps to zap levels; empty defaults to "warn".
// The special value "silent" suppresses all output.
func NewGormLogger(zapLogger *zap.Logger, level string) gormlogger.Interface {
	lvl := parseGormLevel(level)
	silent := level == "silent"
	return &GormLogger{
		logger: zapLogger,
		level:  lvl,
		silent: silent,
	}
}

// LogMode returns a new GormLogger with the given gorm log level.
// It does NOT mutate the receiver.
func (l *GormLogger) LogMode(level gormlogger.LogLevel) gormlogger.Interface {
	newLogger := &GormLogger{
		logger: l.logger,
		level:  l.level,
		silent: l.silent,
	}

	switch level {
	case gormlogger.Silent:
		newLogger.silent = true
	case gormlogger.Error:
		newLogger.level = zapcore.ErrorLevel
		newLogger.silent = false
	case gormlogger.Warn:
		newLogger.level = zapcore.WarnLevel
		newLogger.silent = false
	case gormlogger.Info:
		newLogger.level = zapcore.InfoLevel
		newLogger.silent = false
	}

	return newLogger
}

// Info logs at zap.InfoLevel with structured fields from key-value pairs.
func (l *GormLogger) Info(ctx context.Context, msg string, args ...any) {
	if l.silent || l.level > zapcore.InfoLevel {
		return
	}
	l.logger.Info(msg, argsToFields(args)...)
}

// Warn logs at zap.WarnLevel with structured fields from key-value pairs.
func (l *GormLogger) Warn(ctx context.Context, msg string, args ...any) {
	if l.silent || l.level > zapcore.WarnLevel {
		return
	}
	l.logger.Warn(msg, argsToFields(args)...)
}

// Error logs at zap.ErrorLevel with structured fields from key-value pairs.
func (l *GormLogger) Error(ctx context.Context, msg string, args ...any) {
	if l.silent || l.level > zapcore.ErrorLevel {
		return
	}
	l.logger.Error(msg, argsToFields(args)...)
}

// Trace logs SQL query information based on execution results.
func (l *GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
	if l.silent {
		return
	}

	elapsed := time.Since(begin)
	sql, rows := fc()

	switch {
	case err != nil && !errors.Is(err, gorm.ErrRecordNotFound):
		l.logger.Error("gorm query error",
			zap.Error(err),
			zap.String("sql", sql),
			zap.Int64("rows", rows),
			zap.Duration("elapsed", elapsed),
		)
	case elapsed > slowQueryThreshold:
		l.logger.Warn("gorm slow query",
			zap.String("sql", sql),
			zap.Int64("rows", rows),
			zap.Float64("elapsed_ms", float64(elapsed.Nanoseconds())/1e6),
		)
	default:
		if l.level > zapcore.InfoLevel {
			return
		}
		l.logger.Info("gorm query",
			zap.String("sql", sql),
			zap.Int64("rows", rows),
			zap.Duration("elapsed", elapsed),
		)
	}
}

// parseGormLevel parses a level string into a zapcore.Level.
// Defaults to zapcore.WarnLevel.
func parseGormLevel(level string) zapcore.Level {
	if level == "" || level == "silent" {
		return zapcore.WarnLevel
	}
	var lvl zapcore.Level
	if err := lvl.UnmarshalText([]byte(level)); err != nil {
		return zapcore.WarnLevel
	}
	return lvl
}

// argsToFields converts alternating key-value pairs into zap fields.
func argsToFields(args []any) []zap.Field {
	fields := make([]zap.Field, 0, len(args)/2)
	for i := 0; i+1 < len(args); i += 2 {
		key, ok := args[i].(string)
		if !ok {
			continue
		}
		fields = append(fields, zap.Any(key, args[i+1]))
	}
	return fields
}