Query splitting: add key to merged response (#78522)

* Query splitting: add key to response

* Query splitting: use queryGroupId as key

* Add unit test

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

@@ -12,6 +12,9 @@ import { trackGroupedQueries } from './tracking';
 import { LokiQuery, LokiQueryType } from './types';
 
 jest.mock('./tracking');
+jest.mock('uuid', () => ({
+  v4: jest.fn().mockReturnValue('uuid'),
+}));
 
 describe('runSplitQuery()', () => {
   let datasource: LokiDatasource;
@@ -39,6 +42,14 @@ describe('runSplitQuery()', () => {
     });
   });
 
+  test('Returns a DataQueryResponse with the expected attributes', async () => {
+    await expect(runSplitQuery(datasource, request)).toEmitValuesWith((response) => {
+      expect(response[0].data).toBeDefined();
+      expect(response[0].state).toBe(LoadingState.Done);
+      expect(response[0].key).toBeDefined();
+    });
+  });
+
   test('Correctly splits queries without step', async () => {
     await expect(runSplitQuery(datasource, request)).toEmitValuesWith(() => {
       expect(datasource.runQuery).toHaveBeenNthCalledWith(
@@ -201,6 +212,7 @@ describe('runSplitQuery()', () => {
           {
             data: [],
             state: LoadingState.Done,
+            key: 'uuid',
           },
           [
             {
@@ -233,6 +245,7 @@ describe('runSplitQuery()', () => {
           {
             data: [],
             state: LoadingState.Done,
+            key: 'uuid',
           },
           [
             {

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

@@ -76,7 +76,8 @@ function adjustTargetsFromResponseState(targets: LokiQuery[], response: DataQuer
     .filter((target) => target.maxLines === undefined || target.maxLines > 0);
 }
 export function runSplitGroupedQueries(datasource: LokiDatasource, requests: LokiGroupedRequest[]) {
-  let mergedResponse: DataQueryResponse = { data: [], state: LoadingState.Streaming };
+  const responseKey = requests.length ? requests[0].request.queryGroupId : uuidv4();
+  let mergedResponse: DataQueryResponse = { data: [], state: LoadingState.Streaming, key: responseKey };
   const totalRequests = Math.max(...requests.map(({ partition }) => partition.length));
   const longestPartition = requests.filter(({ partition }) => partition.length === totalRequests)[0].partition;