Add query-frontend config option to log request headers in query logs (…

…#5030)

CHANGELOG.md

@@ -16,6 +16,7 @@
   * `-blocks-storage.bucket-store.chunk-pool-max-bucket-size-bytes`
 * [CHANGE] Store-gateway: remove metrics `cortex_bucket_store_chunk_pool_requested_bytes_total` and `cortex_bucket_store_chunk_pool_returned_bytes_total`. #4996
 * [CHANGE] Compactor: change default of `-compactor.partial-block-deletion-delay` to `1d`. This will automatically clean up partial blocks that were a result of failed block upload or deletion. #5026
+* [FEATURE] Query-frontend: add `-query-frontend.log-query-request-headers` to enable logging of request headers in query logs. #5030
 * [ENHANCEMENT] Add per-tenant limit `-validation.max-native-histogram-buckets` to be able to ignore native histogram samples that have too many buckets. #4765
 * [ENHANCEMENT] Store-gateway: reduce memory usage in some LabelValues calls. #4789
 * [ENHANCEMENT] Store-gateway: add a `stage` label to the metric `cortex_bucket_store_series_data_touched`. This label now applies to `data_type="chunks"` and `data_type="series"`. The `stage` label has 2 values: `processed` - the number of series that parsed - and `returned` - the number of series selected from the processed bytes to satisfy the query. #4797 #4830

cmd/mimir/config-descriptor.json

@@ -3684,6 +3684,17 @@
           "fieldFlag": "query-frontend.log-queries-longer-than",
           "fieldType": "duration"
         },
+        {
+          "kind": "field",
+          "name": "log_query_request_headers",
+          "required": false,
+          "desc": "Comma-separated list of request header names to include in query logs. Applies to both query stats and slow queries logs.",
+          "fieldValue": null,
+          "fieldDefaultValue": "",
+          "fieldFlag": "query-frontend.log-query-request-headers",
+          "fieldType": "string",
+          "fieldCategory": "advanced"
+        },
         {
           "kind": "field",
           "name": "max_body_size",

cmd/mimir/help-all.txt.tmpl

@@ -1633,6 +1633,8 @@ Usage of ./cmd/mimir/mimir:
     	Port to advertise to querier (via scheduler) (defaults to server.grpc-listen-port).
   -query-frontend.log-queries-longer-than duration
     	Log queries that are slower than the specified duration. Set to 0 to disable. Set to < 0 to enable on all queries.
+  -query-frontend.log-query-request-headers comma-separated-list-of-strings
+    	Comma-separated list of request header names to include in query logs. Applies to both query stats and slow queries logs.
   -query-frontend.max-body-size int
     	Max body size for downstream prometheus. (default 10485760)
   -query-frontend.max-cache-freshness duration

docs/sources/mimir/references/configuration-parameters/index.md

@@ -1103,6 +1103,11 @@ The `frontend` block configures the query-frontend.
 # CLI flag: -query-frontend.log-queries-longer-than
 [log_queries_longer_than: <duration> | default = 0s]
 
+# (advanced) Comma-separated list of request header names to include in query
+# logs. Applies to both query stats and slow queries logs.
+# CLI flag: -query-frontend.log-query-request-headers
+[log_query_request_headers: <string> | default = ""]
+
 # (advanced) Max body size for downstream prometheus.
 # CLI flag: -query-frontend.max-body-size
 [max_body_size: <int> | default = 10485760]

pkg/frontend/transport/handler.go

@@ -26,6 +26,7 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/httpgrpc/server"
 
+	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/tenant"
 
 	apierror "github.com/grafana/mimir/pkg/api/error"
@@ -49,13 +50,15 @@ var (
 
 // Config for a Handler.
 type HandlerConfig struct {
-	LogQueriesLongerThan time.Duration `yaml:"log_queries_longer_than"`
-	MaxBodySize          int64         `yaml:"max_body_size" category:"advanced"`
-	QueryStatsEnabled    bool          `yaml:"query_stats_enabled" category:"advanced"`
+	LogQueriesLongerThan   time.Duration          `yaml:"log_queries_longer_than"`
+	LogQueryRequestHeaders flagext.StringSliceCSV `yaml:"log_query_request_headers" category:"advanced"`
+	MaxBodySize            int64                  `yaml:"max_body_size" category:"advanced"`
+	QueryStatsEnabled      bool                   `yaml:"query_stats_enabled" category:"advanced"`
 }
 
 func (cfg *HandlerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.LogQueriesLongerThan, "query-frontend.log-queries-longer-than", 0, "Log queries that are slower than the specified duration. Set to 0 to disable. Set to < 0 to enable on all queries.")
+	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.")
 	f.Int64Var(&cfg.MaxBodySize, "query-frontend.max-body-size", 10*1024*1024, "Max body size for downstream prometheus.")
 	f.BoolVar(&cfg.QueryStatsEnabled, "query-frontend.query-stats-enabled", true, "False to disable query statistics tracking. When enabled, a message with some statistics is logged for every query.")
 }
@@ -240,6 +243,10 @@ func (f *Handler) reportSlowQuery(r *http.Request, queryString url.Values, query
 		"time_taken", queryResponseTime.String(),
 	}, formatQueryString(queryString)...)
 
+	if len(f.cfg.LogQueryRequestHeaders) != 0 {
+		logMessage = append(logMessage, formatRequestHeaders(&r.Header, f.cfg.LogQueryRequestHeaders)...)
+	}
+
 	level.Info(util_log.WithContext(r.Context(), f.log)).Log(logMessage...)
 }
 
@@ -284,6 +291,10 @@ func (f *Handler) reportQueryStats(r *http.Request, queryString url.Values, quer
 		"estimated_series_count", stats.GetEstimatedSeriesCount(),
 	}, formatQueryString(queryString)...)
 
+	if len(f.cfg.LogQueryRequestHeaders) != 0 {
+		logMessage = append(logMessage, formatRequestHeaders(&r.Header, f.cfg.LogQueryRequestHeaders)...)
+	}
+
 	if queryErr != nil {
 		logStatus := "failed"
 		if errors.Is(queryErr, context.Canceled) {
@@ -310,6 +321,15 @@ func formatQueryString(queryString url.Values) (fields []interface{}) {
 	return fields
 }
 
+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 writeError(w http.ResponseWriter, err error) {
 	switch {
 	case errors.Is(err, context.Canceled):

pkg/frontend/transport/handler_test.go

@@ -378,3 +378,18 @@ func (t *testLogger) Log(keyvals ...interface{}) error {
 	t.logMessages = append(t.logMessages, msg)
 	return nil
 }
+
+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!",
+	}
+
+	assert.Equal(t, expected, fields)
+}