Loki Autocomplete: Suggest only possible labels for unwrap (#61411)

* feat(loki-autocomplete): filter valid labels for unwrap autocomplete suggestions

* refactor(unwrap): reuse new function in builder

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

@@ -6,7 +6,11 @@ import { TypeaheadInput } from '@grafana/ui';
 import LanguageProvider, { LokiHistoryItem } from './LanguageProvider';
 import { LokiDatasource } from './datasource';
 import { createLokiDatasource, createMetadataRequest } from './mocks';
-import { extractLogParserFromDataFrame, extractLabelKeysFromDataFrame } from './responseUtils';
+import {
+  extractLogParserFromDataFrame,
+  extractLabelKeysFromDataFrame,
+  extractUnwrapLabelKeysFromDataFrame,
+} from './responseUtils';
 import { LokiQueryType } from './types';
 
 jest.mock('./responseUtils');
@@ -304,11 +308,13 @@ describe('Query imports', () => {
     let datasource: LokiDatasource, languageProvider: LanguageProvider;
     const extractLogParserFromDataFrameMock = jest.mocked(extractLogParserFromDataFrame);
     const extractedLabelKeys = ['extracted', 'label'];
+    const unwrapLabelKeys = ['unwrap', 'labels'];
 
     beforeEach(() => {
       datasource = createLokiDatasource();
       languageProvider = new LanguageProvider(datasource);
       jest.mocked(extractLabelKeysFromDataFrame).mockReturnValue(extractedLabelKeys);
+      jest.mocked(extractUnwrapLabelKeysFromDataFrame).mockReturnValue(unwrapLabelKeys);
     });
 
     it('identifies selectors with JSON parser data', async () => {
@@ -317,6 +323,7 @@ describe('Query imports', () => {
 
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys,
+        unwrapLabelKeys,
         hasJSON: true,
         hasLogfmt: false,
       });
@@ -328,6 +335,7 @@ describe('Query imports', () => {
 
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys,
+        unwrapLabelKeys,
         hasJSON: false,
         hasLogfmt: true,
       });
@@ -339,6 +347,7 @@ describe('Query imports', () => {
 
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys: [],
+        unwrapLabelKeys: [],
         hasJSON: false,
         hasLogfmt: false,
       });

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

@@ -12,7 +12,11 @@ import {
 } from 'app/plugins/datasource/prometheus/language_utils';
 
 import { LokiDatasource } from './datasource';
-import { extractLabelKeysFromDataFrame, extractLogParserFromDataFrame } from './responseUtils';
+import {
+  extractLabelKeysFromDataFrame,
+  extractLogParserFromDataFrame,
+  extractUnwrapLabelKeysFromDataFrame,
+} from './responseUtils';
 import syntax, { FUNCTIONS, PIPE_PARSERS, PIPE_OPERATORS } from './syntax';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -465,15 +469,20 @@ export default class LokiLanguageProvider extends LanguageProvider {
 
   async getParserAndLabelKeys(
     selector: string
-  ): Promise<{ extractedLabelKeys: string[]; hasJSON: boolean; hasLogfmt: boolean }> {
+  ): Promise<{ extractedLabelKeys: string[]; hasJSON: boolean; hasLogfmt: boolean; unwrapLabelKeys: string[] }> {
     const series = await this.datasource.getDataSamples({ expr: selector, refId: 'data-samples' });
 
     if (!series.length) {
-      return { extractedLabelKeys: [], hasJSON: false, hasLogfmt: false };
+      return { extractedLabelKeys: [], unwrapLabelKeys: [], hasJSON: false, hasLogfmt: false };
     }
 
     const { hasLogfmt, hasJSON } = extractLogParserFromDataFrame(series[0]);
 
-    return { extractedLabelKeys: extractLabelKeysFromDataFrame(series[0]), hasJSON, hasLogfmt };
+    return {
+      extractedLabelKeys: extractLabelKeysFromDataFrame(series[0]),
+      unwrapLabelKeys: extractUnwrapLabelKeysFromDataFrame(series[0]),
+      hasJSON,
+      hasLogfmt,
+    };
   }
 }

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

@@ -50,6 +50,7 @@ const otherLabels: Label[] = [
 const seriesLabels = { place: ['series', 'labels'], source: [], other: [] };
 const parserAndLabelKeys = {
   extractedLabelKeys: ['extracted', 'label', 'keys'],
+  unwrapLabelKeys: ['unwrap', 'labels'],
   hasJSON: true,
   hasLogfmt: false,
 };

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

@@ -31,6 +31,7 @@ const labelNames = ['place', 'source'];
 const labelValues = ['moon', 'luna', 'server\\1'];
 // Source is duplicated to test handling duplicated labels
 const extractedLabelKeys = ['extracted', 'place', 'source'];
+const unwrapLabelKeys = ['unwrap', 'labels'];
 const otherLabels: Label[] = [
   {
     name: 'place',
@@ -195,6 +196,7 @@ describe('getCompletions', () => {
     jest.spyOn(completionProvider, 'getLabelValues').mockResolvedValue(labelValues);
     jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
       extractedLabelKeys,
+      unwrapLabelKeys,
       hasJSON: false,
       hasLogfmt: false,
     });
@@ -327,6 +329,7 @@ describe('getCompletions', () => {
     async (afterPipe: boolean, hasSpace: boolean) => {
       jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
         extractedLabelKeys,
+        unwrapLabelKeys,
         hasJSON: true,
         hasLogfmt: false,
       });
@@ -343,6 +346,7 @@ describe('getCompletions', () => {
     async (afterPipe: boolean) => {
       jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
         extractedLabelKeys,
+        unwrapLabelKeys,
         hasJSON: false,
         hasLogfmt: true,
       });
@@ -368,7 +372,20 @@ describe('getCompletions', () => {
     const extractedCompletions = completions.filter((completion) => completion.type === 'LABEL_NAME');
     const functionCompletions = completions.filter((completion) => completion.type === 'FUNCTION');
 
-    expect(extractedCompletions).toHaveLength(3);
+    expect(extractedCompletions).toEqual([
+      {
+        insertText: 'unwrap',
+        label: 'unwrap',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
+      {
+        insertText: 'labels',
+        label: 'labels',
+        triggerOnInsert: false,
+        type: 'LABEL_NAME',
+      },
+    ]);
     expect(functionCompletions).toHaveLength(3);
   });
 });

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

@@ -273,9 +273,9 @@ async function getAfterUnwrapCompletions(
   logQuery: string,
   dataProvider: CompletionDataProvider
 ): Promise<Completion[]> {
-  const { extractedLabelKeys } = await dataProvider.getParserAndLabelKeys(logQuery);
+  const { unwrapLabelKeys } = await dataProvider.getParserAndLabelKeys(logQuery);
 
-  const labelCompletions: Completion[] = extractedLabelKeys.map((label) => ({
+  const labelCompletions: Completion[] = unwrapLabelKeys.map((label) => ({
     type: 'LABEL_NAME',
     label,
     insertText: label,

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

@@ -1,14 +1,13 @@
-import { isNaN } from 'lodash';
 import React, { useState } from 'react';
 
-import { isValidGoDuration, SelectableValue, toOption } from '@grafana/data';
+import { SelectableValue, toOption } from '@grafana/data';
 import { Select } from '@grafana/ui';
 
 import { getOperationParamId } from '../../../prometheus/querybuilder/shared/operationUtils';
 import { QueryBuilderOperationParamEditorProps } from '../../../prometheus/querybuilder/shared/types';
 import { LokiDatasource } from '../../datasource';
-import { isBytesString } from '../../languageUtils';
 import { getLogQueryFromMetricsQuery, isValidQuery } from '../../queryUtils';
+import { extractUnwrapLabelKeysFromDataFrame } from '../../responseUtils';
 import { lokiQueryModeller } from '../LokiQueryModeller';
 import { LokiVisualQuery } from '../types';
 
@@ -62,30 +61,7 @@ async function loadUnwrapOptions(
   }
 
   const samples = await datasource.getDataSamples({ expr: logExpr, refId: 'unwrap_samples' });
-  const labelsArray: Array<{ [key: string]: string }> | undefined =
-    samples[0]?.fields?.find((field) => field.name === 'labels')?.values.toArray() ?? [];
-
-  if (!labelsArray || labelsArray.length === 0) {
-    return [];
-  }
-
-  // We do this only for first label object, because we want to consider only labels that are present in all log lines
-  // possibleUnwrapLabels are labels with 1. number value OR 2. value that is valid go duration OR 3. bytes string value
-  const possibleUnwrapLabels = Object.keys(labelsArray[0]).filter((key) => {
-    const value = labelsArray[0][key];
-    if (!value) {
-      return false;
-    }
-    return !isNaN(Number(value)) || isValidGoDuration(value) || isBytesString(value);
-  });
-
-  const unwrapLabels: string[] = [];
-  for (const label of possibleUnwrapLabels) {
-    // Add only labels that are present in every line to unwrapLabels
-    if (labelsArray.every((obj) => obj[label])) {
-      unwrapLabels.push(label);
-    }
-  }
+  const unwrapLabels = extractUnwrapLabelKeysFromDataFrame(samples[0]);
 
   const labelOptions = unwrapLabels.map((label) => ({
     label,

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

@@ -8,6 +8,7 @@ import {
   extractLevelLikeLabelFromDataFrame,
   extractLogParserFromDataFrame,
   extractLabelKeysFromDataFrame,
+  extractUnwrapLabelKeysFromDataFrame,
 } from './responseUtils';
 
 const frame: DataFrame = {
@@ -105,3 +106,16 @@ describe('extractLabelKeysFromDataFrame', () => {
     expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);
   });
 });
+
+describe('extractUnwrapLabelKeysFromDataFrame', () => {
+  it('returns empty by default', () => {
+    const input = cloneDeep(frame);
+    input.fields[1].values = new ArrayVector([]);
+    expect(extractUnwrapLabelKeysFromDataFrame(input)).toEqual([]);
+  });
+  it('extracts possible unwrap label keys', () => {
+    const input = cloneDeep(frame);
+    input.fields[1].values = new ArrayVector([{ number: 13 }]);
+    expect(extractUnwrapLabelKeysFromDataFrame(input)).toEqual(['number']);
+  });
+});

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

@@ -1,7 +1,9 @@
-import { DataFrame, FieldType, Labels } from '@grafana/data';
+import { DataFrame, FieldType, isValidGoDuration, Labels } from '@grafana/data';
 
 import { isLogLineJSON, isLogLineLogfmt } from './lineParser';
 
+import { isBytesString } from './languageUtils';
+
 export function dataFrameHasLokiError(frame: DataFrame): boolean {
   const labelSets: Labels[] = frame.fields.find((f) => f.name === 'labels')?.values.toArray() ?? [];
   return labelSets.some((labels) => labels.__error__ !== undefined);
@@ -46,6 +48,28 @@ export function extractLabelKeysFromDataFrame(frame: DataFrame): string[] {
   return Object.keys(labelsArray[0]);
 }
 
+export function extractUnwrapLabelKeysFromDataFrame(frame: DataFrame): string[] {
+  const labelsArray: Array<{ [key: string]: string }> | undefined =
+    frame?.fields?.find((field) => field.name === 'labels')?.values.toArray() ?? [];
+
+  if (!labelsArray?.length) {
+    return [];
+  }
+
+  // We do this only for first label object, because we want to consider only labels that are present in all log lines
+  // possibleUnwrapLabels are labels with 1. number value OR 2. value that is valid go duration OR 3. bytes string value
+  const possibleUnwrapLabels = Object.keys(labelsArray[0]).filter((key) => {
+    const value = labelsArray[0][key];
+    if (!value) {
+      return false;
+    }
+    return !isNaN(Number(value)) || isValidGoDuration(value) || isBytesString(value);
+  });
+
+  // Add only labels that are present in every line to unwrapLabels
+  return possibleUnwrapLabels.filter((label) => labelsArray.every((obj) => obj[label]));
+}
+
 export function extractHasErrorLabelFromDataFrame(frame: DataFrame): boolean {
   const labelField = frame.fields.find((field) => field.name === 'labels' && field.type === FieldType.other);
   if (labelField == null) {