Logs: Add experimental support to display a datasource custom UI in LogContext (#62189)

* add loki contextfilter component

* add `getLogRowContextUi` support to DataSourceAPI

* add `runContextQuery` to LogRowContextProvider

* pass `getRowContextUi` to `LogRowContext`

* adapt LogRowContext to show datasource ui

* implement LogRowContextUi in Loki

* add `logsContextDatasourceUi` feature flag

* change state to `Alpha`

* disable the feature if `logsContextDatasourceUi` is not set

* don't fetch labels in the constructor

* adjust to right height

* remove unnecessary eslint disable

* add test for LokiContextUi

* move code down in datasource.ts

* rename `refresh` to `runContextQuery`

* update datasource tests

* don't update if `updateFilter` fn changes

* organized imports in datasource.test.ts

* don't trigger on intialization changes

* change tag to `experimental`

* move `getLogRowContextUi` to props

.betterer.results

@@ -523,7 +523,8 @@ exports[`better eslint`] = {
     ],
     "packages/grafana-data/src/types/logs.ts:5381": [
       [0, 0, 0, "Do not use any type assertions.", "0"],
-      [0, 0, 0, "Do not use any type assertions.", "1"]
+      [0, 0, 0, "Do not use any type assertions.", "1"],
+      [0, 0, 0, "Do not use any type assertions.", "2"]
     ],
     "packages/grafana-data/src/types/logsVolume.ts:5381": [
       [0, 0, 0, "Do not use any type assertions.", "0"]
@@ -6097,9 +6098,10 @@ exports[`better eslint`] = {
       [0, 0, 0, "Unexpected any. Specify a different type.", "1"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "2"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "3"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "4"],
+      [0, 0, 0, "Do not use any type assertions.", "4"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "5"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "6"]
+      [0, 0, 0, "Unexpected any. Specify a different type.", "6"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "7"]
     ],
     "public/app/plugins/datasource/loki/getDerivedFields.ts:5381": [
       [0, 0, 0, "Unexpected any. Specify a different type.", "0"],

docs/sources/setup-grafana/configure-grafana/feature-toggles/index.md

@@ -98,6 +98,7 @@ Alpha features might be changed or removed without prior notice.
 | `alertingBacktesting`              | Rule backtesting API for alerting                                                                                                                                            |
 | `editPanelCSVDragAndDrop`          | Enables drag and drop for CSV and Excel files                                                                                                                                |
 | `azureMultipleResourcePicker`      | Azure multiple resource picker                                                                                                                                               |
+| `logsContextDatasourceUi`          | Allow datasource to provide custom UI for context view                                                                                                                       |
 
 ## Development feature toggles
 

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -92,4 +92,5 @@ export interface FeatureToggles {
   alertingNoNormalState?: boolean;
   azureMultipleResourcePicker?: boolean;
   topNavCommandPalette?: boolean;
+  logsContextDatasourceUi?: boolean;
 }

packages/grafana-data/src/types/logs.ts

@@ -167,6 +167,13 @@ export interface DataSourceWithLogsContextSupport<TQuery extends DataQuery = Dat
    * This method can be used to show "context" button based on runtime conditions (for example row model data or plugin settings, etc.)
    */
   showContextToggle(row?: LogRowModel): boolean;
+
+  /**
+   * This method can be used to display a custom UI in the context view.
+   * @alpha
+   * @internal
+   */
+  getLogRowContextUi?(row: LogRowModel, runContextQuery?: () => void): React.ReactNode;
 }
 
 export const hasLogsContextSupport = (datasource: unknown): datasource is DataSourceWithLogsContextSupport => {
@@ -230,3 +237,13 @@ export const hasSupplementaryQuerySupport = <TQuery extends DataQuery>(
     withSupplementaryQueriesSupport.getSupportedSupplementaryQueryTypes().includes(type)
   );
 };
+
+export const hasLogsContextUiSupport = (datasource: unknown): datasource is DataSourceWithLogsContextSupport => {
+  if (!datasource) {
+    return false;
+  }
+
+  const withLogsSupport = datasource as DataSourceWithLogsContextSupport;
+
+  return withLogsSupport.getLogRowContextUi !== undefined;
+};

pkg/services/featuremgmt/registry.go

@@ -426,5 +426,11 @@ var (
 			State:        FeatureStateBeta,
 			FrontendOnly: true,
 		},
+		{
+			Name:         "logsContextDatasourceUi",
+			Description:  "Allow datasource to provide custom UI for context view",
+			State:        FeatureStateAlpha,
+			FrontendOnly: true,
+		},
 	}
 )

pkg/services/featuremgmt/toggles_gen.go

@@ -310,4 +310,8 @@ const (
 	// FlagTopNavCommandPalette
 	// Launch the Command Palette from the top navigation search box
 	FlagTopNavCommandPalette = "topNavCommandPalette"
+
+	// FlagLogsContextDatasourceUi
+	// Allow datasource to provide custom UI for context view
+	FlagLogsContextDatasourceUi = "logsContextDatasourceUi"
 )

public/app/features/explore/Logs.tsx

@@ -25,6 +25,7 @@ import {
   DataHoverEvent,
   DataHoverClearEvent,
   EventBus,
+  DataSourceWithLogsContextSupport,
 } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
 import { DataQuery } from '@grafana/schema';
