From 3783d87576d4fae8345b1d043a8ee913e10a2753 Mon Sep 17 00:00:00 2001
From: Fabrizio <135109076+fabrizio-grafana@users.noreply.github.com>
Date: Tue, 12 Dec 2023 12:55:21 +0100
Subject: [PATCH] Tempo: Fix autocompletion with strings (#79370)

---
 .../tempo/traceql/autocomplete.test.ts        | 22 ++++++++
 .../datasource/tempo/traceql/situation.ts     | 51 +++++++++++--------
 2 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts b/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
index 347027f5f35..ad15cc0c95b 100644
--- a/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
+++ b/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
@@ -105,6 +105,28 @@ describe('CompletionProvider', () => {
     ]);
   });
 
+  it('suggests options when inside quotes', async () => {
+    const { provider, model } = setup('{.foo=""}', 7, undefined, v2Tags);
+
+    jest.spyOn(provider.languageProvider, 'getOptionsV2').mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          resolve([
+            {
+              type: 'string',
+              value: 'foobar',
+              label: 'foobar',
+            },
+          ]);
+        })
+    );
+
+    const result = await provider.provideCompletionItems(model, emptyPosition);
+    expect((result! as monacoTypes.languages.CompletionList).suggestions).toEqual([
+      expect.objectContaining({ label: 'foobar', insertText: 'foobar' }),
+    ]);
+  });
+
   it('suggests nothing without tags', async () => {
     const { provider, model } = setup('{.foo="}', 8, emptyTags);
     const result = await provider.provideCompletionItems(model, emptyPosition);
diff --git a/public/app/plugins/datasource/tempo/traceql/situation.ts b/public/app/plugins/datasource/tempo/traceql/situation.ts
index 9102bd742a0..6d8681cd6df 100644
--- a/public/app/plugins/datasource/tempo/traceql/situation.ts
+++ b/public/app/plugins/datasource/tempo/traceql/situation.ts
@@ -18,6 +18,7 @@ import {
   SpansetPipeline,
   SpansetPipelineExpression,
   Static,
+  String as StringNode,
   TraceQL,
 } from '@grafana/lezer-traceql';
 
@@ -86,7 +87,7 @@ type Path = Array<[Direction, NodeType[]]>;
 
 type Resolver = {
   path: NodeType[];
-  fun: (node: SyntaxNode, text: string, pos: number, originalPos: number) => SituationType | null;
+  fun: (node: SyntaxNode, text: string, pos: number, originalPos: number) => SituationType | void;
 };
 
 function getErrorNode(tree: Tree, cursorPos: number): SyntaxNode | null {
@@ -178,7 +179,7 @@ export function getSituation(text: string, offset: number): Situation | null {
     ids.push(cur.type.id);
   }
 
-  let situationType: SituationType | null = null;
+  let situationType: SituationType | void = undefined;
   for (let resolver of RESOLVERS) {
     if (isPathMatch(resolver.path, ids)) {
       situationType = resolver.fun(currentNode, text, shiftedOffset, offset);
@@ -220,14 +221,6 @@ const RESOLVERS: Resolver[] = [
     path: [ERROR_NODE_ID, ScalarFilter, SpansetPipeline],
     fun: resolveArithmeticOperator,
   },
-  {
-    path: [ERROR_NODE_ID, TraceQL],
-    fun: () => {
-      return {
-        type: 'UNKNOWN',
-      };
-    },
-  },
   // Curson on valid node cases (the whole query could contain errors nevertheless)
   {
     path: [FieldExpression],
@@ -245,6 +238,10 @@ const RESOLVERS: Resolver[] = [
     path: [TraceQL],
     fun: resolveNewSpansetExpression,
   },
+  {
+    path: [StringNode, Static],
+    fun: resolveExpression,
+  },
 ];
 
 const resolveAttributeCompletion = (node: SyntaxNode, text: string, pos: number): SituationType | void => {
@@ -353,8 +350,25 @@ function resolveExpression(node: SyntaxNode, text: string, _: number, originalPo
     return situation;
   }
 
+  if (
+    walk(node, [
+      ['parent', [Static]],
+      ['parent', [FieldExpression]],
+      ['prevSibling', [FieldOp]],
+    ])
+  ) {
+    let attributeField = node.parent?.parent?.prevSibling?.prevSibling;
+    if (attributeField) {
+      return {
+        type: 'SPANSET_IN_VALUE',
+        tagName: getNodeText(attributeField, text),
+        betweenQuotes: true,
+      };
+    }
+  }
+
   if (node.prevSibling?.type.id === FieldOp) {
-    let attributeField = node.prevSibling.prevSibling;
+    let attributeField = node.prevSibling?.prevSibling;
     if (attributeField) {
       return {
         type: 'SPANSET_IN_VALUE',
@@ -375,16 +389,12 @@ function resolveExpression(node: SyntaxNode, text: string, _: number, originalPo
   };
 }
 
-function resolveArithmeticOperator(node: SyntaxNode, _0: string, _1: number): SituationType {
-  if (node.prevSibling?.type.id === ComparisonOp) {
+function resolveArithmeticOperator(node: SyntaxNode, _0: string, _1: number): SituationType | void {
+  if (node.prevSibling?.type.id !== ComparisonOp) {
     return {
-      type: 'UNKNOWN',
+      type: 'SPANSET_COMPARISON_OPERATORS',
     };
   }
-
-  return {
-    type: 'SPANSET_COMPARISON_OPERATORS',
-  };
 }
 
 function resolveNewSpansetExpression(node: SyntaxNode, text: string, offset: number): SituationType {
@@ -410,16 +420,13 @@ function resolveNewSpansetExpression(node: SyntaxNode, text: string, offset: num
   };
 }
 
-function resolveAttributeForFunction(node: SyntaxNode, _0: string, _1: number): SituationType {
+function resolveAttributeForFunction(node: SyntaxNode, _0: string, _1: number): SituationType | void {
   const parent = node?.parent;
   if (!!parent && [IntrinsicField, Aggregate, GroupOperation, SelectArgs].includes(parent.type.id)) {
     return {
       type: 'ATTRIBUTE_FOR_FUNCTION',
     };
   }
-  return {
-    type: 'UNKNOWN',
-  };
 }
 
 function resolveSpansetPipeline(node: SyntaxNode, _1: string, _2: number): SituationType {