Loki: Support for template variable queries (#20697)

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

@@ -5,6 +5,7 @@ import { AnnotationQueryRequest, DataSourceApi, DataFrame, dateTime, TimeRange }
 import { BackendSrv } from 'app/core/services/backend_srv';
 import { TemplateSrv } from 'app/features/templating/template_srv';
 import { CustomVariable } from 'app/features/templating/custom_variable';
+import { makeMockLokiDatasource } from './mocks';
 import { ExploreMode } from 'app/types';
 import { of } from 'rxjs';
 import omit from 'lodash/omit';
@@ -357,6 +358,57 @@ describe('LokiDatasource', () => {
       expect(res[1].tags).toEqual(['value2']);
     });
   });
+
+  describe('metricFindQuery', () => {
+    const ds = new LokiDatasource(instanceSettings, backendSrv, templateSrvMock);
+    const mocks = makeMetadataAndVersionsMocks();
+
+    mocks.forEach((mock, index) => {
+      it(`should return label names for Loki v${index}`, async () => {
+        ds.getVersion = mock.getVersion;
+        ds.metadataRequest = mock.metadataRequest;
+        const query = 'label_names()';
+        const res = await ds.metricFindQuery(query);
+        expect(res[0].text).toEqual('label1');
+        expect(res[1].text).toEqual('label2');
+        expect(res.length).toBe(2);
+      });
+    });
+
+    mocks.forEach((mock, index) => {
+      it(`should return label names for Loki v${index}`, async () => {
+        ds.getVersion = mock.getVersion;
+        ds.metadataRequest = mock.metadataRequest;
+        const query = 'label_names()';
+        const res = await ds.metricFindQuery(query);
+        expect(res[0].text).toEqual('label1');
+        expect(res[1].text).toEqual('label2');
+        expect(res.length).toBe(2);
+      });
+    });
+
+    mocks.forEach((mock, index) => {
+      it(`should return label values for Loki v${index}`, async () => {
+        ds.getVersion = mock.getVersion;
+        ds.metadataRequest = mock.metadataRequest;
+        const query = 'label_values(label1)';
+        const res = await ds.metricFindQuery(query);
+        expect(res[0].text).toEqual('value1');
+        expect(res[1].text).toEqual('value2');
+        expect(res.length).toBe(2);
+      });
+    });
+
+    mocks.forEach((mock, index) => {
+      it(`should return empty array when incorrect query for Loki v${index}`, async () => {
+        ds.getVersion = mock.getVersion;
+        ds.metadataRequest = mock.metadataRequest;
+        const query = 'incorrect_query';
+        const res = await ds.metricFindQuery(query);
+        expect(res.length).toBe(0);
+      });
+    });
+  });
 });
 
 type LimitTestArgs = {
@@ -418,3 +470,13 @@ function makeAnnotationQueryRequest(): AnnotationQueryRequest<LokiQuery> {
     rangeRaw: timeRange,
   };
 }
+
+function makeMetadataAndVersionsMocks() {
+  const mocks = [];
+  for (let i = 0; i <= 1; i++) {
+    const mock: LokiDatasource = makeMockLokiDatasource({ label1: ['value1', 'value2'], label2: ['value3', 'value4'] });
+    mock.getVersion = jest.fn().mockReturnValue(`v${i}`);
+    mocks.push(mock);
+  }
+  return mocks;
+}

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

@@ -52,9 +52,12 @@ import LanguageProvider from './language_provider';
 
 export type RangeQueryOptions = Pick<DataQueryRequest<LokiQuery>, 'range' | 'intervalMs' | 'maxDataPoints' | 'reverse'>;
 export const DEFAULT_MAX_LINES = 1000;
