Prometheus: Support more binary operations in visual builder (#46241)

* Add definitions and code to parse more binary ops

* Add alternativesKey

* Add signs to nested query

* Change label

public/app/plugins/datasource/prometheus/querybuilder/binaryScalarOperations.ts

@@ -0,0 +1,87 @@
+import { QueryBuilderOperation, QueryBuilderOperationDef } from './shared/types';
+import { PromOperationId, PromVisualQueryOperationCategory } from './types';
+import { defaultAddOperationHandler } from './shared/operationUtils';
+
+export const binaryScalarDefs = [
+  {
+    id: PromOperationId.Addition,
+    name: 'Add scalar',
+    sign: '+',
+  },
+  {
+    id: PromOperationId.Subtraction,
+    name: 'Subtract scalar',
+    sign: '-',
+  },
+  {
+    id: PromOperationId.MultiplyBy,
+    name: 'Multiply by scalar',
+    sign: '*',
+  },
+  {
+    id: PromOperationId.DivideBy,
+    name: 'Divide by scalar',
+    sign: '/',
+  },
+  {
+    id: PromOperationId.Modulo,
+    name: 'Modulo by scalar',
+    sign: '%',
+  },
+  {
+    id: PromOperationId.Exponent,
+    name: 'Exponent',
+    sign: '^',
+  },
+  {
+    id: PromOperationId.EqualTo,
+    name: 'Equal to',
+    sign: '==',
+  },
+  {
+    id: PromOperationId.NotEqualTo,
+    name: 'Not equal to',
+    sign: '!=',
+  },
+  {
+    id: PromOperationId.GreaterThan,
+    name: 'Greater than',
+    sign: '>',
+  },
+  {
+    id: PromOperationId.LessThan,
+    name: 'Less than',
+    sign: '<',
+  },
+  {
+    id: PromOperationId.GreaterOrEqual,
+    name: 'Greater or equal to',
+    sign: '>=',
+  },
+  {
+    id: PromOperationId.LessOrEqual,
+    name: 'Less or equal to',
+    sign: '<=',
+  },
+];
+
+// Not sure about this one. It could also be a more generic 'Simple math operation' where user specifies
+// both the operator and the operand in a single input
+export const binaryScalarOperations = binaryScalarDefs.map((opDef) => {
+  return {
+    id: opDef.id,
+    name: opDef.name,
+    params: [{ name: 'Value', type: 'number' }],
+    defaultParams: [2],
+    alternativesKey: 'binary scalar operations',
+    category: PromVisualQueryOperationCategory.BinaryOps,
+    renderer: getSimpleBinaryRenderer(opDef.sign),
+    addOperationHandler: defaultAddOperationHandler,
+  };
+});
+
+function getSimpleBinaryRenderer(operator: string) {
+  return function binaryRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
+    return `${innerExpr} ${operator} ${model.params[0]}`;
+  };
+}

public/app/plugins/datasource/prometheus/querybuilder/components/NestedQuery.tsx

@@ -7,6 +7,7 @@ import { PrometheusDatasource } from '../../datasource';
 import { AutoSizeInput } from '../shared/AutoSizeInput';
 import { PromVisualQueryBinary } from '../types';
 import { PromQueryBuilder } from './PromQueryBuilder';
