Loki Query Splitting: Enable tracking for split queries (#68645)

* Loki datasource: move tracking out of runQuery

* Types: move LokiGroupedRequest to types file

* Tracking: add function to track grouped queries

* Query splitting: add tracking to split queries runner

* Remove unnecessary types

* Add unit test

* Tracking: include test case with hidden query

* Update public/app/plugins/datasource/loki/tracking.ts

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* Update public/app/plugins/datasource/loki/tracking.ts

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* Update public/app/plugins/datasource/loki/tracking.ts

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* Update public/app/plugins/datasource/loki/tracking.ts

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* Tracking: add is_split dimension

---------

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

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

@@ -287,16 +287,17 @@ export class LokiDatasource
       return runSplitQuery(this, fixedRequest);
     }
 
-    return this.runQuery(fixedRequest);
+    const startTime = new Date();
+    return this.runQuery(fixedRequest).pipe(tap((response) => trackQuery(response, fixedRequest, startTime)));
   }
 
-  runQuery(fixedRequest: DataQueryRequest<LokiQuery> & { targets: LokiQuery[] }) {
-    const startTime = new Date();
-    return super.query(fixedRequest).pipe(
+  runQuery(fixedRequest: DataQueryRequest<LokiQuery>) {
+    return super
+      .query(fixedRequest)
+      .pipe(
         map((response) =>
           transformBackendResult(response, fixedRequest.targets, this.instanceSettings.jsonData.derivedFields ?? [])
-      ),
-      tap((response) => trackQuery(response, fixedRequest, startTime))
+        )
       );
   }
 

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

@@ -9,8 +9,11 @@ import * as logsTimeSplit from './logsTimeSplitting';
 import * as metricTimeSplit from './metricTimeSplitting';
 import { createLokiDatasource, getMockFrames } from './mocks';
 import { runSplitQuery } from './querySplitting';
+import { trackGroupedQueries } from './tracking';
 import { LokiQuery, LokiQueryType } from './types';
 
+jest.mock('./tracking');
+
 describe('runSplitQuery()', () => {
   let datasource: LokiDatasource;
   const range = {
@@ -57,15 +60,36 @@ describe('runSplitQuery()', () => {
     beforeAll(() => {
       jest.spyOn(logsTimeSplit, 'splitTimeRange').mockReturnValue([]);
       jest.spyOn(metricTimeSplit, 'splitTimeRange').mockReturnValue([]);
+      jest.mocked(trackGroupedQueries).mockClear();
+      jest.useFakeTimers().setSystemTime(new Date('Wed May 17 2023 17:20:12 GMT+0200'));
     });
     afterAll(() => {
       jest.mocked(logsTimeSplit.splitTimeRange).mockRestore();
       jest.mocked(metricTimeSplit.splitTimeRange).mockRestore();
+      jest.useRealTimers();
     });
     test('Ignores hidden queries', async () => {
       await expect(runSplitQuery(datasource, request)).toEmitValuesWith(() => {
         expect(logsTimeSplit.splitTimeRange).toHaveBeenCalled();
         expect(metricTimeSplit.splitTimeRange).not.toHaveBeenCalled();
+        expect(trackGroupedQueries).toHaveBeenCalledTimes(1);
+        expect(trackGroupedQueries).toHaveBeenCalledWith(
+          {
+            data: [],
+            state: LoadingState.Done,
+          },
+          [
+            {
+              partition: [],
+              request: {
+                ...request,
+                targets: request.targets.filter((query) => !query.hide),
+              },
+            },
+          ],
+          request,
+          new Date()
+        );
       });
     });
   });

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

@@ -1,5 +1,5 @@
 import { groupBy, partition } from 'lodash';
-import { Observable, Subscriber, Subscription } from 'rxjs';
+import { Observable, Subscriber, Subscription, tap } from 'rxjs';
 import { v4 as uuidv4 } from 'uuid';
 
 import {
@@ -17,7 +17,8 @@ import { splitTimeRange as splitLogsTimeRange } from './logsTimeSplitting';
 import { splitTimeRange as splitMetricTimeRange } from './metricTimeSplitting';
 import { isLogsQuery } from './queryUtils';
 import { combineResponses } from './responseUtils';
-import { LokiQuery, LokiQueryType } from './types';
+import { trackGroupedQueries } from './tracking';
+import { LokiGroupedRequest, LokiQuery, LokiQueryType } from './types';
 
 export function partitionTimeRange(
   isLogsQuery: boolean,
@@ -80,10 +81,7 @@ function adjustTargetsFromResponseState(targets: LokiQuery[], response: DataQuer
     })
     .filter((target) => target.maxLines === undefined || target.maxLines > 0);
 }
-
-type LokiGroupedRequest = Array<{ request: DataQueryRequest<LokiQuery>; partition: TimeRange[] }>;
-
-export function runSplitGroupedQueries(datasource: LokiDatasource, requests: LokiGroupedRequest) {
+export function runSplitGroupedQueries(datasource: LokiDatasource, requests: LokiGroupedRequest[]) {
   let mergedResponse: DataQueryResponse = { data: [], state: LoadingState.Streaming };
   const totalRequests = Math.max(...requests.map(({ partition }) => partition.length));
 
@@ -155,7 +153,7 @@ export function runSplitGroupedQueries(datasource: LokiDatasource, requests: Lok
   return response;
 }
 
-function getNextRequestPointers(requests: LokiGroupedRequest, requestGroup: number, requestN: number) {
+function getNextRequestPointers(requests: LokiGroupedRequest[], requestGroup: number, requestN: number) {
   // There's a pending request from the next group:
   for (let i = requestGroup + 1; i < requests.length; i++) {
     const group = requests[i];
@@ -187,7 +185,7 @@ export function runSplitQuery(datasource: LokiDatasource, request: DataQueryRequ
     query.splitDuration ? durationToMilliseconds(parseDuration(query.splitDuration)) : oneDayMs
   );
 
-  const requests: LokiGroupedRequest = [];
+  const requests: LokiGroupedRequest[] = [];
   for (const [chunkRangeMs, queries] of Object.entries(rangePartitionedLogQueries)) {
     const resolutionPartition = groupBy(queries, (query) => query.resolution || 1);
     for (const resolution in resolutionPartition) {
@@ -227,5 +225,12 @@ export function runSplitQuery(datasource: LokiDatasource, request: DataQueryRequ
     });
   }
 
-  return runSplitGroupedQueries(datasource, requests);
+  const startTime = new Date();
+  return runSplitGroupedQueries(datasource, requests).pipe(
+    tap((response) => {
+      if (response.state === LoadingState.Done) {
+        trackGroupedQueries(response, requests, request, startTime);
+      }
+    })
+  );
 }

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

@@ -0,0 +1,153 @@
+import { getQueryOptions } from 'test/helpers/getQueryOptions';
+
+import { dateTime } from '@grafana/data';
+import { reportInteraction } from '@grafana/runtime';
+
+import { QueryEditorMode } from '../prometheus/querybuilder/shared/types';
+
+import { partitionTimeRange } from './querySplitting';
+import { trackGroupedQueries, trackQuery } from './tracking';
+import { LokiGroupedRequest, LokiQuery } from './types';
+
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  reportInteraction: jest.fn(),
+}));
+
+const baseTarget = {
+  resolution: 1,
+  editorMode: QueryEditorMode.Builder,
+};
+const range = {
+  from: dateTime('2023-02-08T05:00:00.000Z'),
+  to: dateTime('2023-02-10T06:00:00.000Z'),
+  raw: {
+    from: dateTime('2023-02-08T05:00:00.000Z'),
+    to: dateTime('2023-02-10T06:00:00.000Z'),
+  },
+};
+const originalRequest = getQueryOptions<LokiQuery>({
+  targets: [
+    { expr: 'count_over_time({a="b"}[1m])', refId: 'A', ...baseTarget },
+    { expr: '{a="b"}', refId: 'B', maxLines: 10, ...baseTarget },
+    { expr: 'count_over_time({hidden="true"}[1m])', refId: 'C', ...baseTarget, hide: true },
+  ],
+  range,
+  app: 'explore',
+});
+const requests: LokiGroupedRequest[] = [
+  {
+    request: {
+      ...getQueryOptions<LokiQuery>({
+        targets: [{ expr: 'count_over_time({a="b"}[1m])', refId: 'A', ...baseTarget }],
+        range,
+      }),
+      app: 'explore',
+    },
+    partition: partitionTimeRange(true, range, 60000, 1, 24 * 60 * 60 * 1000),
+  },
+  {
+    request: {
+      ...getQueryOptions<LokiQuery>({
+        targets: [{ expr: '{a="b"}', refId: 'B', maxLines: 10, ...baseTarget }],
+        range,
+      }),
+      app: 'explore',
+    },
+    partition: partitionTimeRange(false, range, 60000, 1, 24 * 60 * 60 * 1000),
+  },
+];
+
+beforeAll(() => {
+  jest.useFakeTimers().setSystemTime(new Date('Wed May 17 2023 17:20:12 GMT+0200'));
+});
+afterAll(() => {
+  jest.useRealTimers();
+});
+beforeEach(() => {
+  jest.mocked(reportInteraction).mockClear();
+});
+
+test('Tracks queries', () => {
+  trackQuery({ data: [] }, originalRequest, new Date());
+
+  expect(reportInteraction).toHaveBeenCalledWith('grafana_loki_query_executed', {
+    app: 'explore',
+    bytes_processed: 0,
+    editor_mode: 'builder',
+    grafana_version: '1.0',
+    has_data: false,
+    has_error: false,
+    is_split: false,
+    legend: undefined,
+    line_limit: undefined,
+    obfuscated_query: 'count_over_time({Identifier=String}[1m])',
+    parsed_query:
+      'LogQL,Expr,MetricExpr,RangeAggregationExpr,RangeOp,CountOverTime,LogRangeExpr,Selector,Matchers,Matcher,Identifier,Eq,String,Range,Duration',
+    query_type: 'metric',
+    query_vector_type: undefined,
+    resolution: 1,
+    simultaneously_executed_query_count: 2,
+    simultaneously_hidden_query_count: 1,
+    time_range_from: '2023-02-08T05:00:00.000Z',
+    time_range_to: '2023-02-10T06:00:00.000Z',
+    time_taken: 0,
+  });
+});
+
+test('Tracks grouped queries', () => {
+  trackGroupedQueries({ data: [] }, requests, originalRequest, new Date());
+
+  expect(reportInteraction).toHaveBeenCalledWith('grafana_loki_query_executed', {
+    app: 'explore',
+    bytes_processed: 0,
+    editor_mode: 'builder',
+    grafana_version: '1.0',
+    has_data: false,
+    has_error: false,
+    is_split: true,
+    legend: undefined,
+    line_limit: undefined,
+    obfuscated_query: 'count_over_time({Identifier=String}[1m])',
+    parsed_query:
+      'LogQL,Expr,MetricExpr,RangeAggregationExpr,RangeOp,CountOverTime,LogRangeExpr,Selector,Matchers,Matcher,Identifier,Eq,String,Range,Duration',
+    query_type: 'metric',
+    query_vector_type: undefined,
+    resolution: 1,
+    simultaneously_executed_query_count: 2,
+    simultaneously_hidden_query_count: 1,
+    split_query_group_count: 2,
+    split_query_largest_partition_size: 3,
+    split_query_partition_size: 3,
+    split_query_total_request_count: 6,
+    time_range_from: '2023-02-08T05:00:00.000Z',
+    time_range_to: '2023-02-10T06:00:00.000Z',
+    time_taken: 0,
+  });
+
+  expect(reportInteraction).toHaveBeenCalledWith('grafana_loki_query_executed', {
+    app: 'explore',
+    bytes_processed: 0,
+    editor_mode: 'builder',
+    grafana_version: '1.0',
+    has_data: false,
+    has_error: false,
+    is_split: true,
+    legend: undefined,
+    line_limit: 10,
+    obfuscated_query: '{Identifier=String}',
+    parsed_query: 'LogQL,Expr,LogExpr,Selector,Matchers,Matcher,Identifier,Eq,String',
+    query_type: 'logs',
+    query_vector_type: undefined,
+    resolution: 1,
+    simultaneously_executed_query_count: 2,
+    simultaneously_hidden_query_count: 1,
+    split_query_group_count: 2,
+    split_query_largest_partition_size: 3,
+    split_query_partition_size: 3,
+    split_query_total_request_count: 6,
+    time_range_from: '2023-02-08T05:00:00.000Z',
+    time_range_to: '2023-02-10T06:00:00.000Z',
+    time_taken: 0,
+  });
+});

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

@@ -12,7 +12,7 @@ import {
 } from './datasource';
 import pluginJson from './plugin.json';
 import { getNormalizedLokiQuery, isLogsQuery, obfuscate, parseToNodeNamesArray } from './queryUtils';
-import { LokiQuery, LokiQueryType } from './types';
+import { LokiGroupedRequest, LokiQuery, LokiQueryType } from './types';
 
 type LokiOnDashboardLoadedTrackingEvent = {
   grafana_version?: string;
@@ -132,23 +132,7 @@ const shouldNotReportBasedOnRefId = (refId: string): boolean => {
   return false;
 };
 
-export function trackQuery(
-  response: DataQueryResponse,
-  request: DataQueryRequest<LokiQuery> & { targets: LokiQuery[] },
-  startTime: Date
-): void {
-  // We only want to track usage for these specific apps
-  const { app, targets: queries } = request;
-
-  if (app === CoreApp.Dashboard || app === CoreApp.PanelViewer) {
-    return;
-  }
-
-  // TODO: We need to re-think this for split queries
-  if (config.featureToggles.lokiQuerySplitting) {
-    return;
-  }
-
+const calculateTotalBytes = (response: DataQueryResponse): number => {
   let totalBytes = 0;
   for (const frame of response.data) {
     const byteKey = frame.meta?.custom?.lokiQueryStatKey;
@@ -157,6 +141,23 @@ export function trackQuery(
         frame.meta?.stats?.find((stat: { displayName: string }) => stat.displayName === byteKey)?.value ?? 0;
     }
   }
+  return totalBytes;
+};
+
+export function trackQuery(
+  response: DataQueryResponse,
+  request: DataQueryRequest<LokiQuery>,
+  startTime: Date,
+  extraPayload: Record<string, unknown> = {}
+): void {
+  // We only want to track usage for these specific apps
+  const { app, targets: queries } = request;
+
+  if (app === CoreApp.Dashboard || app === CoreApp.PanelViewer) {
+    return;
+  }
+
+  let totalBytes = calculateTotalBytes(response);
 
   for (const query of queries) {
     if (shouldNotReportBasedOnRefId(query.refId)) {
@@ -182,6 +183,32 @@ export function trackQuery(
       time_range_to: request?.range?.to?.toISOString(),
       time_taken: Date.now() - startTime.getTime(),
       bytes_processed: totalBytes,
+      is_split: false,
+      ...extraPayload,
+    });
+  }
+}
+
+export function trackGroupedQueries(
+  response: DataQueryResponse,
+  groupedRequests: LokiGroupedRequest[],
+  originalRequest: DataQueryRequest<LokiQuery>,
+  startTime: Date
+): void {
+  const splittingPayload = {
+    split_query_group_count: groupedRequests.length,
+    split_query_largest_partition_size: Math.max(...groupedRequests.map(({ partition }) => partition.length)),
+    split_query_total_request_count: groupedRequests.reduce((total, { partition }) => total + partition.length, 0),
+    is_split: true,
+    simultaneously_executed_query_count: originalRequest.targets.filter((query) => !query.hide).length,
+    simultaneously_hidden_query_count: originalRequest.targets.filter((query) => query.hide).length,
+  };
+
+  for (const group of groupedRequests) {
+    const split_query_partition_size = group.partition.length;
+    trackQuery(response, group.request, startTime, {
+      ...splittingPayload,
+      split_query_partition_size,
     });
   }
 }

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

@@ -1,4 +1,4 @@
-import { DataQuery, DataSourceJsonData, QueryResultMeta, ScopedVars } from '@grafana/data';
+import { DataQuery, DataQueryRequest, DataSourceJsonData, QueryResultMeta, ScopedVars, TimeRange } from '@grafana/data';
 
 import { Loki as LokiQueryFromSchema, LokiQueryType, SupportingQueryType, LokiQueryDirection } from './dataquery.gen';
 
@@ -160,3 +160,5 @@ export interface ContextFilter {
   fromParser: boolean;
   description?: string;
 }
+
+export type LokiGroupedRequest = { request: DataQueryRequest<LokiQuery>; partition: TimeRange[] };