Loki: Show label options for unwrap operation (#52810)

* Loki: Show options for unwrap operation

* Add comment for check

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

@@ -0,0 +1,20 @@
+import { isBytesString } from './language_utils';
+
+describe('isBytesString', () => {
+  it('correctly matches bytes string with integers', () => {
+    expect(isBytesString('500b')).toBe(true);
+    expect(isBytesString('2TB')).toBe(true);
+  });
+  it('correctly matches bytes string with float', () => {
+    expect(isBytesString('500.4kib')).toBe(true);
+    expect(isBytesString('10.4654Mib')).toBe(true);
+  });
+  it('does not match integer without unit', () => {
+    expect(isBytesString('500')).toBe(false);
+    expect(isBytesString('10')).toBe(false);
+  });
+  it('does not match float without unit', () => {
+    expect(isBytesString('50.047')).toBe(false);
+    expect(isBytesString('1.234')).toBe(false);
+  });
+});

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

@@ -51,3 +51,36 @@ export function isRegexSelector(selector?: string) {
   }
   return false;
 }
+
+export function isBytesString(string: string) {
+  const BYTES_KEYWORDS = [
+    'b',
+    'kib',
+    'Kib',
+    'kb',
+    'KB',
+    'mib',
+    'Mib',
+    'mb',
+    'MB',
+    'gib',
+    'Gib',
+    'gb',
+    'GB',
+    'tib',
+    'Tib',
+    'tb',
+    'TB',
+    'pib',
+    'Pib',
+    'pb',
+    'PB',
+    'eib',
+    'Eib',
+    'eb',
+    'EB',
+  ];
+  const regex = new RegExp(`^(?:-?\\d+(?:\\.\\d+)?)(?:${BYTES_KEYWORDS.join('|')})$`);
+  const match = string.match(regex);
+  return !!match;
+}

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

@@ -172,3 +172,34 @@ export function isQueryWithLabelFormat(query: string): boolean {
   });
   return queryWithLabelFormat;
 }
+
+export function getLogQueryFromMetricsQuery(query: string): string {
+  if (isLogsQuery(query)) {
+    return query;
+  }
+
+  const tree = parser.parse(query);
+
+  // Log query in metrics query composes of Selector & PipelineExpr
+  let selector = '';
+  tree.iterate({
+    enter: (type, from, to): false | void => {
+      if (type.name === 'Selector') {
+        selector = query.substring(from, to);
+        return false;
+      }
+    },
+  });
+
+  let pipelineExpr = '';
+  tree.iterate({
+    enter: (type, from, to): false | void => {
+      if (type.name === 'PipelineExpr') {
+        pipelineExpr = query.substring(from, to);
+        return false;
+      }
+    },
+  });
+
+  return selector + pipelineExpr;
+}

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

@@ -0,0 +1,147 @@
+import { screen, render } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import React, { ComponentProps } from 'react';
+
+import { DataFrame, DataSourceApi, DataSourcePluginMeta, FieldType, toDataFrame } from '@grafana/data';
+import {
+  QueryBuilderOperation,
+  QueryBuilderOperationParamDef,
+} from 'app/plugins/datasource/prometheus/querybuilder/shared/types';
+
+import { LokiDatasource } from '../../datasource';
+import { LokiOperationId } from '../types';
+
+import { UnwrapParamEditor } from './UnwrapParamEditor';
+
+describe('UnwrapParamEditor', () => {
+  it('shows value if value present', () => {
+    const props = createProps({ value: 'unique' });
+    render(<UnwrapParamEditor {...props} />);
+    expect(screen.getByText('unique')).toBeInTheDocument();
+  });
+
+  it('shows no label options if no samples are returned', async () => {
+    const props = createProps();
+    render(<UnwrapParamEditor {...props} />);
+    const input = screen.getByRole('combobox');
+    await userEvent.click(input);
+    expect(screen.getByText('No labels found')).toBeInTheDocument();
+  });
+
+  it('shows no label options for non-metric query', async () => {
+    const props = createProps({
+      query: {
+        labels: [{ op: '=', label: 'foo', value: 'bar' }],
+        operations: [
+          { id: LokiOperationId.Logfmt, params: [] },
+          { id: LokiOperationId.Unwrap, params: ['', ''] },
+        ],
+      },
+    });
+    render(<UnwrapParamEditor {...props} />);
+    const input = screen.getByRole('combobox');
+    await userEvent.click(input);
+    expect(screen.getByText('No labels found')).toBeInTheDocument();
+  });
+
+  it('shows labels with unwrap-friendly values', async () => {
+    const props = createProps({}, frames);
+    render(<UnwrapParamEditor {...props} />);
+    const input = screen.getByRole('combobox');
+    await userEvent.click(input);
+    expect(await screen.findByText('status')).toBeInTheDocument();
+    expect(await screen.findByText('duration')).toBeInTheDocument();
+  });
+});
+
+const createProps = (
+  propsOverrides?: Partial<ComponentProps<typeof UnwrapParamEditor>>,
+  mockedSample?: DataFrame[]
+) => {
+  const propsDefault = {
+    value: undefined,
+    onChange: jest.fn(),
+    onRunQuery: jest.fn(),
+    index: 1,
+    operationIndex: 1,
+    query: {
+      labels: [{ op: '=', label: 'foo', value: 'bar' }],
+      operations: [
+        { id: LokiOperationId.Logfmt, params: [] },
+        { id: LokiOperationId.Unwrap, params: ['', ''] },
+        { id: LokiOperationId.SumOverTime, params: ['5m'] },
+        { id: '__sum_by', params: ['job'] },
+      ],
+    },
+    paramDef: {} as QueryBuilderOperationParamDef,
+    operation: {} as QueryBuilderOperation,
+    datasource: new LokiDatasource(
+      {
+        id: 1,
+        uid: '',
+        type: 'loki',
+        name: 'loki-test',
+        access: 'proxy',
+        url: '',
+        jsonData: {},
+        meta: {} as DataSourcePluginMeta,
+      },
+      undefined,
+      undefined
+    ) as DataSourceApi,
+  };
+  const props = { ...propsDefault, ...propsOverrides };
+
+  if (props.datasource instanceof LokiDatasource) {
+    const resolvedValue = mockedSample ?? [];
+    props.datasource.getDataSamples = jest.fn().mockResolvedValue(resolvedValue);
+  }
+  return props;
+};
+
+const frames = [
+  toDataFrame({
+    fields: [
+      {
+        name: 'labels',
+        type: FieldType.other,
+        values: [
+          {
+            compose_project: 'docker-compose',
+            compose_service: 'app',
+            container_name: 'docker-compose_app_1',
+            duration: '2.807709ms',
+            filename: '/var/log/docker/37c87fe98cbfa28327c1de10c4aff72c58154d8e4d129118ff2024692360b677/json.log',
+            host: 'docker-desktop',
+            instance: 'docker-compose_app_1',
+            job: 'tns/app',
+            level: 'info',
+            msg: 'HTTP client success',
+            namespace: 'tns',
+            source: 'stdout',
+            status: '200',
+            traceID: '6a3d34c4225776f6',
+            url: 'http://db',
+          },
+          {
+            compose_project: 'docker-compose',
+            compose_service: 'app',
+            container_name: 'docker-compose_app_1',
+            duration: '7.432542ms',
+            filename: '/var/log/docker/37c87fe98cbfa28327c1de10c4aff72c58154d8e4d129118ff2024692360b677/json.log',
+            host: 'docker-desktop',
+            instance: 'docker-compose_app_1',
+            job: 'tns/app',
+            level: 'info',
+            msg: 'HTTP client success',
+            namespace: 'tns',
+            source: 'stdout',
+            status: '200',
+            traceID: '18e99189831471f6',
+            url: 'http://db',
+          },
+        ],
+      },
+    ],
+  }),
+];

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

