Improve logging

cmd/MQTTLogger/main.go

@@ -1,12 +1,15 @@
 package main
 
 import (
+	"os"
+
 	"MQTTLogger/config"
 	"MQTTLogger/internal/csv"
 	"MQTTLogger/internal/mqtt"
 
 	"go.uber.org/fx"
 	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 )
 
 func main() {
@@ -27,7 +30,21 @@ func main() {
 }
 
 func NewLogger() (*zap.Logger, error) {
-	return zap.NewDevelopment()
+	logLevel := os.Getenv("LOG_LEVEL")
+	if logLevel == "" {
+		logLevel = "info"
+	}
+
+	var level zapcore.Level
+	err := level.Set(logLevel)
+	if err != nil {
+		level = zapcore.InfoLevel
+	}
+
+	config := zap.NewDevelopmentConfig()
+	config.Level = zap.NewAtomicLevelAt(level)
+
+	return config.Build()
 }
 
 // Simple logger for fx
@@ -36,7 +53,21 @@ type FXLogger struct {
 }
 
 func NewFXLogger() *FXLogger {
-	logger, _ := zap.NewDevelopment()
+	logLevel := os.Getenv("LOG_LEVEL")
+	if logLevel == "" {
+		logLevel = "info"
+	}
+
+	var level zapcore.Level
+	err := level.Set(logLevel)
+	if err != nil {
+		level = zapcore.InfoLevel
+	}
+
+	config := zap.NewDevelopmentConfig()
+	config.Level = zap.NewAtomicLevelAt(level)
+
+	logger, _ := config.Build()
 	return &FXLogger{logger: logger}
 }
 

internal/mqtt/logger.go

@@ -8,17 +8,27 @@ import (
 	"go.uber.org/zap"
 )
 
-type ZapWriter struct {
-    sugar *zap.SugaredLogger
+type DebugWriter struct {
+    logger *zap.Logger
 }
 
-func (z *ZapWriter) Write(p []byte) (n int, err error) {
-    z.sugar.Debug(strings.TrimSpace(string(p)))
+func (d *DebugWriter) Write(p []byte) (n int, err error) {
+    message := strings.TrimSpace(string(p))
+    d.logger.Debug(message, zap.String("source", "mqtt-debug"))
+    return len(p), nil
+}
+
+type ErrorWriter struct {
+    logger *zap.Logger
+}
+
+func (e *ErrorWriter) Write(p []byte) (n int, err error) {
+    message := strings.TrimSpace(string(p))
+    e.logger.Error(message, zap.String("source", "mqtt-error"))
     return len(p), nil
 }
 
 func SetupMQTTLogging(logger *zap.Logger) {
-    sugar := logger.Sugar()
-    mqtt.DEBUG = log.New(&ZapWriter{sugar: sugar}, "[DEBUG] ", 0)
-    mqtt.ERROR = log.New(&ZapWriter{sugar: sugar}, "[ERROR] ", 0)
+    mqtt.DEBUG = log.New(&DebugWriter{logger: logger}, "[MQTT] ", 0)
+    mqtt.ERROR = log.New(&ErrorWriter{logger: logger}, "[MQTT] ", 0)
 }