Loki: Add structured metadata keys to autocomplete (#78584)

* wip

* remove import

* scope to monaco completions for now

* use `LabelType` enum

* change strucutred metadata documentation

* fix import

* add `responseUtils` tests

* update tests

* fix completions.ts tests

* fix LabelType enum

* fix CompletionDataProvider test

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

@@ -8,7 +8,7 @@ import {
   extractLabelKeysFromDataFrame,
   extractUnwrapLabelKeysFromDataFrame,
 } from './responseUtils';
-import { LokiQueryType } from './types';
+import { LabelType, LokiQueryType } from './types';
 
 jest.mock('./responseUtils');
 
@@ -331,12 +331,19 @@ describe('Query imports', () => {
     let datasource: LokiDatasource, languageProvider: LanguageProvider;
     const extractLogParserFromDataFrameMock = jest.mocked(extractLogParserFromDataFrame);
     const extractedLabelKeys = ['extracted', 'label'];
+    const structuredMetadataKeys = ['structured', 'metadata'];
     const unwrapLabelKeys = ['unwrap', 'labels'];
 
     beforeEach(() => {
       datasource = createLokiDatasource();
       languageProvider = new LanguageProvider(datasource);
-      jest.mocked(extractLabelKeysFromDataFrame).mockReturnValue(extractedLabelKeys);
+      jest.mocked(extractLabelKeysFromDataFrame).mockImplementation((_, type) => {
+        if (type === LabelType.Indexed || !type) {
+          return extractedLabelKeys;
+        } else {
+          return structuredMetadataKeys;
+        }
+      });
       jest.mocked(extractUnwrapLabelKeysFromDataFrame).mockReturnValue(unwrapLabelKeys);
     });
 
@@ -347,6 +354,7 @@ describe('Query imports', () => {
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys,
         unwrapLabelKeys,
+        structuredMetadataKeys,
         hasJSON: true,
         hasLogfmt: false,
         hasPack: false,
@@ -360,6 +368,7 @@ describe('Query imports', () => {
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys,
         unwrapLabelKeys,
+        structuredMetadataKeys,
         hasJSON: false,
         hasLogfmt: true,
         hasPack: false,
@@ -373,6 +382,7 @@ describe('Query imports', () => {
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys: [],
         unwrapLabelKeys: [],
+        structuredMetadataKeys: [],
         hasJSON: false,
         hasLogfmt: false,
         hasPack: false,
@@ -386,6 +396,7 @@ describe('Query imports', () => {
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}')).toEqual({
         extractedLabelKeys: [],
         unwrapLabelKeys: [],
+        structuredMetadataKeys: [],
         hasJSON: false,
         hasLogfmt: false,
         hasPack: false,
@@ -406,6 +417,7 @@ describe('Query imports', () => {
       expect(await languageProvider.getParserAndLabelKeys('{place="luna"}', { maxLines: 5 })).toEqual({
         extractedLabelKeys: [],
         unwrapLabelKeys: [],
+        structuredMetadataKeys: [],
         hasJSON: false,
         hasLogfmt: false,
         hasPack: false,

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

@@ -11,7 +11,7 @@ import {
   extractUnwrapLabelKeysFromDataFrame,
 } from './responseUtils';
 import syntax from './syntax';
-import { ParserAndLabelKeysResult, LokiQuery, LokiQueryType } from './types';
+import { ParserAndLabelKeysResult, LokiQuery, LokiQueryType, LabelType } from './types';
 
 const NS_IN_MS = 1000000;
 
@@ -271,13 +271,21 @@ export default class LokiLanguageProvider extends LanguageProvider {
     );
 
     if (!series.length) {
-      return { extractedLabelKeys: [], unwrapLabelKeys: [], hasJSON: false, hasLogfmt: false, hasPack: false };
+      return {
+        extractedLabelKeys: [],
+        structuredMetadataKeys: [],
+        unwrapLabelKeys: [],
+        hasJSON: false,
+        hasLogfmt: false,
+        hasPack: false,
+      };
     }
 
     const { hasLogfmt, hasJSON, hasPack } = extractLogParserFromDataFrame(series[0]);
 
     return {
       extractedLabelKeys: extractLabelKeysFromDataFrame(series[0]),
+      structuredMetadataKeys: extractLabelKeysFromDataFrame(series[0], LabelType.StructuredMetadata),
       unwrapLabelKeys: extractUnwrapLabelKeysFromDataFrame(series[0]),
       hasJSON,
       hasPack,

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

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

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

@@ -45,6 +45,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 structuredMetadataKeys = ['structured', 'metadata'];
 const unwrapLabelKeys = ['unwrap', 'labels'];
 const otherLabels: Label[] = [
   {
@@ -156,7 +157,8 @@ function buildAfterSelectorCompletions(
   detectedParser: string,
   otherParser: string,
   afterPipe: boolean,
-  hasSpace: boolean
+  hasSpace: boolean,
+  structuredMetadataKeys?: string[]
 ) {
   const explanation = '(detected)';
   let expectedCompletions = afterSelectorCompletions.map((completion) => {
@@ -197,6 +199,20 @@ function buildAfterSelectorCompletions(
     }
   });
 
+  structuredMetadataKeys?.forEach((key) => {
+    let text = `${afterPipe ? ' ' : ' | '}${key}`;
+    if (hasSpace) {
+      text = text.trimStart();
+    }
+
+    expectedCompletions.push({
+      insertText: text,
+      label: `${key} ${explanation}`,
+      documentation: `"${key}" was suggested based on structured metadata attached to your loglines.`,
+      type: 'LABEL_NAME',
+    });
+  });
+
   return expectedCompletions;
 }
 
@@ -218,6 +234,7 @@ describe('getCompletions', () => {
     jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
       extractedLabelKeys,
       unwrapLabelKeys,
+      structuredMetadataKeys,
       hasJSON: false,
       hasLogfmt: false,
       hasPack: false,
@@ -351,6 +368,7 @@ describe('getCompletions', () => {
       jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
         extractedLabelKeys,
         unwrapLabelKeys,
+        structuredMetadataKeys,
         hasJSON: true,
         hasLogfmt: false,
         hasPack: false,
@@ -358,7 +376,7 @@ describe('getCompletions', () => {
       const situation: Situation = { type: 'AFTER_SELECTOR', logQuery: '{job="grafana"}', afterPipe, hasSpace };
       const completions = await getCompletions(situation, completionProvider);
 
-      const expected = buildAfterSelectorCompletions('json', 'logfmt', afterPipe, hasSpace);
+      const expected = buildAfterSelectorCompletions('json', 'logfmt', afterPipe, hasSpace, structuredMetadataKeys);
       expect(completions).toEqual(expected);
     }
   );
@@ -369,6 +387,7 @@ describe('getCompletions', () => {
       jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
         extractedLabelKeys,
         unwrapLabelKeys,
+        structuredMetadataKeys,
         hasJSON: false,
         hasLogfmt: true,
         hasPack: false,
@@ -376,7 +395,7 @@ describe('getCompletions', () => {
       const situation: Situation = { type: 'AFTER_SELECTOR', logQuery: '', afterPipe, hasSpace: true };
       const completions = await getCompletions(situation, completionProvider);
 
-      const expected = buildAfterSelectorCompletions('logfmt', 'json', afterPipe, true);
+      const expected = buildAfterSelectorCompletions('logfmt', 'json', afterPipe, true, structuredMetadataKeys);
       expect(completions).toEqual(expected);
     }
   );
@@ -458,6 +477,7 @@ describe('getAfterSelectorCompletions', () => {
     jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
       extractedLabelKeys: ['abc', 'def'],
       unwrapLabelKeys: [],
+      structuredMetadataKeys: [],
       hasJSON: true,
       hasLogfmt: false,
       hasPack: false,
@@ -480,6 +500,7 @@ describe('getAfterSelectorCompletions', () => {
     jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
       extractedLabelKeys: ['abc', 'def'],
       unwrapLabelKeys: [],
+      structuredMetadataKeys: [],
       hasJSON: true,
       hasLogfmt: false,
       hasPack: true,
@@ -553,6 +574,7 @@ describe('IN_LOGFMT completions', () => {
     jest.spyOn(completionProvider, 'getParserAndLabelKeys').mockResolvedValue({
       extractedLabelKeys: ['label1', 'label2'],
       unwrapLabelKeys: [],
+      structuredMetadataKeys: [],
       hasJSON: true,
       hasLogfmt: false,
       hasPack: false,

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

@@ -309,7 +309,8 @@ export async function getAfterSelectorCompletions(
     query = trimEnd(logQuery, '| ');
   }
 
-  const { extractedLabelKeys, hasJSON, hasLogfmt, hasPack } = await dataProvider.getParserAndLabelKeys(query);
+  const { extractedLabelKeys, structuredMetadataKeys, hasJSON, hasLogfmt, hasPack } =
+    await dataProvider.getParserAndLabelKeys(query);
   const hasQueryParser = isQueryWithParser(query).queryWithParser;
 
   const prefix = `${hasSpace ? '' : ' '}${afterPipe ? '' : '| '}`;
@@ -326,6 +327,15 @@ export async function getAfterSelectorCompletions(
 
   const completions = [...parserCompletions, ...pipeOperations];
 
+  structuredMetadataKeys.forEach((key) => {
+    completions.push({
+      type: 'LABEL_NAME',
+      label: `${key} (detected)`,
+      insertText: `${prefix}${key}`,
+      documentation: `"${key}" was suggested based on structured metadata attached to your loglines.`,
+    });
+  });
+
   // Let's show label options only if query has parser
   if (hasQueryParser) {
     extractedLabelKeys.forEach((key) => {

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

@@ -13,6 +13,7 @@ import {
   cloneQueryResponse,
   combineResponses,
 } from './responseUtils';
+import { LabelType } from './types';
 
 const frame: DataFrame = {
   length: 1,
@@ -38,6 +39,36 @@ const frame: DataFrame = {
   ],
 };
 
+const frameWithTypes: DataFrame = {
+  length: 1,
+  fields: [
+    {
+      name: 'Time',
+      config: {},
+      type: FieldType.time,
+      values: [1],
+    },
+    {
+      name: 'labels',
+      config: {},
+      type: FieldType.other,
+      values: [{ level: 'info', structured: 'foo' }],
+    },
+    {
+      name: 'Line',
+      config: {},
+      type: FieldType.string,
+      values: ['line1'],
+    },
+    {
+      name: 'labelTypes',
+      config: {},
+      type: FieldType.other,
+      values: [{ level: 'I', structured: 'S' }],
+    },
+  ],
+};
+
 describe('dataFrameHasParsingError', () => {
   it('handles frame with parsing error', () => {
     const input = cloneDeep(frame);
@@ -104,10 +135,26 @@ describe('extractLabelKeysFromDataFrame', () => {
     input.fields[1].values = [];
     expect(extractLabelKeysFromDataFrame(input)).toEqual([]);
   });
+
   it('extracts label keys', () => {
     const input = cloneDeep(frame);
     expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);
   });
+
+  it('extracts indexed label keys', () => {
+    const input = cloneDeep(frameWithTypes);
+    expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);
+  });
+
+  it('extracts structured metadata label keys', () => {
+    const input = cloneDeep(frameWithTypes);
+    expect(extractLabelKeysFromDataFrame(input, LabelType.StructuredMetadata)).toEqual(['structured']);
+  });
+
+  it('does not extract structured metadata label keys from non-typed frame', () => {
+    const input = cloneDeep(frame);
+    expect(extractLabelKeysFromDataFrame(input, LabelType.StructuredMetadata)).toEqual([]);
+  });
 });
 
 describe('extractUnwrapLabelKeysFromDataFrame', () => {

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

@@ -13,6 +13,7 @@ import {
 
 import { isBytesString } from './languageUtils';
 import { isLogLineJSON, isLogLineLogfmt, isLogLinePacked } from './lineParser';
+import { LabelType } from './types';
 
 export function dataFrameHasLokiError(frame: DataFrame): boolean {
   const labelSets: Labels[] = frame.fields.find((f) => f.name === 'labels')?.values ?? [];
@@ -54,15 +55,29 @@ export function extractLogParserFromDataFrame(frame: DataFrame): {
   return { hasLogfmt, hasJSON, hasPack };
 }
 
-export function extractLabelKeysFromDataFrame(frame: DataFrame): string[] {
+export function extractLabelKeysFromDataFrame(frame: DataFrame, type: LabelType = LabelType.Indexed): string[] {
   const labelsArray: Array<{ [key: string]: string }> | undefined =
     frame?.fields?.find((field) => field.name === 'labels')?.values ?? [];
+  const labelTypeArray: Array<{ [key: string]: string }> | undefined =
+    frame?.fields?.find((field) => field.name === 'labelTypes')?.values ?? [];
 
   if (!labelsArray?.length) {
     return [];
   }
 
-  return Object.keys(labelsArray[0]);
+  // if there are no label types, only return indexed labels if requested
+  if (!labelTypeArray?.length) {
+    if (type === LabelType.Indexed) {
+      return Object.keys(labelsArray[0]);
+    }
+    return [];
+  }
+
+  const labelTypes = labelTypeArray[0];
+
+  const allLabelKeys = Object.keys(labelsArray[0]).filter((k) => labelTypes[k] === type);
+
+  return allLabelKeys;
 }
 
 export function extractUnwrapLabelKeysFromDataFrame(frame: DataFrame): string[] {

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

@@ -10,6 +10,12 @@ export enum LokiResultType {
   Matrix = 'matrix',
 }
 
+export enum LabelType {
+  Indexed = 'I',
+  StructuredMetadata = 'S',
+  Parsed = 'P',
+}
+
 export interface LokiQuery extends LokiQueryFromSchema {
   direction?: LokiQueryDirection;
   /** Used only to identify supporting queries, e.g. logs volume, logs sample and data sample */
@@ -87,6 +93,7 @@ export interface ContextFilter {
 
 export interface ParserAndLabelKeysResult {
   extractedLabelKeys: string[];
+  structuredMetadataKeys: string[];
   hasJSON: boolean;
   hasLogfmt: boolean;
   hasPack: boolean;