From c6840cfdec6c61d52d3ed613ec9eb8764aa1f4be Mon Sep 17 00:00:00 2001
From: Andrej Ocenas <mr.ocenas@gmail.com>
Date: Mon, 5 Dec 2022 17:05:22 +0100
Subject: [PATCH] Tempo: Fix traceQL suggestions for tag names with dots
 (#59736)

---
 .../datasource/tempo/traceql/autocomplete.ts  | 58 ++++++++++++++++---
 1 file changed, 51 insertions(+), 7 deletions(-)

diff --git a/public/app/plugins/datasource/tempo/traceql/autocomplete.ts b/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
index 96ec31a79d7..374ee42541e 100644
--- a/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
+++ b/public/app/plugins/datasource/tempo/traceql/autocomplete.ts
@@ -56,13 +56,17 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
       // to stop it, we use a number-as-string sortkey,
       // so that monaco keeps the order we use
       const maxIndexDigits = items.length.toString().length;
-      const suggestions: monacoTypes.languages.CompletionItem[] = items.map((item, index) => ({
-        kind: getMonacoCompletionItemKind(item.type, this.monaco!),
-        label: item.label,
-        insertText: item.insertText,
-        sortText: index.toString().padStart(maxIndexDigits, '0'), // to force the order we have
-        range,
-      }));
+      const suggestions: monacoTypes.languages.CompletionItem[] = items.map((item, index) => {
+        const suggestion = {
+          kind: getMonacoCompletionItemKind(item.type, this.monaco!),
+          label: item.label,
+          insertText: item.insertText,
+          sortText: index.toString().padStart(maxIndexDigits, '0'), // to force the order we have
+          range,
+        };
+        fixSuggestion(suggestion, item.type, model, offset, this.monaco!);
+        return suggestion;
+      });
       return { suggestions };
     });
   }
@@ -379,3 +383,43 @@ function getRangeAndOffset(monaco: Monaco, model: monacoTypes.editor.ITextModel,
   const offset = model.getOffsetAt(positionClone);
   return { offset, range };
 }
+
+/**
+ * Fix the suggestions range and insert text. For the range we have to adjust because monaco by default replaces just
+ * the last word which stops at dot while traceQL tags contain dots themselves and we want to replace the whole tag
+ * name when suggesting. The insert text needs to be adjusted for scope (leading dot) if scope is currently missing.
+ * This may be doable also when creating the suggestions but for a particular situation this seems to be easier to do
+ * here.
+ */
+function fixSuggestion(
+  suggestion: monacoTypes.languages.CompletionItem & { range: monacoTypes.IRange },
+  itemType: CompletionType,
+  model: monacoTypes.editor.ITextModel,
+  offset: number,
+  monaco: Monaco
+) {
+  if (itemType === 'TAG_NAME') {
+    const match = model
+      .getValue()
+      .substring(0, offset)
+      .match(/(span\.|resource\.|\.)?([\w./-]*)$/);
+
+    if (match) {
+      const scope = match[1];
+      const tag = match[2];
+
+      if (tag) {
+        // Add the default scope if needed.
+        if (!scope && suggestion.insertText[0] !== '.') {
+          suggestion.insertText = '.' + suggestion.insertText;
+        }
+
+        // Adjust the range, so that we will replace the whole tag.
+        suggestion.range = monaco.Range.lift({
+          ...suggestion.range,
+          startColumn: offset - tag.length + 1,
+        });
+      }
+    }
+  }
+}