Loki: Use better query validation before requesting stats (#70800)

* fix: use correct validate funtion

* rename isValidQuery to isQueryWithError

* lets pretend i didnt break things again

* Loki validation: add missing comments and refactor validation function

* isValidQuery: add unit test

* Loki datasource: add tests for getQueryStats

* Remove isValidQuery function

* UnwrapParamEditor: interpolate query before testing validity

* Stats: trim queries when evaluating change

---------

Co-authored-by: Matias Chomicki <matyax@gmail.com>

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/validation.ts

@@ -16,6 +16,13 @@ interface ParseError {
   node: SyntaxNode;
 }
 
+/**
+ * Conceived to work in combination with the MonacoQueryField component.
+ * Given an original query, and it's interpolated version, it will return an array of ParserErrorBoundary
+ * objects containing nodes which are actual errors. The interpolated version (even with placeholder variables)
+ * is required because variables look like errors for Lezer.
+ * @internal
+ */
 export function validateQuery(
   query: string,
   interpolatedQuery: string,

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

@@ -5,44 +5,43 @@ import { createLokiDatasource } from '../mocks';
 import { getStats, shouldUpdateStats } from './stats';
 
 describe('shouldUpdateStats', () => {
+  const timerange = getDefaultTimeRange();
   it('should return true if the query has changed', () => {
     const query = '{job="grafana"}';
     const prevQuery = '{job="not-grafana"}';
-    const timerange = getDefaultTimeRange();
-    const prevTimerange = timerange;
-    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(true);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(true);
   });
 
   it('should return true if the timerange has changed', () => {
     const query = '{job="grafana"}';
     const prevQuery = '{job="grafana"}';
-    const timerange = getDefaultTimeRange();
     timerange.raw.from = 'now-14h';
     const prevTimerange = getDefaultTimeRange();
     expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(true);
   });
 
-  it('should return false if the query and timerange have not changed', () => {
+  it('should return true if the previous query was undefined', () => {
     const query = '{job="grafana"}';
+    const prevQuery = undefined;
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(true);
+  });
+
+  it('should return true if the query really changed, otherwise false', () => {
     const prevQuery = '{job="grafana"}';
-    const timerange = getDefaultTimeRange();
-    const prevTimerange = timerange;
-    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(false);
+    const query = `${prevQuery} `;
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(false);
   });
 
   it('should return false if the query and timerange have not changed', () => {
     const query = '{job="grafana"}';
     const prevQuery = '{job="grafana"}';
-    const timerange = getDefaultTimeRange();
-    const prevTimerange = getDefaultTimeRange();
-    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(false);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(false);
   });
 
   it('should return false if the query and timerange with absolute and relative mixed have not changed', () => {
     const query = '{job="grafana"}';
     const prevQuery = '{job="grafana"}';
     const now = dateTime(Date.now());
-    const timerange = getDefaultTimeRange();
     timerange.raw.from = now;
 
     const prevTimerange = getDefaultTimeRange();
@@ -83,4 +82,21 @@ describe('makeStatsRequest', () => {
       entries: 78344,
     });
   });
+
+  it('should support queries with variables', () => {
+    const query = 'count_over_time({job="grafana"}[$__interval])';
+
+    datasource.interpolateString = jest
+      .fn()
+      .mockImplementationOnce((value: string) => value.replace('$__interval', '1h'));
+    datasource.getQueryStats = jest
+      .fn()
+      .mockResolvedValue({ streams: 1, chunks: 12611, bytes: 12913664, entries: 78344 });
+    expect(getStats(datasource, query)).resolves.toEqual({
+      streams: 1,
+      chunks: 12611,
+      bytes: 12913664,
+      entries: 78344,
+    });
+  });
 });

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

@@ -39,7 +39,7 @@ export function shouldUpdateStats(
   timerange: TimeRange,
   prevTimerange: TimeRange | undefined
 ): boolean {
-  if (query !== prevQuery) {
+  if (prevQuery === undefined || query.trim() !== prevQuery.trim()) {
     return true;
   }
 

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

@@ -32,7 +32,7 @@ import { LokiDatasource, REF_ID_DATA_SAMPLES } from './datasource';
 import { createLokiDatasource, createMetadataRequest } from './mocks';
 import { runSplitQuery } from './querySplitting';
 import { parseToNodeNamesArray } from './queryUtils';
-import { LokiOptions, LokiQuery, LokiQueryType, LokiVariableQueryType, QueryStats, SupportingQueryType } from './types';
+import { LokiOptions, LokiQuery, LokiQueryType, LokiVariableQueryType, SupportingQueryType } from './types';
 import { LokiVariableSupport } from './variables';
 
 jest.mock('@grafana/runtime', () => {
@@ -1198,11 +1198,43 @@ describe('LokiDatasource', () => {
   });
 
   describe('getQueryStats', () => {
+    let ds: LokiDatasource;
+    beforeEach(() => {
+      ds = createLokiDatasource(templateSrvStub);
+      ds.statsMetadataRequest = jest.fn().mockResolvedValue({ streams: 1, chunks: 1, bytes: 1, entries: 1 });
+      ds.interpolateString = jest.fn().mockImplementation((value: string) => value.replace('$__interval', '1m'));
+    });
+
     it('uses statsMetadataRequest', async () => {
-      const ds = createLokiDatasource(templateSrvStub);
-      const spy = jest.spyOn(ds, 'statsMetadataRequest').mockResolvedValue({} as QueryStats);
-      ds.getQueryStats('{foo="bar"}');
-      expect(spy).toHaveBeenCalled();
+      const result = await ds.getQueryStats('{foo="bar"}');
+
+      expect(ds.statsMetadataRequest).toHaveBeenCalled();
+      expect(result).toEqual({ streams: 1, chunks: 1, bytes: 1, entries: 1 });
+    });
+
+    it('supports queries with template variables', async () => {
+      const result = await ds.getQueryStats('rate({instance="server\\1"}[$__interval])');
+
+      expect(result).toEqual({
+        streams: 1,
+        chunks: 1,
+        bytes: 1,
+        entries: 1,
+      });
+    });
+
+    it('does not call stats if the query is invalid', async () => {
+      const result = await ds.getQueryStats('rate({label="value"}');
+
+      expect(ds.statsMetadataRequest).not.toHaveBeenCalled();
+      expect(result).toBe(undefined);
+    });
+
+    it('combines the stats of each label matcher', async () => {
+      const result = await ds.getQueryStats('count_over_time({foo="bar"}[1m]) + count_over_time({test="test"}[1m])');
+
+      expect(ds.statsMetadataRequest).toHaveBeenCalled();
+      expect(result).toEqual({ streams: 2, chunks: 2, bytes: 2, entries: 2 });
     });
   });
 

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

@@ -51,6 +51,7 @@ import { LiveStreams, LokiLiveTarget } from './LiveStreams';
 import { LogContextProvider } from './LogContextProvider';
 import { transformBackendResult } from './backendResultTransformer';
 import { LokiAnnotationsQueryEditor } from './components/AnnotationsQueryEditor';
+import { placeHolderScopedVars } from './components/monaco-query-field/monaco-completion-provider/validation';
 import { escapeLabelValueInSelector, isRegexSelector } from './languageUtils';
 import { labelNamesRegex, labelValuesRegex } from './migrations/variableQueryMigrations';
 import {
@@ -76,7 +77,7 @@ import {
   getNormalizedLokiQuery,
   getStreamSelectorsFromQuery,
   isLogsQuery,
-  isValidQuery,
+  isQueryWithError,
   requestSupportsSplitting,
 } from './queryUtils';
 import { doLokiChannelStream } from './streaming';
@@ -451,7 +452,7 @@ export class LokiDatasource
 
   async getQueryStats(query: string): Promise<QueryStats | undefined> {
     // if query is invalid, clear stats, and don't request
-    if (!isValidQuery(query)) {
+    if (isQueryWithError(this.interpolateString(query, placeHolderScopedVars))) {
       return undefined;
     }
 
@@ -570,7 +571,7 @@ export class LokiDatasource
 
   async getDataSamples(query: LokiQuery): Promise<DataFrame[]> {
     // Currently works only for logs sample
-    if (!isValidQuery(query.expr) || !isLogsQuery(query.expr)) {
+    if (!isLogsQuery(query.expr) || isQueryWithError(this.interpolateString(query.expr, placeHolderScopedVars))) {
       return [];
     }
 

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

@@ -6,7 +6,7 @@ import {
   isLogsQuery,
   isQueryWithLabelFormat,
   isQueryWithParser,
-  isValidQuery,
+  isQueryWithError,
   parseToNodeNamesArray,
   getParserFromQuery,
   obfuscate,
@@ -188,12 +188,12 @@ describe('getLokiQueryType', () => {
   });
 });
 
-describe('isValidQuery', () => {
+describe('isQueryWithError', () => {
   it('returns false if invalid query', () => {
-    expect(isValidQuery('{job="grafana')).toBe(false);
+    expect(isQueryWithError('{job="grafana')).toBe(true);
   });
   it('returns true if valid query', () => {
-    expect(isValidQuery('{job="grafana"}')).toBe(true);
+    expect(isQueryWithError('{job="grafana"}')).toBe(false);
   });
 });
 

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

@@ -176,8 +176,13 @@ export function getNodeFromQuery(query: string, nodeType: number): SyntaxNode |
   return nodes.length > 0 ? nodes[0] : undefined;
 }
 
-export function isValidQuery(query: string): boolean {
-  return !isQueryWithNode(query, ErrorId);
+/**
+ * Parses the query and looks for error nodes. If there is at least one, it returns false.
+ * Grafana variables are considered errors, so if you need to validate a query
+ * with variables you should interpolate it first.
+ */
+export function isQueryWithError(query: string): boolean {
+  return isQueryWithNode(query, ErrorId);
 }
 
 export function isLogsQuery(query: string): boolean {

public/app/plugins/datasource/loki/querybuilder/components/UnwrapParamEditor.tsx

@@ -5,8 +5,9 @@ import { Select } from '@grafana/ui';
 
 import { getOperationParamId } from '../../../prometheus/querybuilder/shared/operationUtils';
 import { QueryBuilderOperationParamEditorProps } from '../../../prometheus/querybuilder/shared/types';
+import { placeHolderScopedVars } from '../../components/monaco-query-field/monaco-completion-provider/validation';
 import { LokiDatasource } from '../../datasource';
-import { getLogQueryFromMetricsQuery, isValidQuery } from '../../queryUtils';
+import { getLogQueryFromMetricsQuery, isQueryWithError } from '../../queryUtils';
 import { extractUnwrapLabelKeysFromDataFrame } from '../../responseUtils';
 import { lokiQueryModeller } from '../LokiQueryModeller';
 import { LokiVisualQuery } from '../types';
@@ -56,7 +57,7 @@ async function loadUnwrapOptions(
 ): Promise<Array<SelectableValue<string>>> {
   const queryExpr = lokiQueryModeller.renderQuery(query);
   const logExpr = getLogQueryFromMetricsQuery(queryExpr);
-  if (!isValidQuery(logExpr)) {
+  if (isQueryWithError(datasource.interpolateString(logExpr, placeHolderScopedVars))) {
     return [];
   }