Add ability to add request ids to logs

Author: fobilow

logger.go

@@ -1,6 +1,7 @@
 package logger
 
 import (
+	"context"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -14,15 +15,31 @@ const (
 	FormatGoogleCloud
 )
 
+type correlationIdType int
+
+const (
+	requestIdKey correlationIdType = iota
+)
+
 // Default format is format lines
 var (
 	logger        *zap.Logger
 	isInitialized = false
 
 	currentFormat = FormatLines
 	debugLogging  = false
+	logCtx        context.Context
 )
 
+func SetContext(ctx context.Context, requestID string) {
+	ctxRequestID, ok := ctx.Value(requestIdKey).(string)
+	if !ok || len(ctxRequestID) == 0 {
+		ctx = context.WithValue(ctx, requestIdKey, requestID)
+	}
+
+	logCtx = ctx
+}
+
 // SetFormat sets the log output format
 func SetFormat(format Format) {
 	if isInitialized {
@@ -72,6 +89,13 @@ func Instance() *zap.Logger {
 		logger, _ = cfg.Build()
 		isInitialized = true
 	}
+
+	if logCtx != nil {
+		if ctxRequestID, ok := logCtx.Value(requestIdKey).(string); ok {
+			logger = logger.With(zap.String("REQUEST_ID", ctxRequestID))
+		}
+	}
+
 	return logger
 }
 

logger_test.go

@@ -0,0 +1,17 @@
+package logger_test
+
+import (
+	"context"
+	"github.com/lucidcube/logger-go"
+	"go.uber.org/zap"
+	"testing"
+)
+
+//This is more of a usage example than a test
+func TestSetContext(t *testing.T) {
+	ctx := context.Background()
+	logger.SetContext(ctx, "request-1234")
+	logger.SetFormat(logger.FormatGoogleCloud)
+
+	logger.Instance().Info("debug message", zap.String("key", "key-1"))
+}