From f6404b9589839609d2c313de32e4495265ae2fa2 Mon Sep 17 00:00:00 2001
From: Alexa Vargas <239999+axelavargas@users.noreply.github.com>
Date: Wed, 20 Aug 2025 10:18:12 +0200
Subject: [PATCH] Query Library: Connect QueryLibraryEditingHeader in
 QueryEditorRow (#109818)

* Query Library: Connect QueryLibraryEditingHeader in QueryEditorRow

* Add unit test to queryn editor row

* Remove logic of "update query" save disk and add extra condition to prevent dragable action
---
 .../QueryLibrary/QueryLibraryContext.tsx      |  24 ++-
 .../features/explore/QueryLibrary/mocks.tsx   |   1 +
 .../query/components/QueryEditorRow.test.tsx  |  69 +++++---
 .../query/components/QueryEditorRow.tsx       | 161 +++++++++---------
 .../QueryLibraryEditingContainer.tsx          |  25 +++
 public/locales/en-US/grafana.json             |   7 +-
 6 files changed, 179 insertions(+), 108 deletions(-)
 create mode 100644 public/app/features/query/components/QueryLibraryEditingContainer.tsx

diff --git a/public/app/features/explore/QueryLibrary/QueryLibraryContext.tsx b/public/app/features/explore/QueryLibrary/QueryLibraryContext.tsx
index 1ac7ef7cd97..1930c56d81a 100644
--- a/public/app/features/explore/QueryLibrary/QueryLibraryContext.tsx
+++ b/public/app/features/explore/QueryLibrary/QueryLibraryContext.tsx
@@ -40,10 +40,28 @@ export type QueryLibraryContextType = {
   renderSaveQueryButton: (
     query: DataQuery,
     app?: CoreApp,
-    queryLibraryRef?: string,
     onUpdateSuccess?: () => void,
     onSelectQuery?: (query: DataQuery) => void
   ) => ReactNode;
+
+  /**
+   * Returns a header component for editing queries from the library.
+   * used in places like Explore
+   * @param query
+   * @param app
+   * @param queryLibraryRef
+   * @param onCancelEdit
+   * @param onUpdateSuccess
+   */
+  renderQueryLibraryEditingHeader: (
+    query: DataQuery,
+    app?: CoreApp,
+    queryLibraryRef?: string,
+    onCancelEdit?: () => void,
+    onUpdateSuccess?: () => void,
+    onSelectQuery?: (query: DataQuery) => void
+  ) => ReactNode;
+
   queryLibraryEnabled: boolean;
   context: string;
   triggerAnalyticsEvent: (
@@ -66,6 +84,10 @@ export const QueryLibraryContext = createContext<QueryLibraryContextType>({
     return null;
   },
 
+  renderQueryLibraryEditingHeader: () => {
+    return null;
+  },
+
   queryLibraryEnabled: false,
   context: 'unknown',
   triggerAnalyticsEvent: () => {},
diff --git a/public/app/features/explore/QueryLibrary/mocks.tsx b/public/app/features/explore/QueryLibrary/mocks.tsx
index 2c25a64e0bc..ecef7d95870 100644
--- a/public/app/features/explore/QueryLibrary/mocks.tsx
+++ b/public/app/features/explore/QueryLibrary/mocks.tsx
@@ -14,6 +14,7 @@ export function QueryLibraryContextProviderMock(props: PropsWithChildren<Props>)
         closeDrawer: jest.fn(),
         isDrawerOpen: false,
         renderSaveQueryButton: jest.fn(),
+        renderQueryLibraryEditingHeader: jest.fn(),
         queryLibraryEnabled: Boolean(props.queryLibraryEnabled),
         context: 'explore',
         triggerAnalyticsEvent: jest.fn(),
diff --git a/public/app/features/query/components/QueryEditorRow.test.tsx b/public/app/features/query/components/QueryEditorRow.test.tsx
index 8cc33cd4529..80e3ef12919 100644
--- a/public/app/features/query/components/QueryEditorRow.test.tsx
+++ b/public/app/features/query/components/QueryEditorRow.test.tsx
@@ -5,7 +5,7 @@ import { DataQueryRequest, dateTime, LoadingState, PanelData, toDataFrame } from
 import { DataQuery } from '@grafana/schema';
 import { mockDataSource } from 'app/features/alerting/unified/mocks';
 
-import { filterPanelDataToQuery, Props, QueryEditorRow, QueryLibraryEditingBadge } from './QueryEditorRow';
+import { filterPanelDataToQuery, Props, QueryEditorRow } from './QueryEditorRow';
 
 const mockDS = mockDataSource({
   name: 'test',
@@ -15,6 +15,12 @@ const mockDS = mockDataSource({
 // Mock the QueryLibraryContext
 const mockQueryLibraryContext = {
   queryLibraryEnabled: true,
+  renderQueryLibraryEditingHeader: jest.fn(),
+  renderSaveQueryButton: jest.fn(() => null),
+  openDrawer: jest.fn(),
+  closeDrawer: jest.fn(),
+  isDrawerOpen: false,
+  context: 'test',
 };
 
 jest.mock('app/features/explore/QueryLibrary/QueryLibraryContext', () => ({
@@ -404,31 +410,50 @@ describe('QueryEditorRow', () => {
       expect(screen.queryByText('Error!!')).not.toBeInTheDocument();
     });
   });
-});
 
-describe('QueryLibraryBadge', () => {
-  beforeEach(() => {
-    mockQueryLibraryContext.queryLibraryEnabled = true;
-  });
+  describe('Query Library Integration', () => {
+    let testData: PanelData;
+    let mockOnCancelEdit: jest.MockedFunction<() => void>;
 
-  it('should display badge when queryLibraryEnabled is true and queryLibraryRef is provided', () => {
-    render(<QueryLibraryEditingBadge queryLibraryRef="library-query-123" />);
-    expect(screen.getByText('Updating query from library')).toBeInTheDocument();
-  });
+    beforeEach(() => {
+      jest.clearAllMocks();
+      mockQueryLibraryContext.renderQueryLibraryEditingHeader.mockReturnValue(null);
+      mockOnCancelEdit = jest.fn();
 
-  it('should not display badge when queryLibraryEnabled is false', () => {
-    mockQueryLibraryContext.queryLibraryEnabled = false;
-    render(<QueryLibraryEditingBadge queryLibraryRef="library-query-123" />);
-    expect(screen.queryByText('Updating query from library')).not.toBeInTheDocument();
-  });
+      // Standard test data for QueryEditorRow
+      testData = {
+        series: [],
+        timeRange: { from: dateTime(), to: dateTime(), raw: { from: 'now-1d', to: 'now' } },
+        state: LoadingState.Done,
+      };
+    });
 
-  it('should not display badge when queryLibraryRef is not provided', () => {
-    render(<QueryLibraryEditingBadge />);
-    expect(screen.queryByText('Updating query from library')).not.toBeInTheDocument();
-  });
+    it('should render query library editing header when queryLibraryRef is provided', async () => {
+      render(
+        <QueryEditorRow {...props(testData)} queryLibraryRef="test-ref" onCancelQueryLibraryEdit={mockOnCancelEdit} />
+      );
 
-  it('should not display badge when queryLibraryRef is empty string', () => {
-    render(<QueryLibraryEditingBadge queryLibraryRef="" />);
-    expect(screen.queryByText('Updating query from library')).not.toBeInTheDocument();
+      // Wait for async datasource loading and component rendering
+      await waitFor(() => {
+        expect(mockQueryLibraryContext.renderQueryLibraryEditingHeader).toHaveBeenCalledWith(
+          expect.objectContaining({ refId: 'B' }),
+          undefined, // app
+          'test-ref', // queryLibraryRef
+          mockOnCancelEdit, // onCancelEdit
+          expect.any(Function), // onUpdateSuccess
+          expect.any(Function) // onSelectQuery
+        );
+      });
+    });
+
+    it('should not render query library editing header when queryLibraryRef is not provided', async () => {
+      render(<QueryEditorRow {...props(testData)} />);
+
+      await waitFor(() => {
+        expect(screen.getByTestId('query-editor-row')).toBeInTheDocument();
+      });
+
+      expect(mockQueryLibraryContext.renderQueryLibraryEditingHeader).not.toHaveBeenCalled();
+    });
   });
 });
diff --git a/public/app/features/query/components/QueryEditorRow.tsx b/public/app/features/query/components/QueryEditorRow.tsx
index a8bda63f862..d6e58fd8038 100644
--- a/public/app/features/query/components/QueryEditorRow.tsx
+++ b/public/app/features/query/components/QueryEditorRow.tsx
@@ -22,7 +22,7 @@ import { selectors } from '@grafana/e2e-selectors';
 import { Trans, t } from '@grafana/i18n';
 import { getDataSourceSrv, renderLimitedComponents, reportInteraction, usePluginComponents } from '@grafana/runtime';
 import { DataQuery } from '@grafana/schema';
-import { Badge, Divider, ErrorBoundaryAlert, List } from '@grafana/ui';
+import { Badge, ErrorBoundaryAlert, List } from '@grafana/ui';
 import { OperationRowHelp } from 'app/core/components/QueryOperationRow/OperationRowHelp';
 import {
   QueryOperationAction,
@@ -38,6 +38,7 @@ import { useQueryLibraryContext } from '../../explore/QueryLibrary/QueryLibraryC
 import { QueryActionComponent, RowActionComponents } from './QueryActionComponent';
 import { QueryEditorRowHeader } from './QueryEditorRowHeader';
 import { QueryErrorAlert } from './QueryErrorAlert';
+import { QueryLibraryEditingContainer } from './QueryLibraryEditingContainer';
 
 export interface Props<TQuery extends DataQuery> {
   data: PanelData;
@@ -346,11 +347,6 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
     );
   };
 
-  renderQueryLibraryEditingBadge = () => {
-    const { queryLibraryRef } = this.props;
-    return <QueryLibraryEditingBadge key="query-library-editing-badge" queryLibraryRef={queryLibraryRef} />;
-  };
-
   renderExtraActions = () => {
     const { query, queries, data, onAddQuery, dataSource, app } = this.props;
 
@@ -392,14 +388,14 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
 
     return (
       <>
-        {isEditingQueryLibrary && this.renderQueryLibraryEditingBadge()}
-        <MaybeQueryLibrarySaveButton
-          query={query}
-          queryLibraryRef={queryLibraryRef}
-          app={this.props.app}
-          onSelectQuery={this.onSelectQueryFromLibrary}
-          onUpdateSuccess={this.onExitQueryLibraryEditingMode}
-        />
+        {!isEditingQueryLibrary && (
+          <MaybeQueryLibrarySaveButton
+            query={query}
+            app={this.props.app}
+            onSelectQuery={this.onSelectQueryFromLibrary}
+            onUpdateSuccess={this.onExitQueryLibraryEditingMode}
+          />
+        )}
 
         {!isEditingQueryLibrary && (
           <ReplaceQueryFromLibrary
@@ -409,17 +405,6 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
           />
         )}
 
-        {isEditingQueryLibrary && (
-          <>
-            <QueryOperationAction
-              title={t('query-operation.header.cancel-query-library-edit', 'Discard changes')}
-              icon="times"
-              onClick={this.onCancelQueryLibraryEdit}
-            />
-            <Divider direction="vertical" spacing={0} />
-          </>
-        )}
-
         {hasEditorHelp && (
           <QueryOperationToggleAction
             title={t('query-operation.header.datasource-help', 'Show data source help')}
@@ -481,7 +466,18 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
   };
 
   render() {
-    const { query, index, visualization, collapsable, hideActionButtons, isOpen, onQueryOpenChanged } = this.props;
+    const {
+      query,
+      index,
+      visualization,
+      collapsable,
+      hideActionButtons,
+      isOpen,
+      onQueryOpenChanged,
+      app,
+      queryLibraryRef,
+      onCancelQueryLibraryEdit,
+    } = this.props;
     const { datasource, showingHelp, data } = this.state;
     const isHidden = query.hide;
     const error =
@@ -498,64 +494,58 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
     const editor = this.renderPluginEditor();
     const DatasourceCheatsheet = datasource.components?.QueryEditorHelp;
 
+    const queryOperationRow = (
+      <QueryOperationRow
+        id={this.id}
+        draggable={!hideActionButtons && !queryLibraryRef}
+        collapsable={collapsable}
+        index={index}
+        headerElement={this.renderHeader}
+        actions={hideActionButtons ? undefined : this.renderActions}
+        isOpen={isOpen}
+        onOpen={onQueryOpenChanged}
+      >
+        <div className={rowClasses} id={this.id}>
+          <ErrorBoundaryAlert>
+            {showingHelp && DatasourceCheatsheet && (
+              <OperationRowHelp>
+                <DatasourceCheatsheet
+                  onClickExample={(query) => this.onClickExample(query)}
+                  query={this.props.query}
+                  datasource={datasource}
+                />
+              </OperationRowHelp>
+            )}
+            {editor}
+          </ErrorBoundaryAlert>
+          {error && <QueryErrorAlert error={error} />}
+          {visualization}
+        </div>
+      </QueryOperationRow>
+    );
+
     return (
       <div data-testid="query-editor-row" aria-label={selectors.components.QueryEditorRows.rows}>
-        <QueryOperationRow
-          id={this.id}
-          draggable={!hideActionButtons}
-          collapsable={collapsable}
-          index={index}
-          headerElement={this.renderHeader}
-          actions={hideActionButtons ? undefined : this.renderActions}
-          isOpen={isOpen}
-          onOpen={onQueryOpenChanged}
-        >
-          <div className={rowClasses} id={this.id}>
-            <ErrorBoundaryAlert>
-              {showingHelp && DatasourceCheatsheet && (
-                <OperationRowHelp>
-                  <DatasourceCheatsheet
-                    onClickExample={(query) => this.onClickExample(query)}
-                    query={this.props.query}
-                    datasource={datasource}
-                  />
-                </OperationRowHelp>
-              )}
-              {editor}
-            </ErrorBoundaryAlert>
-            {error && <QueryErrorAlert error={error} />}
-            {visualization}
-          </div>
-        </QueryOperationRow>
+        {queryLibraryRef && (
+          <MaybeQueryLibraryEditingHeader
+            query={query}
+            app={app}
+            queryLibraryRef={queryLibraryRef}
+            onCancelEdit={onCancelQueryLibraryEdit}
+            onUpdateSuccess={this.onExitQueryLibraryEditingMode}
+            onSelectQuery={this.onSelectQueryFromLibrary}
+          />
+        )}
+        {queryLibraryRef ? (
+          <QueryLibraryEditingContainer>{queryOperationRow}</QueryLibraryEditingContainer>
+        ) : (
+          queryOperationRow
+        )}
       </div>
     );
   }
 }
 
-export function QueryLibraryEditingBadge(props: { queryLibraryRef?: string }) {
-  const { queryLibraryEnabled } = useQueryLibraryContext();
-  const { queryLibraryRef } = props;
-
-  if (!queryLibraryEnabled || !queryLibraryRef) {
-    return null;
-  }
-
-  return (
-    <Badge
-      color="blue"
-      icon="book"
-      text={t('query-operation.query-library.from-library', 'Updating query from library')}
-      tooltip={t(
-        'query-operation.query-library.editing-tooltip',
-        'Updating query from library\nUID: {{queryLibraryRef}}',
-        {
-          queryLibraryRef,
-        }
-      )}
-    />
-  );
-}
-
 /**
  * Get a version of the PanelData limited to the query we are looking at
  */
@@ -601,15 +591,28 @@ export function filterPanelDataToQuery(data: PanelData, refId: string): PanelDat
 function MaybeQueryLibrarySaveButton(props: {
   query: DataQuery;
   app?: CoreApp;
-  queryLibraryRef?: string;
   onUpdateSuccess?: () => void;
   onSelectQuery: (query: DataQuery) => void;
 }) {
   const { renderSaveQueryButton } = useQueryLibraryContext();
-  return renderSaveQueryButton(
+  return renderSaveQueryButton(props.query, props.app, props.onUpdateSuccess, props.onSelectQuery);
+}
+
+// Will render editing header only if query library is enabled
+function MaybeQueryLibraryEditingHeader(props: {
+  query: DataQuery;
+  app?: CoreApp;
+  queryLibraryRef?: string;
+  onCancelEdit?: () => void;
+  onUpdateSuccess?: () => void;
+  onSelectQuery?: (query: DataQuery) => void;
+}) {
+  const { renderQueryLibraryEditingHeader } = useQueryLibraryContext();
+  return renderQueryLibraryEditingHeader(
     props.query,
     props.app,
     props.queryLibraryRef,
+    props.onCancelEdit,
     props.onUpdateSuccess,
     props.onSelectQuery
   );
diff --git a/public/app/features/query/components/QueryLibraryEditingContainer.tsx b/public/app/features/query/components/QueryLibraryEditingContainer.tsx
new file mode 100644
index 00000000000..97e1c794c7f
--- /dev/null
+++ b/public/app/features/query/components/QueryLibraryEditingContainer.tsx
@@ -0,0 +1,25 @@
+import { css } from '@emotion/css';
+import { ReactNode } from 'react';
+
+import { GrafanaTheme2 } from '@grafana/data';
+import { useStyles2 } from '@grafana/ui';
+
+interface QueryLibraryEditingContainerProps {
+  children: ReactNode;
+}
+
+export function QueryLibraryEditingContainer({ children }: QueryLibraryEditingContainerProps) {
+  const styles = useStyles2(getStyles);
+  return <div className={styles.container}>{children}</div>;
+}
+
+const getStyles = (theme: GrafanaTheme2) => ({
+  container: css({
+    border: `2px solid ${theme.colors.primary.main}`,
+    borderTopLeftRadius: 'unset',
+    borderTopRightRadius: 'unset',
+    borderBottomLeftRadius: theme.shape.radius.default,
+    borderBottomRightRadius: theme.shape.radius.default,
+    overflow: 'hidden',
+  }),
+});
diff --git a/public/locales/en-US/grafana.json b/public/locales/en-US/grafana.json
index 7423a3f2437..e8ad31317dc 100644
--- a/public/locales/en-US/grafana.json
+++ b/public/locales/en-US/grafana.json
@@ -11886,7 +11886,6 @@
   },
   "query-operation": {
     "header": {
-      "cancel-query-library-edit": "Discard changes",
       "collapse-row": "Collapse query row",
       "datasource-help": "Show data source help",
       "drag-and-drop": "Drag and drop to reorder",
@@ -11897,11 +11896,7 @@
       "replace-query-from-library": "Replace with query from library",
       "show-response": "Show response"
     },
-    "query-editor-not-exported": "Data source plugin does not export any Query Editor component",
-    "query-library": {
-      "editing-tooltip": "Updating query from library\nUID: {{queryLibraryRef}}",
-      "from-library": "Updating query from library"
-    }
+    "query-editor-not-exported": "Data source plugin does not export any Query Editor component"
   },
   "recently-deleted": {
     "buttons": {