Add query-frontend option to select request headers in query logs (#11499)

**What this PR does / why we need it**:

Adding feature present in mimir, specifically
https://github.com/grafana/mimir/pull/5030.

Adds a config option to the query-frontend to specify a list of request
headers to include in query logs.

For example, setting
-frontend.log-query-request-headers="X-Grafana-Org-Id" and sending a
query with X-Grafana-Org-Id:1 results in query log lines that include
header_x_grafana_org_id=1.

**Which issue(s) this PR fixes**:
Fixes #11422

CHANGELOG.md

@@ -51,6 +51,7 @@
 * [11654](https://github.com/grafana/loki/pull/11654) **dannykopping** Cache: atomically check background cache size limit correctly.
 * [11682](https://github.com/grafana/loki/pull/11682) **ashwanthgoli** Metadata cache: Adds `frontend.max-metadata-cache-freshness` to configure the time window for which metadata results are not cached. This helps avoid returning inaccurate results by not caching recent results.
 * [11679](https://github.com/grafana/loki/pull/11679) **dannykopping** Cache: extending #11535 to align custom ingester query split with cache keys for correct caching of results.
+* [11499](https://github.com/grafana/loki/pull/11284) **jmichalek132** Config: Adds `frontend.log-query-request-headers` to enable logging of request headers in query logs.
 
 ##### Fixes
 * [11074](https://github.com/grafana/loki/pull/11074) **hainenber** Fix panic in lambda-promtail due to mishandling of empty DROP_LABELS env var.

docs/sources/configure/_index.md

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

pkg/lokifrontend/frontend/transport/handler.go

@@ -12,6 +12,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/grafana/dskit/flagext"
+
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/httpgrpc"
@@ -42,13 +44,15 @@ var (
 
 // Config for a Handler.
 type HandlerConfig struct {
-	LogQueriesLongerThan time.Duration `yaml:"log_queries_longer_than"`
-	MaxBodySize          int64         `yaml:"max_body_size"`
-	QueryStatsEnabled    bool          `yaml:"query_stats_enabled"`
+	LogQueriesLongerThan   time.Duration          `yaml:"log_queries_longer_than"`
+	LogQueryRequestHeaders flagext.StringSliceCSV `yaml:"log_query_request_headers"`
+	MaxBodySize            int64                  `yaml:"max_body_size"`
+	QueryStatsEnabled      bool                   `yaml:"query_stats_enabled"`
 }
 
 func (cfg *HandlerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.LogQueriesLongerThan, "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, "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, "frontend.max-body-size", 10*1024*1024, "Max body size for downstream prometheus.")
 	f.BoolVar(&cfg.QueryStatsEnabled, "frontend.query-stats-enabled", false, "True to enable query statistics tracking. When enabled, a message with some statistics is logged for every query.")
 }
@@ -206,9 +210,22 @@ func (f *Handler) reportQueryStats(r *http.Request, queryString url.Values, quer
 		"fetched_chunks_bytes", numBytes,
 	}, 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...)
 }
 
+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 (f *Handler) parseRequestQueryString(r *http.Request, bodyBuf bytes.Buffer) url.Values {
 	// Use previously buffered body.
 	r.Body = io.NopCloser(&bodyBuf)

pkg/lokifrontend/frontend/transport/handler_test.go

@@ -0,0 +1,23 @@
+package transport
+
+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)
+}