Loki: Support custom X-Query-Tags header (#85123)

* Loki: Support custom `X-Query-Tags` header

* add comment

pkg/tsdb/loki/api.go

@@ -99,7 +99,7 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery, cat
 	}
 
 	if query.SupportingQueryType != SupportingQueryNone {
-		value := getSupportingQueryHeaderValue(req, query.SupportingQueryType)
+		value := getSupportingQueryHeaderValue(query.SupportingQueryType)
 		if value != "" {
 			req.Header.Set("X-Query-Tags", "Source="+value)
 		}
@@ -324,8 +324,13 @@ func (api *LokiAPI) RawQuery(ctx context.Context, resourcePath string) (RawLokiR
 	return rawLokiResponse, nil
 }
 
-func getSupportingQueryHeaderValue(req *http.Request, supportingQueryType SupportingQueryType) string {
+func getSupportingQueryHeaderValue(supportingQueryType SupportingQueryType) string {
 	value := ""
+
+	// we need to map the SupportingQueryType to the actual header value. For
+	// legacy reasons we defined each value, such as "logsVolume" maps to the
+	// "logvolhist" header value to Loki. With #85123, even the value set in the
+	// frontend query can be passed as is to Loki.
 	switch supportingQueryType {
 	case SupportingQueryLogsVolume:
 		value = "logvolhist"
@@ -335,7 +340,8 @@ func getSupportingQueryHeaderValue(req *http.Request, supportingQueryType Suppor
 		value = "datasample"
 	case SupportingQueryInfiniteScroll:
 		value = "infinitescroll"
-	default: // ignore
+	default:
+		value = string(supportingQueryType)
 	}
 
 	return value

pkg/tsdb/loki/api_test.go

@@ -58,7 +58,7 @@ func TestApiLogVolume(t *testing.T) {
 		require.True(t, called)
 	})
 
-	t.Run("non-log-volume queries should not set log-volume http header", func(t *testing.T) {
+	t.Run("none queries should not set X-Query-Tags http header", func(t *testing.T) {
 		called := false
 		api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
 			called = true
@@ -70,6 +70,18 @@ func TestApiLogVolume(t *testing.T) {
 		require.True(t, called)
 	})
 
+	t.Run("any defined supporting query should not set X-Query-Tags http header", func(t *testing.T) {
+		called := false
+		api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
+			called = true
+			require.Equal(t, "Source=foo", req.Header.Get("X-Query-Tags"))
+		}, false)
+
+		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryType("foo"), QueryType: QueryTypeRange}, ResponseOpts{})
+		require.NoError(t, err)
+		require.True(t, called)
+	})
+
 	t.Run("with `structuredMetadata` should set correct http header", func(t *testing.T) {
 		called := false
 		api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {

pkg/tsdb/loki/parse_query.go

@@ -100,23 +100,26 @@ func parseDirection(jsonPointerValue *string) (Direction, error) {
 	}
 }
 
-func parseSupportingQueryType(jsonPointerValue *string) (SupportingQueryType, error) {
+func parseSupportingQueryType(jsonPointerValue *string) SupportingQueryType {
 	if jsonPointerValue == nil {
-		return SupportingQueryNone, nil
-	} else {
-		jsonValue := *jsonPointerValue
-		switch jsonValue {
-		case "logsVolume":
-			return SupportingQueryLogsVolume, nil
-		case "logsSample":
-			return SupportingQueryLogsSample, nil
-		case "dataSample":
-			return SupportingQueryDataSample, nil
-		case "infiniteScroll":
-			return SupportingQueryInfiniteScroll, nil
-		default:
-			return SupportingQueryNone, fmt.Errorf("invalid supportingQueryType: %s", jsonValue)
-		}
+		return SupportingQueryNone
+	}
+
+	jsonValue := *jsonPointerValue
+	switch jsonValue {
+	case "logsVolume":
+		return SupportingQueryLogsVolume
+	case "logsSample":
+		return SupportingQueryLogsSample
+	case "dataSample":
+		return SupportingQueryDataSample
+	case "infiniteScroll":
+		return SupportingQueryInfiniteScroll
+	case "":
+		return SupportingQueryNone
+	default:
+		// `SupportingQueryType` is just a `string` in the schema, so we can just parse this as a string
+		return SupportingQueryType(jsonValue)
 	}
 }
 
@@ -166,10 +169,7 @@ func parseQuery(queryContext *backend.QueryDataRequest) ([]*lokiQuery, error) {
 			legendFormat = *model.LegendFormat
 		}
 
-		supportingQueryType, err := parseSupportingQueryType(model.SupportingQueryType)
-		if err != nil {
-			return nil, err
-		}
+		supportingQueryType := parseSupportingQueryType(model.SupportingQueryType)
 
 		qs = append(qs, &lokiQuery{
 			Expr:                expr,

pkg/tsdb/loki/parse_query_test.go

@@ -35,6 +35,85 @@ func TestParseQuery(t *testing.T) {
 		require.Equal(t, time.Second*15, models[0].Step)
 		require.Equal(t, "go_goroutines 15s 15000 3000s 3000 3000000", models[0].Expr)
 	})
+
+	t.Run("parsing query model with logsVolume supporting query type", func(t *testing.T) {
+		queryContext := &backend.QueryDataRequest{
+			Queries: []backend.DataQuery{
+				{
+					JSON: []byte(`
+					{
+						"expr": "go_goroutines $__interval $__interval_ms $__range $__range_s $__range_ms",
+						"format": "time_series",
+						"refId": "A",
+						"supportingQueryType": "logsVolume"
+					}`,
+					),
+					TimeRange: backend.TimeRange{
+						From: time.Now().Add(-3000 * time.Second),
+						To:   time.Now(),
+					},
+					Interval:      time.Second * 15,
+					MaxDataPoints: 200,
+				},
+			},
+		}
+		models, err := parseQuery(queryContext)
+		require.NoError(t, err)
+		require.Equal(t, SupportingQueryLogsVolume, models[0].SupportingQueryType)
+	})
+
+	t.Run("parsing query model with any supporting query type", func(t *testing.T) {
+		queryContext := &backend.QueryDataRequest{
+			Queries: []backend.DataQuery{
+				{
+					JSON: []byte(`
+					{
+						"expr": "go_goroutines $__interval $__interval_ms $__range $__range_s $__range_ms",
+						"format": "time_series",
+						"refId": "A",
+						"supportingQueryType": "foo"
+					}`,
+					),
+					TimeRange: backend.TimeRange{
+						From: time.Now().Add(-3000 * time.Second),
+						To:   time.Now(),
+					},
+					Interval:      time.Second * 15,
+					MaxDataPoints: 200,
+				},
+			},
+		}
+		models, err := parseQuery(queryContext)
+		require.NoError(t, err)
+		require.Equal(t, SupportingQueryType("foo"), models[0].SupportingQueryType)
+	})
+
+	t.Run("parsing query model with any empty query type", func(t *testing.T) {
+		queryContext := &backend.QueryDataRequest{
+			Queries: []backend.DataQuery{
+				{
+					JSON: []byte(`
+					{
+						"expr": "go_goroutines $__interval $__interval_ms $__range $__range_s $__range_ms",
+						"format": "time_series",
+						"refId": "A",
+						"supportingQueryType": ""
+					}`,
+					),
+					TimeRange: backend.TimeRange{
+						From: time.Now().Add(-3000 * time.Second),
+						To:   time.Now(),
+					},
+					Interval:      time.Second * 15,
+					MaxDataPoints: 200,
+				},
+			},
+		}
+		models, err := parseQuery(queryContext)
+		require.NoError(t, err)
+		require.Equal(t, SupportingQueryNone, models[0].SupportingQueryType)
+	})
+
 	t.Run("interpolate variables, range between 1s and 0.5s", func(t *testing.T) {
 		expr := "go_goroutines $__interval $__interval_ms $__range $__range_s $__range_ms"
 		queryType := dataquery.LokiQueryTypeRange