Loki: Fix incorrect evaluation of real and extracted labels in context (#67112)

* Loki: Fix incorrect evaluation of real/extracted labels in context

* Add tests

* Improve caching and add more tests

* Use mockResolvedValue

* Flip logic for getInitContextFiltersFromLabels

* Update to logic to use fetchSeriesLabels for queries with more than 1 parser

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

@@ -399,8 +399,6 @@ export default class LokiLanguageProvider extends LanguageProvider {
     const cacheKey = this.generateCacheKey(url, start, end, interpolatedMatch);
     let value = this.seriesCache.get(cacheKey);
     if (!value) {
-      // Clear value when requesting new one. Empty object being truthy also makes sure we don't request twice.
-      this.seriesCache.set(cacheKey, {});
       const params = { 'match[]': interpolatedMatch, start, end };
       const data = await this.request(url, params);
       const { values } = processLabels(data);

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

@@ -15,6 +15,7 @@ import { LokiQuery } from './types';
 
 const defaultLanguageProviderMock = {
   start: jest.fn(),
+  fetchSeriesLabels: jest.fn(() => ({ bar: ['baz'], xyz: ['abc'] })),
   getLabelKeys: jest.fn(() => ['bar', 'xyz']),
 } as unknown as LokiLanguageProvider;
 
@@ -41,23 +42,38 @@ describe('LogContextProvider', () => {
   let logContextProvider: LogContextProvider;
   beforeEach(() => {
     logContextProvider = new LogContextProvider(defaultDatasourceMock);
-    logContextProvider.getInitContextFiltersFromLabels = jest.fn(() =>
-      Promise.resolve([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }])
-    );
   });
 
   describe('getLogRowContext', () => {
     it('should call getInitContextFilters if no appliedContextFilters', async () => {
+      logContextProvider.getInitContextFiltersFromLabels = jest
+        .fn()
+        .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
+
       expect(logContextProvider.appliedContextFilters).toHaveLength(0);
-      await logContextProvider.getLogRowContext(defaultLogRow, {
+      await logContextProvider.getLogRowContext(
-        limit: 10,
+        defaultLogRow,
-        direction: LogRowContextQueryDirection.Backward,
+        {
-      });
+          limit: 10,
+          direction: LogRowContextQueryDirection.Backward,
+        },
+        {
+          expr: '{bar="baz"}',
+        } as LokiQuery
+      );
       expect(logContextProvider.getInitContextFiltersFromLabels).toBeCalled();
+      expect(logContextProvider.getInitContextFiltersFromLabels).toHaveBeenCalledWith(
+        { bar: 'baz', foo: 'uniqueParsedLabel', xyz: 'abc' },
+        { expr: '{bar="baz"}' }
+      );
       expect(logContextProvider.appliedContextFilters).toHaveLength(1);
     });
 
     it('should not call getInitContextFilters if appliedContextFilters', async () => {
+      logContextProvider.getInitContextFiltersFromLabels = jest
+        .fn()
+        .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
+
       logContextProvider.appliedContextFilters = [
         { value: 'baz', enabled: true, fromParser: false, label: 'bar' },
         { value: 'abc', enabled: true, fromParser: false, label: 'xyz' },
@@ -73,6 +89,10 @@ describe('LogContextProvider', () => {
 
   describe('getLogRowContextQuery', () => {
     it('should call getInitContextFilters if no appliedContextFilters', async () => {
+      logContextProvider.getInitContextFiltersFromLabels = jest
+        .fn()
+        .mockResolvedValue([{ value: 'baz', enabled: true, fromParser: false, label: 'bar' }]);
+
       const query = await logContextProvider.getLogRowContextQuery(defaultLogRow, {
         limit: 10,
         direction: LogRowContextQueryDirection.Backward,
@@ -180,4 +200,59 @@ describe('LogContextProvider', () => {
       expect(contextQuery.query.expr).toEqual(`{bar="baz"}`);
     });
   });
+
+  describe('getInitContextFiltersFromLabels', () => {
+    describe('query with no parser', () => {
+      const queryWithoutParser = {
+        expr: '{bar="baz"}',
+      } as LokiQuery;
+
+      it('should correctly create contextFilters', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels(
+          defaultLogRow.labels,
+          queryWithoutParser
+        );
+        expect(filters).toEqual([
+          { enabled: true, fromParser: false, label: 'bar', value: 'baz' },
+          { enabled: false, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
+          { enabled: true, fromParser: false, label: 'xyz', value: 'abc' },
+        ]);
+      });
+
+      it('should return empty contextFilters if no query', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, undefined);
+        expect(filters).toEqual([]);
+      });
+
+      it('should return empty contextFilters if no labels', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels({}, queryWithoutParser);
+        expect(filters).toEqual([]);
+      });
+    });
+
+    describe('query with parser', () => {
+      const queryWithParser = {
+        expr: '{bar="baz"} | logfmt',
+      } as LokiQuery;
+
+      it('should correctly create contextFilters', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, queryWithParser);
+        expect(filters).toEqual([
+          { enabled: true, fromParser: false, label: 'bar', value: 'baz' },
+          { enabled: false, fromParser: true, label: 'foo', value: 'uniqueParsedLabel' },
+          { enabled: true, fromParser: false, label: 'xyz', value: 'abc' },
+        ]);
+      });
+
+      it('should return empty contextFilters if no query', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels(defaultLogRow.labels, undefined);
+        expect(filters).toEqual([]);
+      });
+
+      it('should return empty contextFilters if no labels', async () => {
+        const filters = await logContextProvider.getInitContextFiltersFromLabels({}, queryWithParser);
+        expect(filters).toEqual([]);
+      });
+    });
+  });
 });

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

@@ -1,3 +1,4 @@
+import { isEmpty } from 'lodash';
 import { catchError, lastValueFrom, of, switchMap } from 'rxjs';
 
 import {
@@ -13,13 +14,13 @@ import {
   LogRowContextQueryDirection,
   LogRowContextOptions,
 } from '@grafana/data';
-import { DataQuery, Labels } from '@grafana/schema';
+import { Labels } from '@grafana/schema';
 
 import { LokiContextUi } from './components/LokiContextUi';
 import { LokiDatasource, makeRequest, REF_ID_STARTER_LOG_ROW_CONTEXT } from './datasource';
 import { escapeLabelValueInExactSelector } from './languageUtils';
 import { addLabelToQuery, addParserToQuery } from './modifyQuery';
-import { getParserFromQuery, isLokiQuery, isQueryWithParser } from './queryUtils';
+import { getParserFromQuery, getStreamSelectorsFromQuery, isQueryWithParser } from './queryUtils';
 import { sortDataFrameByTime, SortDirection } from './sortDataFrame';
 import { ContextFilter, LokiQuery, LokiQueryDirection, LokiQueryType } from './types';
 
@@ -33,14 +34,15 @@ export class LogContextProvider {
     this.appliedContextFilters = [];
   }
 
-  private async getQueryAndRange(row: LogRowModel, options?: LogRowContextOptions, origQuery?: DataQuery) {
+  private async getQueryAndRange(row: LogRowModel, options?: LogRowContextOptions, origQuery?: LokiQuery) {
     const direction = (options && options.direction) || LogRowContextQueryDirection.Backward;
     const limit = (options && options.limit) || this.datasource.maxLines;
-
     // 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)).filter((filter) => filter.enabled);
+      const filters = (await this.getInitContextFiltersFromLabels(row.labels, origQuery)).filter(
+        (filter) => filter.enabled
+      );
       this.appliedContextFilters = filters;
     }
 
@@ -50,7 +52,7 @@ export class LogContextProvider {
   getLogRowContextQuery = async (
     row: LogRowModel,
     options?: LogRowContextOptions,
-    origQuery?: DataQuery
+    origQuery?: LokiQuery
   ): Promise<LokiQuery> => {
     const { query } = await this.getQueryAndRange(row, options, origQuery);
 
@@ -60,10 +62,9 @@ export class LogContextProvider {
   getLogRowContext = async (
     row: LogRowModel,
     options?: LogRowContextOptions,
-    origQuery?: DataQuery
+    origQuery?: LokiQuery
   ): Promise<{ data: DataFrame[] }> => {
     const direction = (options && options.direction) || LogRowContextQueryDirection.Backward;
-
     const { query, range } = await this.getQueryAndRange(row, options, origQuery);
 
     const processResults = (result: DataQueryResponse): DataQueryResponse => {
@@ -98,14 +99,9 @@ export class LogContextProvider {
     row: LogRowModel,
     limit: number,
     direction: LogRowContextQueryDirection,
-    origQuery?: DataQuery
+    origQuery?: LokiQuery
   ): Promise<{ query: LokiQuery; range: TimeRange }> {
-    let originalLokiQuery: LokiQuery | undefined = undefined;
+    const expr = this.processContextFiltersToExpr(row, this.appliedContextFilters, origQuery);
-    // Type guard for LokiQuery
-    if (origQuery && isLokiQuery(origQuery)) {
-      originalLokiQuery = origQuery;
-    }
-    const expr = this.processContextFiltersToExpr(row, this.appliedContextFilters, originalLokiQuery);
     const contextTimeBuffer = 2 * 60 * 60 * 1000; // 2h buffer
 
     const queryDirection =
@@ -154,7 +150,7 @@ export class LogContextProvider {
     };
   }
 
-  getLogRowContextUi(row: LogRowModel, runContextQuery?: () => void, originalQuery?: DataQuery): React.ReactNode {
+  getLogRowContextUi(row: LogRowModel, runContextQuery?: () => void, origQuery?: LokiQuery): React.ReactNode {
     const updateFilter = (contextFilters: ContextFilter[]) => {
       this.appliedContextFilters = contextFilters;
 
@@ -170,12 +166,6 @@ export class LogContextProvider {
         this.appliedContextFilters = [];
       });
 
-    let origQuery: LokiQuery | undefined = undefined;
-    // Type guard for LokiQuery
-    if (originalQuery && isLokiQuery(originalQuery)) {
-      origQuery = originalQuery;
-    }
-
     return LokiContextUi({
       row,
       origQuery,
@@ -218,9 +208,25 @@ export class LogContextProvider {
     return expr;
   };
 
-  getInitContextFiltersFromLabels = async (labels: Labels) => {
+  getInitContextFiltersFromLabels = async (labels: Labels, query?: LokiQuery) => {
-    await this.datasource.languageProvider.start();
+    if (!query || isEmpty(labels)) {
-    const allLabels = this.datasource.languageProvider.getLabelKeys();
+      return [];
+    }
+
+    let allLabels: string[] = [];
+    if (!isQueryWithParser(query.expr).queryWithParser) {
+      // If there is no parser, we use getLabelKeys because it has better caching
+      // and all labels should already be fetched
+      await this.datasource.languageProvider.start();
+      allLabels = this.datasource.languageProvider.getLabelKeys();
+    } else {
+      // If we have parser, we use fetchSeriesLabels to fetch actual labels for selected stream
+      const stream = getStreamSelectorsFromQuery(query.expr);
+      // We are using stream[0] as log query can always have just 1 stream selector
+      const series = await this.datasource.languageProvider.fetchSeriesLabels(stream[0]);
+      allLabels = Object.keys(series);
+    }
+
     const contextFilters: ContextFilter[] = [];
     Object.entries(labels).forEach(([label, value]) => {
       const filter: ContextFilter = {
@@ -229,6 +235,7 @@ export class LogContextProvider {
         enabled: allLabels.includes(label),
         fromParser: !allLabels.includes(label),
       };
+
       contextFilters.push(filter);
     });
 

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

@@ -123,7 +123,7 @@ export function LokiContextUi(props: LokiContextUiProps) {
 
   useAsync(async () => {
     setLoading(true);
-    const contextFilters = await logContextProvider.getInitContextFiltersFromLabels(row.labels);
+    const contextFilters = await logContextProvider.getInitContextFiltersFromLabels(row.labels, origQuery);
     setContextFilters(contextFilters);
     setInitialized(true);
     setLoading(false);

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

@@ -70,6 +70,7 @@ import { getQueryHints } from './queryHints';
 import { runSplitQuery } from './querySplitting';
 import {
   getLogQueryFromMetricsQuery,
+  getLokiQueryFromDataQuery,
   getNormalizedLokiQuery,
   getStreamSelectorsFromQuery,
   isLogsQuery,
@@ -654,7 +655,7 @@ export class LokiDatasource
     options?: LogRowContextOptions,
     origQuery?: DataQuery
   ): Promise<{ data: DataFrame[] }> => {
-    return await this.logContextProvider.getLogRowContext(row, options, origQuery);
+    return await this.logContextProvider.getLogRowContext(row, options, getLokiQueryFromDataQuery(origQuery));
   };
 
   getLogRowContextQuery = async (
@@ -662,11 +663,11 @@ export class LokiDatasource
     options?: LogRowContextOptions,
     origQuery?: DataQuery
   ): Promise<DataQuery> => {
-    return await this.logContextProvider.getLogRowContextQuery(row, options, origQuery);
+    return await this.logContextProvider.getLogRowContextQuery(row, options, getLokiQueryFromDataQuery(origQuery));
   };
 
   getLogRowContextUi(row: LogRowModel, runContextQuery: () => void, origQuery: DataQuery): React.ReactNode {
-    return this.logContextProvider.getLogRowContextUi(row, runContextQuery, origQuery);
+    return this.logContextProvider.getLogRowContextUi(row, runContextQuery, getLokiQueryFromDataQuery(origQuery));
   }
 
   testDatasource(): Promise<{ status: string; message: string }> {

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

@@ -314,3 +314,11 @@ export const isLokiQuery = (query: DataQuery): query is LokiQuery => {
   const lokiQuery = query as LokiQuery;
   return lokiQuery.expr !== undefined;
 };
+
+export const getLokiQueryFromDataQuery = (query?: DataQuery): LokiQuery | undefined => {
+  if (!query || !isLokiQuery(query)) {
+    return undefined;
+  }
+
+  return query;
+};