@@ -0,0 +1,96 @@
+import { isNaN } from 'lodash';
+import React, { useState } from 'react';
+
+import { isValidGoDuration, 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 '../../language_utils';
+import { getLogQueryFromMetricsQuery, isValidQuery } from '../../query_utils';
+import { lokiQueryModeller } from '../LokiQueryModeller';
+import { LokiVisualQuery } from '../types';
+
+export function UnwrapParamEditor({
+  onChange,
+  index,
+  operationIndex,
+  value,
+  query,
+  datasource,
+}: QueryBuilderOperationParamEditorProps) {
+  const [state, setState] = useState<{
+    options?: Array<SelectableValue<string>>;
+    isLoading?: boolean;
+  }>({});
+
+  return (
+    <Select
+      inputId={getOperationParamId(operationIndex, index)}
+      onOpenMenu={async () => {
+        // This check is always true, we do it to make typescript happy
+        if (datasource instanceof LokiDatasource) {
+          setState({ isLoading: true });
+          const options = await loadUnwrapOptions(query, datasource);
+          setState({ options, isLoading: undefined });
+        }
+      }}
+      isLoading={state.isLoading}
+      allowCustomValue
+      noOptionsMessage="No labels found"
+      loadingMessage="Loading labels"
+      options={state.options}
+      value={value ? toOption(value.toString()) : null}
+      onChange={(value) => {
+        if (value.value) {
+          onChange(index, value.value);
+        }
+      }}
+    />
+  );
+}
+
+async function loadUnwrapOptions(
+  query: LokiVisualQuery,
+  datasource: LokiDatasource
+): Promise<Array<SelectableValue<string>>> {
+  const queryExpr = lokiQueryModeller.renderQuery(query);
+  const logExpr = getLogQueryFromMetricsQuery(queryExpr);
+  if (!isValidQuery(logExpr)) {
+    return [];
+  }
+
+  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 labelOptions = unwrapLabels.map((label) => ({
+    label,
+    value: label,
+  }));
+
+  return labelOptions;
+}

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

@@ -12,6 +12,7 @@ import {
 import { FUNCTIONS } from '../syntax';
 
 import { binaryScalarOperations } from './binaryScalarOperations';
+import { UnwrapParamEditor } from './components/UnwrapParamEditor';
 import { LokiOperationId, LokiOperationOrder, LokiVisualQuery, LokiVisualQueryOperationCategory } from './types';
 
 export function getOperationDefinitions(): QueryBuilderOperationDef[] {
@@ -362,7 +363,14 @@ export function getOperationDefinitions(): QueryBuilderOperationDef[] {
       id: LokiOperationId.Unwrap,
       name: 'Unwrap',
       params: [
-        { name: 'Identifier', type: 'string', hideName: true, minWidth: 16, placeholder: 'Label key' },
+        {
+          name: 'Identifier',
+          type: 'string',
+          hideName: true,
+          minWidth: 16,
+          placeholder: 'Label key',
+          editor: UnwrapParamEditor,
+        },
         {
           name: 'Conversion function',
           hideName: true,