Loki: Create Variable Query Editor for Loki. (#54102)

* feat(loki-query-editor): create base editor component

* feat(loki-query-editor): update editor to use loki query type

* feat(loki-query-editor): add custom variable support to datasource

* feat(loki-query-editor): prevent errors when no label is present

* Add unit test for LokiMetricFindQuery

* Update datasource test

* Add component test

* Add variable query migration support

* feat(loki-query-editor): add migration support for older-format variables

* Fix enum capitalization for consistency

* Move attribute to the top of the class

* Remove unnecessary from()

* Update capitalization of new enum

* Fix enum capitalization in component

* feat(loki-query-editor): replace unnecessary class with class method

public/app/plugins/datasource/loki/components/VariableQueryEditor.test.tsx

@@ -0,0 +1,74 @@
+import { render, screen, waitFor } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import React from 'react';
+import { selectOptionInTest } from 'test/helpers/selectOptionInTest';
+
+import { TemplateSrv } from '@grafana/runtime';
+
+import { createLokiDatasource } from '../mocks';
+import { LokiVariableQueryType } from '../types';
+
+import { LokiVariableQueryEditor, Props } from './VariableQueryEditor';
+
+const props: Props = {
+  datasource: createLokiDatasource({} as unknown as TemplateSrv),
+  query: {
+    refId: 'test',
+    type: LokiVariableQueryType.LabelNames,
+  },
+  onRunQuery: () => {},
+  onChange: () => {},
+};
+
+describe('LokiVariableQueryEditor', () => {
+  test('Allows to create a Label names variable', async () => {
+    const onChange = jest.fn();
+
+    render(<LokiVariableQueryEditor {...props} onChange={onChange} />);
+
+    expect(onChange).not.toHaveBeenCalled();
+
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label names');
+
+    expect(onChange).toHaveBeenCalledWith({
+      type: LokiVariableQueryType.LabelNames,
+      label: '',
+      stream: '',
+      refId: 'LokiVariableQueryEditor-VariableQuery',
+    });
+  });
+
+  test('Allows to create a Label values variable', async () => {
+    const onChange = jest.fn();
+
+    render(<LokiVariableQueryEditor {...props} onChange={onChange} />);
+
+    expect(onChange).not.toHaveBeenCalled();
+
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label values');
+    await userEvent.type(screen.getByLabelText('Label'), 'label');
+    await userEvent.type(screen.getByLabelText('Stream selector'), 'stream');
+
+    await waitFor(() => expect(screen.getByDisplayValue('stream')).toBeInTheDocument());
+
+    await userEvent.click(document.body);
+
+    expect(onChange).toHaveBeenCalledWith({
+      type: LokiVariableQueryType.LabelValues,
+      label: 'label',
+      stream: 'stream',
+      refId: 'LokiVariableQueryEditor-VariableQuery',
+    });
+  });
+
+  test('Migrates legacy string queries to LokiVariableQuery instances', async () => {
+    const query = 'label_values(log stream selector, label)';
+
+    // @ts-expect-error
+    render(<LokiVariableQueryEditor {...props} onChange={() => {}} query={query} />);
+
+    await waitFor(() => expect(screen.getByText('Label values')).toBeInTheDocument());
+    await waitFor(() => expect(screen.getByDisplayValue('label')).toBeInTheDocument());
+    await waitFor(() => expect(screen.getByDisplayValue('log stream selector')).toBeInTheDocument());
+  });
+});

public/app/plugins/datasource/loki/components/VariableQueryEditor.tsx

@@ -0,0 +1,89 @@
+import React, { FC, FormEvent, useState, useEffect } from 'react';
+
+import { QueryEditorProps, SelectableValue } from '@grafana/data';
+import { InlineField, InlineFieldRow, Input, Select } from '@grafana/ui';
+
+import { LokiDatasource } from '../datasource';
+import { migrateVariableQuery } from '../migrations/variableQueryMigrations';
+import { LokiOptions, LokiQuery, LokiVariableQuery, LokiVariableQueryType as QueryType } from '../types';
+
+const variableOptions = [
+  { label: 'Label names', value: QueryType.LabelNames },
+  { label: 'Label values', value: QueryType.LabelValues },
+];
+
+export type Props = QueryEditorProps<LokiDatasource, LokiQuery, LokiOptions, LokiVariableQuery>;
+
+export const LokiVariableQueryEditor: FC<Props> = ({ onChange, query }) => {
+  const [type, setType] = useState<number | undefined>(undefined);
+  const [label, setLabel] = useState('');
+  const [stream, setStream] = useState('');
+
+  useEffect(() => {
+    if (!query || typeof query !== 'string') {
+      return;
+    }
+
+    const variableQuery = migrateVariableQuery(query);
+    setType(variableQuery.type);
+    setLabel(variableQuery.label || '');
+    setStream(variableQuery.stream || '');
+  }, [query]);
+
+  const onQueryTypeChange = (newType: SelectableValue<QueryType>) => {
+    setType(newType.value);
+    if (newType.value !== undefined) {
+      onChange({
+        type: newType.value,
+        label,
+        stream,
+        refId: 'LokiVariableQueryEditor-VariableQuery',
+      });
+    }
+  };
+
+  const onLabelChange = (e: FormEvent<HTMLInputElement>) => {
+    setLabel(e.currentTarget.value);
+  };
+
+  const onStreamChange = (e: FormEvent<HTMLInputElement>) => {
+    setStream(e.currentTarget.value);
+  };
+
+  const handleBlur = () => {
+    if (type !== undefined) {
+      onChange({ type, label, stream, refId: 'LokiVariableQueryEditor-VariableQuery' });
+    }
+  };
+
+  return (
+    <InlineFieldRow>
+      <InlineField label="Query type" labelWidth={20}>
+        <Select
+          aria-label="Query type"
+          onChange={onQueryTypeChange}
+          onBlur={handleBlur}
+          value={type}
+          options={variableOptions}
+          width={16}
+        />
+      </InlineField>
+      {type === QueryType.LabelValues && (
+        <>
+          <InlineField label="Label" labelWidth={20}>
+            <Input type="text" aria-label="Label" value={label} onChange={onLabelChange} onBlur={handleBlur} />
+          </InlineField>
+          <InlineField label="Stream selector" labelWidth={20}>
+            <Input
+              type="text"
+              aria-label="Stream selector"
+              value={stream}
+              onChange={onStreamChange}
+              onBlur={handleBlur}
+            />
+          </InlineField>
+        </>
+      )}
+    </InlineFieldRow>
+  );
+};

public/app/plugins/datasource/loki/components/__snapshots__/LokiExploreQueryEditor.test.tsx.snap

@@ -321,6 +321,11 @@ exports[`LokiExploreQueryEditor should render component 1`] = `
       },
       "type": "",
       "uid": "",
+      "variables": LokiVariableSupport {
+        "datasource": [Circular],
+        "editor": [Function],
+        "query": [Function],
+      },
     }
   }
   history={Array []}

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

@@ -24,6 +24,7 @@ import { CustomVariableModel } from '../../../features/variables/types';
 import { LokiDatasource } from './datasource';
 import { createMetadataRequest, createLokiDatasource } from './mocks';
 import { LokiOptions, LokiQuery, LokiQueryType } from './types';
+import { LokiVariableSupport } from './variables';
 
 const templateSrvStub = {
   getAdhocFilters: jest.fn(() => [] as unknown[]),
@@ -823,6 +824,14 @@ describe('applyTemplateVariables', () => {
   });
 });
 
+describe('Variable support', () => {
+  it('has Loki variable support', () => {
+    const ds = createLokiDatasource(templateSrvStub);
+
+    expect(ds.variables).toBeInstanceOf(LokiVariableSupport);
+  });
+});
+
 function assertAdHocFilters(query: string, expectedResults: string, ds: LokiDatasource) {
   const lokiQuery: LokiQuery = { refId: 'A', expr: query };
   const result = ds.addAdHocFilters(lokiQuery.expr);

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

@@ -49,6 +49,7 @@ import { LokiAnnotationsQueryEditor } from './components/AnnotationsQueryEditor'
 import LanguageProvider from './language_provider';
 import { escapeLabelValueInSelector } from './language_utils';
 import { LiveStreams, LokiLiveTarget } from './live_streams';
+import { labelNamesRegex, labelValuesRegex } from './migrations/variableQueryMigrations';
 import {
   addLabelFormatToQuery,
   addLabelToQuery,
@@ -61,6 +62,7 @@ import { getNormalizedLokiQuery, isLogsQuery, isValidQuery } from './query_utils
 import { sortDataFrameByTime } from './sortDataFrame';
 import { doLokiChannelStream } from './streaming';
 import { LokiOptions, LokiQuery, LokiQueryDirection, LokiQueryType } from './types';
+import { LokiVariableSupport } from './variables';
 
 export type RangeQueryOptions = DataQueryRequest<LokiQuery> | AnnotationQueryRequest<LokiQuery>;
 export const DEFAULT_MAX_LINES = 1000;
@@ -107,6 +109,7 @@ export class LokiDatasource
     this.annotations = {
       QueryEditor: LokiAnnotationsQueryEditor,
     };
+    this.variables = new LokiVariableSupport(this);
   }
 
   getLogsVolumeDataProvider(request: DataQueryRequest<LokiQuery>): Observable<DataQueryResponse> | undefined {
@@ -311,9 +314,6 @@ export class LokiDatasource
   }
 
   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();

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

@@ -0,0 +1,48 @@
+import { LokiVariableQuery, LokiVariableQueryType } from '../types';
+
+import { migrateVariableQuery } from './variableQueryMigrations';
+
+describe('Loki migrateVariableQuery()', () => {
+  it('Does not migrate LokiVariableQuery instances', () => {
+    const query: LokiVariableQuery = {
+      refId: 'test',
+      type: LokiVariableQueryType.LabelValues,
+      label: 'label',
+      stream: 'stream',
+    };
+
+    expect(migrateVariableQuery(query)).toBe(query);
+    expect(migrateVariableQuery(query)).toStrictEqual(query);
+  });
+
+  it('Migrates label_names() queries', () => {
+    const query = 'label_names()';
+
+    expect(migrateVariableQuery(query)).toStrictEqual({
+      refId: 'LokiVariableQueryEditor-VariableQuery',
+      type: LokiVariableQueryType.LabelNames,
+    });
+  });
+
+  it('Migrates label_values(label) queries', () => {
+    const query = 'label_values(label)';
+
+    expect(migrateVariableQuery(query)).toStrictEqual({
+      refId: 'LokiVariableQueryEditor-VariableQuery',
+      type: LokiVariableQueryType.LabelValues,
+      label: 'label',
+      stream: undefined,
+    });
+  });
+
+  it('Migrates label_values(log stream selector, label) queries', () => {
+    const query = 'label_values(log stream selector, label)';
+
+    expect(migrateVariableQuery(query)).toStrictEqual({
+      refId: 'LokiVariableQueryEditor-VariableQuery',
+      type: LokiVariableQueryType.LabelValues,
+      label: 'label',
+      stream: 'log stream selector',
+    });
+  });
+});

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

@@ -0,0 +1,36 @@
+import { LokiVariableQuery, LokiVariableQueryType } from '../types';
+
+export const labelNamesRegex = /^label_names\(\)\s*$/;
+export const labelValuesRegex = /^label_values\((?:(.+),\s*)?([a-zA-Z_][a-zA-Z0-9_]*)\)\s*$/;
+
+export function migrateVariableQuery(rawQuery: string | LokiVariableQuery): LokiVariableQuery {
+  // If not string, we assume LokiVariableQuery
+  if (typeof rawQuery !== 'string') {
+    return rawQuery;
+  }
+
+  const queryBase = {
+    refId: 'LokiVariableQueryEditor-VariableQuery',
+    type: LokiVariableQueryType.LabelNames,
+  };
+
+  const labelNames = rawQuery.match(labelNamesRegex);
+  if (labelNames) {
+    return {
+      ...queryBase,
+      type: LokiVariableQueryType.LabelNames,
+    };
+  }
+
+  const labelValues = rawQuery.match(labelValuesRegex);
+  if (labelValues) {
+    return {
+      ...queryBase,
+      type: LokiVariableQueryType.LabelValues,
+      label: labelValues[2] ? labelValues[2] : labelValues[1],
+      stream: labelValues[2] ? labelValues[1] : undefined,
+    };
+  }
+
+  return queryBase;
+}

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

@@ -142,3 +142,14 @@ export interface TransformerOptions {
   scopedVars: ScopedVars;
   meta?: QueryResultMeta;
 }
+
+export enum LokiVariableQueryType {
+  LabelNames,
+  LabelValues,
+}
+
+export interface LokiVariableQuery extends DataQuery {
+  type: LokiVariableQueryType;
+  label?: string;
+  stream?: string;
+}

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

@@ -0,0 +1,53 @@
+import { TemplateSrv } from 'app/features/templating/template_srv';
+
+import { createLokiDatasource, createMetadataRequest } from './mocks';
+import { LokiVariableQueryType } from './types';
+import { LokiVariableSupport } from './variables';
+
+describe('LokiVariableSupport', () => {
+  let lokiVariableSupport: LokiVariableSupport;
+
+  beforeEach(() => {
+    const datasource = createLokiDatasource({} as unknown as TemplateSrv);
+    jest
+      .spyOn(datasource, 'metadataRequest')
+      .mockImplementation(
+        createMetadataRequest(
+          { label1: ['value1', 'value2'], label2: ['value3', 'value4'] },
+          { '{label1="value1", label2="value2"}': [{ label5: 'value5' }] }
+        )
+      );
+
+    lokiVariableSupport = new LokiVariableSupport(datasource);
+  });
+
+  it('should return label names for Loki', async () => {
+    // label_names()
+    const response = await lokiVariableSupport.execute({ refId: 'test', type: LokiVariableQueryType.LabelNames });
+
+    expect(response).toEqual([{ text: 'label1' }, { text: 'label2' }]);
+  });
+
+  it('should return label values for Loki when no matcher', async () => {
+    // label_values(label1)
+    const response = await lokiVariableSupport.execute({
+      refId: 'test',
+      type: LokiVariableQueryType.LabelValues,
+      label: 'label1',
+    });
+
+    expect(response).toEqual([{ text: 'value1' }, { text: 'value2' }]);
+  });
+
+  it('should return label values for Loki with matcher', async () => {
+    // label_values({label1="value1", label2="value2"},label5)
+    const response = await lokiVariableSupport.execute({
+      refId: 'test',
+      type: LokiVariableQueryType.LabelValues,
+      stream: '{label1="value1", label2="value2"}',
+      label: 'label5',
+    });
+
+    expect(response).toEqual([{ text: 'value5' }]);
+  });
+});

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

@@ -0,0 +1,40 @@
+import { from, Observable } from 'rxjs';
+import { map } from 'rxjs/operators';
+
+import { CustomVariableSupport, DataQueryRequest, DataQueryResponse } from '@grafana/data';
+
+import { LokiVariableQueryEditor } from './components/VariableQueryEditor';
+import { LokiDatasource } from './datasource';
+import { LokiVariableQuery, LokiVariableQueryType } from './types';
+
+export class LokiVariableSupport extends CustomVariableSupport<LokiDatasource, LokiVariableQuery> {
+  editor = LokiVariableQueryEditor;
+
+  constructor(private datasource: LokiDatasource) {
+    super();
+    this.query = this.query.bind(this);
+  }
+
+  async execute(query: LokiVariableQuery) {
+    if (query.type === LokiVariableQueryType.LabelNames) {
+      return this.datasource.labelNamesQuery();
+    }
+
+    if (!query.label) {
+      return [];
+    }
+
+    // If we have query expr, use /series endpoint
+    if (query.stream) {
+      return this.datasource.labelValuesSeriesQuery(query.stream, query.label);
+    }
+
+    return this.datasource.labelValuesQuery(query.label);
+  }
+
+  query(request: DataQueryRequest<LokiVariableQuery>): Observable<DataQueryResponse> {
+    const result = this.execute(request.targets[0]);
+
+    return from(result).pipe(map((data) => ({ data })));
+  }
+}