Loki: Fix processing of all lines labels in getParserAndLabelKeys (#81483)

* Loki: Fix processing of all lines labels in getParserAndLabelKeys

* Refactor

* Update comment

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

@@ -66,6 +66,34 @@ const frameWithTypes: DataFrame = {
   ],
 };
 
+const frameWithMultipleLabels: DataFrame = {
+  length: 1,
+  fields: [
+    {
+      name: 'Time',
+      config: {},
+      type: FieldType.time,
+      values: [1, 2, 3],
+    },
+    {
+      name: 'labels',
+      config: {},
+      type: FieldType.other,
+      values: [
+        { level: 'info', foo: 'bar' },
+        { level: 'info', foo: 'baz', new: 'yes' },
+        { level: 'error', foo: 'baz' },
+      ],
+    },
+    {
+      name: 'Line',
+      config: {},
+      type: FieldType.string,
+      values: ['line1', 'line2', 'line3'],
+    },
+  ],
+};
+
 describe('dataFrameHasParsingError', () => {
   it('handles frame with parsing error', () => {
     const input = cloneDeep(frame);
@@ -138,6 +166,11 @@ describe('extractLabelKeysFromDataFrame', () => {
     expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);
   });
 
+  it('extracts label keys from all logs', () => {
+    const input = cloneDeep(frameWithMultipleLabels);
+    expect(extractLabelKeysFromDataFrame(input)).toEqual(['level', 'foo', 'new']);
+  });
+
   it('extracts indexed label keys', () => {
     const input = cloneDeep(frameWithTypes);
     expect(extractLabelKeysFromDataFrame(input)).toEqual(['level']);

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

@@ -1,6 +1,6 @@
 import { DataFrame, FieldType, isValidGoDuration, Labels } from '@grafana/data';
 
-import { isBytesString } from './languageUtils';
+import { isBytesString, processLabels } from './languageUtils';
 import { isLogLineJSON, isLogLineLogfmt, isLogLinePacked } from './lineParser';
 import { LabelType } from './types';
 
@@ -54,19 +54,26 @@ export function extractLabelKeysFromDataFrame(frame: DataFrame, type: LabelType
     return [];
   }
 
-  // if there are no label types, only return indexed labels if requested
+  // if there are no label types and type is LabelType.Indexed return all label keys
   if (!labelTypeArray?.length) {
     if (type === LabelType.Indexed) {
-      return Object.keys(labelsArray[0]);
+      const { keys: labelKeys } = processLabels(labelsArray);
+      return labelKeys;
     }
     return [];
   }
 
-  const labelTypes = labelTypeArray[0];
+  // If we have label types, we can return only label keys that match type
+  let labelsSet = new Set<string>();
+  for (let i = 0; i < labelsArray.length; i++) {
+    const labels = labelsArray[i];
+    const labelsType = labelTypeArray[i];
 
-  const allLabelKeys = Object.keys(labelsArray[0]).filter((k) => labelTypes[k] === type);
+    const allLabelKeys = Object.keys(labels).filter((key) => labelsType[key] === type);
+    labelsSet = new Set([...labelsSet, ...allLabelKeys]);
+  }
 
-  return allLabelKeys;
+  return Array.from(labelsSet);
 }
 
 export function extractUnwrapLabelKeysFromDataFrame(frame: DataFrame): string[] {