Loki: Add X-Query-Tags header for logs sample and data sample (#62333)

pkg/tsdb/loki/api.go

@@ -85,8 +85,11 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery) (*h
 		return nil, err
 	}
 
-	if query.VolumeQuery {
+	if query.SupportingQueryType != SupportingQueryNone {
-		req.Header.Set("X-Query-Tags", "Source=logvolhist")
+		value := getSupportingQueryHeaderValue(req, query.SupportingQueryType)
+		if value != "" {
+			req.Header.Set("X-Query-Tags", "Source="+value)
+		}
 	}
 
 	return req, nil
@@ -223,3 +226,18 @@ func (api *LokiAPI) RawQuery(ctx context.Context, resourcePath string) (RawLokiR
 
 	return encodedBytes, nil
 }
+
+func getSupportingQueryHeaderValue(req *http.Request, supportingQueryType SupportingQueryType) string {
+	value := ""
+	switch supportingQueryType {
+	case SupportingQueryLogsVolume:
+		value = "logvolhist"
+	case SupportingQueryLogsSample:
+		value = "logsample"
+	case SupportingQueryDataSample:
+		value = "datasample"
+	default: //ignore
+	}
+
+	return value
+}

pkg/tsdb/loki/api_test.go

@@ -28,7 +28,31 @@ func TestApiLogVolume(t *testing.T) {
 			require.Equal(t, "Source=logvolhist", req.Header.Get("X-Query-Tags"))
 		})
 
-		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", VolumeQuery: true, QueryType: QueryTypeRange})
+		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryLogsVolume, QueryType: QueryTypeRange})
+		require.NoError(t, err)
+		require.True(t, called)
+	})
+
+	t.Run("logs sample queries should set logs sample http header", func(t *testing.T) {
+		called := false
+		api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
+			called = true
+			require.Equal(t, "Source=logsample", req.Header.Get("X-Query-Tags"))
+		})
+
+		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryLogsSample, QueryType: QueryTypeRange})
+		require.NoError(t, err)
+		require.True(t, called)
+	})
+
+	t.Run("data sample queries should set data sample http header", func(t *testing.T) {
+		called := false
+		api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
+			called = true
+			require.Equal(t, "Source=datasample", req.Header.Get("X-Query-Tags"))
+		})
+
+		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryDataSample, QueryType: QueryTypeRange})
 		require.NoError(t, err)
 		require.True(t, called)
 	})
@@ -40,7 +64,7 @@ func TestApiLogVolume(t *testing.T) {
 			require.Equal(t, "", req.Header.Get("X-Query-Tags"))
 		})
 
-		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", VolumeQuery: false, QueryType: QueryTypeRange})
+		_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryNone, QueryType: QueryTypeRange})
 		require.NoError(t, err)
 		require.True(t, called)
 	})

pkg/tsdb/loki/loki.go

@@ -65,7 +65,7 @@ type QueryJSONModel struct {
 	IntervalMS          int    `json:"intervalMS"`
 	Resolution          int64  `json:"resolution"`
 	MaxLines            int    `json:"maxLines"`
-	VolumeQuery  bool   `json:"volumeQuery"`
+	SupportingQueryType string `json:"supportingQueryType"`
 }
 
 func parseQueryModel(raw json.RawMessage) (*QueryJSONModel, error) {

pkg/tsdb/loki/parse_query.go

@@ -82,6 +82,21 @@ func parseDirection(jsonValue string) (Direction, error) {
 	}
 }
 
