Prometheus: Preserve custom variables as function parameters after parsing the expression (#106661)

* preserve custom variables as function parameters after parsing the expression

* replaced variable check

packages/grafana-prometheus/src/querybuilder/parsing.test.ts

@@ -948,6 +948,25 @@ describe('buildVisualQueryFromString', () => {
       })
     );
   });
+
+  it('parses query with custom variable', () => {
+    expect(buildVisualQueryFromString('topk($custom, rate(metric_name[$__rate_interval]))')).toEqual(
+      noErrors({
+        metric: 'metric_name',
+        labels: [],
+        operations: [
+          {
+            id: 'rate',
+            params: ['$__rate_interval'],
+          },
+          {
+            id: 'topk',
+            params: ['$custom'],
+          },
+        ],
+      })
+    );
+  });
 });
 
 function noErrors(query: PromVisualQuery) {

packages/grafana-prometheus/src/querybuilder/parsing.ts

@@ -49,9 +49,9 @@ import { PromVisualQuery, PromVisualQueryBinary } from './types';
  * It traverses the tree and uses sort of state machine to update the query model.
  * The query model is modified during the traversal and sent to each handler as context.
  */
-export function buildVisualQueryFromString(expr: string): Context {
+export function buildVisualQueryFromString(expr: string): Omit<Context, 'replacements'> {
   expr = replaceBuiltInVariable(expr);
-  const replacedExpr = replaceVariables(expr);
+  const { replacedExpr, replacedVariables } = replaceVariables(expr);
   const tree = parser.parse(replacedExpr);
   const node = tree.topNode;
 
@@ -64,6 +64,7 @@ export function buildVisualQueryFromString(expr: string): Context {
   const context: Context = {
     query: visQuery,
     errors: [],
+    replacements: replacedVariables,
   };
 
   try {
@@ -83,6 +84,9 @@ export function buildVisualQueryFromString(expr: string): Context {
     context.errors = [];
   }
 
+  // No need to return replaced variables
+  delete context.replacements;
+
   return context;
 }
 
@@ -96,6 +100,7 @@ interface ParsingError {
 interface Context {
   query: PromVisualQuery;
   errors: ParsingError[];
+  replacements?: Record<string, string>;
 }
 
 /**
@@ -359,6 +364,19 @@ function updateFunctionArgs(expr: string, node: SyntaxNode | null, context: Cont
       break;
     }
 
+    case VectorSelector: {
+      // When we replace a custom variable to prevent errors during parsing we receive VectorSelector and Identifier in it.
+      // But this is also a normal case for a normal function body. i.e. topk(5, http_requests_total{})
+      // In such cases we got identifier as http_requests_total. So we shouldn't push this as param.
+      // So we check whether the given VectorSelector is something we replaced earlier.
+      if (context.replacements?.[expr.substring(node.from, node.to)]) {
+        const identifierNode = node.getChild(Identifier);
+        const customVarName = getString(expr, identifierNode);
+        op.params.push(customVarName);
+        break;
+      }
+    }
+
     default: {
       // Means we get to something that does not seem like simple function arg and is probably nested query so jump
       // back to main context
@@ -429,6 +447,7 @@ function handleBinary(expr: string, node: SyntaxNode, context: Context) {
     handleExpression(expr, right, {
       query: binQuery.query,
       errors: context.errors,
+      replacements: context.replacements,
     });
   }
 }

packages/grafana-prometheus/src/querybuilder/parsingUtils.test.ts

@@ -20,9 +20,15 @@ describe('getLeftMostChild', () => {
 
 describe('replaceVariables', () => {
   it('should replace variables', () => {
-    expect(replaceVariables('sum_over_time([[metric_var]]{bar="${app}"}[$__interval])')).toBe(
-      'sum_over_time(__V_1__metric_var__V__{bar="__V_2__app__V__"}[__V_0____interval__V__])'
+    const { replacedExpr, replacedVariables } = replaceVariables(
+      'sum_over_time([[metric_var]]{bar="${app}"}[$__interval])'
     );
+    expect(replacedExpr).toBe('sum_over_time(__V_1__metric_var__V__{bar="__V_2__app__V__"}[__V_0____interval__V__])');
+    expect(replacedVariables).toEqual({
+      __V_1__metric_var__V__: '[[metric_var]]',
+      __V_2__app__V__: '${app}',
+      __V_0____interval__V__: '$__interval',
+    });
   });
 });
 
@@ -42,9 +48,14 @@ describe('getString', () => {
 
   it('is symmetrical with replaceVariables', () => {
     const expr = 'sum_over_time([[metric_var]]{bar="${app}"}[$__interval])';
-    const replaced = replaceVariables(expr);
-    const tree = parser.parse(replaced);
-    expect(getString(replaced, tree.topNode)).toBe(expr);
+    const { replacedExpr, replacedVariables } = replaceVariables(expr);
+    const tree = parser.parse(replacedExpr);
+    expect(getString(replacedExpr, tree.topNode)).toBe(expr);
+    expect(replacedVariables).toEqual({
+      __V_1__metric_var__V__: '[[metric_var]]',
+      __V_2__app__V__: '${app}',
+      __V_0____interval__V__: '$__interval',
+    });
   });
 });
 

packages/grafana-prometheus/src/querybuilder/parsingUtils.ts

@@ -36,7 +36,8 @@ const variableRegex = /\$(\w+)|\[\[([\s\S]+?)(?::(\w+))?\]\]|\${(\w+)(?:\.([^:^\
  * parsable and at the same time we can get the variable and its format back from it.
  */
 export function replaceVariables(expr: string) {
-  return expr.replace(variableRegex, (match, var1, var2, fmt2, var3, fieldPath, fmt3) => {
+  const replacedVariables: Record<string, string> = {};
+  const replacedExpr = expr.replace(variableRegex, (match, var1, var2, fmt2, var3, fieldPath, fmt3) => {
     const fmt = fmt2 || fmt3;
     let variable = var1;
     let varType = '0';
@@ -51,8 +52,12 @@ export function replaceVariables(expr: string) {
       varType = '2';
     }
 
-    return `__V_${varType}__` + variable + '__V__' + (fmt ? '__F__' + fmt + '__F__' : '');
+    const replacement = `__V_${varType}__` + variable + '__V__' + (fmt ? '__F__' + fmt + '__F__' : '');
+    replacedVariables[replacement] = match;
+    return replacement;
   });
+
+  return { replacedExpr, replacedVariables };
 }
 
 const varTypeFunc = [