fix: add logger options to add metadata into ctx

internal/server/server.go

@@ -3,6 +3,9 @@ package server
 import (
 	"context"
 	"fmt"
+	"github.com/google/uuid"
+	"github.com/goto/guardian/core"
+	"google.golang.org/grpc/metadata"
 	"net/http"
 	"runtime/debug"
 	"strings"
@@ -47,7 +50,10 @@ const (
 
 // RunServer runs the application server
 func RunServer(config *Config) error {
-	logger := log.NewCtxLogger(config.LogLevel, []string{domain.TraceIDKey})
+
+	logger := log.NewCtxLogger(config.LogLevel, []string{domain.TraceIDKey},
+		log.WithMetadataExtractor(defaultMetadataExtractor(config)),
+	)
 	crypto := crypto.NewAES(config.EncryptionSecretKeyKey)
 	validator := validator.New()
 	notifier, err := notifiers.NewClient(&config.Notifier, logger)
@@ -231,3 +237,26 @@ func getAuthInterceptor(config *Config) (grpc.UnaryServerInterceptor, error) {
 
 	return authInterceptor, nil
 }
+
+func defaultMetadataExtractor(config *Config) func(context.Context) map[string]interface{} {
+	return func(ctx context.Context) map[string]interface{} {
+		md := map[string]interface{}{
+			"app_name":    "guardian",
+			"app_version": core.Version,
+		}
+
+		// trace id
+		var traceID string
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			if rawTraceID := md.Get(config.AuditLogTraceIDHeaderKey); len(rawTraceID) > 0 {
+				traceID = rawTraceID[0]
+			}
+		}
+		if traceID == "" {
+			traceID = uuid.New().String()
+		}
+		md[domain.TraceIDKey] = traceID
+
+		return md
+	}
+}

internal/server/services.go

@@ -6,8 +6,6 @@ import (
 	"github.com/goto/guardian/plugins/providers/dataplex"
 
 	"github.com/go-playground/validator/v10"
-	"github.com/google/uuid"
-	"github.com/goto/guardian/core"
 	"github.com/goto/guardian/core/activity"
 	"github.com/goto/guardian/core/appeal"
 	"github.com/goto/guardian/core/approval"
@@ -31,7 +29,6 @@ import (
 	"github.com/goto/guardian/plugins/providers/tableau"
 	"github.com/goto/salt/audit"
 	audit_repos "github.com/goto/salt/audit/repositories"
-	"google.golang.org/grpc/metadata"
 )
 
 type Services struct {
@@ -71,26 +68,7 @@ func InitServices(deps ServiceDeps) (*Services, error) {
 
 	auditLogger := audit.New(
 		audit.WithRepository(auditRepository),
-		audit.WithMetadataExtractor(func(ctx context.Context) map[string]interface{} {
-			md := map[string]interface{}{
-				"app_name":    "guardian",
-				"app_version": core.Version,
-			}
-
-			// trace id
-			var traceID string
-			if md, ok := metadata.FromIncomingContext(ctx); ok {
-				if rawTraceID := md.Get(deps.Config.AuditLogTraceIDHeaderKey); len(rawTraceID) > 0 {
-					traceID = rawTraceID[0]
-				}
-			}
-			if traceID == "" {
-				traceID = uuid.New().String()
-			}
-			md[domain.TraceIDKey] = traceID
-
-			return md
-		}),
+		audit.WithMetadataExtractor(defaultMetadataExtractor(deps.Config)),
 		actorExtractor,
 	)
 

pkg/log/logger.go

@@ -2,6 +2,7 @@ package log
 
 import (
 	"context"
+	"errors"
 	"io"
 
 	saltLog "github.com/goto/salt/log"
@@ -36,20 +37,30 @@ type Logger interface {
 	Writer() io.Writer
 }
 
+type LoggerOption func(*CtxLogger)
+type metadataContextKey struct{}
+
 type CtxLogger struct {
-	log  saltLog.Logger
-	keys []string
+	log          saltLog.Logger
+	keys         []string
+	withMetadata func(context.Context) (context.Context, error)
 }
 
 // NewCtxLoggerWithSaltLogger returns a logger that will add context value to the log message, wrapped with saltLog.Logger
-func NewCtxLoggerWithSaltLogger(log saltLog.Logger, ctxKeys []string) *CtxLogger {
-	return &CtxLogger{log: log, keys: ctxKeys}
+func NewCtxLoggerWithSaltLogger(log saltLog.Logger, ctxKeys []string, opts ...LoggerOption) *CtxLogger {
+	ctxLogger := &CtxLogger{log: log, keys: ctxKeys}
+	for _, o := range opts {
+		o(ctxLogger)
+	}
+
+	return ctxLogger
 }
 
 // NewCtxLogger returns a logger that will add context value to the log message
-func NewCtxLogger(logLevel string, ctxKeys []string) *CtxLogger {
+func NewCtxLogger(logLevel string, ctxKeys []string, opts ...LoggerOption) *CtxLogger {
 	saltLogger := saltLog.NewLogrus(saltLog.LogrusWithLevel(logLevel))
-	return NewCtxLoggerWithSaltLogger(saltLogger, ctxKeys)
+	ctxLogger := NewCtxLoggerWithSaltLogger(saltLogger, ctxKeys, opts...)
+	return ctxLogger
 }
 
 func (l *CtxLogger) Debug(ctx context.Context, msg string, args ...interface{}) {
@@ -94,3 +105,30 @@ func (l *CtxLogger) addCtxToArgs(ctx context.Context, args []interface{}) []inte
 
 	return args
 }
+
+func WithMetadata(ctx context.Context, md map[string]interface{}) (context.Context, error) {
+	existingMetadata := ctx.Value(metadataContextKey{})
+	if existingMetadata == nil {
+		return context.WithValue(ctx, metadataContextKey{}, md), nil
+	}
+
+	// append new metadata
+	mapMd, ok := existingMetadata.(map[string]interface{})
+	if !ok {
+		return nil, errors.New("failed to cast existing metadata to map[string]interface{} type")
+	}
+	for k, v := range md {
+		mapMd[k] = v
+	}
+
+	return context.WithValue(ctx, metadataContextKey{}, mapMd), nil
+}
+
+func WithMetadataExtractor(fn func(context.Context) map[string]interface{}) LoggerOption {
+	return func(s *CtxLogger) {
+		s.withMetadata = func(ctx context.Context) (context.Context, error) {
+			md := fn(ctx)
+			return WithMetadata(ctx, md)
+		}
+	}
+}