From 0e9d737e4ca2deb86930237ee32a39cf3f7e8157 Mon Sep 17 00:00:00 2001
From: Thomas Nguyen <thomasvn.dev@gmail.com>
Date: Tue, 2 Apr 2024 00:42:44 -0700
Subject: [PATCH] feat: Allow overiding of Warn and Deprecated loggers (#952)

Signed-off-by: thomasvn <thomasnguyen96@gmail.com>
---
 bun.go             |  2 +-
 internal/logger.go | 39 +++++++++++++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/bun.go b/bun.go
index 4694a4c75..8f71db8fc 100644
--- a/bun.go
+++ b/bun.go
@@ -72,7 +72,7 @@ type AfterDropTableHook interface {
 
 // SetLogger overwriters default Bun logger.
 func SetLogger(logger internal.Logging) {
-	internal.Logger = logger
+	internal.SetLogger(logger)
 }
 
 func In(slice interface{}) schema.QueryAppender {
diff --git a/internal/logger.go b/internal/logger.go
index 2e22a0893..9233dfcff 100644
--- a/internal/logger.go
+++ b/internal/logger.go
@@ -6,14 +6,26 @@ import (
 	"os"
 )
 
-var Warn = log.New(os.Stderr, "WARN: bun: ", log.LstdFlags)
-
-var Deprecated = log.New(os.Stderr, "DEPRECATED: bun: ", log.LstdFlags)
-
 type Logging interface {
 	Printf(format string, v ...interface{})
 }
 
+var defaultLogger = log.New(os.Stderr, "", log.LstdFlags)
+
+var Logger Logging = &logger{
+	log: defaultLogger,
+}
+
+var Warn = &wrapper{
+	prefix: "WARN: bun: ",
+	logger: Logger,
+}
+
+var Deprecated = &wrapper{
+	prefix: "DEPRECATED: bun: ",
+	logger: Logger,
+}
+
 type logger struct {
 	log *log.Logger
 }
@@ -22,6 +34,21 @@ func (l *logger) Printf(format string, v ...interface{}) {
 	_ = l.log.Output(2, fmt.Sprintf(format, v...))
 }
 
-var Logger Logging = &logger{
-	log: log.New(os.Stderr, "bun: ", log.LstdFlags|log.Lshortfile),
+type wrapper struct {
+	prefix string
+	logger Logging
+}
+
+func (w *wrapper) Printf(format string, v ...interface{}) {
+	w.logger.Printf(w.prefix+format, v...)
+}
+
+func SetLogger(newLogger Logging) {
+	if newLogger == nil {
+		Logger = &logger{log: defaultLogger}
+	} else {
+		Logger = newLogger
+	}
+	Warn.logger = Logger
+	Deprecated.logger = Logger
 }