Loki: Add functionality to revert to initial query in log context

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

@@ -206,4 +206,33 @@ describe('LokiContextUi', () => {
       expect(screen.queryByText('Refine the search')).not.toBeInTheDocument();
     });
   });
+
+  it('should revert to original query when revert button clicked', async () => {
+    const props = setupProps();
+    const newProps = {
+      ...props,
+      origQuery: {
+        expr: '{label1="value1"} | logfmt',
+        refId: 'A',
+      },
+    };
+    render(<LokiContextUi {...newProps} />);
+    // In initial query, label3 is not selected
+    await waitFor(() => {
+      expect(screen.queryByText('label3="value3"')).not.toBeInTheDocument();
+    });
+
+    // We select parsed label and label3="value3" should appear
+    const parsedLabelsInput = screen.getAllByRole('combobox')[1];
+    await userEvent.click(parsedLabelsInput);
+    await userEvent.type(parsedLabelsInput, '{enter}');
+    expect(screen.getByText('label3="value3"')).toBeInTheDocument();
+
+    // We click on revert button and label3="value3" should disappear
+    const revertButton = screen.getByTestId('revert-button');
+    await userEvent.click(revertButton);
+    await waitFor(() => {
+      expect(screen.queryByText('label3="value3"')).not.toBeInTheDocument();
+    });
+  });
 });

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

@@ -1,10 +1,10 @@
 import { css } from '@emotion/css';
-import React, { useCallback, useEffect, useState } from 'react';
+import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import { useAsync } from 'react-use';
 
 import { GrafanaTheme2, LogRowModel, SelectableValue } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
-import { Collapse, Icon, Label, MultiSelect, Tooltip, useStyles2 } from '@grafana/ui';
+import { Button, Collapse, Icon, Label, MultiSelect, Spinner, Tooltip, useStyles2 } from '@grafana/ui';
 import store from 'app/core/store';
 
 import { RawQuery } from '../../prometheus/querybuilder/shared/RawQuery';
@@ -33,6 +33,7 @@ function getStyles(theme: GrafanaTheme2) {
       flex-direction: column;
       flex: 1;
       gap: ${theme.spacing(0.5)};
+      position: relative;
     `,
     textWrapper: css`
       display: flex;
@@ -50,10 +51,11 @@ function getStyles(theme: GrafanaTheme2) {
         margin: ${theme.spacing(2)} 0;
       }
     `,
-    query: css`
+    rawQueryContainer: css`
       text-align: start;
       line-break: anywhere;
       margin-top: -${theme.spacing(0.25)};
+      width: calc(100% - 20px);
     `,
     ui: css`
       background-color: ${theme.colors.background.secondary};
@@ -65,6 +67,12 @@ function getStyles(theme: GrafanaTheme2) {
     queryDescription: css`
       margin-left: ${theme.spacing(0.5)};
     `,
+    iconButton: css`
+      position: absolute;
+      top: ${theme.spacing(1)};
+      right: ${theme.spacing(1)};
+      z-index: ${theme.zIndex.navbarFixed};
+    `,
   };
 }
 
@@ -83,6 +91,16 @@ export function LokiContextUi(props: LokiContextUiProps) {
   const timerHandle = React.useRef<number>();
   const previousInitialized = React.useRef<boolean>(false);
   const previousContextFilters = React.useRef<ContextFilter[]>([]);
+
+  const isInitialQuery = useMemo(() => {
+    // Initial query has all regular labels enabled and all parsed labels disabled
+    if (initialized && contextFilters.some((filter) => filter.fromParser === filter.enabled)) {
+      return false;
+    }
+
+    return true;
+  }, [contextFilters, initialized]);
+
   useEffect(() => {
     if (!initialized) {
       return;
@@ -163,6 +181,29 @@ export function LokiContextUi(props: LokiContextUiProps) {
 
   return (
     <div className={styles.wrapper}>
+      <Tooltip content={'Revert to initial log context query.'}>
+        <div className={styles.iconButton}>
+          <Button
+            data-testid="revert-button"
+            icon="history-alt"
+            variant="secondary"
+            disabled={isInitialQuery}
+            onClick={(e) => {
+              reportInteraction('grafana_explore_logs_loki_log_context_reverted', {
+                logRowUid: row.uid,
+              });
+              setContextFilters((contextFilters) => {
+                return contextFilters.map((contextFilter) => ({
+                  ...contextFilter,
+                  // For revert to initial query we need to enable all labels and disable all parsed labels
+                  enabled: !contextFilter.fromParser,
+                }));
+              });
+            }}
+          />
+        </div>
+      </Tooltip>
+
       <Collapse
         collapsible={true}
         isOpen={isOpen}
@@ -175,19 +216,25 @@ export function LokiContextUi(props: LokiContextUiProps) {
           });
         }}
         label={
-          <div className={styles.query}>
-            <RawQuery
-              lang={{ grammar: lokiGrammar, name: 'loki' }}
-              query={logContextProvider.processContextFiltersToExpr(
-                row,
-                contextFilters.filter(({ enabled }) => enabled),
-                origQuery
-              )}
-              className={styles.rawQuery}
-            />
-            <Tooltip content="The initial log context query is created from all labels defining the stream for the selected log line. Use the editor below to customize the log context query.">
-              <Icon name="info-circle" size="sm" className={styles.queryDescription} />
-            </Tooltip>
+          <div className={styles.rawQueryContainer}>
+            {initialized ? (
+              <>
+                <RawQuery
+                  lang={{ grammar: lokiGrammar, name: 'loki' }}
+                  query={logContextProvider.processContextFiltersToExpr(
+                    row,
+                    contextFilters.filter(({ enabled }) => enabled),
+                    origQuery
+                  )}
+                  className={styles.rawQuery}
+                />
+                <Tooltip content="The initial log context query is created from all labels defining the stream for the selected log line. Use the editor below to customize the log context query.">
+                  <Icon name="info-circle" size="sm" className={styles.queryDescription} />
+                </Tooltip>
+              </>
+            ) : (
+              <Spinner />
+            )}
           </div>
         }
       >