+func parseSupportingQueryType(jsonValue string) (SupportingQueryType, error) {
+	switch jsonValue {
+	case "logsVolume":
+		return SupportingQueryLogsVolume, nil
+	case "logsSample":
+		return SupportingQueryLogsSample, nil
+	case "dataSample":
+		return SupportingQueryDataSample, nil
+	case "":
+		return SupportingQueryNone, nil
+	default:
+		return SupportingQueryNone, fmt.Errorf("invalid supportingQueryType: %s", jsonValue)
+	}
+}
+
 func parseQuery(queryContext *backend.QueryDataRequest) ([]*lokiQuery, error) {
 	qs := []*lokiQuery{}
 	for _, query := range queryContext.Queries {
@@ -115,6 +130,11 @@ func parseQuery(queryContext *backend.QueryDataRequest) ([]*lokiQuery, error) {
 			return nil, err
 		}
 
+		supportingQueryType, err := parseSupportingQueryType(model.SupportingQueryType)
+		if err != nil {
+			return nil, err
+		}
+
 		qs = append(qs, &lokiQuery{
 			Expr:                expr,
 			QueryType:           queryType,
@@ -125,7 +145,7 @@ func parseQuery(queryContext *backend.QueryDataRequest) ([]*lokiQuery, error) {
 			Start:               start,
 			End:                 end,
 			RefID:               query.RefID,
-			VolumeQuery:  model.VolumeQuery,
+			SupportingQueryType: supportingQueryType,
 		})
 	}
 

pkg/tsdb/loki/types.go

@@ -3,12 +3,20 @@ package loki
 import "time"
 
 type QueryType string
+type SupportingQueryType string
 
 const (
 	QueryTypeRange   QueryType = "range"
 	QueryTypeInstant QueryType = "instant"
 )
 
+const (
+	SupportingQueryLogsVolume SupportingQueryType = "logsVolume"
+	SupportingQueryLogsSample SupportingQueryType = "logsSample"
+	SupportingQueryDataSample SupportingQueryType = "dataSample"
+	SupportingQueryNone       SupportingQueryType = "none"
+)
+
 type Direction string
 
 const (
@@ -26,5 +34,5 @@ type lokiQuery struct {
 	Start               time.Time
 	End                 time.Time
 	RefID               string
-	VolumeQuery  bool
+	SupportingQueryType SupportingQueryType
 }

public/app/plugins/datasource/loki/datasource.test.ts

@@ -33,7 +33,7 @@ import { CustomVariableModel } from '../../../features/variables/types';
 import { LokiDatasource, REF_ID_DATA_SAMPLES } from './datasource';
 import { createLokiDatasource, createMetadataRequest } from './mocks';
 import { parseToNodeNamesArray } from './queryUtils';
-import { LokiOptions, LokiQuery, LokiQueryType, LokiVariableQueryType } from './types';
+import { LokiOptions, LokiQuery, LokiQueryType, LokiVariableQueryType, SupportingQueryType } from './types';
 import { LokiVariableSupport } from './variables';
 
 jest.mock('@grafana/runtime', () => {
@@ -981,7 +981,7 @@ describe('LokiDatasource', () => {
           instant: false,
           queryType: 'range',
           refId: 'log-volume-A',
-          volumeQuery: true,
+          supportingQueryType: SupportingQueryType.LogsVolume,
         });
       });
 

public/app/plugins/datasource/loki/datasource.ts

@@ -86,6 +86,7 @@ import {
   LokiQueryType,
   LokiVariableQuery,
   LokiVariableQueryType,
+  SupportingQueryType,
 } from './types';
 import { LokiVariableSupport } from './variables';
 
@@ -191,7 +192,7 @@ export class LokiDatasource
           ...normalizedQuery,
           refId: `${REF_ID_STARTER_LOG_VOLUME}${normalizedQuery.refId}`,
           instant: false,
-          volumeQuery: true,
+          supportingQueryType: SupportingQueryType.LogsVolume,
           expr: `sum by (level) (count_over_time(${expr}[$__interval]))`,
         };
 

public/app/plugins/datasource/loki/types.ts

@@ -42,8 +42,8 @@ export interface LokiQuery extends DataQuery {
   legendFormat?: string;
   maxLines?: number;
   resolution?: number;
-  /** Used in range queries */
+  /** Used only to identify supporting queries, e.g. logs volume, logs sample and data sample */
-  volumeQuery?: boolean;
+  supportingQueryType?: SupportingQueryType;
   /* @deprecated now use queryType */
   range?: boolean;
   /* @deprecated now use queryType */
@@ -154,6 +154,12 @@ export interface LokiVariableQuery extends DataQuery {
   stream?: string;
 }
 
+export enum SupportingQueryType {
+  LogsVolume = 'logsVolume',
+  LogsSample = 'logsSample',
+  DataSample = 'dataSample',
+}
+
 export interface ContextFilter {
   enabled: boolean;
   label: string;