From f258d9a3263060a8daa2013d15bc30e207c2c889 Mon Sep 17 00:00:00 2001
From: Karsten Jeschkies <karsten.jeschkies@grafana.com>
Date: Tue, 7 Nov 2023 14:43:44 +0100
Subject: [PATCH] Extract query tags from HTTP requests. (#11147)

**What this PR does / why we need it**:
https://github.com/grafana/loki/pull/10858 removed the code that would
inject query tags into the context of the querier. This change adds an
extraction in the decode method.
---
 pkg/querier/queryrange/codec.go |  5 +++++
 pkg/util/httpreq/tags.go        | 16 ++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/pkg/querier/queryrange/codec.go b/pkg/querier/queryrange/codec.go
index 2f37aaecf7530..bc34b6477b60e 100644
--- a/pkg/querier/queryrange/codec.go
+++ b/pkg/querier/queryrange/codec.go
@@ -378,6 +378,11 @@ func (Codec) DecodeHTTPGrpcRequest(ctx context.Context, r *httpgrpc.HTTPRequest)
 		}
 	}
 
+	// Add query tags
+	if queryTags := httpreq.ExtractQueryTagsFromHTTP(httpReq); queryTags != "" {
+		ctx = httpreq.InjectQueryTags(ctx, queryTags)
+	}
+
 	// If there is not encoding flags in the context, we try the HTTP request.
 	if encFlags := httpreq.ExtractEncodingFlagsFromCtx(ctx); encFlags == nil {
 		encFlags = httpreq.ExtractEncodingFlagsFromProto(r)
diff --git a/pkg/util/httpreq/tags.go b/pkg/util/httpreq/tags.go
index 9295a04ae31a9..a0ed6c0d05386 100644
--- a/pkg/util/httpreq/tags.go
+++ b/pkg/util/httpreq/tags.go
@@ -25,11 +25,9 @@ func ExtractQueryTagsMiddleware() middleware.Interface {
 	return middleware.Func(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 			ctx := req.Context()
-			tags := req.Header.Get(string(QueryTagsHTTPHeader))
-			tags = safeQueryTags.ReplaceAllString(tags, "_")
 
-			if tags != "" {
-				ctx = context.WithValue(ctx, QueryTagsHTTPHeader, tags)
+			if tags := ExtractQueryTagsFromHTTP(req); tags != "" {
+				ctx = InjectQueryTags(ctx, tags)
 				req = req.WithContext(ctx)
 			}
 			next.ServeHTTP(w, req)
@@ -37,6 +35,16 @@ func ExtractQueryTagsMiddleware() middleware.Interface {
 	})
 }
 
+func ExtractQueryTagsFromHTTP(req *http.Request) string {
+	tags := req.Header.Get(string(QueryTagsHTTPHeader))
+	return safeQueryTags.ReplaceAllString(tags, "_")
+}
+
+func InjectQueryTags(ctx context.Context, tags string) context.Context {
+	tags = safeQueryTags.ReplaceAllString(tags, "_")
+	return context.WithValue(ctx, QueryTagsHTTPHeader, tags)
+}
+
 func ExtractQueryMetricsMiddleware() middleware.Interface {
 	return middleware.Func(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {