Loki: Fix creating correct context query when preserved labels used (#69252)

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

@@ -3,10 +3,18 @@ import { of } from 'rxjs';
 import { DataQueryResponse, FieldType, LogRowContextQueryDirection, LogRowModel, createDataFrame } from '@grafana/data';
 
 import LokiLanguageProvider from './LanguageProvider';
-import { LogContextProvider } from './LogContextProvider';
+import { LogContextProvider, LOKI_LOG_CONTEXT_PRESERVED_LABELS } from './LogContextProvider';
 import { createLokiDatasource } from './mocks';
 import { LokiQuery } from './types';
 
+jest.mock('app/core/store', () => {
+  return {
+    get() {
+      return window.localStorage.getItem(LOKI_LOG_CONTEXT_PRESERVED_LABELS);
+    },
+  };
+});
+
 const defaultLanguageProviderMock = {
   start: jest.fn(),
   fetchSeriesLabels: jest.fn(() => ({ bar: ['baz'], xyz: ['abc'] })),
@@ -38,9 +46,13 @@ describe('LogContextProvider', () => {
     logContextProvider = new LogContextProvider(defaultDatasourceMock);
   });
 
+  afterEach(() => {
+    window.localStorage.clear();
+  });
+
   describe('getLogRowContext', () => {
     it('should call getInitContextFilters if no appliedContextFilters', async () => {
-      logContextProvider.getInitContextFiltersFromLabels = jest
+      logContextProvider.getInitContextFilters = jest
         .fn()
         .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
 
@@ -55,8 +67,8 @@ describe('LogContextProvider', () => {
           expr: '{bar="baz"}',
         } as LokiQuery
       );
-      expect(logContextProvider.getInitContextFiltersFromLabels).toBeCalled();
+      expect(logContextProvider.getInitContextFilters).toBeCalled();
-      expect(logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalledWith(
+      expect(logContextProvider.getInitContextFilters).toHaveBeenCalledWith(
         { bar: 'baz', foo: 'uniqueParsedLabel', xyz: 'abc' },
         { expr: '{bar="baz"}' }
       );
@@ -64,7 +76,7 @@ describe('LogContextProvider', () => {
     });
 
     it('should not call getInitContextFilters if appliedContextFilters', async () => {
-      logContextProvider.getInitContextFiltersFromLabels = jest
+      logContextProvider.getInitContextFilters = jest
         .fn()
         .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
 
@@ -76,14 +88,14 @@ describe('LogContextProvider', () => {
         limit: 10,
         direction: LogRowContextQueryDirection.Backward,
       });
-      expect(logContextProvider.getInitContextFiltersFromLabels).not.toBeCalled();
+      expect(logContextProvider.getInitContextFilters).not.toBeCalled();
       expect(logContextProvider.appliedContextFilters).toHaveLength(2);
     });
   });
 
   describe('getLogRowContextQuery', () => {
     it('should call getInitContextFilters if no appliedContextFilters', async () => {
-      logContextProvider.getInitContextFiltersFromLabels = jest
+      logContextProvider.getInitContextFilters = jest
         .fn()
         .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
 
@@ -202,10 +214,7 @@ describe('LogContextProvider', () => {
       } as LokiQuery;
 
       it('should correctly create contextFilters', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels(
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, queryWithoutParser);
-          defaultLogRow.labels,
-          queryWithoutParser
-        );
         expect(filters).toEqual([
           { enabled: true, fromParser: false, label: 'bar', value: 'baz' },
           { enabled: false, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
@@ -214,12 +223,12 @@ describe('LogContextProvider', () => {
       });
 
       it('should return empty contextFilters if no query', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, undefined);
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, undefined);
         expect(filters).toEqual([]);
       });
 
       it('should return empty contextFilters if no labels', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels({}, queryWithoutParser);
+        const filters = await logContextProvider.getInitContextFilters({}, queryWithoutParser);
         expect(filters).toEqual([]);
       });
     });
@@ -230,7 +239,7 @@ describe('LogContextProvider', () => {
       } as LokiQuery;
 
       it('should correctly create contextFilters', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, queryWithParser);
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, queryWithParser);
         expect(filters).toEqual([
           { enabled: true, fromParser: false, label: 'bar', value: 'baz' },
           { enabled: false, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
@@ -239,14 +248,68 @@ describe('LogContextProvider', () => {
       });
 
       it('should return empty contextFilters if no query', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, undefined);
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, undefined);
         expect(filters).toEqual([]);
       });
 
       it('should return empty contextFilters if no labels', async () => {
-        const filters = await logContextProvider.getInitContextFiltersFromLabels({}, queryWithParser);
+        const filters = await logContextProvider.getInitContextFilters({}, queryWithParser);
         expect(filters).toEqual([]);
       });
     });
+
+    describe('with preserved labels', () => {
+      const queryWithParser = {
+        expr: '{bar="baz"} | logfmt',
+      } as LokiQuery;
+
+      it('should correctly apply preserved labels', async () => {
+        window.localStorage.setItem(
+          LOKI_LOG_CONTEXT_PRESERVED_LABELS,
+          JSON.stringify({
+            removedLabels: ['bar'],
+            selectedExtractedLabels: ['foo'],
+          })
+        );
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, queryWithParser);
+        expect(filters).toEqual([
+          { enabled: false, fromParser: false, label: 'bar', value: 'baz' }, // disabled real label
+          { enabled: true, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' }, // enabled parsed label
+          { enabled: true, fromParser: false, label: 'xyz', value: 'abc' },
+        ]);
+      });
+
+      it('should use contextFilters from row labels if all real labels are disabled', async () => {
+        window.localStorage.setItem(
+          LOKI_LOG_CONTEXT_PRESERVED_LABELS,
+          JSON.stringify({
+            removedLabels: ['bar', 'xyz'],
+            selectedExtractedLabels: ['foo'],
+          })
+        );
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, queryWithParser);
+        expect(filters).toEqual([
+          { enabled: true, fromParser: false, label: 'bar', value: 'baz' }, // enabled real label
+          { enabled: false, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
+          { enabled: true, fromParser: false, label: 'xyz', value: 'abc' }, // enabled real label
+        ]);
+      });
+
+      it('should not introduce new labels as context filters', async () => {
+        window.localStorage.setItem(
+          LOKI_LOG_CONTEXT_PRESERVED_LABELS,
+          JSON.stringify({
+            removedLabels: ['bar'],
+            selectedExtractedLabels: ['foo', 'new'],
+          })
+        );
+        const filters = await logContextProvider.getInitContextFilters(defaultLogRow.labels, queryWithParser);
+        expect(filters).toEqual([
+          { enabled: false, fromParser: false, label: 'bar', value: 'baz' },
+          { enabled: true, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
+          { enabled: true, fromParser: false, label: 'xyz', value: 'abc' },
+        ]);
+      });
+    });
   });
 });

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

@@ -15,6 +15,10 @@ import {
   LogRowContextOptions,
 } from '@grafana/data';
 import { Labels } from '@grafana/schema';
+import { notifyApp } from 'app/core/actions';
+import { createSuccessNotification } from 'app/core/copy/appNotification';
+import store from 'app/core/store';
+import { dispatch } from 'app/store/store';
 
 import { LokiContextUi } from './components/LokiContextUi';
 import { LokiDatasource, makeRequest, REF_ID_STARTER_LOG_ROW_CONTEXT } from './datasource';
@@ -24,6 +28,12 @@ import { getParserFromQuery, getStreamSelectorsFromQuery, isQueryWithParser } fr
 import { sortDataFrameByTime, SortDirection } from './sortDataFrame';
 import { ContextFilter, LokiQuery, LokiQueryDirection, LokiQueryType } from './types';
 
+export const LOKI_LOG_CONTEXT_PRESERVED_LABELS = 'lokiLogContextPreservedLabels';
+export type PreservedLabels = {
+  removedLabels: string[];
+  selectedExtractedLabels: string[];
+};
+
 export class LogContextProvider {
   datasource: LokiDatasource;
   appliedContextFilters: ContextFilter[];
@@ -40,9 +50,7 @@ export class LogContextProvider {
     // This happens only on initial load, when user haven't applied any filters yet
     // We need to get the initial filters from the row labels
     if (this.appliedContextFilters.length === 0) {
-      const filters = (await this.getInitContextFiltersFromLabels(row.labels, origQuery)).filter(
+      const filters = (await this.getInitContextFilters(row.labels, origQuery)).filter((filter) => filter.enabled);
-        (filter) => filter.enabled
-      );
       this.appliedContextFilters = filters;
     }
 
@@ -208,11 +216,13 @@ export class LogContextProvider {
     return expr;
   };
 
-  getInitContextFiltersFromLabels = async (labels: Labels, query?: LokiQuery) => {
+  getInitContextFilters = async (labels: Labels, query?: LokiQuery) => {
     if (!query || isEmpty(labels)) {
       return [];
     }
 
+    // 1. First we need to get all labels from the log row's label
+    // and correctly set parsed and not parsed labels
     let allLabels: string[] = [];
     if (!isQueryWithParser(query.expr).queryWithParser) {
       // If there is no parser, we use getLabelKeys because it has better caching
@@ -239,6 +249,44 @@ export class LogContextProvider {
       contextFilters.push(filter);
     });
 
-    return contextFilters;
+    // Secondly we check for preserved labels and update enabled state of filters based on that
+    let preservedLabels: undefined | PreservedLabels = undefined;
+    try {
+      preservedLabels = JSON.parse(store.get(LOKI_LOG_CONTEXT_PRESERVED_LABELS));
+      // Do nothing when error occurs
+    } catch (e) {}
+
+    if (!preservedLabels) {
+      // If we don't have preservedLabels, we return contextFilters as they are
+      return contextFilters;
+    } else {
+      // Otherwise, we need to update filters based on preserved labels
+      let arePreservedLabelsUsed = false;
+      const newContextFilters = contextFilters.map((contextFilter) => {
+        // We checked for undefined above
+        if (preservedLabels!.removedLabels.includes(contextFilter.label)) {
+          arePreservedLabelsUsed = true;
+          return { ...contextFilter, enabled: false };
+        }
+        // We checked for undefined above
+        if (preservedLabels!.selectedExtractedLabels.includes(contextFilter.label)) {
+          arePreservedLabelsUsed = true;
+          return { ...contextFilter, enabled: true };
+        }
+        return { ...contextFilter };
+      });
+
+      const isAtLeastOneRealLabelEnabled = newContextFilters.some(({ enabled, fromParser }) => enabled && !fromParser);
+      if (!isAtLeastOneRealLabelEnabled) {
+        // If we end up with no real labels enabled, we need to reset the init filters
+        return contextFilters;
+      } else {
+        // Otherwise use new filters
+        if (arePreservedLabelsUsed) {
+          dispatch(notifyApp(createSuccessNotification('Previously used log context filters have been applied.')));
+        }
+        return newContextFilters;
+      }
+    }
   };
 }

public/app/plugins/datasource/loki/components/LokiContextUi.test.tsx

@@ -8,7 +8,7 @@ import { LogRowModel } from '@grafana/data';
 import { LogContextProvider } from '../LogContextProvider';
 import { ContextFilter, LokiQuery } from '../types';
 
-import { LokiContextUi, LokiContextUiProps, LOKI_LOG_CONTEXT_PRESERVED_LABELS } from './LokiContextUi';
+import { LokiContextUi, LokiContextUiProps } from './LokiContextUi';
 
 // we have to mock out reportInteraction, otherwise it crashes the test.
 jest.mock('@grafana/runtime', () => ({
@@ -23,9 +23,6 @@ jest.mock('app/core/store', () => {
       return true;
     },
     delete() {},
-    get() {
-      return window.localStorage.getItem(LOKI_LOG_CONTEXT_PRESERVED_LABELS);
-    },
   };
 });
 
@@ -51,7 +48,7 @@ const setupProps = (): LokiContextUiProps => {
 };
 
 const mockLogContextProvider = {
-  getInitContextFiltersFromLabels: jest.fn().mockImplementation(() =>
+  getInitContextFilters: jest.fn().mockImplementation(() =>
     Promise.resolve([
       { value: 'value1', enabled: true, fromParser: false, label: 'label1' },
       { value: 'value3', enabled: false, fromParser: true, label: 'label3' },
@@ -83,10 +80,6 @@ describe('LokiContextUi', () => {
     global = savedGlobal;
   });
 
-  afterEach(() => {
-    window.localStorage.clear();
-  });
-
   it('renders and shows executed query text', async () => {
     const props = setupProps();
     render(<LokiContextUi {...props} />);
@@ -105,7 +98,7 @@ describe('LokiContextUi', () => {
     render(<LokiContextUi {...props} />);
 
     await waitFor(() => {
-      expect(props.logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalled();
+      expect(props.logContextProvider.getInitContextFilters).toHaveBeenCalled();
     });
   });
 
@@ -113,7 +106,7 @@ describe('LokiContextUi', () => {
     const props = setupProps();
     render(<LokiContextUi {...props} />);
     await waitFor(() => {
-      expect(props.logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalled();
+      expect(props.logContextProvider.getInitContextFilters).toHaveBeenCalled();
     });
     const select = await screen.findAllByRole('combobox');
     await selectOptionInTest(select[0], 'label1="value1"');
@@ -123,7 +116,7 @@ describe('LokiContextUi', () => {
     const props = setupProps();
     render(<LokiContextUi {...props} />);
     await waitFor(() => {
-      expect(props.logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalled();
+      expect(props.logContextProvider.getInitContextFilters).toHaveBeenCalled();
     });
     const select = await screen.findAllByRole('combobox');
     await selectOptionInTest(select[1], 'label3="value3"');
@@ -134,7 +127,7 @@ describe('LokiContextUi', () => {
     const props = setupProps();
     render(<LokiContextUi {...props} />);
     await waitFor(() => {
-      expect(props.logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalled();
+      expect(props.logContextProvider.getInitContextFilters).toHaveBeenCalled();
       expect(screen.getAllByRole('combobox')).toHaveLength(2);
     });
     await selectOptionInTest(screen.getAllByRole('combobox')[1], 'label3="value3"');
@@ -243,74 +236,4 @@ describe('LokiContextUi', () => {
       expect(screen.queryByText('label3="value3"')).not.toBeInTheDocument();
     });
   });
-
-  describe('preserve labels', () => {
-    it('should use init contextFilters if all real labels are disabled', async () => {
-      window.localStorage.setItem(
-        LOKI_LOG_CONTEXT_PRESERVED_LABELS,
-        JSON.stringify({
-          removedLabels: ['label1'],
-          selectedExtractedLabels: ['label3'],
-        })
-      );
-      const props = setupProps();
-      const newProps = {
-        ...props,
-        origQuery: {
-          expr: '{label1="value1"} | logfmt',
-          refId: 'A',
-        },
-      };
-      render(<LokiContextUi {...newProps} />);
-      await waitFor(() => {
-        expect(screen.queryByText('label3="value3"')).not.toBeInTheDocument();
-        expect(screen.getByText('label1="value1"')).toBeInTheDocument();
-      });
-    });
-
-    it('should use preserved contextFilters if all at least 1 real labels is enabled', async () => {
-      window.localStorage.setItem(
-        LOKI_LOG_CONTEXT_PRESERVED_LABELS,
-        JSON.stringify({
-          removedLabels: ['foo'],
-          selectedExtractedLabels: ['label3'],
-        })
-      );
-      const props = setupProps();
-      const newProps = {
-        ...props,
-        origQuery: {
-          expr: '{label1="value1"} | logfmt',
-          refId: 'A',
-        },
-      };
-      render(<LokiContextUi {...newProps} />);
-      await waitFor(() => {
-        expect(screen.getByText('label3="value3"')).toBeInTheDocument();
-        expect(screen.getByText('label1="value1"')).toBeInTheDocument();
-      });
-    });
-
-    it('should not introduce new labels in ui', async () => {
-      window.localStorage.setItem(
-        LOKI_LOG_CONTEXT_PRESERVED_LABELS,
-        JSON.stringify({
-          removedLabels: ['foo'],
-          selectedExtractedLabels: ['bar', 'baz'],
-        })
-      );
-      const props = setupProps();
-      const newProps = {
-        ...props,
-        origQuery: {
-          expr: '{label1="value1"} | logfmt',
-          refId: 'A',
-        },
-      };
-      render(<LokiContextUi {...newProps} />);
-      await waitFor(() => {
-        expect(screen.getByText('label1="value1"')).toBeInTheDocument();
-      });
-    });
-  });
 });

public/app/plugins/datasource/loki/components/LokiContextUi.tsx

@@ -5,13 +5,10 @@ import { useAsync } from 'react-use';
 import { GrafanaTheme2, LogRowModel, SelectableValue } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
 import { Button, Collapse, Icon, Label, MultiSelect, Spinner, Tooltip, useStyles2 } from '@grafana/ui';
-import { notifyApp } from 'app/core/actions';
-import { createSuccessNotification } from 'app/core/copy/appNotification';
 import store from 'app/core/store';
-import { dispatch } from 'app/store/store';
 
 import { RawQuery } from '../../prometheus/querybuilder/shared/RawQuery';
-import { LogContextProvider } from '../LogContextProvider';
+import { LogContextProvider, LOKI_LOG_CONTEXT_PRESERVED_LABELS, PreservedLabels } from '../LogContextProvider';
 import { escapeLabelValueInSelector } from '../languageUtils';
 import { isQueryWithParser } from '../queryUtils';
 import { lokiGrammar } from '../syntax';
@@ -80,12 +77,6 @@ function getStyles(theme: GrafanaTheme2) {
 }
 
 const IS_LOKI_LOG_CONTEXT_UI_OPEN = 'isLogContextQueryUiOpen';
-export const LOKI_LOG_CONTEXT_PRESERVED_LABELS = 'lokiLogContextPreservedLabels';
-
-type PreservedLabels = {
-  removedLabels: string[];
-  selectedExtractedLabels: string[];
-};
 
 export function LokiContextUi(props: LokiContextUiProps) {
   const { row, logContextProvider, updateFilter, onClose, origQuery } = props;
@@ -171,45 +162,9 @@ export function LokiContextUi(props: LokiContextUiProps) {
 
   useAsync(async () => {
     setLoading(true);
-    const initContextFilters = await logContextProvider.getInitContextFiltersFromLabels(row.labels, origQuery);
+    const initContextFilters = await logContextProvider.getInitContextFilters(row.labels, origQuery);
+    setContextFilters(initContextFilters);
 
-    let preservedLabels: undefined | PreservedLabels = undefined;
-    try {
-      preservedLabels = JSON.parse(store.get(LOKI_LOG_CONTEXT_PRESERVED_LABELS));
-      // Do nothing when error occurs
-    } catch (e) {}
-
-    if (!preservedLabels) {
-      setContextFilters(initContextFilters);
-    } else {
-      // We need to update filters based on preserved labels
-      let arePreservedLabelsUsed = false;
-      const newContextFilters = initContextFilters.map((contextFilter) => {
-        // We checked for undefined above
-        if (preservedLabels!.removedLabels.includes(contextFilter.label)) {
-          arePreservedLabelsUsed = true;
-          return { ...contextFilter, enabled: false };
-        }
-        // We checked for undefined above
-        if (preservedLabels!.selectedExtractedLabels.includes(contextFilter.label)) {
-          arePreservedLabelsUsed = true;
-          return { ...contextFilter, enabled: true };
-        }
-        return { ...contextFilter };
-      });
-
-      const isAtLeastOneRealLabelEnabled = newContextFilters.some(({ enabled, fromParser }) => enabled && !fromParser);
-      if (!isAtLeastOneRealLabelEnabled) {
-        // If we end up with no real labels enabled, we need to reset the init filters
-        setContextFilters(initContextFilters);
-      } else {
-        // Otherwise use new filters
-        setContextFilters(newContextFilters);
-        if (arePreservedLabelsUsed) {
-          dispatch(notifyApp(createSuccessNotification('Previously used log context filters have been applied.')));
-        }
-      }
-    }
     setInitialized(true);
     setLoading(false);
   });