+import { binaryScalarDefs } from '../binaryScalarOperations';
 
 export interface Props {
   nestedQuery: PromVisualQueryBinary;
@@ -68,14 +69,7 @@ export const NestedQuery = React.memo<Props>(({ nestedQuery, index, datasource,
   );
 });
 
-const operators = [
-  { label: '/', value: '/' },
-  { label: '*', value: '*' },
-  { label: '+', value: '+' },
-  { label: '==', value: '==' },
-  { label: '>', value: '>' },
-  { label: '<', value: '<' },
-];
+const operators = binaryScalarDefs.map((def) => ({ label: def.sign, value: def.sign }));
 
 NestedQuery.displayName = 'NestedQuery';
 

public/app/plugins/datasource/prometheus/querybuilder/operations.ts

@@ -15,6 +15,7 @@ import {
   VisualQueryModeller,
 } from './shared/types';
 import { PromOperationId, PromVisualQuery, PromVisualQueryOperationCategory } from './types';
+import { binaryScalarOperations } from './binaryScalarOperations';
 
 export function getOperationDefinitions(): QueryBuilderOperationDef[] {
   const list: QueryBuilderOperationDef[] = [
@@ -90,26 +91,7 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
       addOperationHandler: addOperationWithRangeVector,
       changeTypeHandler: operationTypeChangedHandlerForRangeFunction,
     }),
-    // Not sure about this one. It could also be a more generic 'Simple math operation' where user specifies
-    // both the operator and the operand in a single input
-    {
-      id: PromOperationId.MultiplyBy,
-      name: 'Multiply by scalar',
-      params: [{ name: 'Factor', type: 'number' }],
-      defaultParams: [2],
-      category: PromVisualQueryOperationCategory.BinaryOps,
-      renderer: getSimpleBinaryRenderer('*'),
-      addOperationHandler: defaultAddOperationHandler,
-    },
-    {
-      id: PromOperationId.DivideBy,
-      name: 'Divide by scalar',
-      params: [{ name: 'Factor', type: 'number' }],
-      defaultParams: [2],
-      category: PromVisualQueryOperationCategory.BinaryOps,
-      renderer: getSimpleBinaryRenderer('/'),
-      addOperationHandler: defaultAddOperationHandler,
-    },
+    ...binaryScalarOperations,
     {
       id: PromOperationId.NestedQuery,
       name: 'Binary operation with query',
@@ -330,12 +312,6 @@ export function operationWithRangeVectorRenderer(
   return `${def.id}(${innerExpr}[${rangeVector}])`;
 }
 
-function getSimpleBinaryRenderer(operator: string) {
-  return function binaryRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
-    return `${innerExpr} ${operator} ${model.params[0]}`;
-  };
-}
-
 /**
  * Since there can only be one operation with range vector this will replace the current one (if one was added )
  */

public/app/plugins/datasource/prometheus/querybuilder/parsing.test.ts

@@ -376,6 +376,58 @@ describe('buildVisualQueryFromString', () => {
       },
     });
   });
