Loki: Add aggregate by label support (#46965)

* Loki: Add aggregate by label support

* Label lookups now work for loki

* Fixed ts issue

* review comment fixed

public/app/plugins/datasource/loki/components/AnnotationsQueryEditor.tsx

@@ -4,7 +4,7 @@ import React, { memo } from 'react';
 import { LokiQuery } from '../types';
 import { LokiQueryField } from './LokiQueryField';
 import { LokiOptionFields } from './LokiOptionFields';
-import LokiDatasource from '../datasource';
+import { LokiDatasource } from '../datasource';
 
 interface Props {
   expr: string;

public/app/plugins/datasource/loki/components/LokiQueryField.tsx

@@ -16,7 +16,7 @@ import { LokiQuery, LokiOptions } from '../types';
 import { LanguageMap, languages as prismLanguages } from 'prismjs';
 import LokiLanguageProvider from '../language_provider';
 import { shouldRefreshLabels } from '../language_utils';
-import LokiDatasource from '../datasource';
+import { LokiDatasource } from '../datasource';
 import { LocalStorageValueProvider } from 'app/core/components/LocalStorageValueProvider';
 
 const LAST_USED_LABELS_KEY = 'grafana.datasources.loki.browser.labels';

public/app/plugins/datasource/loki/components/types.ts

@@ -1,5 +1,5 @@
 import { QueryEditorProps } from '@grafana/data';
-import LokiDatasource from '../datasource';
+import { LokiDatasource } from '../datasource';
 import { LokiOptions, LokiQuery } from '../types';
 
 export type LokiQueryEditorProps = QueryEditorProps<LokiDatasource, LokiQuery, LokiOptions>;

public/app/plugins/datasource/loki/datasource.test.ts

@@ -13,8 +13,7 @@ import {
   toUtc,
 } from '@grafana/data';
 import { BackendSrvRequest, FetchResponse } from '@grafana/runtime';
-
+import { LokiDatasource, RangeQueryOptions } from './datasource';
-import LokiDatasource, { RangeQueryOptions } from './datasource';
 import { LokiQuery, LokiResponse, LokiResultType } from './types';
 import { getQueryOptions } from 'test/helpers/getQueryOptions';
 import { TemplateSrv } from 'app/features/templating/template_srv';

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

@@ -837,5 +837,3 @@ function getLogLevelFromLabels(labels: Labels): LogLevel {
   }
   return levelLabel ? getLogLevelFromKey(labels[levelLabel]) : LogLevel.unknown;
 }
-
-export default LokiDatasource;

public/app/plugins/datasource/loki/language_provider.test.ts

@@ -4,7 +4,7 @@ import LanguageProvider, { LokiHistoryItem } from './language_provider';
 import { TypeaheadInput } from '@grafana/ui';
 
 import { makeMockLokiDatasource } from './mocks';
-import LokiDatasource from './datasource';
+import { LokiDatasource } from './datasource';
 import { AbstractLabelOperator } from '@grafana/data';
 import { LokiQueryType } from './types';
 

public/app/plugins/datasource/loki/language_provider.ts

@@ -15,7 +15,7 @@ import syntax, { FUNCTIONS, PIPE_PARSERS, PIPE_OPERATORS } from './syntax';
 import { LokiQuery, LokiQueryType } from './types';
 import { dateTime, AbsoluteTimeRange, LanguageProvider, HistoryItem, AbstractQuery } from '@grafana/data';
 
-import LokiDatasource from './datasource';
+import { LokiDatasource } from './datasource';
 import { CompletionItem, TypeaheadInput, TypeaheadOutput, CompletionItemGroup } from '@grafana/ui';
 import Prism, { Grammar } from 'prismjs';
 

public/app/plugins/datasource/loki/module.ts

@@ -1,12 +1,12 @@
 import { DataSourcePlugin } from '@grafana/data';
-import Datasource from './datasource';
+import { LokiDatasource } from './datasource';
 
 import LokiCheatSheet from './components/LokiCheatSheet';
 import LokiQueryEditorByApp from './components/LokiQueryEditorByApp';
 import { LokiAnnotationsQueryCtrl } from './LokiAnnotationsQueryCtrl';
 import { ConfigEditor } from './configuration/ConfigEditor';
 
-export const plugin = new DataSourcePlugin(Datasource)
+export const plugin = new DataSourcePlugin(LokiDatasource)
   .setQueryEditor(LokiQueryEditorByApp)
   .setConfigEditor(ConfigEditor)
   .setQueryEditorHelp(LokiCheatSheet)

public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.tsx

@@ -5,7 +5,6 @@ import { RadioButtonGroup, Select } from '@grafana/ui';
 import { LokiQuery, LokiQueryType } from '../../types';
 import { QueryOptionGroup } from 'app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup';
 import { preprocessMaxLines, queryTypeOptions, RESOLUTION_OPTIONS } from '../../components/LokiOptionFields';
-import { getLegendModeLabel } from 'app/plugins/datasource/prometheus/querybuilder/components/PromQueryLegendEditor';
 import { AutoSizeInput } from 'app/plugins/datasource/prometheus/querybuilder/shared/AutoSizeInput';
 import { isMetricsQuery } from '../../datasource';
 
@@ -99,7 +98,9 @@ function getCollapsedInfo(query: LokiQuery, queryType: LokiQueryType, showMaxLin
 
   const items: string[] = [];
 
-  items.push(`Legend: ${getLegendModeLabel(query.legendFormat)}`);
+  if (query.legendFormat) {
+    items.push(`Legend: ${query.legendFormat}`);
+  }
 
   if (query.resolution) {
     items.push(`Resolution: ${resolutionLabel?.label}`);

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

@@ -1,7 +1,5 @@
-import {
+import { createAggregationOperation } from '../../prometheus/querybuilder/aggregations';
-  functionRendererLeft,
+import { getPromAndLokiOperationDisplayName } from '../../prometheus/querybuilder/shared/operationUtils';
-  getPromAndLokiOperationDisplayName,
-} from '../../prometheus/querybuilder/shared/operationUtils';
 import {
   QueryBuilderOperation,
   QueryBuilderOperationDef,
@@ -12,6 +10,20 @@ import { FUNCTIONS } from '../syntax';
 import { LokiOperationId, LokiOperationOrder, LokiVisualQuery, LokiVisualQueryOperationCategory } from './types';
 
 export function getOperationDefintions(): QueryBuilderOperationDef[] {
+  const aggregations = [
+    LokiOperationId.Sum,
+    LokiOperationId.Min,
+    LokiOperationId.Max,
+    LokiOperationId.Avg,
+    LokiOperationId.TopK,
+    LokiOperationId.BottomK,
+  ].flatMap((opId) =>
+    createAggregationOperation(opId, {
+      addOperationHandler: addLokiOperation,
+      orderRank: LokiOperationOrder.Last,
+    })
+  );
+
   const list: QueryBuilderOperationDef[] = [
     createRangeOperation(LokiOperationId.Rate),
     createRangeOperation(LokiOperationId.CountOverTime),
@@ -19,10 +31,7 @@ export function getOperationDefintions(): QueryBuilderOperationDef[] {
     createRangeOperation(LokiOperationId.BytesRate),
     createRangeOperation(LokiOperationId.BytesOverTime),
     createRangeOperation(LokiOperationId.AbsentOverTime),
-    createAggregationOperation(LokiOperationId.Sum),
+    ...aggregations,
-    createAggregationOperation(LokiOperationId.Avg),
-    createAggregationOperation(LokiOperationId.Min),
-    createAggregationOperation(LokiOperationId.Max),
     {
       id: LokiOperationId.Json,
       name: 'Json',
@@ -199,24 +208,6 @@ function createRangeOperation(name: string): QueryBuilderOperationDef {
   };
 }
 
-function createAggregationOperation(name: string): QueryBuilderOperationDef {
-  return {
-    id: name,
-    name: getPromAndLokiOperationDisplayName(name),
-    params: [],
-    defaultParams: [],
-    alternativesKey: 'plain aggregation',
-    category: LokiVisualQueryOperationCategory.Aggregations,
-    orderRank: LokiOperationOrder.Last,
-    renderer: functionRendererLeft,
-    addOperationHandler: addLokiOperation,
-    explainHandler: (op, def) => {
-      const opDocs = FUNCTIONS.find((x) => x.insertText === op.id);
-      return `${opDocs?.documentation}.`;
-    },
-  };
-}
-
 function getRangeVectorParamDef(): QueryBuilderOperationParamDef {
   return {
     name: 'Range',

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

@@ -41,6 +41,8 @@ export enum LokiOperationId {
   Avg = 'avg',
   Min = 'min',
   Max = 'max',
+  TopK = 'topk',
+  BottomK = 'bottomk',
   LineContains = '__line_contains',
   LineContainsNot = '__line_contains_not',
   LineMatchesRegex = '__line_matches_regex',

public/app/plugins/datasource/loki/streaming.ts

@@ -1,7 +1,7 @@
 import { DataFrameJSON, DataQueryRequest, DataQueryResponse, LiveChannelScope, LoadingState } from '@grafana/data';
 import { getGrafanaLiveSrv } from '@grafana/runtime';
 import { map, Observable, defer, mergeMap } from 'rxjs';
-import LokiDatasource from './datasource';
+import { LokiDatasource } from './datasource';
 import { LokiQuery } from './types';
 import { StreamingDataFrame } from 'app/features/live/data/StreamingDataFrame';
 

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

@@ -7,7 +7,12 @@ import {
   getPromAndLokiOperationDisplayName,
   getRangeVectorParamDef,
 } from './shared/operationUtils';
-import { QueryBuilderOperation, QueryBuilderOperationDef, QueryBuilderOperationParamDef } from './shared/types';
+import {
+  QueryBuilderOperation,
+  QueryBuilderOperationDef,
+  QueryBuilderOperationParamDef,
+  QueryWithOperations,
+} from './shared/types';
 import { PromVisualQueryOperationCategory, PromOperationId } from './types';
 
 export function getAggregationOperations(): QueryBuilderOperationDef[] {
@@ -17,7 +22,7 @@ export function getAggregationOperations(): QueryBuilderOperationDef[] {
     ...createAggregationOperation(PromOperationId.Min),
     ...createAggregationOperation(PromOperationId.Max),
     ...createAggregationOperation(PromOperationId.Count),
-    ...createAggregationOperation(PromOperationId.Topk),
+    ...createAggregationOperation(PromOperationId.TopK),
     ...createAggregationOperation(PromOperationId.BottomK),
     createAggregationOverTime(PromOperationId.SumOverTime),
     createAggregationOverTime(PromOperationId.AvgOverTime),
@@ -31,7 +36,13 @@ export function getAggregationOperations(): QueryBuilderOperationDef[] {
   ];
 }
 
-function createAggregationOperation(name: string): 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,
@@ -48,8 +59,10 @@ function createAggregationOperation(name: string): QueryBuilderOperationDef[] {
       alternativesKey: 'plain aggregations',
       category: PromVisualQueryOperationCategory.Aggregations,
       renderer: functionRendererLeft,
-      addOperationHandler: defaultAddOperationHandler,
       paramChangedHandler: getOnLabelAdddedHandler(`__${name}_by`),
+      explainHandler: getAggregationExplainer(name, ''),
+      addOperationHandler: defaultAddOperationHandler,
+      ...overrides,
     },
     {
       id: `__${name}_by`,
@@ -67,10 +80,11 @@ function createAggregationOperation(name: string): QueryBuilderOperationDef[] {
       alternativesKey: 'aggregations by',
       category: PromVisualQueryOperationCategory.Aggregations,
       renderer: getAggregationByRenderer(name),
-      addOperationHandler: defaultAddOperationHandler,
       paramChangedHandler: getLastLabelRemovedHandler(name),
       explainHandler: getAggregationExplainer(name, 'by'),
+      addOperationHandler: defaultAddOperationHandler,
       hideFromList: true,
+      ...overrides,
     },
     {
       id: `__${name}_without`,
@@ -88,10 +102,11 @@ function createAggregationOperation(name: string): QueryBuilderOperationDef[] {
       alternativesKey: 'aggregations by',
       category: PromVisualQueryOperationCategory.Aggregations,
       renderer: getAggregationWithoutRenderer(name),
-      addOperationHandler: defaultAddOperationHandler,
       paramChangedHandler: getLastLabelRemovedHandler(name),
       explainHandler: getAggregationExplainer(name, 'without'),
+      addOperationHandler: defaultAddOperationHandler,
       hideFromList: true,
+      ...overrides,
     },
   ];
 
@@ -126,14 +141,18 @@ function getAggregationWithoutRenderer(aggregation: string) {
 /**
  * Very simple poc implementation, needs to be modified to support all aggregation operators
  */
-function getAggregationExplainer(aggregationName: string, mode: 'by' | 'without') {
+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);
-    if (mode === 'by') {
+
+    switch (mode) {
+      case 'by':
         return `Calculates ${aggregationName} over dimensions while preserving ${labelWord} ${labels}.`;
-    } else {
+      case 'without':
         return `Calculates ${aggregationName} over the dimensions ${labels}. All other labels are preserved.`;
+      default:
+        return `Calculates ${aggregationName} over the dimensions.`;
     }
   };
 }
@@ -164,7 +183,7 @@ function getLastLabelRemovedHandler(changeToOperartionId: string) {
   };
 }
 
-function getOnLabelAdddedHandler(changeToOperartionId: string) {
+export function getOnLabelAdddedHandler(changeToOperartionId: string) {
   return function onParamChanged(index: number, op: QueryBuilderOperation) {
     return {
       ...op,

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

@@ -1,10 +1,10 @@
-import { SelectableValue, toOption } from '@grafana/data';
+import { DataSourceApi, SelectableValue, toOption } from '@grafana/data';
 import { Select } from '@grafana/ui';
 import React, { useState } from 'react';
 import { PrometheusDatasource } from '../../datasource';
 import { promQueryModeller } from '../PromQueryModeller';
 import { getOperationParamId } from '../shared/operationUtils';
-import { QueryBuilderOperationParamEditorProps } from '../shared/types';
+import { QueryBuilderLabelFilter, QueryBuilderOperationParamEditorProps } from '../shared/types';
 import { PromVisualQuery } from '../types';
 
 export function LabelParamEditor({
@@ -28,7 +28,7 @@ export function LabelParamEditor({
       openMenuOnFocus
       onOpenMenu={async () => {
         setState({ isLoading: true });
-        const options = await loadGroupByLabels(query as PromVisualQuery, datasource as PrometheusDatasource);
+        const options = await loadGroupByLabels(query, datasource);
         setState({ options, isLoading: undefined });
       }}
       isLoading={state.isLoading}
@@ -44,11 +44,16 @@ export function LabelParamEditor({
 
 async function loadGroupByLabels(
   query: PromVisualQuery,
-  datasource: PrometheusDatasource
+  datasource: DataSourceApi
 ): Promise<Array<SelectableValue<any>>> {
-  const labels = [{ label: '__name__', op: '=', value: query.metric }, ...query.labels];
+  let labels: QueryBuilderLabelFilter[] = query.labels;
-  const expr = promQueryModeller.renderLabels(labels);
 
+  // This function is used by both Prometheus and Loki and this the only difference
+  if (datasource instanceof PrometheusDatasource) {
+    labels = [{ label: '__name__', op: '=', value: query.metric }, ...query.labels];
+  }
+
+  const expr = promQueryModeller.renderLabels(labels);
   const result = await datasource.languageProvider.fetchSeriesLabels(expr);
 
   return Object.keys(result).map((x) => ({

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

@@ -32,7 +32,7 @@ export interface QueryBuilderOperationDef<T = any> extends RegistryItem {
   renderer: QueryBuilderOperationRenderer;
   addOperationHandler: QueryBuilderAddOperationHandler<T>;
   paramChangedHandler?: QueryBuilderOnParamChangedHandler;
-  explainHandler?: (op: QueryBuilderOperation, def: QueryBuilderOperationDef<T>) => string;
+  explainHandler?: QueryBuilderExplainOperationHandler;
   changeTypeHandler?: (op: QueryBuilderOperation, newDef: QueryBuilderOperationDef<T>) => QueryBuilderOperation;
 }
 
@@ -42,6 +42,8 @@ export type QueryBuilderAddOperationHandler<T> = (
   modeller: VisualQueryModeller
 ) => T;
 
+export type QueryBuilderExplainOperationHandler = (op: QueryBuilderOperation, def: QueryBuilderOperationDef) => string;
+
 export type QueryBuilderOnParamChangedHandler = (
   index: number,
   operation: QueryBuilderOperation,

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

@@ -100,7 +100,7 @@ export enum PromOperationId {
   Tanh = 'tanh',
   Time = 'time',
   Timestamp = 'timestamp',
-  Topk = 'topk',
+  TopK = 'topk',
   Vector = 'vector',
   Year = 'year',
   // Binary ops

public/app/plugins/datasource/tempo/QueryEditor/QueryField.tsx

@@ -16,7 +16,7 @@ import React from 'react';
 import { LokiQueryField } from '../../loki/components/LokiQueryField';
 import { LokiQuery } from '../../loki/types';
 import { TempoDatasource, TempoQuery, TempoQueryType } from '../datasource';
-import LokiDatasource from '../../loki/datasource';
+import { LokiDatasource } from '../../loki/datasource';
 import useAsync from 'react-use/lib/useAsync';
 import NativeSearch from './NativeSearch';
 import { getDS } from './utils';