QueryLibrary: Add From Library Icon Action (#103317)

* QueryLibrary: AddfromIcon init work

* Removed console log

* updated i18n

* Added unit tests for button

* Fixed linting

* Updates per feedback - new replace method

* Remove run query call when replace

* lint fixes

* Feedback + method renaming to 'replace'

* Fix await

public/app/features/explore/QueriesDrawer/mocks.tsx

@@ -4,7 +4,7 @@ import { QueriesDrawerContext, Tabs } from './QueriesDrawerContext';
 
 type Props = {
   setDrawerOpened?: (value: boolean) => {};
-  queryLibraryAvailable?: boolean;
+  queryLibraryEnabled?: boolean;
 } & PropsWithChildren;
 
 export function QueriesDrawerContextProviderMock(props: Props) {

public/app/features/explore/QueryLibrary/mocks.tsx

@@ -3,7 +3,7 @@ import { PropsWithChildren } from 'react';
 import { QueryLibraryContext } from './QueryLibraryContext';
 
 type Props = {
-  queryLibraryAvailable?: boolean;
+  queryLibraryEnabled?: boolean;
 };
 
 export function QueryLibraryContextProviderMock(props: PropsWithChildren<Props>) {
@@ -16,7 +16,7 @@ export function QueryLibraryContextProviderMock(props: PropsWithChildren<Props>)
         openAddQueryModal: jest.fn(),
         closeAddQueryModal: jest.fn(),
         renderSaveQueryButton: jest.fn(),
-        queryLibraryEnabled: false,
+        queryLibraryEnabled: Boolean(props.queryLibraryEnabled),
       }}
     >
       {props.children}

public/app/features/explore/QueryRows.test.tsx

@@ -9,6 +9,7 @@ import { ExploreState } from 'app/types';
 
 import { UserState } from '../profile/state/reducers';
 
+import { QueryLibraryContextProviderMock } from './QueryLibrary/mocks';
 import { QueryRows } from './QueryRows';
 import { makeExplorePaneState } from './state/utils';
 
@@ -92,4 +93,37 @@ describe('Explore QueryRows', () => {
     // We should have another row with refId B
     expect(await screen.findByLabelText('Query editor row title B')).toBeInTheDocument();
   });
+
+  it('Should contain a select query from library button when query library is enabled', async () => {
+    const { store } = setup([{ refId: 'A' }]);
+
+    render(
+      <Provider store={store}>
+        <QueryLibraryContextProviderMock queryLibraryEnabled={true}>
+          <QueryRows exploreId={'left'} />
+        </QueryLibraryContextProviderMock>
+      </Provider>
+    );
+
+    // waiting for the component to fully render.
+    await screen.findAllByText('someDs query editor');
+
+    expect(screen.getByLabelText(/Replace with query from library/i)).toBeInTheDocument();
+  });
+
+  it('Should not contain a select query from library button when query library is disabled', async () => {
+    const { store } = setup([{ refId: 'A' }]);
+
+    render(
+      <Provider store={store}>
+        <QueryLibraryContextProviderMock queryLibraryEnabled={false}>
+          <QueryRows exploreId={'left'} />
+        </QueryLibraryContextProviderMock>
+      </Provider>
+    );
+
+    await screen.findAllByText('someDs query editor');
+
+    expect(screen.queryByLabelText(/Replace with query from library/i)).not.toBeInTheDocument();
+  });
 });

public/app/features/explore/QueryRows.tsx

@@ -3,13 +3,14 @@ import { useCallback, useMemo } from 'react';
 
 import { CoreApp, getNextRefId } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
-import { DataQuery } from '@grafana/schema';
+import { DataQuery, DataSourceRef } from '@grafana/schema';
 import { useDispatch, useSelector } from 'app/types';
 
 import { getDatasourceSrv } from '../plugins/datasource_srv';
 import { QueryEditorRows } from '../query/components/QueryEditorRows';
 
 import { ContentOutlineItem } from './ContentOutline/ContentOutlineItem';
+import { changeDatasource } from './state/datasource';
 import { changeQueries, runQueries } from './state/query';
 import { getExploreItemSelector } from './state/selectors';
 
@@ -55,6 +56,13 @@ export const QueryRows = ({ exploreId }: Props) => {
     [dispatch, exploreId]
   );
 
+  const onUpdateDatasources = useCallback(
+    (datasource: DataSourceRef) => {
+      dispatch(changeDatasource({ exploreId, datasource }));
+    },
+    [dispatch, exploreId]
+  );
+
   const onAddQuery = useCallback(
     (query: DataQuery) => {
       onChange([...queries, { ...query, refId: getNextRefId(queries) }]);
@@ -66,6 +74,10 @@ export const QueryRows = ({ exploreId }: Props) => {
     reportInteraction('grafana_explore_query_row_copy');
   };
 
+  const onQueryReplacedFromLibrary = () => {
+    reportInteraction('grafana_explore_query_replaced_from_library');
+  };
+
   const onQueryRemoved = () => {
     reportInteraction('grafana_explore_query_row_remove');
   };
@@ -79,11 +91,13 @@ export const QueryRows = ({ exploreId }: Props) => {
       dsSettings={dsSettings}
       queries={queries}
       onQueriesChange={onChange}
+      onUpdateDatasources={onUpdateDatasources}
       onAddQuery={onAddQuery}
       onRunQueries={onRunQueries}
       onQueryCopied={onQueryCopied}
       onQueryRemoved={onQueryRemoved}
       onQueryToggled={onQueryToggled}
+      onQueryReplacedFromLibrary={onQueryReplacedFromLibrary}
       data={queryResponse}
       app={CoreApp.Explore}
       history={history}

public/app/features/explore/SecondaryActions.test.tsx

@@ -35,7 +35,7 @@ describe('SecondaryActions', () => {
 
   it('should not render hidden elements', () => {
     render(
-      <QueriesDrawerContextProviderMock queryLibraryAvailable={false}>
+      <QueriesDrawerContextProviderMock queryLibraryEnabled={false}>
         <SecondaryActions
           addQueryRowButtonHidden={true}
           richHistoryRowButtonHidden={true}

public/app/features/query/components/QueryEditorRow.test.tsx

@@ -354,6 +354,7 @@ describe('QueryEditorRow', () => {
     onRunQuery: jest.fn(),
     onChange: jest.fn(),
     onRemoveQuery: jest.fn(),
+    onReplace: jest.fn(),
     index: 0,
     range: { from: dateTime(), to: dateTime(), raw: { from: 'now-1d', to: 'now' } },
   });

public/app/features/query/components/QueryEditorRow.tsx

@@ -52,6 +52,7 @@ export interface Props<TQuery extends DataQuery> {
   onAddQuery: (query: TQuery) => void;
   onRemoveQuery: (query: TQuery) => void;
   onChange: (query: TQuery) => void;
+  onReplace?: (query: DataQuery) => void;
   onRunQuery: () => void;
   visualization?: ReactNode;
   hideHideQueryButton?: boolean;
@@ -63,6 +64,7 @@ export interface Props<TQuery extends DataQuery> {
   onQueryCopied?: () => void;
   onQueryRemoved?: () => void;
   onQueryToggled?: (queryStatus?: boolean | undefined) => void;
+  onQueryReplacedFromLibrary?: () => void;
   collapsable?: boolean;
   hideRefId?: boolean;
 }
@@ -354,7 +356,12 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
   };
 
   renderActions = (props: QueryOperationRowRenderProps) => {
-    const { query, hideHideQueryButton: hideHideQueryButton = false } = this.props;
+    const {
+      query,
+      hideHideQueryButton: hideHideQueryButton = false,
+      onReplace,
+      onQueryReplacedFromLibrary,
+    } = this.props;
     const { datasource, showingHelp } = this.state;
     const isHidden = !!query.hide;
 
@@ -377,6 +384,13 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
           icon="copy"
           onClick={this.onCopyQuery}
         />
+        <ReplaceQueryFromLibrary
+          datasourceFilters={datasource?.name ? [datasource.name] : []}
+          onSelectQuery={(query) => {
+            onQueryReplacedFromLibrary?.();
+            onReplace?.(query);
+          }}
+        />
         {!hideHideQueryButton ? (
           <QueryOperationToggleAction
             dataTestId={selectors.components.QueryEditorRow.actionButton('Hide response')}
@@ -515,6 +529,30 @@ function MaybeQueryLibrarySaveButton(props: { query: DataQuery }) {
   return renderSaveQueryButton(props.query);
 }
 
+interface ReplaceQueryFromLibraryProps<TQuery extends DataQuery> {
+  datasourceFilters: string[];
+  onSelectQuery: (query: DataQuery) => void;
+}
+
+function ReplaceQueryFromLibrary<TQuery extends DataQuery>({
+  datasourceFilters,
+  onSelectQuery,
+}: ReplaceQueryFromLibraryProps<TQuery>) {
+  const { openDrawer, queryLibraryEnabled } = useQueryLibraryContext();
+
+  const onReplaceQueryFromLibrary = () => {
+    openDrawer(datasourceFilters, onSelectQuery);
+  };
+
+  return queryLibraryEnabled ? (
+    <QueryOperationAction
+      title={t('query-operation.header.replace-query-from-library', 'Replace with query from library')}
+      icon="book"
+      onClick={onReplaceQueryFromLibrary}
+    />
+  ) : null;
+}
+
 function AdaptiveTelemetryQueryActions({ query }: { query: DataQuery }) {
   try {
     const { isLoading, components } = usePluginComponents<PluginExtensionQueryEditorRowAdaptiveTelemetryV1Context>({

public/app/features/query/components/QueryEditorRows.test.tsx

@@ -1,6 +1,6 @@
 import { fireEvent, queryByLabelText, render, screen } from '@testing-library/react';
 
-import { type DataQuery } from '@grafana/schema';
+import { DataSourceRef, type DataQuery } from '@grafana/schema';
 import { mockDataSource } from 'app/features/alerting/unified/mocks';
 import { DataSourceType } from 'app/features/alerting/unified/utils/datasource';
 import createMockPanelData from 'app/plugins/datasource/azuremonitor/__mocks__/panelData';
@@ -54,6 +54,9 @@ const props: Props = {
   onRunQueries: function (): void {
     throw new Error('Function not implemented.');
   },
+  onUpdateDatasources: function (datasource: DataSourceRef): void {
+    throw new Error('Function not implemented.');
+  },
   data: createMockPanelData(),
 };
 

public/app/features/query/components/QueryEditorRows.tsx

@@ -11,7 +11,9 @@ import {
   getDataSourceRef,
 } from '@grafana/data';
 import { getDataSourceSrv, reportInteraction } from '@grafana/runtime';
+import { DataSourceRef } from '@grafana/schema';
 import { getTimeSrv } from 'app/features/dashboard/services/TimeSrv';
+import { MIXED_DATASOURCE_NAME } from 'app/plugins/datasource/mixed/MixedDataSource';
 
 import { QueryEditorRow } from './QueryEditorRow';
 
@@ -35,6 +37,8 @@ export interface Props {
   onQueryCopied?: () => void;
   onQueryRemoved?: () => void;
   onQueryToggled?: (queryStatus?: boolean | undefined) => void;
+  onUpdateDatasources?: (datasource: DataSourceRef) => void;
+  onQueryReplacedFromLibrary?: () => void;
   queryRowWrapper?: (children: ReactNode, refId: string) => ReactNode;
 }
 
@@ -57,6 +61,32 @@ export class QueryEditorRows extends PureComponent<Props> {
     );
   }
 
+  onReplaceQuery(query: DataQuery, index: number) {
+    const { queries, onQueriesChange, onUpdateDatasources, dsSettings } = this.props;
+
+    // Replace old query with new query
+    const newQueries = queries.map((item, itemIndex) => {
+      if (itemIndex === index) {
+        return query;
+      }
+      return item;
+    });
+    onQueriesChange(newQueries);
+
+    // Update datasources based on the new query set
+    if (query.datasource?.uid) {
+      const uniqueDatasources = new Set(newQueries.map((q) => q.datasource?.uid));
+      const isMixed = uniqueDatasources.size > 1;
+      const newDatasourceRef = {
+        uid: isMixed ? MIXED_DATASOURCE_NAME : query.datasource.uid,
+      };
+      const shouldChangeDatasource = dsSettings.uid !== newDatasourceRef.uid;
+      if (shouldChangeDatasource) {
+        onUpdateDatasources?.(newDatasourceRef);
+      }
+    }
+  }
+
   onDataSourceChange(dataSource: DataSourceInstanceSettings, index: number) {
     const { queries, onQueriesChange } = this.props;
 
@@ -143,6 +173,7 @@ export class QueryEditorRows extends PureComponent<Props> {
       onQueryCopied,
       onQueryRemoved,
       onQueryToggled,
+      onQueryReplacedFromLibrary,
       queryRowWrapper,
     } = this.props;
 
@@ -168,12 +199,14 @@ export class QueryEditorRows extends PureComponent<Props> {
                       dataSource={dataSourceSettings}
                       onChangeDataSource={onChangeDataSourceSettings}
                       onChange={(query) => this.onChangeQuery(query, index)}
+                      onReplace={(query) => this.onReplaceQuery(query, index)}
                       onRemoveQuery={this.onRemoveQuery}
                       onAddQuery={onAddQuery}
                       onRunQuery={onRunQueries}
                       onQueryCopied={onQueryCopied}
                       onQueryRemoved={onQueryRemoved}
                       onQueryToggled={onQueryToggled}
+                      onQueryReplacedFromLibrary={onQueryReplacedFromLibrary}
                       queries={queries}
                       app={app}
                       range={getTimeSrv().timeRange()}

public/locales/en-US/grafana.json

@@ -6864,6 +6864,7 @@
       "expand-row": "Expand query row",
       "hide-response": "Hide response",
       "remove-query": "Remove query",
+      "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"