Prometheus: Move count_values from function to aggregation (#47260)

* Move count_values from function to aggregation

* Fix typos

* Fix loki operations

* Fix error that change the aggregation variant on blur

* Fix loki ops

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

@@ -1,5 +1,8 @@
-import { createAggregationOperation } from '../../prometheus/querybuilder/aggregations';
-import { getPromAndLokiOperationDisplayName } from '../../prometheus/querybuilder/shared/operationUtils';
+import {
+  createAggregationOperation,
+  createAggregationOperationWithParam,
+  getPromAndLokiOperationDisplayName,
+} from '../../prometheus/querybuilder/shared/operationUtils';
 import {
   QueryBuilderOperation,
   QueryBuilderOperationDef,
@@ -16,8 +19,6 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
     LokiOperationId.Min,
     LokiOperationId.Max,
     LokiOperationId.Avg,
-    LokiOperationId.TopK,
-    LokiOperationId.BottomK,
     LokiOperationId.Stddev,
     LokiOperationId.Stdvar,
     LokiOperationId.Count,
@@ -28,6 +29,20 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
     })
   );
 
+  const aggregationsWithParam = [LokiOperationId.TopK, LokiOperationId.BottomK].flatMap((opId) => {
+    return createAggregationOperationWithParam(
+      opId,
+      {
+        params: [{ name: 'K-value', type: 'number' }],
+        defaultParams: [5],
+      },
+      {
+        addOperationHandler: addLokiOperation,
+        orderRank: LokiOperationOrder.Last,
+      }
+    );
+  });
+
   const list: QueryBuilderOperationDef[] = [
     createRangeOperation(LokiOperationId.Rate),
     createRangeOperation(LokiOperationId.CountOverTime),
@@ -44,6 +59,7 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
     createRangeOperation(LokiOperationId.StddevOverTime),
     createRangeOperation(LokiOperationId.QuantileOverTime),
     ...aggregations,
+    ...aggregationsWithParam,
     {
       id: LokiOperationId.Json,
       name: 'Json',
@@ -145,7 +161,7 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
       renderer: (model, def, innerExpr) => `${innerExpr} | line_format \`${model.params[0]}\``,
       addOperationHandler: addLokiOperation,
       explainHandler: () =>
-        `This will replace log line using a specified template. The template can refer to stream labels and extracted labels. 
+        `This will replace log line using a specified template. The template can refer to stream labels and extracted labels.
 
         Example: \`{{.status_code}} - {{.message}}\`
 
@@ -166,7 +182,7 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
       renderer: (model, def, innerExpr) => `${innerExpr} | label_format ${model.params[1]}=\`${model.params[0]}\``,
       addOperationHandler: addLokiOperation,
       explainHandler: () =>
-        `This will change name of label to desired new label. In the example below, label "error_level" will be renamed to "level". 
+        `This will change name of label to desired new label. In the example below, label "error_level" will be renamed to "level".
 
         Example: error_level=\`level\`
 

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

@@ -1,18 +1,11 @@
-import pluralize from 'pluralize';
-import { LabelParamEditor } from './components/LabelParamEditor';
 import { addOperationWithRangeVector } from './operations';
 import {
-  defaultAddOperationHandler,
-  functionRendererLeft,
+  createAggregationOperation,
+  createAggregationOperationWithParam,
   getPromAndLokiOperationDisplayName,
   getRangeVectorParamDef,
 } from './shared/operationUtils';
-import {
-  QueryBuilderOperation,
-  QueryBuilderOperationDef,
-  QueryBuilderOperationParamDef,
-  QueryWithOperations,
-} from './shared/types';
+import { QueryBuilderOperation, QueryBuilderOperationDef } from './shared/types';
 import { PromVisualQueryOperationCategory, PromOperationId } from './types';
 
 export function getAggregationOperations(): QueryBuilderOperationDef[] {
@@ -22,8 +15,18 @@ export function getAggregationOperations(): QueryBuilderOperationDef[] {
     ...createAggregationOperation(PromOperationId.Min),
     ...createAggregationOperation(PromOperationId.Max),
     ...createAggregationOperation(PromOperationId.Count),
-    ...createAggregationOperation(PromOperationId.TopK),
-    ...createAggregationOperation(PromOperationId.BottomK),
+    ...createAggregationOperationWithParam(PromOperationId.TopK, {
+      params: [{ name: 'K-value', type: 'number' }],
+      defaultParams: [5],
+    }),
+    ...createAggregationOperationWithParam(PromOperationId.BottomK, {
+      params: [{ name: 'K-value', type: 'number' }],
+      defaultParams: [5],
+    }),
+    ...createAggregationOperationWithParam(PromOperationId.CountValues, {
+      params: [{ name: 'Identifier', type: 'string' }],
+      defaultParams: ['count'],
+    }),
     createAggregationOverTime(PromOperationId.SumOverTime),
     createAggregationOverTime(PromOperationId.AvgOverTime),
     createAggregationOverTime(PromOperationId.MinOverTime),
@@ -36,162 +39,6 @@ export function getAggregationOperations(): QueryBuilderOperationDef[] {
   ];
 }
 
-/**
- * This function is shared between Prometheus and Loki variants
- */
-export function createAggregationOperation<T extends QueryWithOperations>(
-  name: string,
-  overrides: Partial<QueryBuilderOperationDef> = {}
-): QueryBuilderOperationDef[] {
-  const operations: QueryBuilderOperationDef[] = [
-    {
-      id: name,
-      name: getPromAndLokiOperationDisplayName(name),
-      params: [
-        {
-          name: 'By label',
-          type: 'string',
-          restParam: true,
-          optional: true,
-        },
-      ],
-      defaultParams: [],
-      alternativesKey: 'plain aggregations',
-      category: PromVisualQueryOperationCategory.Aggregations,
-      renderer: functionRendererLeft,
-      paramChangedHandler: getOnLabelAdddedHandler(`__${name}_by`),
-      explainHandler: getAggregationExplainer(name, ''),
-      addOperationHandler: defaultAddOperationHandler,
-      ...overrides,
-    },
-    {
-      id: `__${name}_by`,
-      name: `${getPromAndLokiOperationDisplayName(name)} by`,
-      params: [
-        {
-          name: 'Label',
-          type: 'string',
-          restParam: true,
-          optional: true,
-          editor: LabelParamEditor,
-        },
-      ],
-      defaultParams: [''],
-      alternativesKey: 'aggregations by',
-      category: PromVisualQueryOperationCategory.Aggregations,
-      renderer: getAggregationByRenderer(name),
-      paramChangedHandler: getLastLabelRemovedHandler(name),
-      explainHandler: getAggregationExplainer(name, 'by'),
-      addOperationHandler: defaultAddOperationHandler,
-      hideFromList: true,
-      ...overrides,
-    },
-    {
-      id: `__${name}_without`,
-      name: `${getPromAndLokiOperationDisplayName(name)} without`,
-      params: [
-        {
-          name: 'Label',
-          type: 'string',
-          restParam: true,
-          optional: true,
-          editor: LabelParamEditor,
-        },
-      ],
-      defaultParams: [''],
-      alternativesKey: 'aggregations by',
-      category: PromVisualQueryOperationCategory.Aggregations,
-      renderer: getAggregationWithoutRenderer(name),
-      paramChangedHandler: getLastLabelRemovedHandler(name),
-      explainHandler: getAggregationExplainer(name, 'without'),
-      addOperationHandler: defaultAddOperationHandler,
-      hideFromList: true,
-      ...overrides,
-    },
-  ];
-
-  // Handle some special aggregations that have parameters
-  if (name === 'topk' || name === 'bottomk') {
-    const param: QueryBuilderOperationParamDef = {
-      name: 'K-value',
-      type: 'number',
-    };
-    operations[0].params.unshift(param);
-    operations[1].params.unshift(param);
-    operations[0].defaultParams = [5];
-    operations[1].defaultParams = [5, ''];
-    operations[1].renderer = getAggregationByRendererWithParameter(name);
-  }
-
-  return operations;
-}
-
-function getAggregationByRenderer(aggregation: string) {
-  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
-    return `${aggregation} by(${model.params.join(', ')}) (${innerExpr})`;
-  };
-}
-
-function getAggregationWithoutRenderer(aggregation: string) {
-  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
-    return `${aggregation} without(${model.params.join(', ')}) (${innerExpr})`;
-  };
-}
-
-/**
- * Very simple poc implementation, needs to be modified to support all aggregation operators
- */
-function getAggregationExplainer(aggregationName: string, mode: 'by' | 'without' | '') {
-  return function aggregationExplainer(model: QueryBuilderOperation) {
-    const labels = model.params.map((label) => `\`${label}\``).join(' and ');
-    const labelWord = pluralize('label', model.params.length);
-
-    switch (mode) {
-      case 'by':
-        return `Calculates ${aggregationName} over dimensions while preserving ${labelWord} ${labels}.`;
-      case 'without':
-        return `Calculates ${aggregationName} over the dimensions ${labels}. All other labels are preserved.`;
-      default:
-        return `Calculates ${aggregationName} over the dimensions.`;
-    }
-  };
-}
-
-function getAggregationByRendererWithParameter(aggregation: string) {
-  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
-    const firstParam = model.params[0];
-    const restParams = model.params.slice(1);
-    return `${aggregation} by(${restParams.join(', ')}) (${firstParam}, ${innerExpr})`;
-  };
-}
-
-/**
- * This function will transform operations without labels to their plan aggregation operation
- */
-function getLastLabelRemovedHandler(changeToOperartionId: string) {
-  return function onParamChanged(index: number, op: QueryBuilderOperation, def: QueryBuilderOperationDef) {
-    // If definition has more params then is defined there are no optional rest params anymore
-    // We then transform this operation into a different one
-    if (op.params.length < def.params.length) {
-      return {
-        ...op,
-        id: changeToOperartionId,
-      };
-    }
-
-    return op;
-  };
-}
-
-export function getOnLabelAdddedHandler(changeToOperartionId: string) {
-  return function onParamChanged(index: number, op: QueryBuilderOperation) {
-    return {
-      ...op,
-      id: changeToOperartionId,
-    };
-  };
-}
-
 function createAggregationOverTime(name: string): QueryBuilderOperationDef {
   return {
     id: name,

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

@@ -155,12 +155,6 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
       id: PromOperationId.Cosh,
       category: PromVisualQueryOperationCategory.Trigonometric,
     }),
