From 5eca966f4523dd5c8509a17876d5237292990f2a Mon Sep 17 00:00:00 2001
From: Hamas Shafiq <hamas.shafiq@grafana.com>
Date: Mon, 23 Jan 2023 15:39:56 +0000
Subject: [PATCH] Tempo: [TraceQL] Integrate the new tag values API (#61637)

* Tempo: [TraceQL] Integrate the new tag values API

* Use the exact string entered in the spanset

* Update a test & add a new test

* Empty commit
---
 .../datasource/tempo/language_provider.ts     | 11 +++----
 .../tempo/traceql/autocomplete.test.ts        | 30 +++++++++++++++++--
 .../datasource/tempo/traceql/autocomplete.ts  | 12 ++------
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/public/app/plugins/datasource/tempo/language_provider.ts b/public/app/plugins/datasource/tempo/language_provider.ts
index 5394584c110..117118e6e4f 100644
--- a/public/app/plugins/datasource/tempo/language_provider.ts
+++ b/public/app/plugins/datasource/tempo/language_provider.ts
@@ -74,7 +74,7 @@ export default class TempoLanguageProvider extends LanguageProvider {
     let tagName = tags[tags.length - 1] ?? '';
     tagName = tagName.split('=')[0];
 
-    const response = await this.request(`/api/search/tag/${tagName}/values`, []);
+    const response = await this.request(`/api/v2/search/tag/${tagName}/values`, []);
 
     const suggestions: CompletionItemGroup[] = [];
 
@@ -88,13 +88,14 @@ export default class TempoLanguageProvider extends LanguageProvider {
   }
 
   async getOptions(tag: string): Promise<Array<SelectableValue<string>>> {
-    const response = await this.request(`/api/search/tag/${tag}/values`);
+    const response = await this.request(`/api/v2/search/tag/${tag}/values`);
     let options: Array<SelectableValue<string>> = [];
 
     if (response && response.tagValues) {
-      options = response.tagValues.map((v: string) => ({
-        value: v,
-        label: v,
+      options = response.tagValues.map((v: { type: string; value: string }) => ({
+        type: v.type,
+        value: v.value,
+        label: v.value,
       }));
     }
 
diff --git a/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts b/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
index 766fa233f2f..516aed301de 100644
--- a/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
+++ b/public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts
@@ -26,7 +26,7 @@ describe('CompletionProvider', () => {
     ]);
   });
 
-  it('suggests tag names with quotes', async () => {
+  it('does not wrap the tag value in quotes if the type in the response is something other than "string"', async () => {
     const { provider, model } = setup('{foo=}', 5, defaultTags);
 
     jest.spyOn(provider.languageProvider, 'getOptions').mockImplementation(
@@ -34,6 +34,32 @@ describe('CompletionProvider', () => {
         new Promise((resolve) => {
           resolve([
             {
+              type: 'int',
+              value: 'foobar',
+              label: 'foobar',
+            },
+          ]);
+        })
+    );
+
+    const result = await provider.provideCompletionItems(
+      model as unknown as monacoTypes.editor.ITextModel,
+      {} as monacoTypes.Position
+    );
+    expect((result! as monacoTypes.languages.CompletionList).suggestions).toEqual([
+      expect.objectContaining({ label: 'foobar', insertText: 'foobar' }),
+    ]);
+  });
+
+  it('wraps the tag value in quotes if the type in the response is set to "string"', async () => {
+    const { provider, model } = setup('{foo=}', 5, defaultTags);
+
+    jest.spyOn(provider.languageProvider, 'getOptions').mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          resolve([
+            {
+              type: 'string',
               value: 'foobar',
               label: 'foobar',
             },
@@ -50,7 +76,7 @@ describe('CompletionProvider', () => {
     ]);
   });
 
-  it('suggests tag names without quotes', async () => {
+  it('inserts the tag value without quotes if the user has entered quotes', async () => {
     const { provider, model } = setup('{foo="}', 6, defaultTags);
 
     jest.spyOn(provider.languageProvider, 'getOptions').mockImplementation(
diff --git a/public/app/plugins/datasource/tempo/traceql/autocomplete.ts b/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
index 889098bc52d..3ba40c6b7cd 100644
--- a/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
+++ b/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
@@ -149,7 +149,6 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
         }));
       case 'SPANSET_IN_VALUE':
         const tagName = this.overrideTagName(situation.tagName);
-        const tagsNoQuotesAroundValue: string[] = ['status'];
         const tagValues = await this.getTagValues(tagName);
         const items: Completion[] = [];
 
@@ -157,7 +156,7 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
           if (situation.betweenQuotes) {
             return val.label!;
           }
-          return tagsNoQuotesAroundValue.includes(situation.tagName) ? val.label! : `"${val.label}"`;
+          return val.type === 'string' ? `"${val.label}"` : val.label!;
         };
 
         tagValues.forEach((val) => {
@@ -275,18 +274,11 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
         };
       }
 
-      // remove the scopes from the word to get accurate autocompletes
-      // Ex: 'span.host.name' won't resolve to any autocomplete values, but removing 'span.' results in 'host.name' which can have autocomplete values
-      const noScopeWord = CompletionProvider.scopes.reduce(
-        (result, word) => result.replace(`${word}.`, ''),
-        nameMatched?.groups?.word || ''
-      );
-
       // We already have an operator and know that the set isn't complete so let's autocomplete the possible values for the tag name
       // { .http.method = |
       return {
         type: 'SPANSET_IN_VALUE',
-        tagName: noScopeWord,
+        tagName: nameFull,
         betweenQuotes: !!matched.groups?.open_quote,
       };
     }