feat: update log with given context (#147)

Update logging with the given context.

internal/log/handler.go

@@ -21,6 +21,8 @@ import (
 	"log/slog"
 	"sync"
 	"time"
+
+	"go.opentelemetry.io/otel/trace"
 )
 
 // ValueTextHandler is a [Handler] that writes Records to an [io.Writer] with values separated by spaces.
@@ -114,3 +116,34 @@ func (h *ValueTextHandler) appendAttr(buf []byte, a slog.Attr) []byte {
 
 	return buf
 }
+
+// spanContextLogHandler is an slog.Handler which adds attributes from the span
+// context.
+type spanContextLogHandler struct {
+	slog.Handler
+}
+
+// handlerWithSpanContext adds attributes from the span context.
+func handlerWithSpanContext(handler slog.Handler) *spanContextLogHandler {
+	return &spanContextLogHandler{Handler: handler}
+}
+
+// Handle overrides slog.Handler's Handle method. This adds attributes from the
+// span context to the slog.Record.
+func (t *spanContextLogHandler) Handle(ctx context.Context, record slog.Record) error {
+	// Get the SpanContext from the golang Context.
+	if s := trace.SpanContextFromContext(ctx); s.IsValid() {
+		// Add trace context attributes following Cloud Logging structured log format described
+		// in https://cloud.google.com/logging/docs/structured-logging#special-payload-fields
+		record.AddAttrs(
+			slog.Any("logging.googleapis.com/trace", s.TraceID()),
+		)
+		record.AddAttrs(
+			slog.Any("logging.googleapis.com/spanId", s.SpanID()),
+		)
+		record.AddAttrs(
+			slog.Bool("logging.googleapis.com/trace_sampled", s.TraceFlags().IsSampled()),
+		)
+	}
+	return t.Handler.Handle(ctx, record)
+}

internal/log/log.go

@@ -15,6 +15,7 @@
 package log
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log/slog"
@@ -45,24 +46,24 @@ func NewStdLogger(outW, errW io.Writer, logLevel string) (Logger, error) {
 	}, nil
 }
 
-// Debug logs debug messages
-func (sl *StdLogger) Debug(msg string, keysAndValues ...interface{}) {
-	sl.outLogger.Debug(msg, keysAndValues...)
+// DebugContext logs debug messages
+func (sl *StdLogger) DebugContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.outLogger.DebugContext(ctx, msg, keysAndValues...)
 }
 
-// Info logs debug messages
-func (sl *StdLogger) Info(msg string, keysAndValues ...interface{}) {
-	sl.outLogger.Info(msg, keysAndValues...)
+// InfoContext logs debug messages
+func (sl *StdLogger) InfoContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.outLogger.InfoContext(ctx, msg, keysAndValues...)
 }
 
-// Warn logs warning messages
-func (sl *StdLogger) Warn(msg string, keysAndValues ...interface{}) {
-	sl.errLogger.Warn(msg, keysAndValues...)
+// WarnContext logs warning messages
+func (sl *StdLogger) WarnContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.errLogger.WarnContext(ctx, msg, keysAndValues...)
 }
 
-// Error logs error messages
-func (sl *StdLogger) Error(msg string, keysAndValues ...interface{}) {
-	sl.errLogger.Error(msg, keysAndValues...)
+// ErrorContext logs error messages
+func (sl *StdLogger) ErrorContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.errLogger.ErrorContext(ctx, msg, keysAndValues...)
 }
 
 const (
@@ -149,36 +150,36 @@ func NewStructuredLogger(outW, errW io.Writer, logLevel string) (Logger, error)
 
 	// Configure structured logs to adhere to Cloud LogEntry format
 	// https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry
-	outHandler := slog.NewJSONHandler(outW, &slog.HandlerOptions{
+	outHandler := handlerWithSpanContext(slog.NewJSONHandler(outW, &slog.HandlerOptions{
 		AddSource:   true,
 		Level:       programLevel,
 		ReplaceAttr: replace,
-	})
-	errHandler := slog.NewJSONHandler(errW, &slog.HandlerOptions{
+	}))
+	errHandler := handlerWithSpanContext(slog.NewJSONHandler(errW, &slog.HandlerOptions{
 		AddSource:   true,
 		Level:       programLevel,
 		ReplaceAttr: replace,
-	})
+	}))
 
 	return &StructuredLogger{outLogger: slog.New(outHandler), errLogger: slog.New(errHandler)}, nil
 }
 
-// Debug logs debug messages
-func (sl *StructuredLogger) Debug(msg string, keysAndValues ...interface{}) {
-	sl.outLogger.Debug(msg, keysAndValues...)
+// DebugContext logs debug messages
+func (sl *StructuredLogger) DebugContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.outLogger.DebugContext(ctx, msg, keysAndValues...)
 }
 
-// Info logs info messages
-func (sl *StructuredLogger) Info(msg string, keysAndValues ...interface{}) {
-	sl.outLogger.Info(msg, keysAndValues...)
+// InfoContext logs info messages
+func (sl *StructuredLogger) InfoContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.outLogger.InfoContext(ctx, msg, keysAndValues...)
 }
 
