Selected dashboards and collection for a metric

web/src/components/features/buttons/SaveMetricToCollectionButton.tsx

@@ -10,16 +10,23 @@ import {
 
 export const SaveMetricToCollectionButton: React.FC<{
   metricIds: string[];
+  selectedCollections: string[];
   buttonType?: 'ghost' | 'default';
   useText?: boolean;
-}> = ({ metricIds, buttonType = 'ghost', useText = false }) => {
+}> = ({
+  metricIds,
+  selectedCollections: selectedCollectionsProp,
+  buttonType = 'ghost',
+  useText = false
+}) => {
   const { openInfoMessage } = useBusterNotifications();
   const { mutateAsync: saveMetricToCollection } = useSaveMetricToCollections();
   const { mutateAsync: removeMetricFromCollection } = useRemoveMetricFromCollection();
 
-  const [selectedCollections, setSelectedCollections] = useState<
-    Parameters<typeof SaveToCollectionsDropdown>[0]['selectedCollections']
-  >([]);
+  const [selectedCollections, setSelectedCollections] =
+    useState<Parameters<typeof SaveToCollectionsDropdown>[0]['selectedCollections']>(
+      selectedCollectionsProp
+    );
 
   const onSaveToCollection = useMemoizedFn(async (collectionIds: string[]) => {
     setSelectedCollections(collectionIds);
@@ -31,6 +38,7 @@ export const SaveMetricToCollectionButton: React.FC<{
   });
 
   const onRemoveFromCollection = useMemoizedFn(async (collectionId: string) => {
+    setSelectedCollections((prev) => prev.filter((x) => x !== collectionId));
     await removeMetricFromCollection({
       metricIds,
       collectionIds: [collectionId]

web/src/components/features/buttons/SaveMetricToDashboardButton.tsx

@@ -1,5 +1,5 @@
 import { useMemoizedFn } from '@/hooks';
-import React from 'react';
+import React, { useState } from 'react';
 import { SaveToDashboardDropdown } from '../dropdowns/SaveToDashboardDropdown';
 import { Button } from '@/components/ui/buttons';
 import { ASSET_ICONS } from '../config/assetIcons';
@@ -8,25 +8,43 @@ import {
   useAddMetricsToDashboard
 } from '@/api/buster_rest/dashboards';
 import { AppTooltip } from '@/components/ui/tooltip';
+import { useBusterNotifications } from '@/context/BusterNotifications';
 
 export const SaveMetricToDashboardButton: React.FC<{
   metricIds: string[];
   disabled?: boolean;
   selectedDashboards: string[];
-}> = React.memo(({ metricIds, disabled = false, selectedDashboards }) => {
+}> = React.memo(({ metricIds, disabled = false, selectedDashboards: selectedDashboardsProp }) => {
   const { mutateAsync: saveMetricsToDashboard } = useAddMetricsToDashboard();
   const { mutateAsync: removeMetricsFromDashboard } = useRemoveMetricsFromDashboard();
+  const { openConfirmModal } = useBusterNotifications();
+
+  const [selectedDashboards, setSelectedDashboards] =
+    useState<Parameters<typeof SaveToDashboardDropdown>[0]['selectedDashboards']>(
+      selectedDashboardsProp
+    );
 
   const onSaveToDashboard = useMemoizedFn(async (dashboardIds: string[]) => {
+    setSelectedDashboards((prev) => [...prev, ...dashboardIds]);
     await Promise.all(
       dashboardIds.map((dashboardId) => saveMetricsToDashboard({ metricIds, dashboardId }))
     );
   });
 
   const onRemoveFromDashboard = useMemoizedFn(async (dashboardIds: string[]) => {
+    const method = async () => {
+      setSelectedDashboards((prev) => prev.filter((x) => !dashboardIds.includes(x)));
       await Promise.all(
-      dashboardIds.map((dashboardId) => removeMetricsFromDashboard({ metricIds, dashboardId }))
+        dashboardIds.map((dashboardId) =>
+          removeMetricsFromDashboard({ useConfirmModal: false, metricIds, dashboardId })
+        )
       );
+    };
+    return await openConfirmModal({
+      title: 'Remove from dashboard',
+      content: 'Are you sure you want to remove this from the dashboard?',
+      onOk: method
+    });
   });
 
   return (

web/src/layouts/ChatLayout/FileContainer/FileContainerHeader/MetricContainerHeaderButtons/MetricContainerHeaderButtons.tsx

@@ -153,7 +153,14 @@ const EditSQLButton = React.memo(({ metricId }: { metricId: string }) => {
 EditSQLButton.displayName = 'EditSQLButton';
 
 const SaveToCollectionButton = React.memo(({ metricId }: { metricId: string }) => {
-  return <SaveMetricToCollectionButton metricIds={[metricId]} />;
+  const { data: collections } = useGetMetric(
+    { id: metricId },
+    { select: (x) => x.collections?.map((x) => x.id) }
+  );
+
+  return (
+    <SaveMetricToCollectionButton metricIds={[metricId]} selectedCollections={collections || []} />
+  );
 });
 SaveToCollectionButton.displayName = 'SaveToCollectionButton';