query-frontend: Extract org id from from request headers (#3245)

* query-frontend: Extract org id from from request headers Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net> * docs: regenerate docs Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net> * query-frontend: gofmt Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net> * query-frontend: inline function header Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net> * docs: rebuild query-frontend docs Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net> * docs: clarify usage of org-id-header flag Signed-off-by: Neil Fordyce <neil.fordyce@skyscanner.net>
thanos-io · Nov 11, 2020 · 713ff65 · 713ff65
1 parent b0af765
commit 713ff65
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -27,6 +27,7 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 - [#3398](https://github.com/thanos-io/thanos/pull/3398) Query Frontend: Add default config for query frontend memcached config.
 - [#3277](https://github.com/thanos-io/thanos/pull/3277) Thanos Query: Introduce dynamic lookback interval. This allows queries with large step to make use of downsampled data.
 - [#3409](https://github.com/thanos-io/thanos/pull/3409) Compactor: Added support for no-compact-mark.json which excludes the block from compaction.
+- [#3245](https://github.com/thanos-io/thanos/pull/3245) Query Frontend: Add `query-frontend.org-id-header` flag to specify HTTP header(s) to populate slow query log (e.g. X-Grafana-User).
 
 ### Fixed
 

diff --git a/cmd/thanos/query_frontend.go b/cmd/thanos/query_frontend.go
@@ -36,6 +36,7 @@ import (
 type queryFrontendConfig struct {
 	http httpConfig
 	queryfrontend.Config
+	orgIdHeaders []string
 }
 
 func registerQueryFrontend(app *extkingpin.App) {
@@ -115,6 +116,11 @@ func registerQueryFrontend(app *extkingpin.App) {
 	cmd.Flag("query-frontend.log-queries-longer-than", "Log queries that are slower than the specified duration. "+
 		"Set to 0 to disable. Set to < 0 to enable on all queries.").Default("0").DurationVar(&cfg.CortexFrontendConfig.LogQueriesLongerThan)
 
+	cmd.Flag("query-frontend.org-id-header", "Request header names used to identify the source of slow queries (repeated flag). "+
+		"The values of the header will be added to the org id field in the slow query log. "+
+		"If multiple headers match the request, the first matching arg specified will take precedence. "+
+		"If no headers match 'anonymous' will be used.").PlaceHolder("<http-header-name>").StringsVar(&cfg.orgIdHeaders)
+
 	cmd.Flag("log.request.decision", "Request Logging for logging the start and end of requests. LogFinishCall is enabled by default. LogFinishCall : Logs the finish call of the requests. LogStartAndFinishCall : Logs the start and finish call of the requests. NoLogCall : Disable request logging.").Default("LogFinishCall").EnumVar(&cfg.RequestLoggingDecision, "NoLogCall", "LogFinishCall", "LogStartAndFinishCall")
 
 	cmd.Setup(func(g *run.Group, logger log.Logger, reg *prometheus.Registry, tracer opentracing.Tracer, _ <-chan struct{}, _ bool) error {
@@ -199,6 +205,7 @@ func runQueryFrontend(
 
 		instr := func(f http.HandlerFunc) http.HandlerFunc {
 			hf := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				orgId := extractOrgId(cfg, r)
 				name := "query-frontend"
 				api.SetCORS(w)
 				ins.NewHandler(
@@ -213,7 +220,7 @@ func runQueryFrontend(
 						),
 					),
 					// Cortex frontend middlewares require orgID.
-				).ServeHTTP(w, r.WithContext(user.InjectOrgID(r.Context(), "fake")))
+				).ServeHTTP(w, r.WithContext(user.InjectOrgID(r.Context(), orgId)))
 			})
 			return hf
 		}
@@ -235,3 +242,13 @@ func runQueryFrontend(
 	statusProber.Ready()
 	return nil
 }
+
+func extractOrgId(conf *queryFrontendConfig, r *http.Request) string {
+	for _, header := range conf.orgIdHeaders {
+		headerVal := r.Header.Get(header)
+		if headerVal != "" {
+			return headerVal
+		}
+	}
+	return "anonymous"
+}
diff --git a/cmd/thanos/query_frontend_test.go b/cmd/thanos/query_frontend_test.go
@@ -0,0 +1,46 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+package main
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/thanos-io/thanos/pkg/testutil"
+)
+
+func Test_extractOrgId(t *testing.T) {
+	var testData = []struct {
+		configuredHeaders []string
+		requestHeaders    map[string]string
+		expectOrgId       string
+	}{
+		{
+			configuredHeaders: []string{"x-grafana-user", "x-server-id"},
+			requestHeaders:    map[string]string{"x-grafana-user": "imagrafanauser", "x-server-id": "iamserverid"},
+			expectOrgId:       "imagrafanauser",
+		},
+		{
+			configuredHeaders: []string{"x-server-id", "x-grafana-user"},
+			requestHeaders:    map[string]string{"x-grafana-user": "imagrafanauser", "x-server-id": "iamserverid"},
+			expectOrgId:       "iamserverid",
+		},
+		{
+			configuredHeaders: []string{},
+			requestHeaders:    map[string]string{"another-header": "another"},
+			expectOrgId:       "anonymous",
+		},
+	}
+	for _, data := range testData {
+		config := queryFrontendConfig{
+			orgIdHeaders: data.configuredHeaders,
+		}
+		req, _ := http.NewRequest("", "", nil)
+		for k, v := range data.requestHeaders {
+			req.Header.Set(k, v)
+		}
+		id := extractOrgId(&config, req)
+		testutil.Equals(t, data.expectOrgId, id)
+	}
+}
diff --git a/docs/components/query-frontend.md b/docs/components/query-frontend.md
@@ -222,6 +222,14 @@ Flags:
                                  Log queries that are slower than the specified
                                  duration. Set to 0 to disable. Set to < 0 to
                                  enable on all queries.
+      --query-frontend.org-id-header=<http-header-name> ...
+                                 Request header names used to identify the
+                                 source of slow queries (repeated flag). The
+                                 values of the header will be added to the org
+                                 id field in the slow query log. If multiple
+                                 headers match the request, the first matching
+                                 arg specified will take precedence. If no
+                                 headers match 'anonymous' will be used.
       --log.request.decision=LogFinishCall
                                  Request Logging for logging the start and end
                                  of requests. LogFinishCall is enabled by