Loki Monaco Editor: implement extracted label keys (#57368)

* feat(loki-monaco-editor): implement extracted label keys

* Chore: add missing responseUtils tests

* feat(loki-monaco-editor): suggest extracted labels

* Chore: fix test case name

* feat(loki-monaco-editor): dont suggest labels in logs query

* Chore: remove console log

* Chore: remove extracted keyword from suggested label

* feat(loki-monaco-editor): do not suggest duplicated labels

* refactor(loki-monaco-editor): pass query and offset to the completions resolver

* Revert "refactor(loki-monaco-editor): pass query and offset to the completions resolver"

This reverts commit d39464fd1a4624d5cd5420156dd2d1e2dad2eecf.

* refactor(loki-monaco-editor): refactor label completions for grouping

* Chore: remove obsolete function

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

@@ -6,7 +6,7 @@ import { TypeaheadInput } from '@grafana/ui';
 import LanguageProvider, { LokiHistoryItem } from './LanguageProvider';
 import { LokiDatasource } from './datasource';
 import { createLokiDatasource, createMetadataRequest } from './mocks';
-import { extractLogParserFromDataFrame } from './responseUtils';
+import { extractLogParserFromDataFrame, extractLabelKeysFromDataFrame } from './responseUtils';
 import { LokiQueryType } from './types';
 
 jest.mock('./responseUtils');
@@ -302,10 +302,13 @@ describe('Query imports', () => {
 
   describe('getParserAndLabelKeys()', () => {
     let datasource: LokiDatasource, languageProvider: LanguageProvider;
-    const extractLogParserFromDataFrameMock = extractLogParserFromDataFrame as jest.Mock;
+    const extractLogParserFromDataFrameMock = jest.mocked(extractLogParserFromDataFrame);
+    const extractedLabelKeys = ['extracted', 'label'];
+
     beforeEach(() => {
       datasource = createLokiDatasource();
       languageProvider = new LanguageProvider(datasource);
+      jest.mocked(extractLabelKeysFromDataFrame).mockReturnValue(extractedLabelKeys);
     });
 
     it('identifies selectors with JSON parser data', async () => {
@@ -313,7 +316,7 @@ describe('Query imports', () => {
       extractLogParserFromDataFrameMock.mockReturnValueOnce({ hasLogfmt: false, hasJSON: true });
 
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
-        extractedLabelKeys: [],
+        extractedLabelKeys,
         hasJSON: true,
         hasLogfmt: false,
       });
@@ -324,7 +327,7 @@ describe('Query imports', () => {
       extractLogParserFromDataFrameMock.mockReturnValueOnce({ hasLogfmt: true, hasJSON: false });
 
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
-        extractedLabelKeys: [],
+        extractedLabelKeys,
         hasJSON: false,
         hasLogfmt: true,
       });

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

@@ -12,7 +12,7 @@ import {
 } from 'app/plugins/datasource/prometheus/language_utils';
 
 import { LokiDatasource } from './datasource';
-import { extractLogParserFromDataFrame } from './responseUtils';
+import { extractLabelKeysFromDataFrame, extractLogParserFromDataFrame } from './responseUtils';
 import syntax, { FUNCTIONS, PIPE_PARSERS, PIPE_OPERATORS } from './syntax';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -474,7 +474,6 @@ export default class LokiLanguageProvider extends LanguageProvider {
 
     const { hasLogfmt, hasJSON } = extractLogParserFromDataFrame(series[0]);
 
-    // TODO: figure out extractedLabelKeys
-    return { extractedLabelKeys: [], hasJSON, hasLogfmt };
+    return { extractedLabelKeys: extractLabelKeysFromDataFrame(series[0]), hasJSON, hasLogfmt };
   }
 }

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

@@ -29,7 +29,8 @@ const history = [
 
 const labelNames = ['place', 'source'];
 const labelValues = ['moon', 'luna', 'server\\1'];
-const extractedLabelKeys = ['extracted', 'label'];
+// Source is duplicated to test handling duplicated labels
+const extractedLabelKeys = ['extracted', 'place', 'source'];
 const otherLabels: Label[] = [
   {
     name: 'place',
@@ -98,12 +99,17 @@ const afterSelectorCompletions = [
   },
   {
     insertText: '| unwrap extracted',
-    label: 'unwrap extracted (detected)',
+    label: 'unwrap extracted',
     type: 'LINE_FILTER',
   },
   {
-    insertText: '| unwrap label',
-    label: 'unwrap label (detected)',
+    insertText: '| unwrap place',
+    label: 'unwrap place',
+    type: 'LINE_FILTER',
+  },
+  {
+    insertText: '| unwrap source',
+    label: 'unwrap source',
     type: 'LINE_FILTER',
   },
   {
@@ -215,6 +221,12 @@ describe('getCompletions', () => {
     const completions = await getCompletions(situation, completionProvider);
 
     expect(completions).toEqual([
+      {
+        insertText: 'extracted',
+        label: 'extracted',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
       {
         insertText: 'place',
         label: 'place',
@@ -227,18 +239,6 @@ describe('getCompletions', () => {
         triggerOnInsert: false,
         type: 'LABEL_NAME',
       },
-      {
-        insertText: 'extracted',
-        label: 'extracted (parsed)',
-        triggerOnInsert: false,
-        type: 'LABEL_NAME',
-      },
-      {
-        insertText: 'label',
-        label: 'label (parsed)',
-        triggerOnInsert: false,
-        type: 'LABEL_NAME',
-      },
     ]);
   });
 

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

@@ -109,48 +109,32 @@ async function getAllHistoryCompletions(dataProvider: CompletionDataProvider): P
   }));
 }
 
-async function getLabelNamesForCompletions(
-  suffix: string,
-  triggerOnInsert: boolean,
-  addExtractedLabels: boolean,
-  otherLabels: Label[],
-  dataProvider: CompletionDataProvider
-): Promise<Completion[]> {
-  const labelNames = await dataProvider.getLabelNames(otherLabels);
-  const result: Completion[] = labelNames.map((text) => ({
-    type: 'LABEL_NAME',
-    label: text,
-    insertText: `${text}${suffix}`,
-    triggerOnInsert,
-  }));
-
-  if (addExtractedLabels) {
-    const { extractedLabelKeys } = await dataProvider.getParserAndLabelKeys(otherLabels);
-    extractedLabelKeys.forEach((key) => {
-      result.push({
-        type: 'LABEL_NAME',
-        label: `${key} (parsed)`,
-        insertText: `${key}${suffix}`,
-        triggerOnInsert,
-      });
-    });
-  }
-
-  return result;
-}
-
 async function getLabelNamesForSelectorCompletions(
   otherLabels: Label[],
   dataProvider: CompletionDataProvider
 ): Promise<Completion[]> {
-  return getLabelNamesForCompletions('=', true, false, otherLabels, dataProvider);
+  const labelNames = await dataProvider.getLabelNames(otherLabels);
+
+  return labelNames.map((label) => ({
+    type: 'LABEL_NAME',
+    label,
+    insertText: `${label}=`,
+    triggerOnInsert: true,
+  }));
 }
 
 async function getInGroupingCompletions(
   otherLabels: Label[],
   dataProvider: CompletionDataProvider
 ): Promise<Completion[]> {
-  return getLabelNamesForCompletions('', false, true, otherLabels, dataProvider);
+  const { extractedLabelKeys } = await dataProvider.getParserAndLabelKeys(otherLabels);
+
+  return extractedLabelKeys.map((label) => ({
+    type: 'LABEL_NAME',
+    label,
+    insertText: label,
+    triggerOnInsert: false,
+  }));
 }
 
 const PARSERS = ['json', 'logfmt', 'pattern', 'regexp', 'unpack'];
@@ -204,7 +188,7 @@ async function getAfterSelectorCompletions(
   extractedLabelKeys.forEach((key) => {
     completions.push({
       type: 'LINE_FILTER',
-      label: `unwrap ${key} (detected)`,
+      label: `unwrap ${key}`,
       insertText: `${prefix}unwrap ${key}`,
     });
   });

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

@@ -2,7 +2,13 @@ import { cloneDeep } from 'lodash';
 
 import { ArrayVector, DataFrame, FieldType } from '@grafana/data';
 
-import { dataFrameHasLevelLabel, dataFrameHasLokiError, extractLevelLikeLabelFromDataFrame } from './responseUtils';
+import {
+  dataFrameHasLevelLabel,
+  dataFrameHasLokiError,
+  extractLevelLikeLabelFromDataFrame,
+  extractLogParserFromDataFrame,
+  extractLabelKeysFromDataFrame,
+} from './responseUtils';
 
 const frame: DataFrame = {
   length: 1,
@@ -70,3 +76,32 @@ describe('extractLevelLikeLabelFromDataFrame', () => {
     expect(extractLevelLikeLabelFromDataFrame(input)).toBe(null);
   });
 });
+
+describe('extractLogParserFromDataFrame', () => {
+  it('returns false by default', () => {
+    const input = cloneDeep(frame);
+    expect(extractLogParserFromDataFrame(input)).toEqual({ hasJSON: false, hasLogfmt: false });
+  });
+  it('identifies JSON', () => {
+    const input = cloneDeep(frame);
+    input.fields[2].values = new ArrayVector(['{"a":"b"}']);
+    expect(extractLogParserFromDataFrame(input)).toEqual({ hasJSON: true, hasLogfmt: false });
+  });
+  it('identifies logfmt', () => {
+    const input = cloneDeep(frame);
+    input.fields[2].values = new ArrayVector(['a=b']);
+    expect(extractLogParserFromDataFrame(input)).toEqual({ hasJSON: false, hasLogfmt: true });
+  });
+});
+
+describe('extractLabelKeysFromDataFrame', () => {
+  it('returns empty by default', () => {
+    const input = cloneDeep(frame);
+    input.fields[1].values = new ArrayVector([]);
+    expect(extractLabelKeysFromDataFrame(input)).toEqual([]);
+  });
+  it('extracts label keys', () => {
+    const input = cloneDeep(frame);
+    expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);
+  });
+});

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

@@ -36,6 +36,17 @@ export function extractLogParserFromDataFrame(frame: DataFrame): { hasLogfmt: bo
   return { hasLogfmt, hasJSON };
 }
 
+export function extractLabelKeysFromDataFrame(frame: DataFrame): string[] {
+  const labelsArray: Array<{ [key: string]: string }> | undefined =
+    frame?.fields?.find((field) => field.name === 'labels')?.values.toArray() ?? [];
+
+  if (!labelsArray?.length) {
+    return [];
+  }
+
+  return Object.keys(labelsArray[0]);
+}
+
 export function extractHasErrorLabelFromDataFrame(frame: DataFrame): boolean {
   const labelField = frame.fields.find((field) => field.name === 'labels' && field.type === FieldType.other);
   if (labelField == null) {