From f6f9a7f97acf4fdd076ff60e1319cc26dc9a9ae1 Mon Sep 17 00:00:00 2001
From: Derrick Hammer <derrick@derrickhammer.com>
Date: Sat, 24 Feb 2024 08:41:40 -0500
Subject: [PATCH] feat: add custom logger wrapper to use zap logger for
 database

---
 cmd/portal/main.go |  4 +--
 db/db.go           |  9 ++++--
 db/logger.go       | 75 ++++++++++++++++++++++++++++++++++++++++++++++
 logger/logger.go   |  7 +++--
 4 files changed, 87 insertions(+), 8 deletions(-)
 create mode 100644 db/logger.go

diff --git a/cmd/portal/main.go b/cmd/portal/main.go
index 28ac3fe..119f5e4 100644
--- a/cmd/portal/main.go
+++ b/cmd/portal/main.go
@@ -23,7 +23,7 @@ import (
 
 func main() {
 
-	logger := _logger.NewLogger()
+	logger, logLevel := _logger.NewLogger()
 	cfg, err := config.NewManager(logger)
 
 	if err != nil {
@@ -50,7 +50,7 @@ func main() {
 
 	fx.New(
 		fx.Supply(cfg),
-		fx.Supply(logger),
+		fx.Supply(logger, logLevel),
 		fxLogger,
 		fx.Invoke(initCheckRequiredConfig),
 		fx.Provide(NewIdentity),
diff --git a/db/db.go b/db/db.go
index 3826a12..e3d1de2 100644
--- a/db/db.go
+++ b/db/db.go
@@ -19,8 +19,9 @@ import (
 
 type DatabaseParams struct {
 	fx.In
-	Config *config.Manager
-	Logger *zap.Logger
+	Config      *config.Manager
+	Logger      *zap.Logger
+	LoggerLevel *zap.AtomicLevel
 }
 
 var Module = fx.Module("db",
@@ -39,7 +40,9 @@ func NewDatabase(lc fx.Lifecycle, params DatabaseParams) *gorm.DB {
 
 	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", username, password, host, port, dbname, charset)
 
-	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
+	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
+		Logger: newLogger(params.Logger, params.LoggerLevel),
+	})
 	if err != nil {
 		panic(err)
 	}
diff --git a/db/logger.go b/db/logger.go
new file mode 100644
index 0000000..83ac703
--- /dev/null
+++ b/db/logger.go
@@ -0,0 +1,75 @@
+package db
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"go.uber.org/zap"
+	dbLogger "gorm.io/gorm/logger"
+)
+
+var _ dbLogger.Interface = (*logger)(nil)
+
+var (
+	levels = map[dbLogger.LogLevel]zap.AtomicLevel{
+		dbLogger.Silent: zap.NewAtomicLevelAt(zap.InfoLevel),
+		dbLogger.Error:  zap.NewAtomicLevelAt(zap.ErrorLevel),
+		dbLogger.Warn:   zap.NewAtomicLevelAt(zap.WarnLevel),
+		dbLogger.Info:   zap.NewAtomicLevelAt(zap.InfoLevel),
+	}
+)
+
+type logger struct {
+	logger *zap.Logger
+	level  *zap.AtomicLevel
+}
+
+func (l logger) LogMode(level dbLogger.LogLevel) dbLogger.Interface {
+	if atomicLevel, ok := levels[level]; ok {
+		l.level.SetLevel(atomicLevel.Level())
+		return l
+	}
+
+	l.logger.Fatal("invalid log level", zap.Int("level", int(level)))
+	return nil
+}
+
+func (l logger) Info(ctx context.Context, s string, i ...interface{}) {
+	l.logger.Info(s, interfacesToFields(i...)...)
+}
+
+func (l logger) Warn(ctx context.Context, s string, i ...interface{}) {
+	l.logger.Warn(s, interfacesToFields(i...)...)
+}
+
+func (l logger) Error(ctx context.Context, s string, i ...interface{}) {
+	l.logger.Error(s, interfacesToFields(i...)...)
+}
+
+func (l logger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
+	if l.level.Level() <= zap.DebugLevel {
+		sql, rowsAffected := fc()
+		fields := []zap.Field{
+			zap.String("sql", sql),
+			zap.Int64("rows_affected", rowsAffected),
+			zap.Duration("elapsed", time.Since(begin)),
+		}
+		if err != nil {
+			fields = append(fields, zap.Error(err))
+		}
+		l.logger.Debug("trace", fields...)
+	}
+}
+
+func newLogger(zlog *zap.Logger, zlogLevel *zap.AtomicLevel) *logger {
+	return &logger{logger: zlog, level: zlogLevel}
+}
+
+func interfacesToFields(i ...interface{}) []zap.Field {
+	fields := make([]zap.Field, 0)
+	for idx, v := range i {
+		fields = append(fields, zap.Any(strconv.Itoa(idx), v))
+	}
+	return fields
+}
diff --git a/logger/logger.go b/logger/logger.go
index a570042..4518f36 100644
--- a/logger/logger.go
+++ b/logger/logger.go
@@ -1,10 +1,11 @@
 package logger
 
 import (
+	"os"
+
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
-	"os"
 )
 
 func NewFallbackLogger() *zap.Logger {
@@ -13,7 +14,7 @@ func NewFallbackLogger() *zap.Logger {
 	return logger
 }
 
-func NewLogger() *zap.Logger {
+func NewLogger() (*zap.Logger, *zap.AtomicLevel) {
 
 	// Create a new atomic level
 	atomicLevel := zap.NewAtomicLevel()
@@ -28,7 +29,7 @@ func NewLogger() *zap.Logger {
 		atomicLevel,
 	))
 
-	return logger
+	return logger, &atomicLevel
 }
 
 func mapLogLevel(level string) zapcore.Level {