From 3f6f5d3357e232abe53f1775f893fdf8f842712c Mon Sep 17 00:00:00 2001
From: Shan Aminzadeh <shan.aminzadeh@chronosphere.io>
Date: Fri, 17 Feb 2023 04:23:16 -0500
Subject: [PATCH] Scope GroupBy labels to metric (#11914)

Signed-off-by: Shan Aminzadeh <shan.aminzadeh@chronosphere.io>
---
 .../src/complete/hybrid.test.ts               | 16 ++++++++++--
 .../codemirror-promql/src/complete/hybrid.ts  | 26 ++++++++++++++++---
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/web/ui/module/codemirror-promql/src/complete/hybrid.test.ts b/web/ui/module/codemirror-promql/src/complete/hybrid.test.ts
index a09a118a1..0f1a8b80a 100644
--- a/web/ui/module/codemirror-promql/src/complete/hybrid.test.ts
+++ b/web/ui/module/codemirror-promql/src/complete/hybrid.test.ts
@@ -231,7 +231,19 @@ describe('analyzeCompletion test', () => {
       title: 'starting to autocomplete labelName in aggregate modifier',
       expr: 'sum by ()',
       pos: 8, // cursor is between the bracket
-      expectedContext: [{ kind: ContextKind.LabelName }],
+      expectedContext: [{ kind: ContextKind.LabelName, metricName: '' }],
+    },
+    {
+      title: 'starting to autocomplete labelName in aggregate modifier with metric name',
+      expr: 'sum(up) by ()',
+      pos: 12, // cursor is between ()
+      expectedContext: [{ kind: ContextKind.LabelName, metricName: 'up' }],
+    },
+    {
+      title: 'starting to autocomplete labelName in aggregate modifier with metric name in front',
+      expr: 'sum by ()(up)',
+      pos: 8, // cursor is between ()
+      expectedContext: [{ kind: ContextKind.LabelName, metricName: 'up' }],
     },
     {
       title: 'continue to autocomplete labelName in aggregate modifier',
@@ -243,7 +255,7 @@ describe('analyzeCompletion test', () => {
       title: 'autocomplete labelName in a list',
       expr: 'sum by (myLabel1,)',
       pos: 17, // cursor is between the bracket after the string myLab
-      expectedContext: [{ kind: ContextKind.LabelName }],
+      expectedContext: [{ kind: ContextKind.LabelName, metricName: '' }],
     },
     {
       title: 'autocomplete labelName in a list 2',
diff --git a/web/ui/module/codemirror-promql/src/complete/hybrid.ts b/web/ui/module/codemirror-promql/src/complete/hybrid.ts
index e3c02f890..cf23aa11a 100644
--- a/web/ui/module/codemirror-promql/src/complete/hybrid.ts
+++ b/web/ui/module/codemirror-promql/src/complete/hybrid.ts
@@ -110,6 +110,26 @@ export interface Context {
   matchers?: Matcher[];
 }
 
+function getMetricNameInGroupBy(tree: SyntaxNode, state: EditorState): string {
+  // There should be an AggregateExpr as parent of the GroupingLabels.
+  // Then we should find the VectorSelector child to be able to find the metric name.
+  const currentNode: SyntaxNode | null = walkBackward(tree, AggregateExpr);
+  if (!currentNode) {
+    return '';
+  }
+  let metricName = '';
+  currentNode.cursor().iterate((node) => {
+    // Continue until we find the VectorSelector, then look up the metric name.
+    if (node.type.id === VectorSelector) {
+      metricName = getMetricNameInVectorSelector(node.node, state);
+      if (metricName) {
+        return false;
+      }
+    }
+  });
+  return metricName;
+}
+
 function getMetricNameInVectorSelector(tree: SyntaxNode, state: EditorState): string {
   // Find if there is a defined metric name. Should be used to autocomplete a labelValue or a labelName
   // First find the parent "VectorSelector" to be able to find then the subChild "Identifier" if it exists.
@@ -344,9 +364,9 @@ export function analyzeCompletion(state: EditorState, node: SyntaxNode): Context
       break;
     case GroupingLabels:
       // In this case we are in the given situation:
-      //      sum by ()
-      // So we have to autocomplete any labelName
-      result.push({ kind: ContextKind.LabelName });
+      //      sum by () or sum (metric_name) by ()
+      // so we have or to autocomplete any kind of labelName or to autocomplete only the labelName associated to the metric
+      result.push({ kind: ContextKind.LabelName, metricName: getMetricNameInGroupBy(node, state) });
       break;
     case LabelMatchers:
       // In that case we are in the given situation: