[Tempo] - Random API response and other improvements (#54905)

* Moved the SearchResponse type to types.ts.
Created a mockSearchResponse function to generate search responses.

* Generate spans in mocked response. Extend results table to accomodate spans

* Show the first spanset attributes in the table

* Added a shortcut to run the query in TraceQL editor. Added a label and link to docs above the editor

* Improved autocomplete list sorting. Improved value regex.

* Rename column to "Span"

* Tests are great!

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

@@ -49,7 +49,9 @@ import {
   transformTraceList,
   transformFromOTLP as transformFromOTEL,
   createTableFrameFromSearch,
+  createTableFrameFromTraceQlQuery,
 } from './resultTransformer';
+import { mockedSearchResponse } from './traceql/mockedSearchResponse';
 import { SearchQueryParams, TempoQuery, TempoJsonData } from './types';
 
 export const DEFAULT_LIMIT = 20;
@@ -164,6 +166,22 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
         return of({ error: { message: error instanceof Error ? error.message : 'Unknown error occurred' }, data: [] });
       }
     }
+    if (targets.traceql?.length) {
+      try {
+        reportInteraction('grafana_traces_traceql_queried', {
+          datasourceType: 'tempo',
+          app: options.app ?? '',
+        });
+
+        subQueries.push(
+          of({
+            data: [createTableFrameFromTraceQlQuery(mockedSearchResponse().traces, this.instanceSettings)],
+          })
+        );
+      } catch (error) {
+        return of({ error: { message: error instanceof Error ? error.message : 'Unknown error occurred' }, data: [] });
+      }
+    }
 
     if (targets.upload?.length) {
       if (this.uploadedJson) {

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

@@ -2,13 +2,7 @@ import { collectorTypes } from '@opentelemetry/exporter-collector';
 
 import { FieldType, MutableDataFrame, PluginType, DataSourceInstanceSettings, dateTime } from '@grafana/data';
 
-import {
-  SearchResponse,
-  createTableFrame,
-  transformToOTLP,
-  transformFromOTLP,
-  createTableFrameFromSearch,
-} from './resultTransformer';
+import { createTableFrame, transformToOTLP, transformFromOTLP, createTableFrameFromSearch } from './resultTransformer';
 import {
   badOTLPResponse,
   otlpDataFrameToResponse,
@@ -16,6 +10,7 @@ import {
   otlpResponse,
   tempoSearchResponse,
 } from './testResponse';
+import { TraceSearchMetadata } from './types';
 
 const defaultSettings: DataSourceInstanceSettings = {
   id: 0,
@@ -94,7 +89,7 @@ describe('createTableFrameFromSearch()', () => {
   const mockTimeUnix = dateTime(1643357709095).valueOf();
   global.Date.now = jest.fn(() => mockTimeUnix);
   test('transforms search response to dataFrame', () => {
-    const frame = createTableFrameFromSearch(tempoSearchResponse.traces as SearchResponse[], defaultSettings);
+    const frame = createTableFrameFromSearch(tempoSearchResponse.traces as TraceSearchMetadata[], defaultSettings);
     expect(frame.fields[0].name).toBe('traceID');
     expect(frame.fields[0].values.get(0)).toBe('e641dcac1c3a0565');
 

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

@@ -18,6 +18,7 @@ import {
 } from '@grafana/data';
 
 import { createGraphFrames } from './graphTransform';
+import { Span, TraceSearchMetadata } from './types';
 
 export function createTableFrame(
   logsFrame: DataFrame,
@@ -549,15 +550,7 @@ export function transformTrace(response: DataQueryResponse, nodeGraph = false):
   };
 }
 
-export type SearchResponse = {
-  traceID: string;
-  rootServiceName: string;
-  rootTraceName: string;
-  startTimeUnixNano: string;
-  durationMs: number;
-};
-
-export function createTableFrameFromSearch(data: SearchResponse[], instanceSettings: DataSourceInstanceSettings) {
+export function createTableFrameFromSearch(data: TraceSearchMetadata[], instanceSettings: DataSourceInstanceSettings) {
   const frame = new MutableDataFrame({
     fields: [
       {
@@ -605,7 +598,7 @@ export function createTableFrameFromSearch(data: SearchResponse[], instanceSetti
   return frame;
 }
 
-function transformToTraceData(data: SearchResponse) {
+function transformToTraceData(data: TraceSearchMetadata) {
   let traceName = '';
   if (data.rootServiceName) {
     traceName += data.rootServiceName + ' ';
@@ -633,6 +626,121 @@ function transformToTraceData(data: SearchResponse) {
   };
 }
 
+export function createTableFrameFromTraceQlQuery(
+  data: TraceSearchMetadata[],
+  instanceSettings: DataSourceInstanceSettings
+) {
+  const frame = new MutableDataFrame({
+    fields: [
+      {
+        name: 'traceID',
+        type: FieldType.string,
+        config: {
+          unit: 'string',
+          displayNameFromDS: 'Trace ID',
+          links: [
+            {
+              title: 'Trace: ${__value.raw}',
+              url: '',
+              internal: {
+                datasourceUid: instanceSettings.uid,
+                datasourceName: instanceSettings.name,
+                query: {
+                  query: '${__value.raw}',
+                  queryType: 'traceId',
+                },
+              },
+            },
+          ],
+        },
+      },
+      {
+        name: 'spanID',
+        type: FieldType.string,
+        config: {
+          unit: 'string',
+          displayNameFromDS: 'Span',
+          links: [
+            {
+              title: 'Span: ${__value.raw}',
+              url: '',
+              internal: {
+                datasourceUid: instanceSettings.uid,
+                datasourceName: instanceSettings.name,
+                query: {
+                  query: '${__value.raw}',
+                  queryType: 'spanId',
+                },
+              },
+            },
+          ],
+        },
+      },
+      { name: 'traceName', type: FieldType.string, config: { displayNameFromDS: 'Name' } },
+      { name: 'attributes', type: FieldType.string, config: { displayNameFromDS: 'Attributes' } },
+      { name: 'startTime', type: FieldType.string, 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
+    .sort((a, b) => parseInt(b?.startTimeUnixNano!, 10) / 1000000 - parseInt(a?.startTimeUnixNano!, 10) / 1000000)
+    .reduce((list: TraceTableData[], t) => {
+      const firstSpanSet = t.spanSets?.[0];
+      const trace: TraceTableData = transformToTraceData(t);
+      trace.attributes = firstSpanSet?.attributes
+        .map((atr) => Object.keys(atr).map((key) => `${key} = ${atr[key]}`))
+        .join(', ');
+      list.push(trace);
+      firstSpanSet?.spans.forEach((span) => list.push(transformSpanToTraceData(span)));
+      return list;
+    }, []);
+
+  for (const trace of traceData) {
+    frame.add(trace);
+  }
+
+  return frame;
+}
+
+interface TraceTableData {
+  traceID?: string;
+  traceName: string;
+  spanID?: string;
+  attributes?: string;
+  startTime: string;
+  duration: number;
+}
+
+function transformSpanToTraceData(data: Span): TraceTableData {
+  const traceStartTime = data.startTimeUnixNano / 1000000;
+  const traceEndTime = data.endTimeUnixNano / 1000000;
+
+  let startTime = dateTimeFormat(traceStartTime);
+
+  if (Math.abs(differenceInHours(new Date(traceStartTime), Date.now())) <= 1) {
+    startTime = formatDistance(new Date(traceStartTime), Date.now(), {
+      addSuffix: true,
+      includeSeconds: true,
+    });
+  }
+
+  return {
+    traceID: undefined,
+    traceName: data.name,
+    spanID: data.spanId,
+    attributes: data.attributes?.map((atr) => Object.keys(atr).map((key) => `${key} = ${atr[key]}`)).join(', '),
+    startTime,
+    duration: traceEndTime - traceStartTime,
+  };
+}
+
 const emptyDataQueryResponse = {
   data: [
     new MutableDataFrame({

public/app/plugins/datasource/tempo/traceql/QueryEditor.tsx

@@ -3,7 +3,7 @@ import { defaults } from 'lodash';
 import React from 'react';
 
 import { QueryEditorProps } from '@grafana/data';
-import { useStyles2 } from '@grafana/ui';
+import { InlineLabel, useStyles2 } from '@grafana/ui';
 
 import { TempoDatasource } from '../datasource';
 import { defaultQuery, MyDataSourceOptions, TempoQuery } from '../types';
@@ -23,7 +23,22 @@ export function QueryEditor(props: Props) {
 
   return (
     <>
-      <TraceQLEditor value={query.query} onChange={onEditorChange} datasource={props.datasource} />
+      <InlineLabel>
+        Build complex queries using TraceQL to select a list of traces.{' '}
+        <a
+          rel="noreferrer"
+          target="_blank"
+          href="https://github.com/grafana/tempo/blob/main/docs/design-proposals/2022-04%20TraceQL%20Concepts.md"
+        >
+          Documentation
+        </a>
+      </InlineLabel>
+      <TraceQLEditor
+        value={query.query}
+        onChange={onEditorChange}
+        datasource={props.datasource}
+        onRunQuery={props.onRunQuery}
+      />
       <div className={styles.optionsContainer}>
         <TempoQueryBuilderOptions query={query} onChange={props.onChange} />
       </div>

public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx

@@ -16,10 +16,12 @@ import { languageDefinition } from './traceql';
 interface Props {
   value: string;
   onChange: (val: string) => void;
+  onRunQuery: () => void;
   datasource: TempoDatasource;
 }
 
 export function TraceQLEditor(props: Props) {
+  const { onRunQuery } = props;
   const setupAutocompleteFn = useAutocomplete(props.datasource);
   const styles = useStyles2(getStyles);
   return (
@@ -47,11 +49,29 @@ export function TraceQLEditor(props: Props) {
       onBeforeEditorMount={ensureTraceQL}
       onEditorDidMount={(editor, monaco) => {
         setupAutocompleteFn(editor, monaco);
+        setupActions(editor, monaco, onRunQuery);
       }}
     />
   );
 }
 
+function setupActions(editor: monacoTypes.editor.IStandaloneCodeEditor, monaco: Monaco, onRunQuery: () => void) {
+  editor.addAction({
+    id: 'run-query',
+    label: 'Run Query',
+
+    keybindings: [monaco.KeyMod.Shift | monaco.KeyCode.Enter],
+
+    contextMenuGroupId: 'navigation',
+
+    contextMenuOrder: 1.5,
+
+    run: function () {
+      onRunQuery();
+    },
+  });
+}
+
 /**
  * Hook that returns function that will set up monaco autocomplete for the label selector
  * @param datasource

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

@@ -17,10 +17,10 @@ describe('CompletionProvider', () => {
     const { provider, model } = setup('{}', 1, defaultTags);
     const result = await provider.provideCompletionItems(model as any, {} as any);
     expect((result! as monacoTypes.languages.CompletionList).suggestions).toEqual([
-      expect.objectContaining({ label: 'foo', insertText: '.foo' }),
-      expect.objectContaining({ label: 'bar', insertText: '.bar' }),
-      ...CompletionProvider.intrinsics.map((s) => expect.objectContaining({ label: s, insertText: s })),
       ...CompletionProvider.scopes.map((s) => expect.objectContaining({ label: s, insertText: s })),
+      ...CompletionProvider.intrinsics.map((s) => expect.objectContaining({ label: s, insertText: s })),
+      expect.objectContaining({ label: 'bar', insertText: '.bar' }),
+      expect.objectContaining({ label: 'foo', insertText: '.foo' }),
     ]);
   });
 
@@ -76,10 +76,10 @@ describe('CompletionProvider', () => {
     const { provider, model } = setup('', 0, defaultTags);
     const result = await provider.provideCompletionItems(model as any, {} as any);
     expect((result! as monacoTypes.languages.CompletionList).suggestions).toEqual([
-      expect.objectContaining({ label: 'foo', insertText: '{ .foo' }),
-      expect.objectContaining({ label: 'bar', insertText: '{ .bar' }),
-      ...CompletionProvider.intrinsics.map((s) => expect.objectContaining({ label: s, insertText: `{ ${s}` })),
       ...CompletionProvider.scopes.map((s) => expect.objectContaining({ label: s, insertText: `{ ${s}` })),
+      ...CompletionProvider.intrinsics.map((s) => expect.objectContaining({ label: s, insertText: `{ ${s}` })),
+      expect.objectContaining({ label: 'bar', insertText: '{ .bar' }),
+      expect.objectContaining({ label: 'foo', insertText: '{ .foo' }),
     ]);
   });
 
@@ -95,8 +95,8 @@ describe('CompletionProvider', () => {
     const { provider, model } = setup('{ resource. }', 11, defaultTags);
     const result = await provider.provideCompletionItems(model as any, {} as any);
     expect((result! as monacoTypes.languages.CompletionList).suggestions).toEqual([
-      ...defaultTags.map((s) => expect.objectContaining({ label: s, insertText: s })),
       ...CompletionProvider.intrinsics.map((s) => expect.objectContaining({ label: s, insertText: s })),
+      ...defaultTags.map((s) => expect.objectContaining({ label: s, insertText: s })),
     ]);
   });
 
@@ -130,7 +130,7 @@ describe('CompletionProvider', () => {
   });
 });
 
-const defaultTags = ['foo', 'bar'];
+const defaultTags = ['bar', 'foo'];
 
 function setup(value: string, offset: number, tags?: string[]) {
   const ds = new TempoDatasource(defaultSettings);

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

@@ -20,9 +20,9 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
 
   triggerCharacters = ['{', '.', '[', '(', '=', '~', ' ', '"'];
 
-  static readonly intrinsics: string[] = ['name', 'status', 'duration'];
-  static readonly scopes: string[] = ['span', 'resource'];
-  static readonly operators: string[] = ['=', '-', '+', '<', '>', '>=', '<='];
+  static readonly intrinsics: string[] = ['duration', 'name', 'status'];
+  static readonly scopes: string[] = ['resource', 'span'];
+  static readonly operators: string[] = ['=', '-', '+', '<', '>', '>=', '<=', '=~'];
   static readonly logicalOps: string[] = ['&&', '||'];
 
   // We set these directly and ae required for the provider to function.
@@ -110,16 +110,16 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
         return [];
       }
       case 'EMPTY': {
-        return this.getTagsCompletions('{ .')
+        return this.getScopesCompletions('{ ')
           .concat(this.getIntrinsicsCompletions('{ '))
-          .concat(this.getScopesCompletions('{ '));
+          .concat(this.getTagsCompletions('{ .'));
       }
       case 'SPANSET_EMPTY':
-        return this.getTagsCompletions('.').concat(this.getIntrinsicsCompletions()).concat(this.getScopesCompletions());
+        return this.getScopesCompletions().concat(this.getIntrinsicsCompletions()).concat(this.getTagsCompletions('.'));
       case 'SPANSET_IN_NAME':
-        return this.getTagsCompletions().concat(this.getIntrinsicsCompletions()).concat(this.getScopesCompletions());
+        return this.getScopesCompletions().concat(this.getIntrinsicsCompletions()).concat(this.getTagsCompletions());
       case 'SPANSET_IN_NAME_SCOPE':
-        return this.getTagsCompletions().concat(this.getIntrinsicsCompletions());
+        return this.getIntrinsicsCompletions().concat(this.getTagsCompletions());
       case 'SPANSET_AFTER_NAME':
         return CompletionProvider.operators.map((key) => ({
           label: key,
@@ -152,11 +152,13 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
   }
 
   private getTagsCompletions(prepend?: string): Completion[] {
-    return Object.keys(this.tags).map((key) => ({
-      label: key,
-      insertText: (prepend || '') + key,
-      type: 'TAG_NAME' as CompletionType,
-    }));
+    return Object.keys(this.tags)
+      .sort((a, b) => a.localeCompare(b, undefined, { sensitivity: 'accent' }))
+      .map((key) => ({
+        label: key,
+        insertText: (prepend || '') + key,
+        type: 'TAG_NAME' as CompletionType,
+      }));
   }
 
   private getIntrinsicsCompletions(prepend?: string): Completion[] {
@@ -178,11 +180,12 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
   private getSituationInSpanSet(textUntilCaret: string): Situation {
     const nameRegex = /(?<name>[\w./-]+)?/;
     const opRegex = /(?<op>[!=+\-<>]+)/;
-    const valueRegex = /(?<value>(?<open_quote>")?(\w[^"\n&|]*\w)?(?<close_quote>")?)?/;
+    // only allow spaces in the value if it's enclosed by quotes
+    const valueRegex = /(?<value>(?<open_quote>")([^"\n&|]+)?(?<close_quote>")?|([^"\n\s&|]+))?/;
 
     // prettier-ignore
     const fullRegex = new RegExp(
-      '([\\s{])' +      // Space(s) or initial opening bracket {
+        '([\\s{])' +      // Space(s) or initial opening bracket {
         '(' +                   // Open full set group
         nameRegex.source +
         '(?<space1>\\s*)' +     // Optional space(s) between name and operator
@@ -214,7 +217,6 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
       if (!op) {
         // There's no operator so we check if the name is one of the known scopes
         // { resource.|
-
         if (CompletionProvider.scopes.filter((w) => w === nameMatched?.groups?.word) && nameMatched?.groups?.post_dot) {
           return {
             type: 'SPANSET_IN_NAME_SCOPE',
@@ -260,8 +262,6 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
 
   /**
    * Figure out where is the cursor and what kind of suggestions are appropriate.
-   * As currently TraceQL handles just a simple {foo="bar", baz="zyx"} kind of values we can do with simple regex to figure
-   * out where we are with the cursor.
    * @param text
    * @param offset
    */

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

@@ -0,0 +1,62 @@
+import { v4 as uuidv4 } from 'uuid';
+
+import { SearchResponse, Span, SpanKind, TraceSearchMetadata } from '../types';
+
+export const mockedSearchResponse = (): SearchResponse => {
+  const traces: TraceSearchMetadata[] = [];
+  const attributes = [
+    { 'http.status.code': '500' },
+    { 'http.status.code': '200' },
+    { 'http.status.code': '404' },
+    { job: '"test-job"' },
+    { job: '"main-job"' },
+    { job: '"long-job"' },
+    { error: '"lorem ipsum"' },
+    { error: '"something went wrong"' },
+  ];
+
+  const tracesCount = Math.random() * 20 + 20;
+  for (let i = 0; i < tracesCount; i++) {
+    const attr = Math.floor(Math.random() * attributes.length);
+    const startTime = (Date.now() - Math.random() * (i + 1) * 100000) * 1000000;
+    const t: TraceSearchMetadata = {
+      traceID: uuidv4().replace(/-/, '').substring(0, 16),
+      rootServiceName: 'service' + i,
+      rootTraceName: 'trace' + i,
+      startTimeUnixNano: startTime.toString(10),
+      durationMs: Math.random() * 1000,
+      spanSets: [],
+    };
+
+    const spanAttributes = [];
+    for (let k = 0; k < Math.random() * 2; k++) {
+      const newAttr = Math.floor(Math.random() * attributes.length);
+      if (newAttr !== attr) {
+        spanAttributes.push(attributes[newAttr]);
+      }
+    }
+
+    const spans: Span[] = [];
+    for (let j = 0; j < Math.random() * 3 + 1; j++) {
+      spans.push({
+        traceId: t.traceID,
+        spanId: uuidv4().replace(/-/, '').substring(0, 16),
+        name: uuidv4().replace(/-/, '').substring(0, 6),
+        startTimeUnixNano: startTime,
+        endTimeUnixNano: startTime + Math.random() * 10000000,
+        kind: SpanKind.INTERNAL,
+        attributes: spanAttributes,
+      });
+    }
+    t.spanSets!.push({ spans, attributes: [attributes[attr]] });
+    traces.push(t);
+  }
+
+  return {
+    traces,
+    metrics: {
+      inspectedTraces: tracesCount,
+      inspectedBytes: 83720,
+    },
+  };
+};

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

@@ -1,5 +1,5 @@
 import { DataQuery } from '@grafana/data';
-import { DataSourceJsonData } from '@grafana/data/src';
+import { DataSourceJsonData, KeyValue } from '@grafana/data/src';
 import { NodeGraphOptions } from 'app/core/components/NodeGraphSettings';
 import { TraceToLogsOptions } from 'app/core/components/TraceToLogs/TraceToLogsSettings';
 
@@ -51,3 +51,54 @@ export interface TempoQuery extends DataQuery {
 export interface MyDataSourceOptions extends DataSourceJsonData {}
 
 export const defaultQuery: Partial<TempoQuery> = {};
+
+export type TraceSearchMetadata = {
+  traceID: string;
+  rootServiceName: string;
+  rootTraceName: string;
+  startTimeUnixNano: string;
+  durationMs: number;
+  spanSets?: Spanset[];
+};
+
+export type SearchMetrics = {
+  inspectedTraces?: number;
+  inspectedBytes?: number;
+  inspectedBlocks?: number;
+  skippedBlocks?: number;
+  skippedTraces?: number;
+  totalBlockBytes?: number;
+  spanSets?: Spanset[];
+};
+
+export enum SpanKind {
+  UNSPECIFIED,
+  INTERNAL,
+  SERVER,
+  CLIENT,
+  PRODUCER,
+  CONSUMER,
+}
+
+export type Span = {
+  traceId: string;
+  spanId: string;
+  traceState?: string;
+  parentSpanId?: string;
+  name: string;
+  kind: SpanKind;
+  startTimeUnixNano: number;
+  endTimeUnixNano: number;
+  attributes?: KeyValue[];
+  dropped_attributes_count?: number;
+};
+
+export type Spanset = {
+  attributes: KeyValue[];
+  spans: Span[];
+};
+
+export type SearchResponse = {
+  traces: TraceSearchMetadata[];
+  metrics: SearchMetrics;
+};