Tempo: Add Tempo search behind feature flag (#37765)

* Add Tempo search behind feature flag

* Add query fields for Tempo search

* Only show loki search if a logs-to-traces datasource is set up

* Refactor tempo search to use separate fields for service name, span name, and tags

* Add tests to buildSearchQuery

* Move search to separate component and rename type to native search

* Improve Tempo tokenizer syntax

packages/grafana-data/src/types/config.ts

@@ -49,6 +49,7 @@ export interface FeatureToggles {
   trimDefaults: boolean;
   accesscontrol: boolean;
   tempoServiceGraph: boolean;
+  tempoSearch: boolean;
 }
 
 /**

packages/grafana-runtime/src/config.ts

@@ -64,6 +64,7 @@ export class GrafanaBootConfig implements GrafanaConfig {
     accesscontrol: false,
     trimDefaults: false,
     tempoServiceGraph: false,
+    tempoSearch: false,
   };
   licenseInfo: LicenseInfo = {} as LicenseInfo;
   rendererAvailable = false;

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

@@ -0,0 +1,177 @@
+import React, { useState, useEffect, useMemo } from 'react';
+import {
+  InlineFieldRow,
+  InlineField,
+  Input,
+  QueryField,
+  Select,
+  SlatePrism,
+  BracesPlugin,
+  TypeaheadInput,
+  TypeaheadOutput,
+} from '@grafana/ui';
+import { tokenizer } from './syntax';
+import Prism from 'prismjs';
+import { Node } from 'slate';
+import { css } from '@emotion/css';
+import { SelectableValue } from '@grafana/data';
+import TempoLanguageProvider from './language_provider';
+import { TempoDatasource, TempoQuery } from './datasource';
+
+interface Props {
+  datasource: TempoDatasource;
+  query: TempoQuery;
+  onChange: (value: TempoQuery) => void;
+  onBlur?: () => void;
+  onRunQuery: () => void;
+}
+
+const PRISM_LANGUAGE = 'tempo';
+const durationPlaceholder = 'e.g. 1.2s, 100ms, 500us';
+const plugins = [
+  BracesPlugin(),
+  SlatePrism({
+    onlyIn: (node: Node) => node.object === 'block' && node.type === 'code_block',
+    getSyntax: () => PRISM_LANGUAGE,
+  }),
+];
+
+Prism.languages[PRISM_LANGUAGE] = tokenizer;
+
+const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props) => {
+  const languageProvider = useMemo(() => new TempoLanguageProvider(datasource), [datasource]);
+  const [hasSyntaxLoaded, setHasSyntaxLoaded] = useState(false);
+  const [autocomplete, setAutocomplete] = useState<{
+    serviceNameOptions: Array<SelectableValue<string>>;
+    spanNameOptions: Array<SelectableValue<string>>;
+  }>({
+    serviceNameOptions: [],
+    spanNameOptions: [],
+  });
+
+  useEffect(() => {
+    const fetchAutocomplete = async () => {
+      await languageProvider.start();
+      const serviceNameOptions = await languageProvider.getOptions('service.name');
+      const spanNameOptions = await languageProvider.getOptions('name');
+      setHasSyntaxLoaded(true);
+      setAutocomplete({ serviceNameOptions, spanNameOptions });
+    };
+    fetchAutocomplete();
+  }, [languageProvider]);
+
+  const onTypeahead = async (typeahead: TypeaheadInput): Promise<TypeaheadOutput> => {
+    return await languageProvider.provideCompletionItems(typeahead);
+  };
+
+  const cleanText = (text: string) => {
+    const splittedText = text.split(/\s+(?=([^"]*"[^"]*")*[^"]*$)/g);
+    if (splittedText.length > 1) {
+      return splittedText[splittedText.length - 1];
+    }
+    return text;
+  };
+
+  return (
+    <div className={css({ maxWidth: '500px' })}>
+      <InlineFieldRow>
+        <InlineField label="Service Name" labelWidth={14} grow>
+          <Select
+            menuShouldPortal
+            options={autocomplete.serviceNameOptions}
+            value={query.serviceName || ''}
+            onChange={(v) => {
+              onChange({
+                ...query,
+                serviceName: v?.value || undefined,
+              });
+            }}
+            placeholder="Select a service"
+            isClearable
+          />
+        </InlineField>
+      </InlineFieldRow>
+      <InlineFieldRow>
+        <InlineField label="Span Name" labelWidth={14} grow>
+          <Select
+            menuShouldPortal
+            options={autocomplete.spanNameOptions}
+            value={query.spanName || ''}
+            onChange={(v) => {
+              onChange({
+                ...query,
+                spanName: v?.value || undefined,
+              });
+            }}
+            placeholder="Select a span"
+            isClearable
+          />
+        </InlineField>
+      </InlineFieldRow>
+      <InlineFieldRow>
+        <InlineField label="Tags" labelWidth={14} grow tooltip="Values should be in the logfmt format.">
+          <QueryField
+            additionalPlugins={plugins}
+            query={query.search}
+            onTypeahead={onTypeahead}
+            onBlur={onBlur}
+            onChange={(value) => {
+              onChange({
+                ...query,
+                search: value,
+              });
+            }}
+            cleanText={cleanText}
+            onRunQuery={onRunQuery}
+            syntaxLoaded={hasSyntaxLoaded}
+            portalOrigin="tempo"
+          />
+        </InlineField>
+      </InlineFieldRow>
+      <InlineFieldRow>
+        <InlineField label="Min Duration" labelWidth={14} grow>
+          <Input
+            value={query.minDuration || ''}
+            placeholder={durationPlaceholder}
+            onChange={(v) =>
+              onChange({
+                ...query,
+                minDuration: v.currentTarget.value,
+              })
+            }
+          />
+        </InlineField>
+      </InlineFieldRow>
+      <InlineFieldRow>
+        <InlineField label="Max Duration" labelWidth={14} grow>
+          <Input
+            value={query.maxDuration || ''}
+            placeholder={durationPlaceholder}
+            onChange={(v) =>
+              onChange({
+                ...query,
+                maxDuration: v.currentTarget.value,
+              })
+            }
+          />
+        </InlineField>
+      </InlineFieldRow>
+      <InlineFieldRow>
+        <InlineField label="Limit" labelWidth={14} grow tooltip="Maximum numbers of returned results">
+          <Input
+            value={query.limit || ''}
+            type="number"
+            onChange={(v) =>
+              onChange({
+                ...query,
+                limit: v.currentTarget.value ? parseInt(v.currentTarget.value, 10) : undefined,
+              })
+            }
+          />
+        </InlineField>
+      </InlineFieldRow>
+    </div>
+  );
+};
+
+export default NativeSearch;

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

@@ -15,11 +15,12 @@ import {
 import { TraceToLogsOptions } from 'app/core/components/TraceToLogsSettings';
 import React from 'react';
 import { LokiQueryField } from '../loki/components/LokiQueryField';
+import { LokiQuery } from '../loki/types';
 import { TempoDatasource, TempoQuery, TempoQueryType } from './datasource';
 import LokiDatasource from '../loki/datasource';
-import { LokiQuery } from '../loki/types';
 import { PrometheusDatasource } from '../prometheus/datasource';
 import useAsync from 'react-use/lib/useAsync';
+import NativeSearch from './NativeSearch';
 
 interface Props extends ExploreQueryFieldProps<TempoDatasource, TempoQuery>, Themeable2 {}
 
@@ -31,6 +32,7 @@ interface State {
   serviceMapDatasourceUid?: string;
   serviceMapDatasource?: PrometheusDatasource;
 }
+
 class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
   state = {
     linkedDatasourceUid: undefined,
@@ -82,7 +84,6 @@ class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
     const graphDatasourceUid = datasource.serviceMap?.datasourceUid;
 
     const queryTypeOptions: Array<SelectableValue<TempoQueryType>> = [
-      { value: 'search', label: 'Search' },
       { value: 'traceId', label: 'TraceID' },
       { value: 'upload', label: 'JSON file' },
     ];
@@ -91,6 +92,14 @@ class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
       queryTypeOptions.push({ value: 'serviceMap', label: 'Service Map' });
     }
 
+    if (config.featureToggles.tempoSearch) {
+      queryTypeOptions.unshift({ value: 'nativeSearch', label: 'Search' });
+    }
+
+    if (logsDatasourceUid) {
+      queryTypeOptions.push({ value: 'search', label: 'Loki Search' });
+    }
+
     return (
       <>
         <InlineFieldRow>
@@ -116,6 +125,15 @@ class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
             onChange={this.onChangeLinkedQuery}
           />
         )}
+        {query.queryType === 'nativeSearch' && (
+          <NativeSearch
+            datasource={this.props.datasource}
+            query={query}
+            onChange={onChange}
+            onBlur={this.props.onBlur}
+            onRunQuery={this.props.onRunQuery}
+          />
+        )}
         {query.queryType === 'upload' && (
           <div className={css({ padding: this.props.theme.spacing(2) })}>
             <FileDropzone

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

@@ -11,8 +11,8 @@ import {
 } from '@grafana/data';
 
 import { createFetchResponse } from 'test/helpers/createFetchResponse';
-import { TempoDatasource } from './datasource';
 import { BackendDataSourceResponse, FetchResponse, setBackendSrv, setDataSourceSrv } from '@grafana/runtime';
+import { TempoDatasource, TempoQuery } from './datasource';
 import mockJson from './mockJsonResponse.json';
 
 describe('Tempo data source', () => {
@@ -118,6 +118,44 @@ describe('Tempo data source', () => {
     expect(field.values.get(0)).toBe('60ba2abb44f13eae');
     expect(field.values.length).toBe(6);
   });
+
+  it('should build search query correctly', () => {
+    const ds = new TempoDatasource(defaultSettings);
+    const tempoQuery: TempoQuery = {
+      queryType: 'search',
+      refId: 'A',
+      query: '',
+      serviceName: 'frontend',
+      spanName: '/config',
+      search: 'root.http.status_code=500',
+      minDuration: '1ms',
+      maxDuration: '100s',
+      limit: 10,
+    };
+    const builtQuery = ds.buildSearchQuery(tempoQuery);
+    expect(builtQuery).toStrictEqual({
+      'service.name': 'frontend',
+      name: '/config',
+      'root.http.status_code': '500',
+      minDuration: '1ms',
+      maxDuration: '100s',
+      limit: 10,
+    });
+  });
+
+  it('should ignore incomplete tag queries', () => {
+    const ds = new TempoDatasource(defaultSettings);
+    const tempoQuery: TempoQuery = {
+      queryType: 'search',
+      refId: 'A',
+      query: '',
+      search: 'root.ip root.http.status_code=500',
+    };
+    const builtQuery = ds.buildSearchQuery(tempoQuery);
+    expect(builtQuery).toStrictEqual({
+      'root.http.status_code': '500',
+    });
+  });
 });
 
 const backendSrvWithPrometheus = {

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

@@ -1,5 +1,4 @@
-import { groupBy } from 'lodash';
-import { from, lastValueFrom, merge, Observable, of, throwError } from 'rxjs';
+import { lastValueFrom, from, merge, Observable, of, throwError } from 'rxjs';
 import { map, mergeMap, toArray } from 'rxjs/operators';
 import {
   DataQuery,
@@ -10,17 +9,26 @@ import {
   DataSourceJsonData,
   LoadingState,
 } from '@grafana/data';
-import { DataSourceWithBackend } from '@grafana/runtime';
-
 import { TraceToLogsOptions } from 'app/core/components/TraceToLogsSettings';
+import { BackendSrvRequest, DataSourceWithBackend, getBackendSrv } from '@grafana/runtime';
+import { serializeParams } from 'app/core/utils/fetch';
 import { getDatasourceSrv } from 'app/features/plugins/datasource_srv';
+import { identity, pick, pickBy, groupBy } from 'lodash';
+import Prism from 'prismjs';
 import { LokiOptions, LokiQuery } from '../loki/types';
-import { transformTrace, transformTraceList, transformFromOTLP as transformFromOTEL } from './resultTransformer';
 import { PrometheusDatasource } from '../prometheus/datasource';
 import { PromQuery } from '../prometheus/types';
 import { mapPromMetricsToServiceMap, serviceMapMetrics } from './graphTransform';
+import {
+  transformTrace,
+  transformTraceList,
+  transformFromOTLP as transformFromOTEL,
+  createTableFrameFromSearch,
+} from './resultTransformer';
+import { tokenizer } from './syntax';
 
-export type TempoQueryType = 'search' | 'traceId' | 'serviceMap' | 'upload';
+// search = Loki search, nativeSearch = Tempo search for backwards compatibility
+export type TempoQueryType = 'search' | 'traceId' | 'serviceMap' | 'upload' | 'nativeSearch';
 
 export interface TempoJsonData extends DataSourceJsonData {
   tracesToLogs?: TraceToLogsOptions;
@@ -33,7 +41,13 @@ export type TempoQuery = {
   query: string;
   // Query to find list of traces, e.g., via Loki
   linkedQuery?: LokiQuery;
+  search: string;
   queryType: TempoQueryType;
+  serviceName?: string;
+  spanName?: string;
+  minDuration?: string;
+  maxDuration?: string;
+  limit?: number;
 } & DataQuery;
 
 export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJsonData> {
@@ -43,7 +57,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
   };
   uploadedJson?: string | ArrayBuffer | null = null;
 
-  constructor(instanceSettings: DataSourceInstanceSettings<TempoJsonData>) {
+  constructor(private instanceSettings: DataSourceInstanceSettings<TempoJsonData>) {
     super(instanceSettings);
     this.tracesToLogs = instanceSettings.jsonData.tracesToLogs;
     this.serviceMap = instanceSettings.jsonData.serviceMap;
@@ -84,6 +98,19 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
       );
     }
 
+    if (targets.nativeSearch?.length) {
+      const searchQuery = this.buildSearchQuery(targets.nativeSearch[0]);
+      subQueries.push(
+        this._request('/api/search', searchQuery).pipe(
+          map((response) => {
+            return {
+              data: [createTableFrameFromSearch(response.data.traces, this.instanceSettings)],
+            };
+          })
+        )
+      );
+    }
+
     if (targets.upload?.length) {
       if (this.uploadedJson) {
         const otelTraceData = JSON.parse(this.uploadedJson as string);
@@ -118,6 +145,18 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     return merge(...subQueries);
   }
 
+  async metadataRequest(url: string, params = {}) {
+    return await this._request(url, params, { method: 'GET', hideFromInspector: true }).toPromise();
+  }
+
+  private _request(apiUrl: string, data?: any, options?: Partial<BackendSrvRequest>): Observable<Record<string, any>> {
+    const params = data ? serializeParams(data) : '';
+    const url = `${this.instanceSettings.url}${apiUrl}${params.length ? `?${params}` : ''}`;
+    const req = { ...options, url };
+
+    return getBackendSrv().fetch(req);
+  }
+
   async testDatasource(): Promise<any> {
     // to test Tempo we send a dummy traceID and verify Tempo answers with 'trace not found'
     const response = await lastValueFrom(super.query({ targets: [{ query: '0' }] } as any));
@@ -137,6 +176,41 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
   getQueryDisplayText(query: TempoQuery) {
     return query.query;
   }
+
+  buildSearchQuery(query: TempoQuery) {
+    const tokens = query.search ? Prism.tokenize(query.search, tokenizer) : [];
+
+    // Build key value pairs
+    let tagsQuery: Array<{ [key: string]: string }> = [];
+    for (let i = 0; i < tokens.length - 1; i++) {
+      const token = tokens[i];
+      const lookupToken = tokens[i + 2];
+
+      // Ensure there is a valid key value pair with accurate types
+      if (
+        typeof token !== 'string' &&
+        token.type === 'key' &&
+        typeof token.content === 'string' &&
+        typeof lookupToken !== 'string' &&
+        lookupToken.type === 'value' &&
+        typeof lookupToken.content === 'string'
+      ) {
+        tagsQuery.push({ [token.content]: lookupToken.content });
+      }
+    }
+
+    let tempoQuery = pick(query, ['minDuration', 'maxDuration', 'limit']);
+    // Remove empty properties
+    tempoQuery = pickBy(tempoQuery, identity);
+    if (query.serviceName) {
+      tagsQuery.push({ ['service.name']: query.serviceName });
+    }
+    if (query.spanName) {
+      tagsQuery.push({ ['name']: query.spanName });
+    }
+    const tagsQueryObject = tagsQuery.reduce((tagQuery, item) => ({ ...tagQuery, ...item }), {});
+    return { ...tagsQueryObject, ...tempoQuery };
+  }
 }
 
 function queryServiceMapPrometheus(request: DataQueryRequest<PromQuery>, datasourceUid: string) {

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

@@ -0,0 +1,103 @@
+import { HistoryItem, LanguageProvider, SelectableValue } from '@grafana/data';
+import { CompletionItemGroup, TypeaheadInput, TypeaheadOutput } from '@grafana/ui';
+import { Value } from 'slate';
+import { TempoDatasource } from './datasource';
+
+export default class TempoLanguageProvider extends LanguageProvider {
+  datasource: TempoDatasource;
+  tags?: string[];
+  constructor(datasource: TempoDatasource, initialValues?: any) {
+    super();
+
+    this.datasource = datasource;
+    Object.assign(this, initialValues);
+  }
+
+  request = async (url: string, defaultValue: any, params = {}) => {
+    try {
+      const res = await this.datasource.metadataRequest(url, params);
+      return res?.data;
+    } catch (error) {
+      console.error(error);
+    }
+
+    return defaultValue;
+  };
+
+  start = async () => {
+    await this.fetchTags();
+    return [];
+  };
+
+  async fetchTags() {
+    try {
+      const response = await this.request('/api/search/tags', []);
+      this.tags = response.tagNames;
+    } catch (error) {
+      console.error(error);
+    }
+  }
+
+  provideCompletionItems = async (
+    { prefix, text, value, labelKey, wrapperClasses }: TypeaheadInput,
+    context: { history: Array<HistoryItem<any>> } = { history: [] }
+  ): Promise<TypeaheadOutput> => {
+    const emptyResult: TypeaheadOutput = { suggestions: [] };
+
+    if (!value) {
+      return emptyResult;
+    }
+    if (text === '=') {
+      return this.getTagValueCompletionItems(value);
+    }
+    return this.getTagsCompletionItems();
+  };
+
+  getTagsCompletionItems = (): TypeaheadOutput => {
+    const { tags } = this;
+    const suggestions: CompletionItemGroup[] = [];
+
+    if (tags?.length) {
+      suggestions.push({
+        label: `Tag`,
+        items: tags.map((tag) => ({ label: tag })),
+      });
+    }
+
+    return { suggestions };
+  };
+
+  async getTagValueCompletionItems(value: Value) {
+    const tagNames = value.endText.getText().split(' ');
+    let tagName = tagNames[0];
+    // Get last item if multiple tags
+    if (tagNames.length > 1) {
+      tagName = tagNames[tagNames.length - 1];
+    }
+    tagName = tagName.slice(0, -1);
+    const response = await this.request(`/api/search/tag/${tagName}/values`, []);
+    const suggestions: CompletionItemGroup[] = [];
+
+    if (response && response.tagValues) {
+      suggestions.push({
+        label: `TagValues`,
+        items: response.tagValues.map((tagValue: string) => ({ label: tagValue })),
+      });
+    }
+    return { suggestions };
+  }
+
+  async getOptions(tag: string): Promise<Array<SelectableValue<string>>> {
+    const response = await this.request(`/api/search/tag/${tag}/values`, []);
+    let options: Array<SelectableValue<string>> = [];
+
+    if (response && response.tagValues) {
+      options = response.tagValues.map((v: string) => ({
+        value: v,
+        label: v,
+      }));
+    }
+
+    return options;
+  }
+}

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

@@ -2,6 +2,7 @@ import {
   ArrayVector,
   DataFrame,
   DataQueryResponse,
+  DataSourceInstanceSettings,
   Field,
   FieldType,
   MutableDataFrame,
@@ -315,6 +316,75 @@ function parseJsonFields(frame: DataFrame) {
   }
 }
 
+type SearchResponse = {
+  traceID: string;
+  rootServiceName: string;
+  rootTraceName: string;
+  startTimeUnixNano: string;
+  durationMs: number;
+};
+
+export function createTableFrameFromSearch(data: SearchResponse[], instanceSettings: DataSourceInstanceSettings) {
+  const frame = new MutableDataFrame({
+    fields: [
+      {
+        name: 'traceID',
+        type: FieldType.string,
+        config: {
+          displayNameFromDS: 'Trace ID',
+          links: [
+            {
+              title: 'Trace: ${__value.raw}',
+              url: '',
+              internal: {
+                datasourceUid: instanceSettings.uid,
+                datasourceName: instanceSettings.name,
+                query: {
+                  query: '${__value.raw}',
+                  queryType: 'traceId',
+                },
+              },
+            },
+          ],
+        },
+      },
+      { name: 'traceName', type: FieldType.string, config: { displayNameFromDS: 'Trace name' } },
+      { name: 'startTime', type: FieldType.time, config: { displayNameFromDS: 'Start time' } },
+      { name: 'duration', type: FieldType.number, config: { displayNameFromDS: 'Duration', unit: 'ms' } },
+    ],
+    meta: {
+      preferredVisualisationType: 'table',
+    },
+  });
+  if (!data?.length) {
+    return frame;
+  }
+  // Show the most recent traces
+  const traceData = data.map(transformToTraceData).sort((a, b) => b?.startTime! - a?.startTime!);
+
+  for (const trace of traceData) {
+    frame.add(trace);
+  }
+
+  return frame;
+}
+
+function transformToTraceData(data: SearchResponse) {
+  let traceName = '';
+  if (data.rootServiceName) {
+    traceName += data.rootServiceName + ' ';
+  }
+  if (data.rootTraceName) {
+    traceName += data.rootTraceName;
+  }
+  return {
+    traceID: data.traceID,
+    startTime: parseInt(data.startTimeUnixNano, 10) / 1000 / 1000,
+    duration: data.durationMs,
+    traceName,
+  };
+}
+
 const emptyDataQueryResponse = {
   data: [
     new MutableDataFrame({

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

@@ -0,0 +1,19 @@
+import { tokenizer } from './syntax';
+import Prism from 'prismjs';
+
+describe('Loki syntax', () => {
+  it('should highlight Loki query correctly', () => {
+    expect(Prism.highlight('key=value', tokenizer, 'tempo')).toBe(
+      '<span class="token key attr-name">key</span><span class="token operator">=</span><span class="token value">value</span>'
+    );
+    expect(Prism.highlight('root.ip=172.123.0.1', tokenizer, 'tempo')).toBe(
+      '<span class="token key attr-name">root.ip</span><span class="token operator">=</span><span class="token value">172.123.0.1</span>'
+    );
+    expect(Prism.highlight('root.name="http get /config"', tokenizer, 'tempo')).toBe(
+      '<span class="token key attr-name">root.name</span><span class="token operator">=</span><span class="token value">"http get /config"</span>'
+    );
+    expect(Prism.highlight('key=value key2=value2', tokenizer, 'tempo')).toBe(
+      '<span class="token key attr-name">key</span><span class="token operator">=</span><span class="token value">value</span> <span class="token key attr-name">key2</span><span class="token operator">=</span><span class="token value">value2</span>'
+    );
+  });
+});

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

@@ -0,0 +1,17 @@
+import { Grammar } from 'prismjs';
+
+export const tokenizer: Grammar = {
+  key: {
+    pattern: /[^\s]+(?==)/,
+    alias: 'attr-name',
+  },
+  operator: /[=]/,
+  value: [
+    {
+      pattern: /"(.+)"/,
+    },
+    {
+      pattern: /[^\s]+/,
+    },
+  ],
+};