Loki: Add support for distinct operation in autocomplete and query builder (#69003)

* Loki Autocomplete: add support for suggesting distinct

* Loki query builder: add distinct as format operation

* Remove unused import

* Loki visual query: add support to parse distinct filters

* Query builder: use label param editor for distinct

* Loki Autocomplete: Improve distinct label suggestions

* Query Builder: improve distinct parsing

* Fix tests

* Update tests

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/completions.test.ts

@@ -124,6 +124,12 @@ const afterSelectorCompletions = [
     type: 'PIPE_OPERATION',
     documentation: 'Operator docs',
   },
+  {
+    documentation: 'Operator docs',
+    insertText: '| distinct',
+    label: 'distinct',
+    type: 'PIPE_OPERATION',
+  },
 ];
 
 function buildAfterSelectorCompletions(
@@ -382,6 +388,32 @@ describe('getCompletions', () => {
     ]);
     expect(functionCompletions).toHaveLength(3);
   });
+
+  test('Returns completion options when the situation is AFTER_DISTINCT', async () => {
+    const situation: Situation = { type: 'AFTER_DISTINCT', logQuery: '{label="value"}' };
+    const completions = await getCompletions(situation, completionProvider);
+
+    expect(completions).toEqual([
+      {
+        insertText: 'extracted',
+        label: 'extracted',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
+      {
+        insertText: 'place',
+        label: 'place',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
+      {
+        insertText: 'source',
+        label: 'source',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
+    ]);
+  });
 });
 
 describe('getAfterSelectorCompletions', () => {

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/completions.ts

@@ -288,6 +288,13 @@ export async function getAfterSelectorCompletions(
     documentation: explainOperator(LokiOperationId.Decolorize),
   });
 
+  completions.push({
+    type: 'PIPE_OPERATION',
+    label: 'distinct',
+    insertText: `${prefix}distinct`,
+    documentation: explainOperator(LokiOperationId.Distinct),
+  });
+
   // Let's show label options only if query has parser
   if (hasQueryParser) {
     extractedLabelKeys.forEach((key) => {
@@ -340,6 +347,18 @@ async function getAfterUnwrapCompletions(
   return [...labelCompletions, ...UNWRAP_FUNCTION_COMPLETIONS];
 }
 
+async function getAfterDistinctCompletions(logQuery: string, dataProvider: CompletionDataProvider) {
+  const { extractedLabelKeys } = await dataProvider.getParserAndLabelKeys(logQuery);
+  const labelCompletions: Completion[] = extractedLabelKeys.map((label) => ({
+    type: 'LABEL_NAME',
+    label,
+    insertText: label,
+    triggerOnInsert: false,
+  }));
+
+  return [...labelCompletions];
+}
+
 export async function getCompletions(
   situation: Situation,
   dataProvider: CompletionDataProvider
@@ -374,6 +393,8 @@ export async function getCompletions(
       return getAfterUnwrapCompletions(situation.logQuery, dataProvider);
     case 'IN_AGGREGATION':
       return [...FUNCTION_COMPLETIONS, ...AGGREGATION_COMPLETIONS];
+    case 'AFTER_DISTINCT':
+      return getAfterDistinctCompletions(situation.logQuery, dataProvider);
     default:
       throw new NeverCaseError(situation);
   }

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/situation.test.ts

@@ -264,4 +264,16 @@ describe('situation', () => {
       ],
     });
   });
+
+  it('identifies AFTER_DISTINCT autocomplete situations', () => {
+    assertSituation('{label="value"} | logfmt | distinct^', {
+      type: 'AFTER_DISTINCT',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | distinct id,^', {
+      type: 'AFTER_DISTINCT',
+      logQuery: '{label="value"} | logfmt ',
+    });
+  });
 });

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/situation.ts

@@ -20,6 +20,8 @@ import {
   LiteralExpr,
   MetricExpr,
   UnwrapExpr,
+  DistinctFilter,
+  DistinctLabel,
 } from '@grafana/lezer-logql';
 
 import { getLogQueryFromMetricsQuery } from '../../../queryUtils';
@@ -125,6 +127,10 @@ export type Situation =
   | {
       type: 'AFTER_UNWRAP';
       logQuery: string;
+    }
+  | {
+      type: 'AFTER_DISTINCT';
+      logQuery: string;
     };
 
 type Resolver = {
@@ -191,6 +197,14 @@ const RESOLVERS: Resolver[] = [
     path: [UnwrapExpr],
     fun: resolveAfterUnwrap,
   },
+  {
+    path: [ERROR_NODE_ID, DistinctFilter],
+    fun: resolveAfterDistinct,
+  },
+  {
+    path: [ERROR_NODE_ID, DistinctLabel],
+    fun: resolveAfterDistinct,
+  },
 ];
 
 const LABEL_OP_MAP = new Map<string, LabelOperator>([
@@ -495,6 +509,29 @@ function resolveSelector(node: SyntaxNode, text: string, pos: number): Situation
   };
 }
 
+function resolveAfterDistinct(node: SyntaxNode, text: string, pos: number): Situation | null {
+  let logQuery = getLogQueryFromMetricsQuery(text).trim();
+
+  let distinctFilterParent: SyntaxNode | null = null;
+  let parent = node.parent;
+  while (parent !== null) {
+    if (parent.type.id === PipelineStage) {
+      distinctFilterParent = parent;
+      break;
+    }
+    parent = parent.parent;
+  }
+
+  if (distinctFilterParent?.type.id === PipelineStage) {
+    logQuery = logQuery.slice(0, distinctFilterParent.from);
+  }
+
+  return {
+    type: 'AFTER_DISTINCT',
+    logQuery,
+  };
+}
+
 // we find the first error-node in the tree that is at the cursor-position.
 // NOTE: this might be too slow, might need to optimize it
 // (ideas: we do not need to go into every subtree, based on from/to)

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

@@ -1,3 +1,4 @@
+import { LabelParamEditor } from '../../prometheus/querybuilder/components/LabelParamEditor';
 import {
   createAggregationOperation,
   createAggregationOperationWithParam,
@@ -486,6 +487,27 @@ Example: \`\`error_level=\`level\` \`\`
       addOperationHandler: addLokiOperation,
       explainHandler: () => `This will remove ANSI color codes from log lines.`,
     },
+    {
+      id: LokiOperationId.Distinct,
+      name: 'Distinct',
+      params: [
+        {
+          name: 'Label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+          editor: LabelParamEditor,
+        },
+      ],
+      defaultParams: [''],
+      alternativesKey: 'format',
+      category: LokiVisualQueryOperationCategory.Formats,
+      orderRank: LokiOperationOrder.Unwrap,
+      renderer: (op, def, innerExpr) => `${innerExpr} | distinct ${op.params.join(',')}`,
+      addOperationHandler: addLokiOperation,
+      explainHandler: () =>
+        'Allows filtering log lines using their original and extracted labels to filter out duplicate label values. The first line occurrence of a distinct value is returned, and the others are dropped.',
+    },
     ...binaryScalarOperations,
     {
       id: LokiOperationId.NestedQuery,

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

@@ -716,6 +716,36 @@ describe('buildVisualQueryFromString', () => {
       },
     });
   });
+
+  it('parses a log query with distinct and no labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | distinct')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Distinct, params: [] }],
+      })
+    );
+  });
+
+  it('parses a log query with distinct and labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | distinct id, email')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Distinct, params: ['id', 'email'] }],
+      })
+    );
+  });
 });
 
 function noErrors(query: LokiVisualQuery) {

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

@@ -8,6 +8,8 @@ import {
   By,
   ConvOp,
   Decolorize,
+  DistinctFilter,
+  DistinctLabel,
   Filter,
   FilterOp,
   Grouping,
@@ -205,6 +207,11 @@ export function handleExpression(expr: string, node: SyntaxNode, context: Contex
       break;
     }
 
+    case DistinctFilter: {
+      visQuery.operations.push(handleDistinctFilter(expr, node, context));
+      break;
+    }
+
     default: {
       // Any other nodes we just ignore and go to its children. This should be fine as there are lots of wrapper
       // nodes that can be skipped.
@@ -422,6 +429,7 @@ function handleUnwrapExpr(
 
   return {};
 }
+
 function handleRangeAggregation(expr: string, node: SyntaxNode, context: Context) {
   const nameNode = node.getChild(RangeOp);
   const funcName = getString(expr, nameNode);
@@ -632,3 +640,20 @@ function isEmptyQuery(query: LokiVisualQuery) {
   }
   return false;
 }
+
+function handleDistinctFilter(expr: string, node: SyntaxNode, context: Context): QueryBuilderOperation {
+  const labels: string[] = [];
+  let exploringNode = node.getChild(DistinctLabel);
+  while (exploringNode) {
+    const label = getString(expr, exploringNode.getChild(Identifier));
+    if (label) {
+      labels.push(label);
+    }
+    exploringNode = exploringNode?.getChild(DistinctLabel);
+  }
+  labels.reverse();
+  return {
+    id: LokiOperationId.Distinct,
+    params: labels,
+  };
+}

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

@@ -38,6 +38,7 @@ export enum LokiOperationId {
   Regexp = 'regexp',
   Pattern = 'pattern',
   Unpack = 'unpack',
+  Distinct = 'distinct',
   LineFormat = 'line_format',
   LabelFormat = 'label_format',
   Decolorize = 'decolorize',