fix(explore): Metrics disappearing after removing metric from dataset (#17201)

* fix(explore): Metrics disappearing after removing metric from dataset

* fix test

* Apply fix to non-dnd controls

* Make adhoc metrics pick up changes from dataset columns

* Remove console log

* Fix bug in nondnd controls

superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndMetricSelect.tsx

@@ -19,7 +19,6 @@
 
 import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import {
-  DatasourceType,
   ensureIsArray,
   FeatureFlag,
   GenericDataType,
@@ -42,7 +41,6 @@ import { DndItemType } from 'src/explore/components/DndItemType';
 import DndSelectLabel from 'src/explore/components/controls/DndColumnSelectControl/DndSelectLabel';
 import { savedMetricType } from 'src/explore/components/controls/MetricControl/types';
 import { AGGREGATES } from 'src/explore/constants';
-import { DndControlProps } from './types';
 
 const EMPTY_OBJECT = {};
 const DND_ACCEPTED_TYPES = [DndItemType.Column, DndItemType.Metric];
@@ -82,39 +80,48 @@ const getOptionsForSavedMetrics = (
 
 type ValueType = Metric | AdhocMetric | QueryFormMetric;
 
-const columnsContainAllMetrics = (
-  value: ValueType | ValueType[] | null | undefined,
+// TODO: use typeguards to distinguish saved metrics from adhoc metrics
+const getMetricsMatchingCurrentDataset = (
+  values: ValueType[],
   columns: ColumnMeta[],
   savedMetrics: (savedMetricType | Metric)[],
+  prevColumns: ColumnMeta[],
+  prevSavedMetrics: (savedMetricType | Metric)[],
 ) => {
-  const columnNames = new Set(
-    [...(columns || []), ...(savedMetrics || [])]
-      // eslint-disable-next-line camelcase
-      .map(
-        item =>
-          (item as ColumnMeta).column_name ||
-          (item as savedMetricType).metric_name,
-      ),
-  );
+  const areSavedMetricsEqual =
+    !prevSavedMetrics || isEqual(prevSavedMetrics, savedMetrics);
+  const areColsEqual = !prevColumns || isEqual(prevColumns, columns);
 
-  return (
-    ensureIsArray(value)
-      .filter(metric => metric)
-      // find column names
-      .map(metric =>
-        (metric as AdhocMetric).column
-          ? (metric as AdhocMetric).column.column_name
-          : (metric as ColumnMeta).column_name || metric,
-      )
-      .filter(name => name && typeof name === 'string')
-      .every(name => columnNames.has(name))
-  );
-};
+  if (areColsEqual && areSavedMetricsEqual) {
+    return values;
+  }
+  return values.reduce((acc: ValueType[], metric) => {
+    if (
+      (typeof metric === 'string' || (metric as Metric).metric_name) &&
+      (areSavedMetricsEqual ||
+        savedMetrics?.some(
+          savedMetric =>
+            savedMetric.metric_name === metric ||
+            savedMetric.metric_name === (metric as Metric).metric_name,
+        ))
+    ) {
+      acc.push(metric);
+      return acc;
+    }
 
-export type DndMetricSelectProps = DndControlProps<ValueType> & {
-  savedMetrics: savedMetricType[];
-  columns: ColumnMeta[];
-  datasourceType?: DatasourceType;
+    if (!areColsEqual) {
+      const newCol = columns?.find(
+        column =>
+          (metric as AdhocMetric).column?.column_name === column.column_name,
+      );
+      if (newCol) {
+        acc.push({ ...(metric as AdhocMetric), column: newCol });
+      }
+    } else {
+      acc.push(metric);
+    }
+    return acc;
+  }, []);
 };
 
 export const DndMetricSelect = (props: any) => {
@@ -158,25 +165,25 @@ export const DndMetricSelect = (props: any) => {
   }, [JSON.stringify(props.value)]);
 
   useEffect(() => {
-    if (
-      !isEqual(prevColumns, columns) ||
-      !isEqual(prevSavedMetrics, savedMetrics)
-    ) {
-      // Remove all metrics if selected value no longer a valid column
-      // in the dataset. Must use `nextProps` here because Redux reducers may
-      // have already updated the value for this control.
-      if (!columnsContainAllMetrics(props.value, columns, savedMetrics)) {
-        onChange([]);
-      }
+    // Remove selected custom metrics that do not exist in the dataset anymore
+    // Remove selected adhoc metrics that use columns which do not exist in the dataset anymore
+    // Sync adhoc metrics with dataset columns when they are modified by the user
+    if (!props.value) {
+      return;
     }
-  }, [
-    prevColumns,
-    columns,
-    prevSavedMetrics,
-    savedMetrics,
-    props.value,
-    onChange,
-  ]);
+    const propsValues = ensureIsArray(props.value);
+    const matchingMetrics = getMetricsMatchingCurrentDataset(
+      propsValues,
+      columns,
+      savedMetrics,
+      prevColumns,
+      prevSavedMetrics,
+    );
+
+    if (!isEqual(propsValues, matchingMetrics)) {
+      handleChange(matchingMetrics);
+    }
+  }, [columns, savedMetrics, handleChange]);
 
   const canDrop = useCallback(
     (item: DatasourcePanelDndItem) => {
@@ -337,7 +344,7 @@ export const DndMetricSelect = (props: any) => {
     ) {
       const itemValue = droppedItem.value as ColumnMeta;
       const config: Partial<AdhocMetric> = {
-        column: { column_name: itemValue?.column_name },
+        column: itemValue,
       };
       if (isFeatureEnabled(FeatureFlag.UX_BETA)) {
         if (itemValue.type_generic === GenericDataType.NUMERIC) {