Tempo: Highlight errors in TraceQL query (#74697)

* Highlight errors

* Chores

* Refactor

* Address PR comments

* Refactoring

* Fix

* Handle another case

* Handle more use cases

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

@@ -0,0 +1,66 @@
+import { computeErrorMessage, getErrorNodes } from './errorHighlighting';
+
+describe('computeErrorMarkers', () => {
+  it.each([
+    ['{span.http.status_code = }', 'Invalid value after comparison or aritmetic operator.'],
+    ['{span.http.status_code 200}', 'Invalid comparison operator after field expression.'],
+    ['{span.http.status_code ""}', 'Invalid comparison operator after field expression.'],
+    ['{span.http.status_code @ 200}', 'Invalid comparison operator after field expression.'],
+    ['{span.http.status_code span.http.status_code}', 'Invalid operator after field expression.'],
+    [
+      '{span.http.status_code = 200} {span.http.status_code = 200}',
+      'Invalid spanset combining operator after spanset expression.',
+    ],
+    [
+      '{span.http.status_code = 200} + {span.http.status_code = 200}',
+      'Invalid spanset combining operator after spanset expression.',
+    ],
+    ['{span.http.status_code = 200} &&', 'Invalid spanset expression after spanset combining operator.'],
+    [
+      '{span.http.status_code = 200} && {span.http.status_code = 200} | foo() > 3',
+      'Invalid aggregation operator after pipepile operator.',
+    ],
+    [
+      '{span.http.status_code = 200} && {span.http.status_code = 200} | avg() > 3',
+      'Invalid expression for aggregator operator.',
+    ],
+    ['{ 1 + 1 = 2 + }', 'Invalid value after comparison or aritmetic operator.'],
+    ['{ .a && }', 'Invalid value after logical operator.'],
+    ['{ .a || }', 'Invalid value after logical operator.'],
+    ['{ .a + }', 'Invalid value after comparison or aritmetic operator.'],
+    ['{ 200 = 200 200 }', 'Invalid comparison operator after field expression.'],
+    ['{.foo   300}', 'Invalid comparison operator after field expression.'],
+    ['{.foo  300 && .bar = 200}', 'Invalid operator after field expression.'],
+    ['{.foo  300 && .bar  200}', 'Invalid operator after field expression.'],
+    ['{.foo=1}  {.bar=2}', 'Invalid spanset combining operator after spanset expression.'],
+    ['{ span.http.status_code = 200 &&  }', 'Invalid value after logical operator.'],
+    ['{ span.http.status_code = 200 ||  }', 'Invalid value after logical operator.'],
+    ['{ .foo = 200 } && ', 'Invalid spanset expression after spanset combining operator.'],
+    ['{ .foo = 200 } || ', 'Invalid spanset expression after spanset combining operator.'],
+    ['{ .foo = 200 } >> ', 'Invalid spanset expression after spanset combining operator.'],
+    ['{.foo=1} | avg()', 'Invalid expression for aggregator operator.'],
+    ['{.foo=1} | avg(.foo) > ', 'Invalid value after comparison operator.'],
+    ['{.foo=1} | avg() < 1s', 'Invalid expression for aggregator operator.'],
+    ['{.foo=1} | max() = 3', 'Invalid expression for aggregator operator.'],
+    ['{.foo=1} | by()', 'Invalid expression for aggregator operator.'],
+    ['{.foo=1} | select()', 'Invalid expression for aggregator operator.'],
+    ['{foo}', 'Invalid expression for spanset.'],
+    ['{.}', 'Invalid expression for spanset.'],
+    ['{ resource. }', 'Invalid expression for spanset.'],
+    ['{ span. }', 'Invalid expression for spanset.'],
+    ['{.foo=}', 'Invalid value after comparison or aritmetic operator.'],
+    ['{.foo="}', 'Invalid value after comparison or aritmetic operator.'],
+    ['{.foo=300} |', 'Invalid aggregation operator after pipepile operator.'],
+    ['{.foo=300} && {.bar=200} |', 'Invalid aggregation operator after pipepile operator.'],
+    ['{.foo=300} && {.bar=300} && {.foo=300} |', 'Invalid aggregation operator after pipepile operator.'],
+    ['{.foo=300} | avg(.value)', 'Invalid comparison operator after aggregator operator.'],
+    ['{.foo=300} && {.foo=300} | avg(.value)', 'Invalid comparison operator after aggregator operator.'],
+    ['{.foo=300} | avg(.value) =', 'Invalid value after comparison operator.'],
+    ['{.foo=300} && {.foo=300} | avg(.value) =', 'Invalid value after comparison operator.'],
+    ['{.foo=300} | max(duration) > 1hs', 'Invalid value after comparison operator.'],
+    ['{ span.http.status_code', 'Invalid comparison operator after field expression.'],
+  ])('error message for invalid query - %s, %s', (query: string, expectedErrorMessage: string) => {
+    const errorNode = getErrorNodes(query)[0];
+    expect(computeErrorMessage(errorNode)).toBe(expectedErrorMessage);
+  });
+});

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

@@ -11,6 +11,7 @@ import { dispatch } from '../../../../store/store';
 import { TempoDatasource } from '../datasource';
 
 import { CompletionProvider, CompletionType } from './autocomplete';
+import { getErrorNodes, setErrorMarkers } from './errorHighlighting';
 import { languageDefinition } from './traceql';
 
 interface Props {
@@ -32,6 +33,8 @@ export function TraceQLEditor(props: Props) {
   const onRunQueryRef = useRef(onRunQuery);
   onRunQueryRef.current = onRunQuery;
 
+  const errorTimeoutId = useRef<number>();
+
   return (
     <CodeEditor
       value={props.value}
@@ -63,6 +66,33 @@ export function TraceQLEditor(props: Props) {
           setupPlaceholder(editor, monaco, styles);
         }
         setupAutoSize(editor);
+
+        // Register callback for query changes
+        editor.onDidChangeModelContent((changeEvent) => {
+          const model = editor.getModel();
+          if (!model) {
+            return;
+          }
+
+          // Remove previous callback if existing, to prevent squiggles from been shown while the user is still typing
+          window.clearTimeout(errorTimeoutId.current);
+
+          const errorNodes = getErrorNodes(model.getValue());
+          const cursorPosition = changeEvent.changes[0].rangeOffset;
+
+          // Immediately updates the squiggles, in case the user fixed an error,
+          // excluding the error around the cursor position
+          setErrorMarkers(
+            monaco,
+            model,
+            errorNodes.filter((errorNode) => !(errorNode.from <= cursorPosition && cursorPosition <= errorNode.to))
+          );
+
+          // Later on, show all errors
+          errorTimeoutId.current = window.setTimeout(() => {
+            setErrorMarkers(monaco, model, errorNodes);
+          }, 500);
+        });
       }}
     />
   );

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

@@ -0,0 +1,123 @@
+import { SyntaxNode } from '@lezer/common';
+
+import {
+  Aggregate,
+  And,
+  AttributeField,
+  ComparisonOp,
+  FieldExpression,
+  FieldOp,
+  IntrinsicField,
+  Or,
+  parser,
+  Pipe,
+  ScalarExpression,
+  ScalarFilter,
+  SpansetFilter,
+  SpansetPipelineExpression,
+} from '@grafana/lezer-traceql';
+import { monacoTypes } from '@grafana/ui';
+
+/**
+ * Given an error node, generate an error message to be displayed to the user.
+ *
+ * @param errorNode the error node, as returned by the TraceQL Lezer parser
+ * @returns the error message
+ */
+export const computeErrorMessage = (errorNode: SyntaxNode) => {
+  switch (errorNode.parent?.type.id) {
+    case FieldExpression:
+      switch (errorNode.prevSibling?.type.id) {
+        case And:
+        case Or:
+          return 'Invalid value after logical operator.';
+        case FieldOp:
+          return 'Invalid value after comparison or aritmetic operator.';
+        default:
+          return 'Invalid operator after field expression.';
+      }
+    case SpansetFilter:
+      if (errorNode.prevSibling?.type.id === FieldExpression) {
+        return 'Invalid comparison operator after field expression.';
+      }
+      return 'Invalid expression for spanset.';
+    case SpansetPipelineExpression:
+      switch (errorNode.prevSibling?.type.id) {
+        case SpansetPipelineExpression:
+          return 'Invalid spanset combining operator after spanset expression.';
+        case Pipe:
+          return 'Invalid aggregation operator after pipepile operator.';
+        default:
+          return 'Invalid spanset expression after spanset combining operator.';
+      }
+    case IntrinsicField:
+    case Aggregate:
+      return 'Invalid expression for aggregator operator.';
+    case AttributeField:
+      return 'Invalid expression for spanset.';
+    case ScalarFilter:
+      switch (errorNode.prevSibling?.type.id) {
+        case ComparisonOp:
+          return 'Invalid value after comparison operator.';
+        case ScalarExpression:
+          if (errorNode.prevSibling?.firstChild?.type.id === Aggregate) {
+            return 'Invalid comparison operator after aggregator operator.';
+          }
+        default:
+          return 'Invalid value after comparison operator.';
+      }
+    default:
+      return 'Invalid query.';
+  }
+};
+
+/**
+ * Parse the given query and find the error nodes, if any, in the resulting tree.
+ *
+ * @param query the TraceQL query of the user
+ * @returns the error nodes
+ */
+export const getErrorNodes = (query: string): SyntaxNode[] => {
+  const tree = parser.parse(query);
+
+  // Find all error nodes and compute the associated erro boundaries
+  const errorNodes: SyntaxNode[] = [];
+  tree.iterate({
+    enter: (nodeRef) => {
+      if (nodeRef.type.id === 0) {
+        errorNodes.push(nodeRef.node);
+      }
+    },
+  });
+
+  return errorNodes;
+};
+
+/**
+ * Use red markers (squiggles) to highlight syntax errors in queries.
+ *
+ */
+export const setErrorMarkers = (
+  monaco: typeof monacoTypes,
+  model: monacoTypes.editor.ITextModel,
+  errorNodes: SyntaxNode[]
+) => {
+  monaco.editor.setModelMarkers(
+    model,
+    'owner', // default value
+    errorNodes.map((errorNode) => {
+      return {
+        message: computeErrorMessage(errorNode),
+        severity: monaco.MarkerSeverity.Error,
+
+        // As of now, we support only single-line queries
+        startLineNumber: 0,
+        endLineNumber: 0,
+
+        // `+ 1` because squiggles seem shifted by one
+        startColumn: errorNode.from + 1,
+        endColumn: errorNode.to + 1,
+      };
+    })
+  );
+};