DataSourceWithBackend: Add `queryGroupId` to find correlated/related queries (#64587)

* add `correlationId` to queries

* trace correlation id in backend

* add correlation id to loki's span

* add correlation id to query chunks

* fix test

* fix DataSourceWithBackend test

* change to `queryGroupId`

* remove empty line

* fix test in `DataSourceWithBackend`

packages/grafana-data/src/types/datasource.ts

@@ -522,6 +522,9 @@ export interface DataQueryRequest<TQuery extends DataQuery = DataQuery> {
 
   // Make it possible to hide support queries from the inspector
   hideFromInspector?: boolean;
+
+  // Used to correlate multiple related requests
+  queryGroupId?: string;
 }
 
 export interface DataQueryTimings {

packages/grafana-runtime/src/utils/DataSourceWithBackend.test.ts

@@ -49,6 +49,7 @@ describe('DataSourceWithBackend', () => {
       targets: [{ refId: 'A' }, { refId: 'B', datasource: { type: 'sample' } }],
       dashboardUID: 'dashA',
       panelId: 123,
+      queryGroupId: 'abc',
     } as DataQueryRequest);
 
     const args = mock.calls[0][0];
@@ -87,6 +88,7 @@ describe('DataSourceWithBackend', () => {
           "X-Datasource-Uid": "abc, <mockuid>",
           "X-Panel-Id": "123",
           "X-Plugin-Id": "dummy, sample",
+          "X-Query-Group-Id": "abc",
         },
         "hideFromInspector": false,
         "method": "POST",

packages/grafana-runtime/src/utils/DataSourceWithBackend.ts

@@ -77,6 +77,7 @@ enum PluginRequestHeaders {
   DatasourceUID = 'X-Datasource-Uid', // can be used for routing/ load balancing
   DashboardUID = 'X-Dashboard-Uid', // mainly useful for debuging slow queries
   PanelID = 'X-Panel-Id', // mainly useful for debuging slow queries
+  QueryGroupID = 'X-Query-Group-Id', // mainly useful to find related queries with query chunking
 }
 
 /**
@@ -210,6 +211,9 @@ class DataSourceWithBackend<
     if (request.panelId) {
       headers[PluginRequestHeaders.PanelID] = `${request.panelId}`;
     }
+    if (request.queryGroupId) {
+      headers[PluginRequestHeaders.QueryGroupID] = `${request.queryGroupId}`;
+    }
     return getBackendSrv()
       .fetch<BackendDataSourceResponse>({
         url,

pkg/services/pluginsintegration/clientmiddleware/tracing_header_middleware.go

@@ -33,7 +33,7 @@ func (m *TracingHeaderMiddleware) applyHeaders(ctx context.Context, req backend.
 		return
 	}
 
-	var headersList = []string{query.HeaderPanelID, query.HeaderDashboardUID, query.HeaderDatasourceUID, `X-Grafana-Org-Id`}
+	var headersList = []string{query.HeaderQueryGroupID, query.HeaderPanelID, query.HeaderDashboardUID, query.HeaderDatasourceUID, `X-Grafana-Org-Id`}
 
 	for _, headerName := range headersList {
 		gotVal := reqCtx.Req.Header.Get(headerName)

pkg/services/pluginsintegration/clientmiddleware/tracing_header_middleware_test.go

@@ -18,6 +18,7 @@ func TestTracingHeaderMiddleware(t *testing.T) {
 		req.Header[`X-Datasource-Uid`] = []string{}
 		req.Header[`X-Grafana-Org-Id`] = []string{}
 		req.Header[`X-Panel-Id`] = []string{}
+		req.Header[`X-Query-Group-Id`] = []string{}
 
 		pluginCtx := backend.PluginContext{
 			DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
@@ -111,6 +112,7 @@ func TestTracingHeaderMiddleware(t *testing.T) {
 		req.Header[`X-Datasource-Uid`] = []string{"aIyC_OcVz"}
 		req.Header[`X-Grafana-Org-Id`] = []string{"1"}
 		req.Header[`X-Panel-Id`] = []string{"2"}
+		req.Header[`X-Query-Group-Id`] = []string{"d26e337d-cb53-481a-9212-0112537b3c1a"}
 
 		pluginCtx := backend.PluginContext{
 			DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
@@ -131,11 +133,12 @@ func TestTracingHeaderMiddleware(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			require.Len(t, cdt.QueryDataReq.GetHTTPHeaders(), 4)
+			require.Len(t, cdt.QueryDataReq.GetHTTPHeaders(), 5)
 			require.Equal(t, `lN53lOcVk`, cdt.QueryDataReq.GetHTTPHeader(`X-Dashboard-Uid`))
 			require.Equal(t, `aIyC_OcVz`, cdt.QueryDataReq.GetHTTPHeader(`X-Datasource-Uid`))
 			require.Equal(t, `1`, cdt.QueryDataReq.GetHTTPHeader(`X-Grafana-Org-Id`))
 			require.Equal(t, `2`, cdt.QueryDataReq.GetHTTPHeader(`X-Panel-Id`))
+			require.Equal(t, `d26e337d-cb53-481a-9212-0112537b3c1a`, cdt.QueryDataReq.GetHTTPHeader(`X-Query-Group-Id`))
 		})
 
 		t.Run("tracing headers are set for health check", func(t *testing.T) {
@@ -153,11 +156,12 @@ func TestTracingHeaderMiddleware(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			require.Len(t, cdt.CheckHealthReq.GetHTTPHeaders(), 4)
+			require.Len(t, cdt.CheckHealthReq.GetHTTPHeaders(), 5)
 			require.Equal(t, `lN53lOcVk`, cdt.CheckHealthReq.GetHTTPHeader(`X-Dashboard-Uid`))
 			require.Equal(t, `aIyC_OcVz`, cdt.CheckHealthReq.GetHTTPHeader(`X-Datasource-Uid`))
 			require.Equal(t, `1`, cdt.CheckHealthReq.GetHTTPHeader(`X-Grafana-Org-Id`))
 			require.Equal(t, `2`, cdt.CheckHealthReq.GetHTTPHeader(`X-Panel-Id`))
+			require.Equal(t, `d26e337d-cb53-481a-9212-0112537b3c1a`, cdt.CheckHealthReq.GetHTTPHeader(`X-Query-Group-Id`))
 		})
 	})
 }

pkg/services/query/query.go

@@ -28,6 +28,7 @@ const (
 	HeaderDatasourceUID = "X-Datasource-Uid" // can be used for routing/ load balancing
 	HeaderDashboardUID  = "X-Dashboard-Uid"  // mainly useful for debuging slow queries
 	HeaderPanelID       = "X-Panel-Id"       // mainly useful for debuging slow queries
+	HeaderQueryGroupID  = "X-Query-Group-Id" // mainly useful for finding related queries with query chunking
 )
 
 func ProvideService(

pkg/tsdb/loki/loki.go

@@ -163,6 +163,10 @@ func queryData(ctx context.Context, req *backend.QueryDataRequest, dsInfo *datas
 		span.SetAttributes("start_unixnano", query.Start, attribute.Key("start_unixnano").Int64(query.Start.UnixNano()))
 		span.SetAttributes("stop_unixnano", query.End, attribute.Key("stop_unixnano").Int64(query.End.UnixNano()))
 
+		if req.GetHTTPHeader("X-Query-Group-Id") != "" {
+			span.SetAttributes("query_group_id", req.GetHTTPHeader("X-Query-Group-Id"), attribute.Key("query_group_id").String(req.GetHTTPHeader("X-Query-Group-Id")))
+		}
+
 		logger := logger.FromContext(ctx) // get logger with trace-id and other contextual info
 		logger.Debug("Sending query", "start", query.Start, "end", query.End, "step", query.Step, "query", query.Expr)
 

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

@@ -1,5 +1,6 @@
 import { groupBy, partition } from 'lodash';
 import { Observable, Subscriber, Subscription } from 'rxjs';
+import { v4 as uuidv4 } from 'uuid';
 
 import {
   DataQueryRequest,
@@ -172,6 +173,8 @@ export function runPartitionedQueries(datasource: LokiDatasource, request: DataQ
   const [instantQueries, normalQueries] = partition(queries, (query) => query.queryType === LokiQueryType.Instant);
   const [logQueries, metricQueries] = partition(normalQueries, (query) => isLogsQuery(query.expr));
 
+  request.queryGroupId = uuidv4();
+
   const oneDayMs = 24 * 60 * 60 * 1000;
   const rangePartitionedLogQueries = groupBy(logQueries, (query) =>
     query.chunkDuration ? durationToMilliseconds(parseDuration(query.chunkDuration)) : oneDayMs