fix: trace id in logs from custom modules (wundergraph#1299)

router-tests/modules/module_test.go

@@ -4,6 +4,9 @@ import (
 	"encoding/json"
 	"testing"
 
+	"github.com/wundergraph/cosmo/router/pkg/trace/tracetest"
+	"go.uber.org/zap/zapcore"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/wundergraph/cosmo/router-tests/testenv"
@@ -39,3 +42,48 @@ func TestModuleSetCustomHeader(t *testing.T) {
 		assert.JSONEq(t, res.Body, `{"data":{"employees":[{"id":1},{"id":2},{"id":3},{"id":4},{"id":5},{"id":7},{"id":8},{"id":10},{"id":11},{"id":12}]}}`)
 	})
 }
+
+func TestCustomModuleLogs(t *testing.T) {
+	cfg := config.Config{
+		Graph: config.Graph{},
+		Modules: map[string]interface{}{
+			"myModule": module.MyModule{
+				Value: 1,
+			},
+		},
+	}
+
+	exporter := tracetest.NewInMemoryExporter(t)
+
+	testenv.Run(t, &testenv.Config{
+		RouterOptions: []core.Option{
+			core.WithModulesConfig(cfg.Modules),
+		},
+		LogObservation: testenv.LogObservationConfig{
+			Enabled:  true,
+			LogLevel: zapcore.InfoLevel,
+		},
+		TraceExporter: exporter,
+	}, func(t *testing.T, xEnv *testenv.Environment) {
+		res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+			Query:         `query MyQuery { employees { id } }`,
+			OperationName: json.RawMessage(`"MyQuery"`),
+		})
+		require.NoError(t, err)
+
+		assert.Equal(t, 200, res.Response.StatusCode)
+		assert.JSONEq(t, res.Body, `{"data":{"employees":[{"id":1},{"id":2},{"id":3},{"id":4},{"id":5},{"id":7},{"id":8},{"id":10},{"id":11},{"id":12}]}}`)
+
+		requestLog := xEnv.Observer().FilterMessage("Test custom module logs")
+		assert.Equal(t, requestLog.Len(), 1)
+		requestContext := requestLog.All()[0].ContextMap()
+
+		expectedKeys := []string{
+			"trace_id", "request_id", "hostname", "pid",
+		}
+
+		for _, key := range expectedKeys {
+			assert.NotEmpty(t, requestContext[key])
+		}
+	})
+}

router/cmd/custom/module/module.go

@@ -80,6 +80,8 @@ func (m *MyModule) Middleware(ctx core.RequestContext, next http.Handler) {
 
 	operation := ctx.Operation()
 
+	logger := ctx.Logger()
+	logger.Info("Test custom module logs")
 	// Access the GraphQL operation context
 	fmt.Println(
 		operation.Name(),

router/core/graph_server.go

@@ -485,12 +485,16 @@ func (s *graphServer) buildGraphMux(ctx context.Context,
 			rtrace.WithTracePreHandler(
 				func(r *http.Request, w http.ResponseWriter) {
 					reqContext := getRequestContext(r.Context())
+					traceID := rtrace.GetTraceID(r.Context())
+					requestLogger := reqContext.Logger().With(logging.WithTraceID(traceID))
+
+					reqContext.logger = requestLogger
+
 					span := oteltrace.SpanFromContext(r.Context())
 					span.SetAttributes(reqContext.telemetry.CommonAttrs()...)
 
 					// Set the trace ID in the response header
 					if s.traceConfig.ResponseTraceHeader.Enabled {
-						traceID := rtrace.GetTraceID(r.Context())
 						w.Header().Set(s.traceConfig.ResponseTraceHeader.HeaderName, traceID)
 					}
 				}),
@@ -523,7 +527,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context,
 	subgraphResolver := NewSubgraphResolver(subgraphs)
 	httpRouter.Use(func(h http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			requestLogger := s.logger.With(logging.WithRequestID(middleware.GetReqID(r.Context())), logging.WithTraceID(rtrace.GetTraceID(r.Context())))
+			requestLogger := s.logger.With(logging.WithRequestID(middleware.GetReqID(r.Context())))
 			r = r.WithContext(withSubgraphResolver(r.Context(), subgraphResolver))
 
 			reqContext := buildRequestContext(requestContextOptions{