-const LEGACY_QUERY_ENDPOINT = '/api/prom/query';
-const RANGE_QUERY_ENDPOINT = '/loki/api/v1/query_range';
-const INSTANT_QUERY_ENDPOINT = '/loki/api/v1/query';
+export const LEGACY_LOKI_ENDPOINT = '/api/prom';
+export const LOKI_ENDPOINT = '/loki/api/v1';
+
+const LEGACY_QUERY_ENDPOINT = `${LEGACY_LOKI_ENDPOINT}/query`;
+const RANGE_QUERY_ENDPOINT = `${LOKI_ENDPOINT}/query_range`;
+const INSTANT_QUERY_ENDPOINT = `${LOKI_ENDPOINT}/query`;
 
 const DEFAULT_QUERY_PARAMS: Partial<LokiLegacyQueryRequest> = {
   direction: 'BACKWARD',
@@ -76,9 +79,9 @@ interface LokiContextQueryOptions {
 
 export class LokiDatasource extends DataSourceApi<LokiQuery, LokiOptions> {
   private streams = new LiveStreams();
+  private version: string;
   languageProvider: LanguageProvider;
   maxLines: number;
-  version: string;
 
   /** @ngInject */
   constructor(
@@ -374,6 +377,46 @@ export class LokiDatasource extends DataSourceApi<LokiQuery, LokiOptions> {
     };
   }
 
+  async metricFindQuery(query: string) {
+    if (!query) {
+      return Promise.resolve([]);
+    }
+    const interpolated = this.templateSrv.replace(query, {}, this.interpolateQueryExpr);
+    return await this.processMetricFindQuery(interpolated);
+  }
+
+  async processMetricFindQuery(query: string) {
+    const labelNamesRegex = /^label_names\(\)\s*$/;
+    const labelValuesRegex = /^label_values\((?:(.+),\s*)?([a-zA-Z_][a-zA-Z0-9_]*)\)\s*$/;
+
+    const labelNames = query.match(labelNamesRegex);
+    if (labelNames) {
+      return await this.labelNamesQuery();
+    }
+
+    const labelValues = query.match(labelValuesRegex);
+    if (labelValues) {
+      return await this.labelValuesQuery(labelValues[2]);
+    }
+
+    return Promise.resolve([]);
+  }
+
+  async labelNamesQuery() {
+    const url = (await this.getVersion()) === 'v0' ? `${LEGACY_LOKI_ENDPOINT}/label` : `${LOKI_ENDPOINT}/label`;
+    const result = await this.metadataRequest(url);
+    return result.data.data.map((value: string) => ({ text: value }));
+  }
+
+  async labelValuesQuery(label: string) {
+    const url =
+      (await this.getVersion()) === 'v0'
+        ? `${LEGACY_LOKI_ENDPOINT}/label/${label}/values`
+        : `${LOKI_ENDPOINT}/label/${label}/values`;
+    const result = await this.metadataRequest(url);
+    return result.data.data.map((value: string) => ({ text: value }));
+  }
+
   interpolateQueryExpr(value: any, variable: any) {
     // if no multi or include all do not regexEscape
     if (!variable.multi && !variable.includeAll) {

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

@@ -1,17 +1,23 @@
-import LokiDatasource from './datasource';
+import { LokiDatasource, LOKI_ENDPOINT, LEGACY_LOKI_ENDPOINT } from './datasource';
 import { DataSourceSettings } from '@grafana/data';
 import { LokiOptions } from './types';
 import { createDatasourceSettings } from '../../../features/datasources/mocks';
 
 export function makeMockLokiDatasource(labelsAndValues: { [label: string]: string[] }): LokiDatasource {
+  const legacyLokiLabelsAndValuesEndpointRegex = /^\/api\/prom\/label\/(\w*)\/values/;
+  const lokiLabelsAndValuesEndpointRegex = /^\/loki\/api\/v1\/label\/(\w*)\/values/;
+
+  const legacyLokiLabelsEndpoint = `${LEGACY_LOKI_ENDPOINT}/label`;
+  const lokiLabelsEndpoint = `${LOKI_ENDPOINT}/label`;
+
   const labels = Object.keys(labelsAndValues);
   return {
     metadataRequest: (url: string) => {
       let responseData;
-      if (url === '/api/prom/label') {
+      if (url === legacyLokiLabelsEndpoint || url === lokiLabelsEndpoint) {
         responseData = labels;
       } else {
-        const match = url.match(/^\/api\/prom\/label\/(\w*)\/values/);
+        const match = url.match(legacyLokiLabelsAndValuesEndpointRegex) || url.match(lokiLabelsAndValuesEndpointRegex);
         if (match) {
           responseData = labelsAndValues[match[1]];
         }