-    createFunction({
-      id: PromOperationId.CountValues,
-      params: [{ name: 'Identifier', type: 'string' }],
-      defaultParams: ['count'],
-      renderer: functionRendererLeft,
-    }),
     createFunction({
       id: PromOperationId.DayOfMonth,
       category: PromVisualQueryOperationCategory.Time,

public/app/plugins/datasource/prometheus/querybuilder/shared/operationUtils.test.ts

@@ -0,0 +1,127 @@
+import { createAggregationOperation, createAggregationOperationWithParam } from './operationUtils';
+
+describe('createAggregationOperation', () => {
+  it('returns correct aggregation definitions with overrides', () => {
+    expect(createAggregationOperation('test_aggregation', { category: 'test_category' })).toMatchObject([
+      {
+        addOperationHandler: {},
+        alternativesKey: 'plain aggregations',
+        category: 'test_category',
+        defaultParams: [],
+        explainHandler: {},
+        id: 'test_aggregation',
+        name: 'Test aggregation',
+        paramChangedHandler: {},
+        params: [
+          {
+            name: 'By label',
+            optional: true,
+            restParam: true,
+            type: 'string',
+          },
+        ],
+        renderer: {},
+      },
+      {
+        alternativesKey: 'aggregations by',
+        category: 'test_category',
+        defaultParams: [''],
+        explainHandler: {},
+        hideFromList: true,
+        id: '__test_aggregation_by',
+        name: 'Test aggregation by',
+        paramChangedHandler: {},
+        params: [
+          {
+            editor: {},
+            name: 'Label',
+            optional: true,
+            restParam: true,
+            type: 'string',
+          },
+        ],
+        renderer: {},
+      },
+      {
+        alternativesKey: 'aggregations by',
+        category: 'test_category',
+        defaultParams: [''],
+        explainHandler: {},
+        hideFromList: true,
+        id: '__test_aggregation_without',
+        name: 'Test aggregation without',
+        paramChangedHandler: {},
+        params: [
+          {
+            name: 'Label',
+            optional: true,
+            restParam: true,
+            type: 'string',
+          },
+        ],
+        renderer: {},
+      },
+    ]);
+  });
+});
+
+describe('createAggregationOperationWithParams', () => {
+  it('returns correct aggregation definitions with overrides and params', () => {
+    expect(
+      createAggregationOperationWithParam(
+        'test_aggregation',
+        {
+          params: [{ name: 'K-value', type: 'number' }],
+          defaultParams: [5],
+        },
+        { category: 'test_category' }
+      )
+    ).toMatchObject([
+      {
+        addOperationHandler: {},
+        alternativesKey: 'plain aggregations',
+        category: 'test_category',
+        defaultParams: [5],
+        explainHandler: {},
+        id: 'test_aggregation',
+        name: 'Test aggregation',
+        paramChangedHandler: {},
+        params: [
+          { name: 'K-value', type: 'number' },
+          { name: 'By label', optional: true, restParam: true, type: 'string' },
+        ],
+        renderer: {},
+      },
+      {
+        alternativesKey: 'aggregations by',
+        category: 'test_category',
+        defaultParams: [5, ''],
+        explainHandler: {},
+        hideFromList: true,
+        id: '__test_aggregation_by',
+        name: 'Test aggregation by',
+        paramChangedHandler: {},
+        params: [
+          { name: 'K-value', type: 'number' },
+          { editor: {}, name: 'Label', optional: true, restParam: true, type: 'string' },
+        ],
+        renderer: {},
+      },
+      {
+        alternativesKey: 'aggregations by',
+        category: 'test_category',
+        defaultParams: [5, ''],
+        explainHandler: {},
+        hideFromList: true,
+        id: '__test_aggregation_without',
+        name: 'Test aggregation without',
+        paramChangedHandler: {},
+        params: [
+          { name: 'K-value', type: 'number' },
+          { name: 'Label', optional: true, restParam: true, type: 'string' },
+        ],
+        renderer: {},
+      },
+    ]);
+  });
+});

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

@@ -3,9 +3,13 @@ import {
   QueryBuilderOperation,
   QueryBuilderOperationDef,
   QueryBuilderOperationParamDef,
+  QueryBuilderOperationParamValue,
   QueryWithOperations,
 } from './types';
 import { SelectableValue } from '@grafana/data/src';
+import { LabelParamEditor } from '../components/LabelParamEditor';
+import { PromVisualQueryOperationCategory } from '../types';
+import pluralize from 'pluralize';
 
 export function functionRendererLeft(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
   const params = renderParams(model, def, innerExpr);
@@ -144,3 +148,175 @@ export function getRangeVectorParamDef(withRateInterval = false): QueryBuilderOp
 
   return param;
 }
+
+/**
+ * This function is shared between Prometheus and Loki variants
+ */
+export function createAggregationOperation<T extends QueryWithOperations>(
+  name: string,
+  overrides: Partial<QueryBuilderOperationDef> = {}
+): QueryBuilderOperationDef[] {
+  const operations: QueryBuilderOperationDef[] = [
+    {
+      id: name,
+      name: getPromAndLokiOperationDisplayName(name),
+      params: [
+        {
+          name: 'By label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+        },
+      ],
+      defaultParams: [],
+      alternativesKey: 'plain aggregations',
+      category: PromVisualQueryOperationCategory.Aggregations,
+      renderer: functionRendererLeft,
+      paramChangedHandler: getOnLabelAddedHandler(`__${name}_by`),
+      explainHandler: getAggregationExplainer(name, ''),
+      addOperationHandler: defaultAddOperationHandler,
+      ...overrides,
+    },
+    {
+      id: `__${name}_by`,
+      name: `${getPromAndLokiOperationDisplayName(name)} by`,
+      params: [
+        {
+          name: 'Label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+          editor: LabelParamEditor,
+        },
+      ],
+      defaultParams: [''],
+      alternativesKey: 'aggregations by',
+      category: PromVisualQueryOperationCategory.Aggregations,
+      renderer: getAggregationByRenderer(name),
+      paramChangedHandler: getLastLabelRemovedHandler(name),
+      explainHandler: getAggregationExplainer(name, 'by'),
+      addOperationHandler: defaultAddOperationHandler,
+      hideFromList: true,
+      ...overrides,
+    },
+    {
+      id: `__${name}_without`,
+      name: `${getPromAndLokiOperationDisplayName(name)} without`,
+      params: [
+        {
+          name: 'Label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+          editor: LabelParamEditor,
+        },
+      ],
+      defaultParams: [''],
+      alternativesKey: 'aggregations by',
+      category: PromVisualQueryOperationCategory.Aggregations,
+      renderer: getAggregationWithoutRenderer(name),
+      paramChangedHandler: getLastLabelRemovedHandler(name),
+      explainHandler: getAggregationExplainer(name, 'without'),
+      addOperationHandler: defaultAddOperationHandler,
+      hideFromList: true,
+      ...overrides,
+    },
+  ];
+
+  return operations;
+}
+
+export function createAggregationOperationWithParam(
+  name: string,
+  paramsDef: { params: QueryBuilderOperationParamDef[]; defaultParams: QueryBuilderOperationParamValue[] },
+  overrides: Partial<QueryBuilderOperationDef> = {}
+): QueryBuilderOperationDef[] {
+  const operations = createAggregationOperation(name, overrides);
+  operations[0].params.unshift(...paramsDef.params);
+  operations[1].params.unshift(...paramsDef.params);
+  operations[2].params.unshift(...paramsDef.params);
+  operations[0].defaultParams = paramsDef.defaultParams;
+  operations[1].defaultParams = [...paramsDef.defaultParams, ''];
+  operations[2].defaultParams = [...paramsDef.defaultParams, ''];
+  operations[1].renderer = getAggregationByRendererWithParameter(name);
+  operations[2].renderer = getAggregationByRendererWithParameter(name);
+  return operations;
+}
+
+function getAggregationByRenderer(aggregation: string) {
+  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
+    return `${aggregation} by(${model.params.join(', ')}) (${innerExpr})`;
+  };
+}
+
+function getAggregationWithoutRenderer(aggregation: string) {
+  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
+    return `${aggregation} without(${model.params.join(', ')}) (${innerExpr})`;
+  };
+}
+
+/**
+ * Very simple poc implementation, needs to be modified to support all aggregation operators
+ */
+function getAggregationExplainer(aggregationName: string, mode: 'by' | 'without' | '') {
+  return function aggregationExplainer(model: QueryBuilderOperation) {
+    const labels = model.params.map((label) => `\`${label}\``).join(' and ');
+    const labelWord = pluralize('label', model.params.length);
+
+    switch (mode) {
+      case 'by':
+        return `Calculates ${aggregationName} over dimensions while preserving ${labelWord} ${labels}.`;
+      case 'without':
+        return `Calculates ${aggregationName} over the dimensions ${labels}. All other labels are preserved.`;
+      default:
+        return `Calculates ${aggregationName} over the dimensions.`;
+    }
+  };
+}
+
+function getAggregationByRendererWithParameter(aggregation: string) {
+  return function aggregationRenderer(model: QueryBuilderOperation, def: QueryBuilderOperationDef, innerExpr: string) {
+    function mapType(p: QueryBuilderOperationParamValue) {
+      if (typeof p === 'string') {
+        return `\"${p}\"`;
+      }
+      return p;
+    }
+    const params = model.params.slice(0, -1);
+    const restParams = model.params.slice(1);
+    return `${aggregation} by(${restParams.join(', ')}) (${params.map(mapType).join(', ')}, ${innerExpr})`;
+  };
+}
+
+/**
+ * This function will transform operations without labels to their plan aggregation operation
+ */
+function getLastLabelRemovedHandler(changeToOperationId: string) {
+  return function onParamChanged(index: number, op: QueryBuilderOperation, def: QueryBuilderOperationDef) {
+    // If definition has more params then is defined there are no optional rest params anymore.
+    // We then transform this operation into a different one
+    if (op.params.length < def.params.length) {
+      return {
+        ...op,
+        id: changeToOperationId,
+      };
+    }
+
+    return op;
+  };
+}
+
+function getOnLabelAddedHandler(changeToOperationId: string) {
+  return function onParamChanged(index: number, op: QueryBuilderOperation, def: QueryBuilderOperationDef) {
+    // Check if we actually have the label param. As it's optional the aggregation can have one less, which is the
+    // case of just simple aggregation without label. When user adds the label it now has the same number of params
+    // as it's definition, and now we can change it to it's `_by` variant.
+    if (op.params.length === def.params.length) {
+      return {
+        ...op,
+        id: changeToOperationId,
+      };
+    }
+    return op;
+  };
+}