UI/PromQL: autocomplete topk like aggregation function parameters

Signed-off-by: Augustin Husson <husson.augustin@gmail.com>

web/ui/module/codemirror-promql/src/complete/hybrid.test.ts

@@ -583,12 +583,30 @@ describe('analyzeCompletion test', () => {
       pos: 5,
       expectedContext: [{ kind: ContextKind.AtModifiers }],
     },
+    {
+      title: 'autocomplete topk params',
+      expr: 'topk()',
+      pos: 5,
+      expectedContext: [{ kind: ContextKind.Number }],
+    },
+    {
+      title: 'autocomplete topk params 2',
+      expr: 'topk(inf,)',
+      pos: 9,
+      expectedContext: [{ kind: ContextKind.MetricName, metricName: '' }, { kind: ContextKind.Function }, { kind: ContextKind.Aggregation }],
+    },
+    {
+      title: 'autocomplete topk params 3',
+      expr: 'topk(inf,r)',
+      pos: 10,
+      expectedContext: [{ kind: ContextKind.MetricName, metricName: 'r' }, { kind: ContextKind.Function }, { kind: ContextKind.Aggregation }],
+    },
   ];
   testCases.forEach((value) => {
     it(value.title, () => {
       const state = createEditorState(value.expr);
       const node = syntaxTree(state).resolve(value.pos, -1);
-      const result = analyzeCompletion(state, node);
+      const result = analyzeCompletion(state, node, value.pos);
       expect(value.expectedContext).toEqual(result);
     });
   });

web/ui/module/codemirror-promql/src/complete/hybrid.ts

@@ -54,6 +54,12 @@ import {
   QuotedLabelName,
   NumberDurationLiteralInDurationContext,
   NumberDurationLiteral,
+  AggregateOp,
+  Topk,
+  Bottomk,
+  LimitK,
+  LimitRatio,
+  CountValues,
 } from '@prometheus-io/lezer-promql';
 import { Completion, CompletionContext, CompletionResult } from '@codemirror/autocomplete';
 import { EditorState } from '@codemirror/state';
