Loki: Implement keep and drop operations (#73636)

* Update lezer

* Add functionalities for code and builder

* Add comment

package.json

@@ -240,7 +240,7 @@
     "@grafana/faro-core": "1.1.2",
     "@grafana/faro-web-sdk": "1.1.2",
     "@grafana/google-sdk": "0.1.1",
-    "@grafana/lezer-logql": "0.1.8",
+    "@grafana/lezer-logql": "0.1.9",
     "@grafana/lezer-traceql": "0.0.4",
     "@grafana/monaco-logql": "^0.0.7",
     "@grafana/runtime": "workspace:*",

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

@@ -130,6 +130,18 @@ const afterSelectorCompletions = [
     label: 'distinct',
     type: 'PIPE_OPERATION',
   },
+  {
+    documentation: 'Operator docs',
+    insertText: '| drop',
+    label: 'drop',
+    type: 'PIPE_OPERATION',
+  },
+  {
+    documentation: 'Operator docs',
+    insertText: '| keep',
+    label: 'keep',
+    type: 'PIPE_OPERATION',
+  },
 ];
 
 function buildAfterSelectorCompletions(
@@ -413,6 +425,32 @@ describe('getCompletions', () => {
       },
     ]);
   });
+
+  test('Returns completion options when the situation is AFTER_KEEP_AND_DROP', async () => {
+    const situation: Situation = { type: 'AFTER_KEEP_AND_DROP', 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

@@ -293,6 +293,20 @@ export async function getAfterSelectorCompletions(
     documentation: explainOperator(LokiOperationId.Distinct),
   });
 
+  completions.push({
+    type: 'PIPE_OPERATION',
+    label: 'drop',
+    insertText: `${prefix}drop`,
+    documentation: explainOperator(LokiOperationId.Drop),
+  });
+
+  completions.push({
+    type: 'PIPE_OPERATION',
+    label: 'keep',
+    insertText: `${prefix}keep`,
+    documentation: explainOperator(LokiOperationId.Keep),
+  });
+
   // Let's show label options only if query has parser
   if (hasQueryParser) {
     extractedLabelKeys.forEach((key) => {
@@ -357,6 +371,18 @@ async function getAfterDistinctCompletions(logQuery: string, dataProvider: Compl
   return [...labelCompletions];
 }
 
+async function getAfterKeepAndDropCompletions(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
@@ -393,6 +419,8 @@ export async function getCompletions(
       return [...FUNCTION_COMPLETIONS, ...AGGREGATION_COMPLETIONS];
     case 'AFTER_DISTINCT':
       return getAfterDistinctCompletions(situation.logQuery, dataProvider);
+    case 'AFTER_KEEP_AND_DROP':
+      return getAfterKeepAndDropCompletions(situation.logQuery, dataProvider);
     default:
       throw new NeverCaseError(situation);
   }

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

@@ -276,4 +276,36 @@ describe('situation', () => {
       logQuery: '{label="value"} | logfmt ',
     });
   });
+
+  it('identifies AFTER_KEEP_AND_DROP autocomplete situations', () => {
+    assertSituation('{label="value"} | logfmt | drop^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | keep^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | drop id,^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | keep id,^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | drop id, name="test",^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+
+    assertSituation('{label="value"} | logfmt | keep id, name="test",^', {
+      type: 'AFTER_KEEP_AND_DROP',
+      logQuery: '{label="value"} | logfmt ',
+    });
+  });
 });

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

@@ -22,6 +22,10 @@ import {
   UnwrapExpr,
   DistinctFilter,
   DistinctLabel,
+  DropLabelsExpr,
+  KeepLabelsExpr,
+  DropLabels,
+  KeepLabels,
 } from '@grafana/lezer-logql';
 
 import { getLogQueryFromMetricsQuery } from '../../../queryUtils';
@@ -131,6 +135,10 @@ export type Situation =
   | {
       type: 'AFTER_DISTINCT';
       logQuery: string;
+    }
+  | {
+      type: 'AFTER_KEEP_AND_DROP';
+      logQuery: string;
     };
 
 type Resolver = {
@@ -205,6 +213,22 @@ const RESOLVERS: Resolver[] = [
     path: [ERROR_NODE_ID, DistinctLabel],
     fun: resolveAfterDistinct,
   },
+  {
+    path: [ERROR_NODE_ID, DropLabelsExpr],
+    fun: resolveAfterKeepAndDrop,
+  },
+  {
+    path: [ERROR_NODE_ID, DropLabels],
+    fun: resolveAfterKeepAndDrop,
+  },
+  {
+    path: [ERROR_NODE_ID, KeepLabelsExpr],
+    fun: resolveAfterKeepAndDrop,
+  },
+  {
+    path: [ERROR_NODE_ID, KeepLabels],
+    fun: resolveAfterKeepAndDrop,
+  },
 ];
 
 const LABEL_OP_MAP = new Map<string, LabelOperator>([
@@ -532,6 +556,28 @@ function resolveAfterDistinct(node: SyntaxNode, text: string, pos: number): Situ
   };
 }
 
+function resolveAfterKeepAndDrop(node: SyntaxNode, text: string, pos: number): Situation | null {
+  let logQuery = getLogQueryFromMetricsQuery(text).trim();
+  let keepAndDropParent: SyntaxNode | null = null;
+  let parent = node.parent;
+  while (parent !== null) {
+    if (parent.type.id === PipelineStage) {
+      keepAndDropParent = parent;
+      break;
+    }
+    parent = parent.parent;
+  }
+
+  if (keepAndDropParent?.type.id === PipelineStage) {
+    logQuery = logQuery.slice(0, keepAndDropParent.from);
+  }
+
+  return {
+    type: 'AFTER_KEEP_AND_DROP',
+    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

@@ -508,6 +508,55 @@ Example: \`\`error_level=\`level\` \`\`
       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.',
     },
+    {
+      id: LokiOperationId.Drop,
+      name: 'Drop',
+      params: [
+        // As drop can support both labels (e.g. job) and expressions (e.g. job="grafana"), we
+        // use input and not LabelParamEditor.
+        {
+          name: 'Label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+          minWidth: 18,
+          placeholder: 'job="grafana"',
+          description: 'Specify labels or expressions to drop.',
+        },
+      ],
+      defaultParams: [''],
+      alternativesKey: 'format',
+      category: LokiVisualQueryOperationCategory.Formats,
+      orderRank: LokiOperationOrder.PipeOperations,
+      renderer: (op, def, innerExpr) => `${innerExpr} | drop ${op.params.join(',')}`,
+      addOperationHandler: addLokiOperation,
+      explainHandler: () => 'The drop expression will drop the given labels in the pipeline.',
+    },
+    {
+      id: LokiOperationId.Keep,
+      name: 'Keep',
+      params: [
+        // As keep can support both labels (e.g. job) and expressions (e.g. job="grafana"), we
+        // use input and not LabelParamEditor.
+        {
+          name: 'Label',
+          type: 'string',
+          restParam: true,
+          optional: true,
+          minWidth: 18,
+          placeholder: 'job="grafana"',
+          description: 'Specify labels or expressions to keep.',
+        },
+      ],
+      defaultParams: [''],
+      alternativesKey: 'format',
+      category: LokiVisualQueryOperationCategory.Formats,
+      orderRank: LokiOperationOrder.PipeOperations,
+      renderer: (op, def, innerExpr) => `${innerExpr} | keep ${op.params.join(',')}`,
+      addOperationHandler: addLokiOperation,
+      explainHandler: () =>
+        'The keep expression will keep only the specified labels in the pipeline and drop all the other labels.',
+    },
     ...binaryScalarOperations,
     {
       id: LokiOperationId.NestedQuery,

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

@@ -776,6 +776,96 @@ describe('buildVisualQueryFromString', () => {
       })
     );
   });
+
+  it('parses a log query with drop and no labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | drop')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Drop, params: [] }],
+      })
+    );
+  });
+
+  it('parses a log query with drop and labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | drop id, email')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Drop, params: ['id', 'email'] }],
+      })
+    );
+  });
+
+  it('parses a log query with drop, labels and expressions', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | drop id, email, test="test1"')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Drop, params: ['id', 'email', 'test="test1"'] }],
+      })
+    );
+  });
+
+  it('parses a log query with keep and no labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | keep')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Keep, params: [] }],
+      })
+    );
+  });
+
+  it('parses a log query with keep and labels', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | keep id, email')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Keep, params: ['id', 'email'] }],
+      })
+    );
+  });
+
+  it('parses a log query with keep, labels and expressions', () => {
+    expect(buildVisualQueryFromString('{app="frontend"} | keep id, email, test="test1"')).toEqual(
+      noErrors({
+        labels: [
+          {
+            op: '=',
+            value: 'frontend',
+            label: 'app',
+          },
+        ],
+        operations: [{ id: LokiOperationId.Keep, params: ['id', 'email', 'test="test1"'] }],
+      })
+    );
+  });
 });
 
 function noErrors(query: LokiVisualQuery) {

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

@@ -10,6 +10,9 @@ import {
   Decolorize,
   DistinctFilter,
   DistinctLabel,
+  DropLabel,
+  DropLabels,
+  DropLabelsExpr,
   Filter,
   FilterOp,
   Grouping,
@@ -21,6 +24,9 @@ import {
   Json,
   JsonExpression,
   JsonExpressionParser,
+  KeepLabel,
+  KeepLabels,
+  KeepLabelsExpr,
   LabelFilter,
   LabelFormatMatcher,
   LabelParser,
@@ -212,6 +218,16 @@ export function handleExpression(expr: string, node: SyntaxNode, context: Contex
       break;
     }
 
+    case DropLabelsExpr: {
+      visQuery.operations.push(handleDropFilter(expr, node, context));
+      break;
+    }
+
+    case KeepLabelsExpr: {
+      visQuery.operations.push(handleKeepFilter(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.
@@ -660,3 +676,37 @@ function handleDistinctFilter(expr: string, node: SyntaxNode, context: Context):
     params: labels,
   };
 }
+
+function handleDropFilter(expr: string, node: SyntaxNode, context: Context): QueryBuilderOperation {
+  const labels: string[] = [];
+  let exploringNode = node.getChild(DropLabels);
+  while (exploringNode) {
+    const label = getString(expr, exploringNode.getChild(DropLabel));
+    if (label) {
+      labels.push(label);
+    }
+    exploringNode = exploringNode?.getChild(DropLabels);
+  }
+  labels.reverse();
+  return {
+    id: LokiOperationId.Drop,
+    params: labels,
+  };
+}
+
+function handleKeepFilter(expr: string, node: SyntaxNode, context: Context): QueryBuilderOperation {
+  const labels: string[] = [];
+  let exploringNode = node.getChild(KeepLabels);
+  while (exploringNode) {
+    const label = getString(expr, exploringNode.getChild(KeepLabel));
+    if (label) {
+      labels.push(label);
+    }
+    exploringNode = exploringNode?.getChild(KeepLabels);
+  }
+  labels.reverse();
+  return {
+    id: LokiOperationId.Keep,
+    params: labels,
+  };
+}

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

@@ -42,6 +42,8 @@ export enum LokiOperationId {
   LineFormat = 'line_format',
   LabelFormat = 'label_format',
   Decolorize = 'decolorize',
+  Drop = 'drop',
+  Keep = 'keep',
   Rate = 'rate',
   RateCounter = 'rate_counter',
   CountOverTime = 'count_over_time',

yarn.lock

@@ -3857,14 +3857,14 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@grafana/lezer-logql@npm:0.1.8":
-  version: 0.1.8
-  resolution: "@grafana/lezer-logql@npm:0.1.8"
+"@grafana/lezer-logql@npm:0.1.9":
+  version: 0.1.9
+  resolution: "@grafana/lezer-logql@npm:0.1.9"
   dependencies:
     lodash: ^4.17.21
   peerDependencies:
     "@lezer/lr": ^1.0.0
-  checksum: f0f301b6d4fbd2d79563b5b4e34303257be0ea995b2b9fa1f012648654b4afaa9cea91642bc59eddb70e9fa24ec8804489c161f7065b41eef49db68d3a2ca561
+  checksum: 6cceb18586413864137ef2305bbe2fdce054796c61a3fde4864e3f2d30ea3dac853aa701728867bf609a4009334dbf92fddbc06f7291dbca9459a86c9dc29e67
   languageName: node
   linkType: hard
 
@@ -19244,7 +19244,7 @@ __metadata:
     "@grafana/faro-core": 1.1.2
     "@grafana/faro-web-sdk": 1.1.2
     "@grafana/google-sdk": 0.1.1
-    "@grafana/lezer-logql": 0.1.8
+    "@grafana/lezer-logql": 0.1.9
     "@grafana/lezer-traceql": 0.0.4
     "@grafana/monaco-logql": ^0.0.7
     "@grafana/runtime": "workspace:*"