-// Warn logs warning messages
-func (sl *StructuredLogger) Warn(msg string, keysAndValues ...interface{}) {
-	sl.errLogger.Warn(msg, keysAndValues...)
+// WarnContext logs warning messages
+func (sl *StructuredLogger) WarnContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.errLogger.WarnContext(ctx, msg, keysAndValues...)
 }
 
-// Error logs error messages
-func (sl *StructuredLogger) Error(msg string, keysAndValues ...interface{}) {
-	sl.errLogger.Error(msg, keysAndValues...)
+// ErrorContext logs error messages
+func (sl *StructuredLogger) ErrorContext(ctx context.Context, msg string, keysAndValues ...interface{}) {
+	sl.errLogger.ErrorContext(ctx, msg, keysAndValues...)
 }

internal/log/log_test.go

@@ -16,6 +16,7 @@ package log
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"log/slog"
 	"strings"
@@ -121,15 +122,16 @@ func TestLevelToSeverityError(t *testing.T) {
 }
 
 func runLogger(logger Logger, logMsg string) {
+	ctx := context.Background()
 	switch logMsg {
 	case "info":
-		logger.Info("log info")
+		logger.InfoContext(ctx, "log info")
 	case "debug":
-		logger.Debug("log debug")
+		logger.DebugContext(ctx, "log debug")
 	case "warn":
-		logger.Warn("log warn")
+		logger.WarnContext(ctx, "log warn")
 	case "error":
-		logger.Error("log error")
+		logger.ErrorContext(ctx, "log error")
 	}
 }
 

internal/log/logger.go

@@ -14,14 +14,18 @@
 
 package log
 
+import (
+	"context"
+)
+
 // Logger is the interface used throughout the project for logging.
 type Logger interface {
-	// Debug is for reporting additional information about internal operations.
-	Debug(format string, args ...interface{})
-	// Info is for reporting informational messages.
-	Info(format string, args ...interface{})
-	// Warn is for reporting warning messages.
-	Warn(format string, args ...interface{})
-	// Error is for reporting errors.
-	Error(format string, args ...interface{})
+	// DebugContext is for reporting additional information about internal operations.
+	DebugContext(ctx context.Context, format string, args ...interface{})
+	// InfoContext is for reporting informational messages.
+	InfoContext(ctx context.Context, format string, args ...interface{})
+	// WarnContext is for reporting warning messages.
+	WarnContext(ctx context.Context, format string, args ...interface{})
+	// ErrorContext is for reporting errors.
+	ErrorContext(ctx context.Context, format string, args ...interface{})
 }

internal/server/server.go

@@ -44,7 +44,7 @@ type Server struct {
 }
 
 // NewServer returns a Server object based on provided Config.
-func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) {
+func NewServer(ctx context.Context, cfg ServerConfig, log logLib.Logger) (*Server, error) {
 	logLevel, err := logLib.SeverityToLevel(cfg.LogLevel.String())
 	if err != nil {
 		return nil, fmt.Errorf("unable to initialize http log: %w", err)
@@ -89,7 +89,7 @@ func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) {
 		}
 		sourcesMap[name] = s
 	}
-	log.Info(fmt.Sprintf("Initialized %d sources.", len(sourcesMap)))
+	log.InfoContext(ctx, fmt.Sprintf("Initialized %d sources.", len(sourcesMap)))
 
 	// initialize and validate the auth sources
 	authSourcesMap := make(map[string]auth.AuthSource)
@@ -100,7 +100,7 @@ func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) {
 		}
 		authSourcesMap[name] = a
 	}
-	log.Info(fmt.Sprintf("Initialized %d authSources.", len(authSourcesMap)))
+	log.InfoContext(ctx, fmt.Sprintf("Initialized %d authSources.", len(authSourcesMap)))
 
 	// initialize and validate the tools
 	toolsMap := make(map[string]tools.Tool)
@@ -111,7 +111,7 @@ func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) {
 		}
 		toolsMap[name] = t
 	}
-	log.Info(fmt.Sprintf("Initialized %d tools.", len(toolsMap)))
+	log.InfoContext(ctx, fmt.Sprintf("Initialized %d tools.", len(toolsMap)))
 
 	// create a default toolset that contains all tools
 	allToolNames := make([]string, 0, len(toolsMap))
@@ -131,7 +131,7 @@ func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) {
 		}
 		toolsetsMap[name] = t
 	}
-	log.Info(fmt.Sprintf("Initialized %d toolsets.", len(toolsetsMap)))
+	log.InfoContext(ctx, fmt.Sprintf("Initialized %d toolsets.", len(toolsetsMap)))
 
 	s := &Server{
 		conf:        cfg,

internal/server/server_test.go

@@ -55,7 +55,7 @@ func TestServe(t *testing.T) {
 		t.Fatalf("unexpected error: %s", err)
 	}
 
-	s, err := server.NewServer(cfg, testLogger)
+	s, err := server.NewServer(context.Background(), cfg, testLogger)
 	if err != nil {
 		t.Fatalf("unable to initialize server! %v", err)
 	}