+
+  it('handles multiple binary scalar operations', () => {
+    expect(buildVisualQueryFromString('cluster_namespace_slug_dialer_name + 1 - 1 / 1 * 1 % 1 ^ 1')).toEqual({
+      errors: [],
+      query: {
+        metric: 'cluster_namespace_slug_dialer_name',
+        labels: [],
+        operations: [
+          {
+            id: '__addition',
+            params: [1],
+          },
+          {
+            id: '__subtraction',
+            params: [1],
+          },
+          {
+            id: '__divide_by',
+            params: [1],
+          },
+          {
+            id: '__multiply_by',
+            params: [1],
+          },
+          {
+            id: '__modulo',
+            params: [1],
+          },
+          {
+            id: '__exponent',
+            params: [1],
+          },
+        ],
+      },
+    });
+  });
+
+  it('handles scalar comparison operators', () => {
+    expect(buildVisualQueryFromString('cluster_namespace_slug_dialer_name <= 2')).toEqual({
+      errors: [],
+      query: {
+        metric: 'cluster_namespace_slug_dialer_name',
+        labels: [],
+        operations: [
+          {
+            id: '__less_or_equal',
+            params: [2],
+          },
+        ],
+      },
+    });
+  });
 });
 
 function noErrors(query: PromVisualQuery) {

public/app/plugins/datasource/prometheus/querybuilder/parsing.ts

@@ -1,7 +1,8 @@
 import { parser } from 'lezer-promql';
-import { SyntaxNode } from '@lezer/common';
+import { SyntaxNode, TreeCursor } from '@lezer/common';
 import { QueryBuilderLabelFilter, QueryBuilderOperation } from './shared/types';
 import { PromVisualQuery } from './types';
+import { binaryScalarDefs } from './binaryScalarOperations';
 
 // Taken from template_srv, but copied so to not mess with the regex.index which is manipulated in the service
 /*
@@ -322,10 +323,10 @@ function updateFunctionArgs(expr: string, node: SyntaxNode | null, context: Cont
   }
 }
 
-const operatorToOpName: Record<string, string> = {
-  '/': '__divide_by',
-  '*': '__multiply_by',
-};
+const operatorToOpName = binaryScalarDefs.reduce((acc, def) => {
+  acc[def.sign] = def.id;
+  return acc;
+}, {} as Record<string, string>);
 
 /**
  * Right now binary expressions can be represented in 2 way in visual query. As additional operation in case it is
@@ -338,6 +339,7 @@ function handleBinary(expr: string, node: SyntaxNode, context: Context) {
   const visQuery = context.query;
   const left = node.firstChild!;
   const op = getString(expr, left.nextSibling);
+  // TODO: we are skipping BinModifiers
   const right = node.lastChild!;
 
   const opName = operatorToOpName[op];
@@ -347,9 +349,40 @@ function handleBinary(expr: string, node: SyntaxNode, context: Context) {
 
   if (leftNumber || rightNumber) {
     // Scalar case, just add operation.
-    const [num, query] = leftNumber ? [leftNumber, right] : [rightNumber, left];
-    visQuery.operations.push({ id: opName, params: [parseInt(getString(expr, num), 10)] });
-    handleExpression(expr, query, context);
+    if (leftNumber) {
+      // TODO: this should be already handled in case parent is binary expression as it has to be added to parent
+      //  if query starts with a number that isn't handled now.
+    } else {
+      handleExpression(expr, left, context);
+    }
+
+    if (rightNumber) {
+      // TODO: this should be already handled in case parent is binary expression as it has to be added to parent
+      //  if query starts with a number that isn't handled now.
+      visQuery.operations.push({ id: opName, params: [parseInt(getString(expr, right), 10)] });
+    } else {
+      handleExpression(expr, right, context);
+    }
+    return;
+  }
+
+  const leftBinary = left.getChild('BinaryExpr');
+  const rightBinary = right.getChild('BinaryExpr');
+
+  if (leftBinary || rightBinary) {
+    // One of the sides is binary which means we don't really know if there is a query or just chained scalars. So
+    // we have to traverse a bit deeper to know
+    handleExpression(expr, left, context);
+
+    // Due to the way binary ops are parsed we can get a binary operation on the right that starts with a number which
+    // is a factor for a current binary operation. So we have to add it as an operation now.
+    const leftMostChild = getLeftMostChild(right);
+    if (leftMostChild?.name === 'NumberLiteral') {
+      visQuery.operations.push({ id: opName, params: [parseInt(getString(expr, leftMostChild), 10)] });
+    }
+    // If we added the first number literal as operation here we still can continue and handle the rest as the first
+    // number will be just skipped.
+    handleExpression(expr, right, context);
   } else {
     // Two queries case so we create a binary query.
     visQuery.binaryQueries = visQuery.binaryQueries || [];
@@ -377,7 +410,7 @@ function handleBinary(expr: string, node: SyntaxNode, context: Context) {
  * @param expr
  * @param node
  */
-function getString(expr: string, node: SyntaxNode | null) {
+function getString(expr: string, node: SyntaxNode | TreeCursor | null) {
   if (!node) {
     return '';
   }
@@ -407,6 +440,18 @@ function getAllByType(expr: string, cur: SyntaxNode, type: string): string[] {
   return values;
 }
 
+function getLeftMostChild(cur: SyntaxNode): SyntaxNode | null {
+  let child = cur;
+  while (true) {
+    if (child.firstChild) {
+      child = child.firstChild;
+    } else {
+      break;
+    }
+  }
+  return child;
+}
+
 // Debugging function for convenience.
 // @ts-ignore
 function log(expr: string, cur?: SyntaxNode) {

public/app/plugins/datasource/prometheus/querybuilder/shared/LokiAndPromQueryModellerBase.ts

@@ -43,7 +43,10 @@ export abstract class LokiAndPromQueryModellerBase<T extends QueryWithOperations
 
   renderOperations(queryString: string, operations: QueryBuilderOperation[]) {
     for (const operation of operations) {
-      const def = this.operationsRegisty.get(operation.id);
+      const def = this.operationsRegisty.getIfExists(operation.id);
+      if (!def) {
+        throw new Error(`Could not find operation ${operation.id} in the registry`);
+      }
       queryString = def.renderer(operation, def, queryString);
     }
 

public/app/plugins/datasource/prometheus/querybuilder/types.ts

@@ -103,9 +103,20 @@ export enum PromOperationId {
   Topk = 'topk',
   Vector = 'vector',
   Year = 'year',
+  // Binary ops
+  Addition = '__addition',
+  Subtraction = '__subtraction',
   MultiplyBy = '__multiply_by',
   DivideBy = '__divide_by',
+  Modulo = '__modulo',
+  Exponent = '__exponent',
   NestedQuery = '__nested_query',
+  EqualTo = '__equal_to',
+  NotEqualTo = '__not_equal_to',
+  GreaterThan = '__greater_than',
+  LessThan = '__less_than',
+  GreaterOrEqual = '__greater_or_equal',
+  LessOrEqual = '__less_or_equal',
 }
 
 export interface PromQueryPattern {