@@ -185,7 +191,7 @@ export function computeStartCompletePosition(state: EditorState, node: SyntaxNod
   if (node.type.id === LabelMatchers || node.type.id === GroupingLabels) {
     start = computeStartCompleteLabelPositionInLabelMatcherOrInGroupingLabel(node, pos);
   } else if (
-    node.type.id === FunctionCallBody ||
+    (node.type.id === FunctionCallBody && node.firstChild === null) ||
     (node.type.id === StringLiteral && (node.parent?.type.id === UnquotedLabelMatcher || node.parent?.type.id === QuotedLabelMatcher))
   ) {
     // When the cursor is between bracket, quote, we need to increment the starting position to avoid to consider the open bracket/ first string.
@@ -198,6 +204,7 @@ export function computeStartCompletePosition(state: EditorState, node: SyntaxNod
     // So we have to analyze the string about the current node to see if the duration unit is already present or not.
     (node.type.id === NumberDurationLiteralInDurationContext && !durationTerms.map((v) => v.label).includes(currentText[currentText.length - 1])) ||
     (node.type.id === NumberDurationLiteral && node.parent?.type.id === 0 && node.parent.parent?.type.id === SubqueryExpr) ||
+    (node.type.id === FunctionCallBody && isItATopKLikeAggregationFunc(node) && node.firstChild !== null) ||
     (node.type.id === 0 &&
       (node.parent?.type.id === OffsetExpr ||
         node.parent?.type.id === MatrixSelector ||
@@ -208,10 +215,28 @@ export function computeStartCompletePosition(state: EditorState, node: SyntaxNod
   return start;
 }
 
+function isItATopKLikeAggregationFunc(functionCallBody: SyntaxNode): boolean {
+  const prevSibling = functionCallBody.prevSibling;
+  if (prevSibling !== null && prevSibling.type.id === AggregateOp) {
+    const aggregationOpType = prevSibling.firstChild;
+    if (
+      aggregationOpType !== null &&
+      (aggregationOpType.type.id == Topk ||
+        aggregationOpType.type.id === Bottomk ||
+        aggregationOpType.type.id === LimitK ||
+        aggregationOpType.type.id === LimitRatio ||
+        aggregationOpType.type.id === CountValues)
+    ) {
+      return true;
+    }
+  }
+  return false;
+}
+
 // analyzeCompletion is going to determinate what should be autocompleted.
 // The value of the autocompletion is then calculate by the function buildCompletion.
 // Note: this method is exported for testing purpose only. Do not use it directly.
-export function analyzeCompletion(state: EditorState, node: SyntaxNode): Context[] {
+export function analyzeCompletion(state: EditorState, node: SyntaxNode, pos: number): Context[] {
   const result: Context[] = [];
   switch (node.type.id) {
     case 0: // 0 is the id of the error node
@@ -330,7 +355,7 @@ export function analyzeCompletion(state: EditorState, node: SyntaxNode): Context
       }
       // now we have to know if we have two Expr in the direct children of the `parent`
       const containExprTwice = containsChild(parent, 'Expr', 'Expr');
-      if (containExprTwice) {
+      if (containExprTwice && parent.type.id !== FunctionCallBody) {
         if (parent.type.id === BinaryExpr && !containsAtLeastOneChild(parent, 0)) {
           // We are likely in the case 1 or 5
           result.push(
@@ -460,7 +485,23 @@ export function analyzeCompletion(state: EditorState, node: SyntaxNode): Context
       result.push({ kind: ContextKind.Duration });
       break;
     case FunctionCallBody:
-      // In this case we are in the given situation:
+      // For aggregation function such as Topk, the first parameter is a number.
+      // The second one is an expression.
+      // When moving to the second parameter, the node is an error node.
+      // Unfortunately, as a current node, codemirror doesn't give us the error node but instead the FunctionCallBody
+      // The tree looks like that: PromQL(AggregateExpr(AggregateOp(Topk),FunctionCallBody(NumberDurationLiteral,⚠)))
+      // So, we need to figure out if the cursor is on the first parameter or in the second.
+      if (isItATopKLikeAggregationFunc(node)) {
+        if (node.firstChild === null || (node.firstChild.from <= pos && node.firstChild.to >= pos)) {
+          // it means the FunctionCallBody has no child, which means we are autocompleting the first parameter
+          result.push({ kind: ContextKind.Number });
+          break;
+        }
+        // at this point we are necessary autocompleting the second parameter
+        result.push({ kind: ContextKind.MetricName, metricName: '' }, { kind: ContextKind.Function }, { kind: ContextKind.Aggregation });
+        break;
+      }
+      // In all other cases, we are in the given situation:
       //       sum() or in rate()
       // with the cursor between the bracket. So we can autocomplete the metric, the function and the aggregation.
       result.push({ kind: ContextKind.MetricName, metricName: '' }, { kind: ContextKind.Function }, { kind: ContextKind.Aggregation });
@@ -516,7 +557,11 @@ export class HybridComplete implements CompleteStrategy {
   promQL(context: CompletionContext): Promise<CompletionResult | null> | CompletionResult | null {
     const { state, pos } = context;
     const tree = syntaxTree(state).resolve(pos, -1);
-    const contexts = analyzeCompletion(state, tree);
+    // The lines above can help you to print the current lezer tree.
+    // It's useful when you are trying to understand why it doesn't autocomplete.
+    // console.log(syntaxTree(state).topNode.toString());
+    // console.log(`current node: ${tree.type.name}`);
+    const contexts = analyzeCompletion(state, tree, pos);
     let asyncResult: Promise<Completion[]> = Promise.resolve([]);
     let completeSnippet = false;
     let span = true;