chore: add query-frontend option to select request headers in query logs

grafana · Feb 12, 2024 · 04ec072 · 04ec072
1 parent 9d82759
commit 04ec072
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 41 deletions.
diff --git a/docs/sources/tempo/configuration/_index.md b/docs/sources/tempo/configuration/_index.md
@@ -437,6 +437,10 @@ query_frontend:
     # (default: true)
     [multi_tenant_queries_enabled: <bool>]
 
+    # Comma-separated list of request header names to include in query logs. Applies
+    # to both query stats and slow queries logs.
+    [log_query_request_headers: <string> | default = ""]
+
     search:
 
         # The number of concurrent jobs to execute when searching the backend.

diff --git a/modules/frontend/frontend.go b/modules/frontend/frontend.go
@@ -102,16 +102,15 @@ func New(cfg Config, next http.RoundTripper, o overrides.Interface, reader tempo
 
 	metrics := metricsMiddleware.Wrap(next)
 	queryrange := queryRangeMiddleware.Wrap(next)
-
 	return &QueryFrontend{
-		TraceByIDHandler:          newHandler(traces, traceByIDSLOPostHook(cfg.TraceByID.SLO), nil, logger),
-		SearchHandler:             newHandler(search, searchSLOPostHook(cfg.Search.SLO), searchSLOPreHook, logger),
-		SearchTagsHandler:         newHandler(searchTags, nil, nil, logger),
-		SearchTagsV2Handler:       newHandler(searchTagsV2, nil, nil, logger),
-		SearchTagsValuesHandler:   newHandler(searchTagValues, nil, nil, logger),
-		SearchTagsValuesV2Handler: newHandler(searchTagValuesV2, nil, nil, logger),
-		SpanMetricsSummaryHandler: newHandler(metrics, nil, nil, logger),
-		QueryRangeHandler:         newHandler(queryrange, nil, nil, logger),
+		TraceByIDHandler:          newHandler(cfg.Config.LogQueryRequestHeaders, traces, traceByIDSLOPostHook(cfg.TraceByID.SLO), nil, logger),
+		SearchHandler:             newHandler(cfg.Config.LogQueryRequestHeaders, search, searchSLOPostHook(cfg.Search.SLO), searchSLOPreHook, logger),
+		SearchTagsHandler:         newHandler(cfg.Config.LogQueryRequestHeaders, searchTags, nil, nil, logger),
+		SearchTagsV2Handler:       newHandler(cfg.Config.LogQueryRequestHeaders, searchTagsV2, nil, nil, logger),
+		SearchTagsValuesHandler:   newHandler(cfg.Config.LogQueryRequestHeaders, searchTagValues, nil, nil, logger),
+		SearchTagsValuesV2Handler: newHandler(cfg.Config.LogQueryRequestHeaders, searchTagValuesV2, nil, nil, logger),
+		SpanMetricsSummaryHandler: newHandler(cfg.Config.LogQueryRequestHeaders, metrics, nil, nil, logger),
+		QueryRangeHandler:         newHandler(cfg.Config.LogQueryRequestHeaders, queryrange, nil, nil, logger),
 		cacheProvider:             cacheProvider,
 		streamingSearch:           newSearchStreamingGRPCHandler(cfg, o, retryWare.Wrap(next), reader, searchCache, apiPrefix, logger),
 		logger:                    logger,

diff --git a/modules/frontend/handler.go b/modules/frontend/handler.go
@@ -3,11 +3,14 @@ package frontend
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"strings"
 	"time"
 
+	"github.com/grafana/dskit/flagext"
+
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/httpgrpc"
@@ -38,19 +41,21 @@ type (
 // handler exists to wrap a roundtripper with an HTTP handler. It wraps all
 // frontend endpoints and should only contain functionality that is common to all.
 type handler struct {
-	roundTripper http.RoundTripper
-	logger       log.Logger
-	post         handlerPostHook
-	pre          handlerPreHook
+	roundTripper           http.RoundTripper
+	logger                 log.Logger
+	post                   handlerPostHook
+	pre                    handlerPreHook
+	LogQueryRequestHeaders flagext.StringSliceCSV
 }
 
 // newHandler creates a handler
-func newHandler(rt http.RoundTripper, post handlerPostHook, pre handlerPreHook, logger log.Logger) http.Handler {
+func newHandler(LogQueryRequestHeaders flagext.StringSliceCSV, rt http.RoundTripper, post handlerPostHook, pre handlerPreHook, logger log.Logger) http.Handler {
 	return &handler{
-		roundTripper: rt,
-		logger:       logger,
-		post:         post,
-		pre:          pre,
+		LogQueryRequestHeaders: LogQueryRequestHeaders,
+		roundTripper:           rt,
+		logger:                 logger,
+		post:                   post,
+		pre:                    pre,
 	}
 }
 
@@ -81,35 +86,40 @@ func (f *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		f.post(ctx, resp, orgID, elapsed, err)
 	}
 
+	logMessage := append([]interface{}{
+		"tenant", orgID,
+		"method", r.Method,
+		"traceID", traceID,
+		"url", r.URL.RequestURI(),
+		"duration", elapsed.String(),
+	})
+	if len(f.LogQueryRequestHeaders) != 0 {
+		logMessage = append(logMessage, formatRequestHeaders(&r.Header, f.LogQueryRequestHeaders)...)
+	}
+
 	if err != nil {
 		statusCode := http.StatusInternalServerError
 		err = writeError(w, err)
-		level.Info(f.logger).Log(
-			"tenant", orgID,
-			"method", r.Method,
-			"traceID", traceID,
-			"url", r.URL.RequestURI(),
-			"duration", elapsed.String(),
-			"response_size", 0,
+		logMessage = append(
+			logMessage,
 			"status", statusCode,
 			"err", err.Error(),
+			"response_size", 0,
 		)
+		level.Info(f.logger).Log(logMessage...)
 		return
 	}
 
 	if resp == nil {
 		statusCode := http.StatusInternalServerError
 		err = writeError(w, errors.New(NilResponseError))
-		level.Info(f.logger).Log(
-			"tenant", orgID,
-			"method", r.Method,
-			"traceID", traceID,
-			"url", r.URL.RequestURI(),
-			"duration", elapsed.String(),
-			"response_size", 0,
+		logMessage = append(
+			logMessage,
 			"status", statusCode,
 			"err", err.Error(),
+			"response_size", 0,
 		)
+		level.Info(f.logger).Log(logMessage...)
 		return
 	}
 
@@ -131,15 +141,21 @@ func (f *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		contentLength = resp.ContentLength
 	}
 
-	level.Info(f.logger).Log(
-		"tenant", orgID,
-		"method", r.Method,
-		"traceID", traceID,
-		"url", r.URL.RequestURI(),
-		"duration", elapsed.String(),
+	logMessage = append(
+		logMessage,
 		"response_size", contentLength,
 		"status", statusCode,
 	)
+	level.Info(f.logger).Log(logMessage...)
+}
+
+func formatRequestHeaders(h *http.Header, headersToLog []string) (fields []interface{}) {
+	for _, s := range headersToLog {
+		if v := h.Get(s); v != "" {
+			fields = append(fields, fmt.Sprintf("header_%s", strings.ReplaceAll(strings.ToLower(s), "-", "_")), v)
+		}
+	}
+	return fields
 }
 
 func copyHeader(dst, src http.Header) {

diff --git a/modules/frontend/handler_test.go b/modules/frontend/handler_test.go
@@ -0,0 +1,23 @@
+package frontend
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFormatRequestHeaders(t *testing.T) {
+	h := http.Header{}
+	h.Add("X-Header-To-Log", "i should be logged!")
+	h.Add("X-Header-To-Not-Log", "i shouldn't be logged!")
+
+	fields := formatRequestHeaders(&h, []string{"X-Header-To-Log", "X-Header-Not-Present"})
+
+	expected := []interface{}{
+		"header_x_header_to_log",
+		"i should be logged!",
+	}
+
+	require.Equal(t, expected, fields)
+}
diff --git a/modules/frontend/v1/frontend.go b/modules/frontend/v1/frontend.go
@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/grafana/dskit/flagext"
+
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/httpgrpc"
@@ -27,15 +29,17 @@ var errTooManyRequest = httpgrpc.Errorf(http.StatusTooManyRequests, "too many ou
 
 // Config for a Frontend.
 type Config struct {
-	MaxOutstandingPerTenant int           `yaml:"max_outstanding_per_tenant"`
-	QuerierForgetDelay      time.Duration `yaml:"querier_forget_delay"`
-	MaxBatchSize            int           `yaml:"max_batch_size"`
+	MaxOutstandingPerTenant int                    `yaml:"max_outstanding_per_tenant"`
+	QuerierForgetDelay      time.Duration          `yaml:"querier_forget_delay"`
+	MaxBatchSize            int                    `yaml:"max_batch_size"`
+	LogQueryRequestHeaders  flagext.StringSliceCSV `yaml:"log_query_request_headers"`
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet.
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.IntVar(&cfg.MaxOutstandingPerTenant, "querier.max-outstanding-requests-per-tenant", 2000, "Maximum number of outstanding requests per tenant per frontend; requests beyond this error with HTTP 429.")
 	f.DurationVar(&cfg.QuerierForgetDelay, "query-frontend.querier-forget-delay", 0, "If a querier disconnects without sending notification about graceful shutdown, the query-frontend will keep the querier in the tenant's shard until the forget delay has passed. This feature is useful to reduce the blast radius when shuffle-sharding is enabled.")
+	f.Var(&cfg.LogQueryRequestHeaders, "query-frontend.log-query-request-headers", "Comma-separated list of request header names to include in query logs. Applies to both query stats and slow queries logs.")
 }
 
 type Limits interface {