@@ -79,6 +80,7 @@ interface Props extends Themeable2 {
   onStartScanning?: () => void;
   onStopScanning?: () => void;
   getRowContext?: (row: LogRowModel, options?: RowContextOptions) => Promise<any>;
+  getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'];
   getFieldLinks: (field: Field, rowIndex: number, dataFrame: DataFrame) => Array<LinkModel<Field>>;
   addResultsToCache: () => void;
   clearCache: () => void;
@@ -344,6 +346,8 @@ class UnthemedLogs extends PureComponent<Props, State> {
       addResultsToCache,
       exploreId,
       scrollElement,
+      getRowContext,
+      getLogRowContextUi,
     } = this.props;
 
     const {
@@ -487,7 +491,8 @@ class UnthemedLogs extends PureComponent<Props, State> {
                 logRows={logRows}
                 deduplicatedRows={dedupedRows}
                 dedupStrategy={dedupStrategy}
-                getRowContext={this.props.getRowContext}
+                getRowContext={getRowContext}
+                getLogRowContextUi={getLogRowContextUi}
                 onClickFilterLabel={onClickFilterLabel}
                 onClickFilterOutLabel={onClickFilterOutLabel}
                 showContextToggle={showContextToggle}

public/app/features/explore/LogsContainer.tsx

@@ -5,6 +5,7 @@ import {
   AbsoluteTimeRange,
   Field,
   hasLogsContextSupport,
+  hasLogsContextUiSupport,
   LoadingState,
   LogRowModel,
   RawTimeRange,
@@ -63,6 +64,16 @@ class LogsContainer extends PureComponent<LogsContainerProps> {
     return [];
   };
 
+  getLogRowContextUi = (row: LogRowModel, runContextQuery?: () => void): React.ReactNode => {
+    const { datasourceInstance } = this.props;
+
+    if (hasLogsContextUiSupport(datasourceInstance) && datasourceInstance.getLogRowContextUi) {
+      return datasourceInstance.getLogRowContextUi(row, runContextQuery);
+    }
+
+    return <></>;
+  };
+
   showContextToggle = (row?: LogRowModel): boolean => {
     const { datasourceInstance } = this.props;
 
@@ -159,6 +170,7 @@ class LogsContainer extends PureComponent<LogsContainerProps> {
             scanRange={range.raw}
             showContextToggle={this.showContextToggle}
             getRowContext={this.getLogRowContext}
+            getLogRowContextUi={this.getLogRowContextUi}
             getFieldLinks={this.getFieldLinks}
             addResultsToCache={() => addResultsToCache(exploreId)}
             clearCache={() => clearCache(exploreId)}

public/app/features/logs/components/LogRow.tsx

@@ -12,6 +12,7 @@ import {
   GrafanaTheme2,
   CoreApp,
   DataFrame,
+  DataSourceWithLogsContextSupport,
 } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
 import { styleMixins, withTheme2, Themeable2, Icon, Tooltip } from '@grafana/ui';
@@ -53,6 +54,7 @@ interface Props extends Themeable2 {
   onClickFilterOutLabel?: (key: string, value: string) => void;
   onContextClick?: () => void;
   getRowContext: (row: LogRowModel, options?: RowContextOptions) => Promise<DataQueryResponse>;
+  getLogRowContextUi?: (row: LogRowModel) => React.ReactNode;
   getFieldLinks?: (field: Field, rowIndex: number, dataFrame: DataFrame) => Array<LinkModel<Field>>;
   showContextToggle?: (row?: LogRowModel) => boolean;
   onClickShowField?: (key: string) => void;
@@ -143,7 +145,9 @@ class UnThemedLogRow extends PureComponent<Props, State> {
     errors?: LogRowContextQueryErrors,
     hasMoreContextRows?: HasMoreContextRows,
     updateLimit?: () => void,
-    logsSortOrder?: LogsSortOrder | null
+    logsSortOrder?: LogsSortOrder | null,
+    getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'],
+    runContextQuery?: () => void
   ) {
     const {
       getRows,
@@ -230,6 +234,8 @@ class UnThemedLogRow extends PureComponent<Props, State> {
               getRows={getRows}
               errors={errors}
               hasMoreContextRows={hasMoreContextRows}
+              getLogRowContextUi={getLogRowContextUi}
+              runContextQuery={runContextQuery}
               updateLimit={updateLimit}
               context={context}
               contextIsOpen={showContext}
@@ -267,14 +273,26 @@ class UnThemedLogRow extends PureComponent<Props, State> {
 
   render() {
     const { showContext } = this.state;
-    const { logsSortOrder, row, getRowContext } = this.props;
+    const { logsSortOrder, row, getRowContext, getLogRowContextUi } = this.props;
 
     if (showContext) {
       return (
         <>
           <LogRowContextProvider row={row} getRowContext={getRowContext} logsSortOrder={logsSortOrder}>
-            {({ result, errors, hasMoreContextRows, updateLimit, logsSortOrder }) => {
+            {({ result, errors, hasMoreContextRows, updateLimit, runContextQuery, logsSortOrder }) => {
-              return <>{this.renderLogRow(result, errors, hasMoreContextRows, updateLimit, logsSortOrder)}</>;
+              return (
+                <>
+                  {this.renderLogRow(
+                    result,
+                    errors,
+                    hasMoreContextRows,
+                    updateLimit,
+                    logsSortOrder,
+                    getLogRowContextUi,
+                    runContextQuery
+                  )}
+                </>
+              );
             }}
           </LogRowContextProvider>
         </>

public/app/features/logs/components/LogRowContext.tsx

@@ -1,8 +1,16 @@
 import { css, cx } from '@emotion/css';
-import React, { useEffect, useLayoutEffect, useRef, useState } from 'react';
+import React, { useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';
 import usePrevious from 'react-use/lib/usePrevious';
 
-import { DataQueryError, GrafanaTheme2, LogRowModel, LogsSortOrder, textUtil } from '@grafana/data';
+import {
+  DataQueryError,
+  GrafanaTheme2,
+  LogRowModel,
+  LogsSortOrder,
+  textUtil,
+  DataSourceWithLogsContextSupport,
+} from '@grafana/data';
+import { config } from '@grafana/runtime';
 import { Alert, Button, ClickOutsideWrapper, CustomScrollbar, IconButton, List, useStyles2 } from '@grafana/ui';
 
 import { LogMessageAnsi } from './LogMessageAnsi';
@@ -22,9 +30,16 @@ interface LogRowContextProps {
   logsSortOrder?: LogsSortOrder | null;
   onOutsideClick: (method: string) => void;
   onLoadMoreContext: () => void;
+  runContextQuery?: () => void;
+  getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'];
 }
 
-const getLogRowContextStyles = (theme: GrafanaTheme2, wrapLogMessage?: boolean) => {
+const getLogRowContextStyles = (theme: GrafanaTheme2, wrapLogMessage?: boolean, datasourceUiHeight?: number) => {
+  if (config.featureToggles.logsContextDatasourceUi) {
+    datasourceUiHeight = datasourceUiHeight ?? 55;
+  } else {
+    datasourceUiHeight = 0;
+  }
   /**
    * This is workaround for displaying uncropped context when we have unwrapping log messages.
    * We are using margins to correctly position context. Because non-wrapped logs have always 1 line of log
@@ -34,7 +49,8 @@ const getLogRowContextStyles = (theme: GrafanaTheme2, wrapLogMessage?: boolean)
 
   const headerHeight = 40;
   const logsHeight = 220;
-  const contextHeight = headerHeight + logsHeight;
+  const contextHeight = datasourceUiHeight + headerHeight + logsHeight;
+  const bottomContextHeight = headerHeight + logsHeight;
   const width = wrapLogMessage ? '100%' : '75%';
   const afterContext = wrapLogMessage
     ? css`
@@ -55,6 +71,9 @@ const getLogRowContextStyles = (theme: GrafanaTheme2, wrapLogMessage?: boolean)
     width: css`
       width: ${width};
     `,
+    bottomContext: css`
+      height: ${bottomContextHeight}px;
+    `,
     commonStyles: css`
       position: absolute;
       height: ${contextHeight}px;
@@ -73,6 +92,13 @@ const getLogRowContextStyles = (theme: GrafanaTheme2, wrapLogMessage?: boolean)
       align-items: center;
       background: ${theme.colors.background.canvas};
     `,
+    datasourceUi: css`
+      height: ${datasourceUiHeight}px;
+      padding: ${theme.spacing(0, 1.25)};
+      display: flex;
+      align-items: center;
+      background: ${theme.colors.background.canvas};
+    `,
     top: css`
       border-radius: 0 0 ${theme.shape.borderRadius(2)} ${theme.shape.borderRadius(2)};
       box-shadow: 0 0 ${theme.spacing(1.25)} ${theme.v1.palette.black};
@@ -132,13 +158,14 @@ interface LogRowContextGroupHeaderProps {
   shouldScrollToBottom?: boolean;
   canLoadMoreRows?: boolean;
   logsSortOrder?: LogsSortOrder | null;
+  getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'];
+  runContextQuery?: () => void;
 }
 interface LogRowContextGroupProps extends LogRowContextGroupHeaderProps {
   rows: Array<string | DataQueryError>;
   groupPosition: LogGroupPosition;
   className?: string;
   error?: string;
-  logsSortOrder?: LogsSortOrder | null;
 }
 
 const LogRowContextGroupHeader: React.FunctionComponent<LogRowContextGroupHeaderProps> = ({
@@ -148,8 +175,16 @@ const LogRowContextGroupHeader: React.FunctionComponent<LogRowContextGroupHeader
   canLoadMoreRows,
   groupPosition,
   logsSortOrder,
+  getLogRowContextUi,
+  runContextQuery,
 }) => {
-  const { header, headerButton } = useStyles2(getLogRowContextStyles);
+  const [height, setHeight] = useState(50);
+  const datasourceUiRef = React.createRef<HTMLDivElement>();
+  const {
+    datasourceUi: dsUi,
+    header,
+    headerButton,
+  } = useStyles2((theme) => getLogRowContextStyles(theme, undefined, height));
 
   // determine the position in time for this LogGroup by taking the ordering of
   // logs and position of the component itself into account.
@@ -162,21 +197,56 @@ const LogRowContextGroupHeader: React.FunctionComponent<LogRowContextGroupHeader
     logGroupPosition = 'before';
   }
 
+  if (config.featureToggles.logsContextDatasourceUi) {
+    // disabling eslint here, because this condition does not change in runtime
+
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    const resizeObserver = useMemo(
+      () =>
+        new ResizeObserver((entries) => {
+          for (let entry of entries) {
+            setHeight(entry.contentRect.height);
+          }
+        }),
+      []
+    );
+
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    useLayoutEffect(() => {
+      // observe the first child of the ref, which is the datasource controlled component and varies in height
+      // TODO: this is a bit of a hack and we can remove this as soon as we move back from the absolute positioned context
+      const child = datasourceUiRef.current?.children.item(0);
+      if (child) {
+        resizeObserver.observe(child);
+      }
+      return () => {
+        resizeObserver.disconnect();
+      };
+    }, [datasourceUiRef, resizeObserver]);
+  }
+
   return (
-    <div className={header}>
+    <>
-      <span
+      {config.featureToggles.logsContextDatasourceUi && getLogRowContextUi && (
-        className={css`
+        <div ref={datasourceUiRef} className={dsUi}>
-          opacity: 0.6;
+          {getLogRowContextUi(row, runContextQuery)}
-        `}
+        </div>
-      >
-        Showing {rows.length} lines {logGroupPosition} match.
-      </span>
-      {(rows.length >= 10 || (rows.length > 10 && rows.length % 10 !== 0)) && canLoadMoreRows && (
-        <Button className={headerButton} variant="secondary" size="sm" onClick={onLoadMoreContext}>
-          Load 10 more lines
-        </Button>
       )}
-    </div>
+      <div className={header}>
+        <span
+          className={css`
+            opacity: 0.6;
+          `}
+        >
+          Showing {rows.length} lines {logGroupPosition} match.
+        </span>
+        {(rows.length >= 10 || (rows.length > 10 && rows.length % 10 !== 0)) && canLoadMoreRows && (
+          <Button className={headerButton} variant="secondary" size="sm" onClick={onLoadMoreContext}>
+            Load 10 more lines
+          </Button>
+        )}
+      </div>
+    </>
   );
 };
 
@@ -190,8 +260,10 @@ export const LogRowContextGroup: React.FunctionComponent<LogRowContextGroupProps
   onLoadMoreContext,
   groupPosition,
   logsSortOrder,
+  getLogRowContextUi,
+  runContextQuery,
 }) => {
-  const { commonStyles, logs } = useStyles2(getLogRowContextStyles);
+  const { commonStyles, logs, bottomContext } = useStyles2(getLogRowContextStyles);
   const [scrollTop, setScrollTop] = useState(0);
   const [scrollHeight, setScrollHeight] = useState(0);
 
@@ -243,10 +315,12 @@ export const LogRowContextGroup: React.FunctionComponent<LogRowContextGroupProps
     canLoadMoreRows,
     groupPosition,
     logsSortOrder,
+    getLogRowContextUi,
+    runContextQuery,
   };
 
   return (
-    <div className={cx(commonStyles, className)}>
+    <div className={cx(commonStyles, className, groupPosition === LogGroupPosition.Bottom ? bottomContext : '')}>
       {/* When displaying "after" context */}
       {shouldScrollToBottom && !error && <LogRowContextGroupHeader {...headerProps} />}
       <div className={logs}>
@@ -284,9 +358,11 @@ export const LogRowContext: React.FunctionComponent<LogRowContextProps> = ({
   errors,
   onOutsideClick,
   onLoadMoreContext,
+  runContextQuery: runContextQuery,
   hasMoreContextRows,
   wrapLogMessage,
   logsSortOrder,
+  getLogRowContextUi,
 }) => {
   useEffect(() => {
     const handleEscKeyDown = (e: KeyboardEvent): void => {
@@ -321,6 +397,8 @@ export const LogRowContext: React.FunctionComponent<LogRowContextProps> = ({
             onLoadMoreContext={onLoadMoreContext}
             groupPosition={LogGroupPosition.Top}
             logsSortOrder={logsSortOrder}
+            getLogRowContextUi={getLogRowContextUi}
+            runContextQuery={runContextQuery}
           />
         )}
 

public/app/features/logs/components/LogRowContextProvider.tsx

@@ -34,6 +34,7 @@ export interface HasMoreContextRows {
 interface ResultType {
   data: string[][];
   errors: string[];
+  doNotCheckForMore?: boolean;
 }
 
 interface LogRowContextProviderProps {
@@ -45,6 +46,7 @@ interface LogRowContextProviderProps {
     errors: LogRowContextQueryErrors;
     hasMoreContextRows: HasMoreContextRows;
     updateLimit: () => void;
+    runContextQuery: () => void;
     limit: number;
     logsSortOrder?: LogsSortOrder | null;
   }) => JSX.Element;
@@ -55,7 +57,7 @@ export const getRowContexts = async (
   row: LogRowModel,
   limit: number,
   logsSortOrder?: LogsSortOrder | null
-) => {
+): Promise<ResultType> => {
   const promises = [
     getRowContext(row, {
       limit,
@@ -159,6 +161,8 @@ export const LogRowContextProvider: React.FunctionComponent<LogRowContextProvide
     after: true,
   });
 
+  const [results, setResults] = useState<ResultType>();
+
   // React Hook that resolves two promises every time the limit prop changes
   // First promise fetches limit number of rows backwards in time from a specific point in time
   // Second promise fetches limit number of rows forwards in time from a specific point in time
@@ -166,40 +170,46 @@ export const LogRowContextProvider: React.FunctionComponent<LogRowContextProvide
     return await getRowContexts(getRowContext, row, limit, logsSortOrder); // Moved it to a separate function for debugging purposes
   }, [limit]);
 
+  useEffect(() => {
+    setResults(value);
+  }, [value]);
+
   // React Hook that performs a side effect every time the value (from useAsync hook) prop changes
   // The side effect changes the result state with the response from the useAsync hook
   // The side effect changes the hasMoreContextRows state if there are more context rows before or after the current result
   useEffect(() => {
-    if (value) {
+    if (results) {
       setResult((currentResult) => {
-        let hasMoreLogsBefore = true,
+        if (!results.doNotCheckForMore) {
-          hasMoreLogsAfter = true;
+          let hasMoreLogsBefore = true,
+            hasMoreLogsAfter = true;
 
-        const currentResultBefore = currentResult?.data[0];
+          const currentResultBefore = currentResult?.data[0];
-        const currentResultAfter = currentResult?.data[1];
+          const currentResultAfter = currentResult?.data[1];
-        const valueBefore = value.data[0];
+          const valueBefore = results.data[0];
-        const valueAfter = value.data[1];
+          const valueAfter = results.data[1];
 
-        // checks if there are more log rows in a given direction
+          // checks if there are more log rows in a given direction
-        // if after fetching additional rows the length of result is the same,
+          // if after fetching additional rows the length of result is the same,
-        // we can assume there are no logs in that direction within a given time range
+          // we can assume there are no logs in that direction within a given time range
-        if (currentResult && (!valueBefore || currentResultBefore.length === valueBefore.length)) {
+          if (currentResult && (!valueBefore || currentResultBefore.length === valueBefore.length)) {
-          hasMoreLogsBefore = false;
+            hasMoreLogsBefore = false;
+          }
+
+          if (currentResult && (!valueAfter || currentResultAfter.length === valueAfter.length)) {
+            hasMoreLogsAfter = false;
+          }
+
+          setHasMoreContextRows({
+            before: hasMoreLogsBefore,
+            after: hasMoreLogsAfter,
+          });
         }
 
-        if (currentResult && (!valueAfter || currentResultAfter.length === valueAfter.length)) {
+        return results;
-          hasMoreLogsAfter = false;
-        }
-
-        setHasMoreContextRows({
-          before: hasMoreLogsBefore,
-          after: hasMoreLogsAfter,
-        });
-
-        return value;
       });
     }
-  }, [value]);
+  }, [results]);
 
   return children({
     result: {
@@ -221,6 +231,11 @@ export const LogRowContextProvider: React.FunctionComponent<LogRowContextProvide
         newLimit: limit + 10,
       });
     },
+    runContextQuery: async () => {
+      const results = await getRowContexts(getRowContext, row, limit, logsSortOrder);
+      results.doNotCheckForMore = true;
+      setResults(results);
+    },
     limit,
     logsSortOrder,
   });

public/app/features/logs/components/LogRowMessage.tsx

@@ -4,7 +4,14 @@ import React, { PureComponent } from 'react';
 import Highlighter from 'react-highlight-words';
 import tinycolor from 'tinycolor2';
 
-import { LogRowModel, findHighlightChunksInText, GrafanaTheme2, LogsSortOrder, CoreApp } from '@grafana/data';
+import {
+  LogRowModel,
+  findHighlightChunksInText,
+  GrafanaTheme2,
+  LogsSortOrder,
+  CoreApp,
+  DataSourceWithLogsContextSupport,
+} from '@grafana/data';
 import { withTheme2, Themeable2, IconButton, Tooltip } from '@grafana/ui';
 
 import { LogMessageAnsi } from './LogMessageAnsi';
@@ -26,9 +33,11 @@ interface Props extends Themeable2 {
   app?: CoreApp;
   scrollElement?: HTMLDivElement;
   showContextToggle?: (row?: LogRowModel) => boolean;
+  getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'];
   getRows: () => LogRowModel[];
   onToggleContext: (method: string) => void;
   updateLimit?: () => void;
+  runContextQuery?: () => void;
   logsSortOrder?: LogsSortOrder | null;
 }
 
@@ -154,6 +163,7 @@ class UnThemedLogRowMessage extends PureComponent<Props> {
       errors,
       hasMoreContextRows,
       updateLimit,
+      runContextQuery,
       context,
       contextIsOpen,
       showRowMenu,
@@ -163,6 +173,7 @@ class UnThemedLogRowMessage extends PureComponent<Props> {
       app,
       logsSortOrder,
       showContextToggle,
+      getLogRowContextUi,
     } = this.props;
 
     const style = getLogRowStyles(theme, row.logLevel);
@@ -191,6 +202,8 @@ class UnThemedLogRowMessage extends PureComponent<Props> {
             {contextIsOpen && context && (
               <LogRowContext
                 row={row}
+                getLogRowContextUi={getLogRowContextUi}
+                runContextQuery={runContextQuery}
                 context={context}
                 errors={errors}
                 wrapLogMessage={wrapLogMessage}

public/app/features/logs/components/LogRows.tsx

@@ -10,6 +10,7 @@ import {
   LogsSortOrder,
   CoreApp,
   DataFrame,
+  DataSourceWithLogsContextSupport,
 } from '@grafana/data';
 import { withTheme2, Themeable2 } from '@grafana/ui';
 
@@ -42,6 +43,7 @@ export interface Props extends Themeable2 {
   onClickFilterLabel?: (key: string, value: string) => void;
   onClickFilterOutLabel?: (key: string, value: string) => void;
   getRowContext?: (row: LogRowModel, options?: RowContextOptions) => Promise<any>;
+  getLogRowContextUi?: DataSourceWithLogsContextSupport['getLogRowContextUi'];
   getFieldLinks?: (field: Field, rowIndex: number, dataFrame: DataFrame) => Array<LinkModel<Field>>;
   onClickShowField?: (key: string) => void;
   onClickHideField?: (key: string) => void;
@@ -128,6 +130,7 @@ class UnThemedLogRows extends PureComponent<Props, State> {
       onLogRowHover,
       app,
       scrollElement,
+      getLogRowContextUi,
     } = this.props;
     const { renderAll, contextIsOpen } = this.state;
     const { logsRowsTable } = getLogRowStyles(theme);
@@ -156,6 +159,7 @@ class UnThemedLogRows extends PureComponent<Props, State> {
                 key={row.uid}
                 getRows={getRows}
                 getRowContext={getRowContext}
+                getLogRowContextUi={getLogRowContextUi}
                 row={row}
                 showContextToggle={showContextToggle}
                 showRowMenu={!contextIsOpen}
@@ -187,6 +191,7 @@ class UnThemedLogRows extends PureComponent<Props, State> {
                 key={row.uid}
                 getRows={getRows}
                 getRowContext={getRowContext}
+                getLogRowContextUi={getLogRowContextUi}
                 row={row}
                 showContextToggle={showContextToggle}
                 showRowMenu={!contextIsOpen}

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

@@ -0,0 +1,116 @@
+import { act, render, screen, waitFor } from '@testing-library/react';
+import React from 'react';
+import { selectOptionInTest } from 'test/helpers/selectOptionInTest';
+
+import { LogRowModel } from '@grafana/data';
+
+import LokiLanguageProvider from '../LanguageProvider';
+
+import { LokiContextUi, LokiContextUiProps } from './LokiContextUi';
+
+// we have to mock out reportInteraction, otherwise it crashes the test.
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  reportInteraction: () => null,
+}));
+
+describe('LokiContextUi', () => {
+  const setupProps = (): LokiContextUiProps => {
+    const mockLanguageProvider = {
+      start: jest.fn().mockImplementation(() => Promise.resolve()),
+      getLabelValues: (name: string) => {
+        switch (name) {
+          case 'label1':
+            return ['value1-1', 'value1-2'];
+          case 'label2':
+            return ['value2-1', 'value2-2'];
+          case 'label3':
+            return ['value3-1', 'value3-2'];
+        }
+        return [];
+      },
+      fetchSeriesLabels: (selector: string) => {
+        switch (selector) {
+          case '{label1="value1-1"}':
+            return { label1: ['value1-1'], label2: ['value2-1'], label3: ['value3-1'] };
+          case '{label1=~"value1-1|value1-2"}':
+            return { label1: ['value1-1', 'value1-2'], label2: ['value2-1'], label3: ['value3-1', 'value3-2'] };
+        }
+        // Allow full set by default
+        return {
+          label1: ['value1-1', 'value1-2'],
+          label2: ['value2-1', 'value2-2'],
+        };
+      },
+      getLabelKeys: () => ['label1', 'label2'],
+    };
+
+    const defaults: LokiContextUiProps = {
+      languageProvider: mockLanguageProvider as unknown as LokiLanguageProvider,
+      updateFilter: jest.fn(),
+      row: {
+        entry: 'WARN test 1.23 on [xxx]',
+        labels: {
+          label1: 'value1',
+          label3: 'value3',
+        },
+      } as unknown as LogRowModel,
+    };
+
+    return defaults;
+  };
+
+  it('renders and shows basic text', async () => {
+    const props = setupProps();
+    render(<LokiContextUi {...props} />);
+
+    // Initial set of labels is available and not selected
+    expect(await screen.findByText(/Select labels to include in the context query/)).toBeInTheDocument();
+  });
+
+  it('starts the languageProvider', async () => {
+    const props = setupProps();
+    render(<LokiContextUi {...props} />);
+
+    await waitFor(() => {
+      expect(props.languageProvider.start).toHaveBeenCalled();
+    });
+  });
+
+  it('finds label1 as a real label', async () => {
+    const props = setupProps();
+    render(<LokiContextUi {...props} />);
+    await waitFor(() => {
+      expect(props.languageProvider.start).toHaveBeenCalled();
+    });
+    const select = await screen.findAllByRole('combobox');
+    await selectOptionInTest(select[0], 'label1');
+  });
+
+  it('finds label3 as a parsed label', async () => {
+    const props = setupProps();
+    render(<LokiContextUi {...props} />);
+    await waitFor(() => {
+      expect(props.languageProvider.start).toHaveBeenCalled();
+    });
+    const select = await screen.findAllByRole('combobox');
+    await selectOptionInTest(select[1], 'label3');
+  });
+
+  it('calls updateFilter when selecting a label', async () => {
+    jest.useFakeTimers();
+    const props = setupProps();
+    render(<LokiContextUi {...props} />);
+    await waitFor(() => {
+      expect(props.languageProvider.start).toHaveBeenCalled();
+    });
+    const select = await screen.findAllByRole('combobox');
+    await selectOptionInTest(select[1], 'label3');
+    act(() => {
+      jest.runAllTimers();
+    });
+    expect(props.updateFilter).toHaveBeenCalled();
+
+    jest.useRealTimers();
+  });
+});

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

@@ -0,0 +1,178 @@
+import { css } from '@emotion/css';
+import memoizeOne from 'memoize-one';
+import React, { useEffect, useState } from 'react';
+import { useAsync } from 'react-use';
+
+import { GrafanaTheme2, LogRowModel, SelectableValue } from '@grafana/data';
+import { MultiSelect, Tag, Tooltip, useStyles2 } from '@grafana/ui';
+
+import LokiLanguageProvider from '../LanguageProvider';
+import { ContextFilter } from '../types';
+
+export interface LokiContextUiProps {
+  languageProvider: LokiLanguageProvider;
+  row: LogRowModel;
+  updateFilter: (value: ContextFilter[]) => void;
+}
+
+function getStyles(theme: GrafanaTheme2) {
+  return {
+    labels: css`
+      display: flex;
+      gap: 2px;
+    `,
+    multiSelectWrapper: css`
+      display: flex;
+      flex-direction: column;
+      flex: 1;
+      margin-top: ${theme.spacing(1)};
+      gap: ${theme.spacing(0.5)};
+    `,
+    multiSelect: css`
+      & .scrollbar-view {
+        overscroll-behavior: contain;
+      }
+    `,
+  };
+}
+
+const formatOptionLabel = memoizeOne(({ label, description }: SelectableValue<string>) => (
+  <Tooltip content={`${label}="${description}"`} placement="top" interactive={true}>
+    <span>{label}</span>
+  </Tooltip>
+));
+
+export function LokiContextUi(props: LokiContextUiProps) {
+  const { row, languageProvider, updateFilter } = props;
+  const styles = useStyles2(getStyles);
+
+  const [contextFilters, setContextFilters] = useState<ContextFilter[]>([]);
+  const [initialized, setInitialized] = useState(false);
+  const timerHandle = React.useRef<number>();
+  const previousInitialized = React.useRef<boolean>(false);
+  useEffect(() => {
+    if (!initialized) {
+      return;
+    }
+
+    // don't trigger if we initialized, this will be the same query anyways.
+    if (!previousInitialized.current) {
+      previousInitialized.current = initialized;
+      return;
+    }
+
+    if (timerHandle.current) {
+      clearTimeout(timerHandle.current);
+    }
+    timerHandle.current = window.setTimeout(() => {
+      updateFilter(contextFilters);
+    }, 1500);
+
+    return () => {
+      clearTimeout(timerHandle.current);
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [contextFilters, initialized]);
+
+  useAsync(async () => {
+    await languageProvider.start();
+    const allLabels = languageProvider.getLabelKeys();
+    const contextFilters: ContextFilter[] = [];
+
+    Object.entries(row.labels).forEach(([label, value]) => {
+      const filter: ContextFilter = {
+        label,
+        value: label, // this looks weird in the first place, but we need to set the label as value here
+        enabled: allLabels.includes(label),
+        fromParser: !allLabels.includes(label),
+        description: value,
+      };
+      contextFilters.push(filter);
+    });
+
+    setContextFilters(contextFilters);
+    setInitialized(true);
+  });
+
+  const realLabels = contextFilters.filter(({ fromParser }) => !fromParser);
+  const realLabelsEnabled = realLabels.filter(({ enabled }) => enabled);
+
+  const parsedLabels = contextFilters.filter(({ fromParser }) => fromParser);
+  const parsedLabelsEnabled = parsedLabels.filter(({ enabled }) => enabled);
+
+  return (
+    <div className={styles.multiSelectWrapper}>
+      <div>
+        {' '}
+        <Tooltip
+          content={
+            'This feature is experimental and only works on log queries containing no more than 1 parser (logfmt, json).'
+          }
+          placement="top"
+        >
+          <Tag
+            className={css({
+              fontSize: 10,
+              padding: '1px 5px',
+              verticalAlign: 'text-bottom',
+            })}
+            name={'Experimental'}
+            colorIndex={1}
+          />
+        </Tooltip>{' '}
+        Select labels to include in the context query:
+      </div>
+      <div>
+        <MultiSelect
+          className={styles.multiSelect}
+          prefix="Labels"
+          options={realLabels}
+          value={realLabelsEnabled}
+          formatOptionLabel={formatOptionLabel}
+          closeMenuOnSelect={true}
+          maxMenuHeight={200}
+          menuShouldPortal={false}
+          noOptionsMessage="No further labels available"
+          onChange={(keys) => {
+            return setContextFilters(
+              contextFilters.map((filter) => {
+                if (filter.fromParser) {
+                  return filter;
+                }
+                filter.enabled = keys.some((key) => key.value === filter.value);
+                return filter;
+              })
+            );
+          }}
+        />
+      </div>
+      {parsedLabels.length > 0 && (
+        <div>
+          <MultiSelect
+            className={styles.multiSelect}
+            prefix="Parsed Labels"
+            options={parsedLabels}
+            value={parsedLabelsEnabled}
+            formatOptionLabel={formatOptionLabel}
+            closeMenuOnSelect={true}
+            menuShouldPortal={false}
+            maxMenuHeight={200}
+            noOptionsMessage="No further labels available"
+            isClearable={true}
+            onChange={(keys) => {
+              setContextFilters(
+                contextFilters.map((filter) => {
+                  if (!filter.fromParser) {
+                    return filter;
+                  }
+                  filter.enabled = keys.some((key) => key.value === filter.value);
+                  return filter;
+                })
+              );
+            }}
+          />
+        </div>
+      )}
+    </div>
+  );
+}

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

@@ -1168,3 +1168,57 @@ function makeAnnotationQueryRequest(options = {}): AnnotationQueryRequest<LokiQu
     rangeRaw: timeRange,
   };
 }
+
+describe('new context ui', () => {
+  it('returns expression with 1 label', async () => {
+    const ds = createLokiDatasource(templateSrvStub);
+
+    const row: LogRowModel = {
+      rowIndex: 0,
+      dataFrame: new MutableDataFrame({
+        fields: [
+          {
+            name: 'ts',
+            type: FieldType.time,
+            values: [0],
+          },
+        ],
+      }),
+      labels: { bar: 'baz', foo: 'uniqueParsedLabel' },
+      uid: '1',
+    } as unknown as LogRowModel;
+
+    jest.spyOn(ds.languageProvider, 'start').mockImplementation(() => Promise.resolve([]));
+    jest.spyOn(ds.languageProvider, 'getLabelKeys').mockImplementation(() => ['foo']);
+
+    const result = await ds.prepareContextExpr(row);
+
+    expect(result).toEqual('{foo="uniqueParsedLabel"}');
+  });
+
+  it('returns empty expression for parsed labels', async () => {
+    const ds = createLokiDatasource(templateSrvStub);
+
+    const row: LogRowModel = {
+      rowIndex: 0,
+      dataFrame: new MutableDataFrame({
+        fields: [
+          {
+            name: 'ts',
+            type: FieldType.time,
+            values: [0],
+          },
+        ],
+      }),
+      labels: { bar: 'baz', foo: 'uniqueParsedLabel' },
+      uid: '1',
+    } as unknown as LogRowModel;
+
+    jest.spyOn(ds.languageProvider, 'start').mockImplementation(() => Promise.resolve([]));
+    jest.spyOn(ds.languageProvider, 'getLabelKeys').mockImplementation(() => []);
+
+    const result = await ds.prepareContextExpr(row);
+
+    expect(result).toEqual('{}');
+  });
+});

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

@@ -35,6 +35,7 @@ import {
   toUtc,
 } from '@grafana/data';
 import { config, DataSourceWithBackend, FetchError } from '@grafana/runtime';
+import { DataQuery } from '@grafana/schema';
 import { queryLogsSample, queryLogsVolume } from 'app/core/logsModel';
 import { convertToWebSocketUrl } from 'app/core/utils/explore';
 import { getTimeSrv, TimeSrv } from 'app/features/dashboard/services/TimeSrv';
@@ -50,6 +51,7 @@ import LanguageProvider from './LanguageProvider';
 import { LiveStreams, LokiLiveTarget } from './LiveStreams';
 import { transformBackendResult } from './backendResultTransformer';
 import { LokiAnnotationsQueryEditor } from './components/AnnotationsQueryEditor';
+import { LokiContextUi } from './components/LokiContextUi';
 import { escapeLabelValueInExactSelector, escapeLabelValueInSelector, isRegexSelector } from './languageUtils';
 import { labelNamesRegex, labelValuesRegex } from './migrations/variableQueryMigrations';
 import {
@@ -66,11 +68,18 @@ import {
   getLabelFilterPositions,
 } from './modifyQuery';
 import { getQueryHints } from './queryHints';
-import { getLogQueryFromMetricsQuery, getNormalizedLokiQuery, isLogsQuery, isValidQuery } from './queryUtils';
+import {
+  getLogQueryFromMetricsQuery,
+  getNormalizedLokiQuery,
+  getParserFromQuery,
+  isLogsQuery,
+  isValidQuery,
+} from './queryUtils';
 import { sortDataFrameByTime } from './sortDataFrame';
 import { doLokiChannelStream } from './streaming';
 import { trackQuery } from './tracking';
 import {
+  ContextFilter,
   LokiOptions,
   LokiQuery,
   LokiQueryDirection,
@@ -594,10 +603,14 @@ export class LokiDatasource
     return Math.ceil(date.valueOf() * 1e6);
   }
 
-  getLogRowContext = async (row: LogRowModel, options?: RowContextOptions): Promise<{ data: DataFrame[] }> => {
+  getLogRowContext = async (
+    row: LogRowModel,
+    options?: RowContextOptions,
+    origQuery?: DataQuery
+  ): Promise<{ data: DataFrame[] }> => {
     const direction = (options && options.direction) || 'BACKWARD';
     const limit = (options && options.limit) || 10;
-    const { query, range } = await this.prepareLogRowContextQueryTarget(row, limit, direction);
+    const { query, range } = await this.prepareLogRowContextQueryTarget(row, limit, direction, origQuery);
 
     const processDataFrame = (frame: DataFrame): DataFrame => {
       // log-row-context requires specific field-names to work, so we set them here: "ts", "line", "id"
@@ -663,29 +676,17 @@ export class LokiDatasource
   prepareLogRowContextQueryTarget = async (
     row: LogRowModel,
     limit: number,
-    direction: 'BACKWARD' | 'FORWARD'
+    direction: 'BACKWARD' | 'FORWARD',
+    origQuery?: DataQuery
   ): Promise<{ query: LokiQuery; range: TimeRange }> => {
-    // need to await the languageProvider to be started to have all labels. This call is not blocking after it has been called once.
+    let expr = await this.prepareContextExpr(row, origQuery);
-    await this.languageProvider.start();
-    const labels = this.languageProvider.getLabelKeys();
-    const expr = Object.keys(row.labels)
-      .map((label: string) => {
-        if (labels.includes(label)) {
-          // escape backslashes in label as users can't escape them by themselves
-          return `${label}="${escapeLabelValueInExactSelector(row.labels[label])}"`;
-        }
-        return '';
-      })
-      // Filter empty strings
-      .filter((label) => !!label)
-      .join(',');
 
     const contextTimeBuffer = 2 * 60 * 60 * 1000; // 2h buffer
 
     const queryDirection = direction === 'FORWARD' ? LokiQueryDirection.Forward : LokiQueryDirection.Backward;
 
     const query: LokiQuery = {
-      expr: `{${expr}}`,
+      expr,
       queryType: LokiQueryType.Range,
       refId: `${REF_ID_STARTER_LOG_ROW_CONTEXT}${row.dataFrame.refId || ''}`,
       maxLines: limit,
@@ -726,6 +727,71 @@ export class LokiDatasource
     };
   };
 
+  async prepareContextExpr(row: LogRowModel, origQuery?: DataQuery): Promise<string> {
+    await this.languageProvider.start();
+    const labels = this.languageProvider.getLabelKeys();
+    const expr = Object.keys(row.labels)
+      .map((label: string) => {
+        if (labels.includes(label)) {
+          // escape backslashes in label as users can't escape them by themselves
+          return `${label}="${escapeLabelValueInExactSelector(row.labels[label])}"`;
+        }
+        return '';
+      })
+      .filter((label) => !!label)
+      .join(',');
+
+    return `{${expr}}`;
+  }
+
+  getLogRowContextUi(row: LogRowModel, runContextQuery: () => void): React.ReactNode {
+    return LokiContextUi({
+      row,
+      languageProvider: this.languageProvider,
+      updateFilter: (contextFilters: ContextFilter[]) => {
+        this.prepareContextExpr = async (row: LogRowModel, origQuery?: DataQuery) => {
+          await this.languageProvider.start();
+          const labels = this.languageProvider.getLabelKeys();
+
+          let expr = contextFilters
+            .map((filter) => {
+              const label = filter.value;
+              if (filter && !filter.fromParser && filter.enabled && labels.includes(label)) {
+                // escape backslashes in label as users can't escape them by themselves
+                return `${label}="${escapeLabelValueInExactSelector(row.labels[label])}"`;
+              }
+              return '';
+            })
+            // Filter empty strings
+            .filter((label) => !!label)
+            .join(',');
+
+          expr = `{${expr}}`;
+
+          const parserContextFilters = contextFilters.filter((filter) => filter.fromParser && filter.enabled);
+          if (parserContextFilters.length) {
+            // we should also filter for labels from parsers, let's find the right parser
+            if (origQuery) {
+              const parser = getParserFromQuery((origQuery as LokiQuery).expr);
+              if (parser) {
+                expr = addParserToQuery(expr, parser);
+              }
+            }
+            for (const filter of parserContextFilters) {
+              if (filter.enabled) {
+                expr = addLabelToQuery(expr, filter.label, '=', row.labels[filter.label]);
+              }
+            }
+          }
+          return expr;
+        };
+        if (runContextQuery) {
+          runContextQuery();
+        }
+      },
+    });
+  }
+
   testDatasource(): Promise<{ status: string; message: string }> {
     // Consider only last 10 minutes otherwise request takes too long
     const nowMs = Date.now();

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

@@ -153,3 +153,11 @@ export interface LokiVariableQuery extends DataQuery {
   label?: string;
   stream?: string;
 }
+
+export interface ContextFilter {
+  enabled: boolean;
+  label: string;
+  value: string;
+  fromParser: boolean;
+  description?: string;
+}