Dashboard: Add support for Tempo query variables (#72745)

.betterer.results

@@ -4298,22 +4298,24 @@ exports[`better eslint`] = {
     ],
     "public/app/plugins/datasource/tempo/datasource.ts:5381": [
       [0, 0, 0, "Do not use any type assertions.", "0"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "1"],
+      [0, 0, 0, "Do not use any type assertions.", "1"],
       [0, 0, 0, "Do not use any type assertions.", "2"],
-      [0, 0, 0, "Do not use any type assertions.", "3"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "4"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "5"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "3"],
+      [0, 0, 0, "Do not use any type assertions.", "4"],
+      [0, 0, 0, "Do not use any type assertions.", "5"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "6"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "7"],
-      [0, 0, 0, "Do not use any type assertions.", "8"],
-      [0, 0, 0, "Do not use any type assertions.", "9"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "8"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "9"],
       [0, 0, 0, "Do not use any type assertions.", "10"],
       [0, 0, 0, "Do not use any type assertions.", "11"],
       [0, 0, 0, "Do not use any type assertions.", "12"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "13"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "14"],
+      [0, 0, 0, "Do not use any type assertions.", "13"],
+      [0, 0, 0, "Do not use any type assertions.", "14"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "15"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "16"]
+      [0, 0, 0, "Unexpected any. Specify a different type.", "16"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "17"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "18"]
     ],
     "public/app/plugins/datasource/tempo/language_provider.ts:5381": [
       [0, 0, 0, "Unexpected any. Specify a different type.", "0"]

public/app/plugins/datasource/tempo/VariableQueryEditor.test.tsx

@@ -0,0 +1,68 @@
+import { render, screen } 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 {
+  TempoVariableQuery,
+  TempoVariableQueryEditor,
+  TempoVariableQueryEditorProps,
+  TempoVariableQueryType,
+} from './VariableQueryEditor';
+import { createTempoDatasource } from './mocks';
+
+const refId = 'TempoDatasourceVariableQueryEditor-VariableQuery';
+
+describe('TempoVariableQueryEditor', () => {
+  let props: TempoVariableQueryEditorProps;
+  let onChange: (value: TempoVariableQuery) => void;
+
+  beforeEach(() => {
+    props = {
+      datasource: createTempoDatasource({} as unknown as TemplateSrv),
+      query: { type: 0, refId: 'test' },
+      onChange: (_: TempoVariableQuery) => {},
+    };
+
+    onChange = jest.fn();
+  });
+
+  test('Allows to create a Label names variable', async () => {
+    expect(onChange).not.toHaveBeenCalled();
+    render(<TempoVariableQueryEditor {...props} onChange={onChange} />);
+
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label names');
+    await userEvent.click(document.body);
+
+    expect(onChange).toHaveBeenCalledWith({
+      type: TempoVariableQueryType.LabelNames,
+      label: '',
+      refId,
+    });
+  });
+
+  test('Allows to create a Label values variable', async () => {
+    jest.spyOn(props.datasource, 'labelNamesQuery').mockResolvedValue([
+      {
+        text: 'moon',
+      },
+      {
+        text: 'luna',
+      },
+    ]);
+    expect(onChange).not.toHaveBeenCalled();
+    render(<TempoVariableQueryEditor {...props} onChange={onChange} />);
+
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label values');
+    await selectOptionInTest(screen.getByLabelText('Label'), 'luna');
+    await userEvent.click(document.body);
+
+    expect(onChange).toHaveBeenCalledWith({
+      type: TempoVariableQueryType.LabelValues,
+      label: 'luna',
+      refId,
+    });
+  });
+});

public/app/plugins/datasource/tempo/VariableQueryEditor.tsx

@@ -0,0 +1,106 @@
+import React, { useEffect, useState } from 'react';
+
+import { DataQuery, SelectableValue } from '@grafana/data';
+import { InlineField, InlineFieldRow, Select } from '@grafana/ui';
+
+import { TempoDatasource } from './datasource';
+
+export enum TempoVariableQueryType {
+  LabelNames,
+  LabelValues,
+}
+
+export interface TempoVariableQuery extends DataQuery {
+  type: TempoVariableQueryType;
+  label?: string;
+  stream?: string;
+}
+
+const variableOptions = [
+  { label: 'Label names', value: TempoVariableQueryType.LabelNames },
+  { label: 'Label values', value: TempoVariableQueryType.LabelValues },
+];
+
+const refId = 'TempoDatasourceVariableQueryEditor-VariableQuery';
+
+export type TempoVariableQueryEditorProps = {
+  onChange: (value: TempoVariableQuery) => void;
+  query: TempoVariableQuery;
+  datasource: TempoDatasource;
+};
+
+export const TempoVariableQueryEditor = ({ onChange, query, datasource }: TempoVariableQueryEditorProps) => {
+  const [label, setLabel] = useState(query.label || '');
+  const [type, setType] = useState<number | undefined>(query.type);
+  const [labelOptions, setLabelOptions] = useState<Array<SelectableValue<string>>>([]);
+
+  useEffect(() => {
+    if (type === TempoVariableQueryType.LabelValues) {
+      datasource.labelNamesQuery().then((labelNames: Array<{ text: string }>) => {
+        setLabelOptions(labelNames.map(({ text }) => ({ label: text, value: text })));
+      });
+    }
+  }, [datasource, query, type]);
+
+  const onQueryTypeChange = (newType: SelectableValue<TempoVariableQueryType>) => {
+    setType(newType.value);
+    if (newType.value !== undefined) {
+      onChange({
+        type: newType.value,
+        label,
+        refId,
+      });
+    }
+  };
+
+  const onLabelChange = (newLabel: SelectableValue<string>) => {
+    const newLabelValue = newLabel.value || '';
+    setLabel(newLabelValue);
+    if (type !== undefined) {
+      onChange({
+        type,
+        label: newLabelValue,
+        refId,
+      });
+    }
+  };
+
+  const handleBlur = () => {
+    if (type !== undefined) {
+      onChange({ type, label, refId });
+    }
+  };
+
+  return (
+    <>
+      <InlineFieldRow>
+        <InlineField label="Query type" labelWidth={20}>
+          <Select
+            aria-label="Query type"
+            onChange={onQueryTypeChange}
+            onBlur={handleBlur}
+            value={type}
+            options={variableOptions}
+            width={32}
+          />
+        </InlineField>
+      </InlineFieldRow>
+
+      {type === TempoVariableQueryType.LabelValues && (
+        <InlineFieldRow>
+          <InlineField label="Label" labelWidth={20}>
+            <Select
+              aria-label="Label"
+              onChange={onLabelChange}
+              onBlur={handleBlur}
+              value={{ label, value: label }}
+              options={labelOptions}
+              width={32}
+              allowCustomValue
+            />
+          </InlineField>
+        </InlineFieldRow>
+      )}
+    </>
+  );
+};

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

@@ -18,6 +18,7 @@ import { BackendDataSourceResponse, FetchResponse, setBackendSrv, setDataSourceS
 import { BarGaugeDisplayMode, TableCellDisplayMode } from '@grafana/schema';
 import { TemplateSrv } from 'app/features/templating/template_srv';
 
+import { TempoVariableQueryType } from './VariableQueryEditor';
 import { TraceqlSearchScope } from './dataquery.gen';
 import {
   DEFAULT_LIMIT,
@@ -32,6 +33,7 @@ import {
 } from './datasource';
 import mockJson from './mockJsonResponse.json';
 import mockServiceGraph from './mockServiceGraph.json';
+import { createMetadataRequest, createTempoDatasource } from './mocks';
 import { TempoJsonData, TempoQuery } from './types';
 
 let mockObservable: () => Observable<any>;
@@ -769,6 +771,105 @@ describe('Tempo service graph view', () => {
   });
 });
 
+describe('label names - v2 tags', () => {
+  let datasource: TempoDatasource;
+
+  beforeEach(() => {
+    datasource = createTempoDatasource();
+    jest.spyOn(datasource, 'metadataRequest').mockImplementation(
+      createMetadataRequest({
+        data: {
+          scopes: [{ name: 'span', tags: ['label1', 'label2'] }],
+        },
+      })
+    );
+  });
+
+  it('get label names', async () => {
+    // label_names()
+    const response = await datasource.executeVariableQuery({ refId: 'test', type: TempoVariableQueryType.LabelNames });
+
+    expect(response).toEqual([{ text: 'label1' }, { text: 'label2' }]);
+  });
+});
+
+describe('label names - v1 tags', () => {
+  let datasource: TempoDatasource;
+
+  beforeEach(() => {
+    datasource = createTempoDatasource();
+    jest
+      .spyOn(datasource, 'metadataRequest')
+      .mockImplementationOnce(() => {
+        throw Error;
+      })
+      .mockImplementation(
+        createMetadataRequest({
+          data: {
+            tagNames: ['label1', 'label2'],
+          },
+        })
+      );
+  });
+
+  it('get label names', async () => {
+    // label_names()
+    const response = await datasource.executeVariableQuery({ refId: 'test', type: TempoVariableQueryType.LabelNames });
+    expect(response).toEqual([{ text: 'label1' }, { text: 'label2' }, { text: 'status.code' }]);
+  });
+});
+
+describe('label values', () => {
+  let datasource: TempoDatasource;
+
+  beforeEach(() => {
+    datasource = createTempoDatasource();
+    jest.spyOn(datasource, 'metadataRequest').mockImplementation(
+      createMetadataRequest({
+        data: {
+          tagValues: [
+            {
+              type: 'value1',
+              value: 'value1',
+              label: 'value1',
+            },
+            {
+              type: 'value2',
+              value: 'value2',
+              label: 'value2',
+            },
+          ],
+        },
+      })
+    );
+  });
+
+  it('get label values for given label', async () => {
+    // label_values("label")
+    const response = await datasource.executeVariableQuery({
+      refId: 'test',
+      type: TempoVariableQueryType.LabelValues,
+      label: 'label',
+    });
+
+    expect(response).toEqual([
+      { text: { type: 'value1', value: 'value1', label: 'value1' } },
+      { text: { type: 'value2', value: 'value2', label: 'value2' } },
+    ]);
+  });
+
+  it('do not raise error when label is not set', async () => {
+    // label_values()
+    const response = await datasource.executeVariableQuery({
+      refId: 'test',
+      type: TempoVariableQueryType.LabelValues,
+      label: undefined,
+    });
+
+    expect(response).toEqual([]);
+  });
+});
+
 const backendSrvWithPrometheus = {
   async get(uid: string) {
     if (uid === 'prom') {

public/app/plugins/datasource/tempo/datasource.ts

@@ -39,6 +39,7 @@ import { PrometheusDatasource } from '../prometheus/datasource';
 import { PromQuery } from '../prometheus/types';
 
 import { generateQueryFromFilters } from './SearchTraceQLEditor/utils';
+import { TempoVariableQuery, TempoVariableQueryType } from './VariableQueryEditor';
 import { TraceqlFilter, TraceqlSearchScope } from './dataquery.gen';
 import {
   failedMetric,
@@ -63,6 +64,7 @@ import {
 import { doTempoChannelStream } from './streaming';
 import { SearchQueryParams, TempoQuery, TempoJsonData } from './types';
 import { getErrorMessage } from './utils';
+import { TempoVariableSupport } from './variables';
 
 export const DEFAULT_LIMIT = 20;
 
@@ -115,6 +117,66 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
         ],
       };
     }
+
+    this.variables = new TempoVariableSupport(this);
+  }
+
+  async executeVariableQuery(query: TempoVariableQuery) {
+    // Avoid failing if the user did not select the query type (label names, label values, etc.)
+    if (query.type === undefined) {
+      return new Promise<Array<{ text: string }>>(() => []);
+    }
+
+    switch (query.type) {
+      case TempoVariableQueryType.LabelNames: {
+        return await this.labelNamesQuery();
+      }
+      case TempoVariableQueryType.LabelValues: {
+        return this.labelValuesQuery(query.label);
+      }
+      default: {
+        throw Error('Invalid query type', query.type);
+      }
+    }
+  }
+
+  async labelNamesQuery(): Promise<Array<{ text: string }>> {
+    await this.languageProvider.fetchTags();
+    const tags = this.languageProvider.getAutocompleteTags();
+    return tags.filter((tag) => tag !== undefined).map((tag) => ({ text: tag })) as Array<{ text: string }>;
+  }
+
+  async labelValuesQuery(labelName?: string): Promise<Array<{ text: string }>> {
+    if (!labelName) {
+      return [];
+    }
+
+    let options;
+    try {
+      // Retrieve the scope of the tag
+      // Example: given `http.status_code`, we want scope `span`
+      // Note that we ignore possible name clashes, e.g., `http.status_code` in both `span` and `resource`
+      const scope: string | undefined = (this.languageProvider.tagsV2 || [])
+        // flatten the Scope objects
+        .flatMap((tagV2) => tagV2.tags.map((tag) => ({ scope: tagV2.name, name: tag })))
+        // find associated scope
+        .find((tag) => tag.name === labelName)?.scope;
+      if (!scope) {
+        throw Error(`Scope for tag ${labelName} not found`);
+      }
+
+      // For V2, we need to send scope and tag name, e.g. `span.http.status_code`,
+      // unless the tag has intrinsic scope
+      const scopeAndTag = scope === 'intrinsic' ? labelName : `${scope}.${labelName}`;
+      options = await this.languageProvider.getOptionsV2(scopeAndTag);
+    } catch {
+      // For V1, the tag name (e.g. `http.status_code`) is enough
+      options = await this.languageProvider.getOptionsV1(labelName);
+    }
+
+    return options.filter((option) => option.value !== undefined).map((option) => ({ text: option.value })) as Array<{
+      text: string;
+    }>;
   }
 
   query(options: DataQueryRequest<TempoQuery>): Observable<DataQueryResponse> {

public/app/plugins/datasource/tempo/language_provider.ts

@@ -129,7 +129,7 @@ export default class TempoLanguageProvider extends LanguageProvider {
     return options;
   }
 
-  async getOptionsV2(tag: string, query: string): Promise<Array<SelectableValue<string>>> {
+  async getOptionsV2(tag: string, query?: string): Promise<Array<SelectableValue<string>>> {
     const response = await this.request(`/api/v2/search/tag/${tag}/values`, query ? { q: query } : {});
     let options: Array<SelectableValue<string>> = [];
     if (response && response.tagValues) {

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

@@ -0,0 +1,68 @@
+import { DataSourceInstanceSettings, PluginType, toUtc } from '@grafana/data';
+import { TemplateSrv } from '@grafana/runtime';
+
+import { TempoDatasource } from './datasource';
+import { TempoJsonData } from './types';
+
+const rawRange = {
+  from: toUtc('2018-04-25 10:00'),
+  to: toUtc('2018-04-25 11:00'),
+};
+
+const defaultTimeSrvMock = {
+  timeRange: jest.fn().mockReturnValue({
+    from: rawRange.from,
+    to: rawRange.to,
+    raw: rawRange,
+  }),
+};
+
+const defaultTemplateSrvMock = {
+  replace: (input: string) => input,
+};
+
+export function createTempoDatasource(
+  templateSrvMock: Partial<TemplateSrv> = defaultTemplateSrvMock,
+  settings: Partial<DataSourceInstanceSettings<TempoJsonData>> = {},
+  timeSrvStub = defaultTimeSrvMock
+): TempoDatasource {
+  const customSettings: DataSourceInstanceSettings<TempoJsonData> = {
+    url: 'myloggingurl',
+    id: 0,
+    uid: '',
+    type: '',
+    name: '',
+    meta: {
+      id: 'id',
+      name: 'name',
+      type: PluginType.datasource,
+      module: '',
+      baseUrl: '',
+      info: {
+        author: {
+          name: 'Test',
+        },
+        description: '',
+        links: [],
+        logos: {
+          large: '',
+          small: '',
+        },
+        screenshots: [],
+        updated: '',
+        version: '',
+      },
+    },
+    readOnly: false,
+    jsonData: {},
+    access: 'direct',
+    ...settings,
+  };
+
+  // @ts-expect-error
+  return new TempoDatasource(customSettings, templateSrvMock, timeSrvStub);
+}
+
+export function createMetadataRequest(labelsAndValues: Record<string, Record<string, unknown>>) {
+  return async () => labelsAndValues;
+}

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

@@ -0,0 +1,51 @@
+import { lastValueFrom } from 'rxjs';
+
+import { DataQueryRequest, TimeRange } from '@grafana/data';
+
+import { TempoVariableQuery } from './VariableQueryEditor';
+import { createMetadataRequest, createTempoDatasource } from './mocks';
+import { TempoVariableSupport } from './variables';
+
+describe('TempoVariableSupport', () => {
+  let TempoVariableSupportMock: TempoVariableSupport;
+
+  beforeEach(() => {
+    const datasource = createTempoDatasource();
+    jest.spyOn(datasource, 'metadataRequest').mockImplementation(
+      createMetadataRequest({
+        data: {
+          tagNames: ['label1', 'label2'],
+          scopes: [{ name: 'span', tags: ['label1', 'label2'] }],
+        },
+      })
+    );
+    TempoVariableSupportMock = new TempoVariableSupport(datasource);
+  });
+
+  it('should return label names for Tempo', async () => {
+    const response = TempoVariableSupportMock.query({
+      app: 'undefined',
+      startTime: 0,
+      requestId: '1',
+      interval: 'undefined',
+      scopedVars: {},
+      timezone: 'undefined',
+      type: 0,
+      maxDataPoints: 10,
+      intervalMs: 5000,
+      targets: [
+        {
+          refId: 'A',
+          datasource: { uid: 'GRAFANA_DATASOURCE_NAME', type: 'sample' },
+          type: 0,
+        },
+      ],
+      panelId: 1,
+      publicDashboardAccessToken: '',
+      range: { from: new Date().toLocaleString(), to: new Date().toLocaleString() } as unknown as TimeRange,
+    } as DataQueryRequest<TempoVariableQuery>);
+
+    const data = (await lastValueFrom(response)).data;
+    expect(data).toEqual([{ text: 'label1' }, { text: 'label2' }]);
+  });
+});

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

@@ -0,0 +1,25 @@
+import { from, Observable } from 'rxjs';
+import { map } from 'rxjs/operators';
+
+import { DataQueryRequest, DataQueryResponse, CustomVariableSupport } from '@grafana/data';
+
+import { TempoVariableQuery, TempoVariableQueryEditor } from './VariableQueryEditor';
+import { TempoDatasource } from './datasource';
+
+export class TempoVariableSupport extends CustomVariableSupport<TempoDatasource, TempoVariableQuery> {
+  editor = TempoVariableQueryEditor;
+
+  constructor(private datasource: TempoDatasource) {
+    super();
+    this.query = this.query.bind(this);
+  }
+
+  query(request: DataQueryRequest<TempoVariableQuery>): Observable<DataQueryResponse> {
+    if (!this.datasource) {
+      throw new Error('Datasource not initialized');
+    }
+
+    const result = this.datasource.executeVariableQuery(request.targets[0]);
+    return from(result).pipe(map((data) => ({ data })));
+  }
+}