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.test.tsx

@@ -19,14 +19,44 @@
 import React from 'react';
 import { render, screen } from 'spec/helpers/testing-library';
 import { DndMetricSelect } from 'src/explore/components/controls/DndColumnSelectControl/DndMetricSelect';
+import { AGGREGATES } from 'src/explore/constants';
+import { EXPRESSION_TYPES } from '../MetricControl/AdhocMetric';
 
 const defaultProps = {
   savedMetrics: [
     {
-      metric_name: 'Metric A',
-      expression: 'Expression A',
+      metric_name: 'metric_a',
+      expression: 'expression_a',
+    },
+    {
+      metric_name: 'metric_b',
+      expression: 'expression_b',
+      verbose_name: 'Metric B',
     },
   ],
+  columns: [
+    {
+      column_name: 'column_a',
+    },
+    {
+      column_name: 'column_b',
+      verbose_name: 'Column B',
+    },
+  ],
+  onChange: () => {},
+};
+
+const adhocMetricA = {
+  expressionType: EXPRESSION_TYPES.SIMPLE,
+  column: defaultProps.columns[0],
+  aggregate: AGGREGATES.SUM,
+  optionName: 'abc',
+};
+const adhocMetricB = {
+  expressionType: EXPRESSION_TYPES.SIMPLE,
+  column: defaultProps.columns[1],
+  aggregate: AGGREGATES.SUM,
+  optionName: 'def',
 };
 
 test('renders with default props', () => {
@@ -38,3 +68,161 @@ test('renders with default props and multi = true', () => {
   render(<DndMetricSelect {...defaultProps} multi />, { useDnd: true });
   expect(screen.getByText('Drop columns or metrics here')).toBeInTheDocument();
 });
+
+test('render selected metrics correctly', () => {
+  const metricValues = ['metric_a', 'metric_b', adhocMetricB];
+  render(<DndMetricSelect {...defaultProps} value={metricValues} multi />, {
+    useDnd: true,
+  });
+  expect(screen.getByText('metric_a')).toBeVisible();
+  expect(screen.getByText('Metric B')).toBeVisible();
+  expect(screen.getByText('SUM(Column B)')).toBeVisible();
+});
+
+test('remove selected custom metric when metric gets removed from dataset', () => {
+  let metricValues = ['metric_a', 'metric_b', adhocMetricA, adhocMetricB];
+  const onChange = (val: any[]) => {
+    metricValues = val;
+  };
+
+  const { rerender } = render(
+    <DndMetricSelect
+      {...defaultProps}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+    {
+      useDnd: true,
+    },
+  );
+
+  const newPropsWithRemovedMetric = {
+    ...defaultProps,
+    savedMetrics: [
+      {
+        metric_name: 'metric_a',
+        expression: 'expression_a',
+      },
+    ],
+  };
+  rerender(
+    <DndMetricSelect
+      {...newPropsWithRemovedMetric}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+  );
+  expect(screen.getByText('metric_a')).toBeVisible();
+  expect(screen.queryByText('Metric B')).not.toBeInTheDocument();
+  expect(screen.getByText('SUM(column_a)')).toBeVisible();
+  expect(screen.getByText('SUM(Column B)')).toBeVisible();
+});
+
+test('remove selected adhoc metric when column gets removed from dataset', async () => {
+  let metricValues = ['metric_a', 'metric_b', adhocMetricA, adhocMetricB];
+  const onChange = (val: any[]) => {
+    metricValues = val;
+  };
+
+  const { rerender } = render(
+    <DndMetricSelect
+      {...defaultProps}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+    {
+      useDnd: true,
+    },
+  );
+
+  const newPropsWithRemovedColumn = {
+    ...defaultProps,
+    columns: [
+      {
+        column_name: 'column_a',
+      },
+    ],
+  };
+
+  // rerender twice - first to update columns, second to update value
+  rerender(
+    <DndMetricSelect
+      {...newPropsWithRemovedColumn}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+  );
+  rerender(
+    <DndMetricSelect
+      {...newPropsWithRemovedColumn}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+  );
+
+  expect(screen.getByText('metric_a')).toBeVisible();
+  expect(screen.getByText('Metric B')).toBeVisible();
+  expect(screen.getByText('SUM(column_a)')).toBeVisible();
+  expect(screen.queryByText('SUM(Column B)')).not.toBeInTheDocument();
+});
+
+test('update adhoc metric name when column label in dataset changes', () => {
+  let metricValues = ['metric_a', 'metric_b', adhocMetricA, adhocMetricB];
+  const onChange = (val: any[]) => {
+    metricValues = val;
+  };
+
+  const { rerender } = render(
+    <DndMetricSelect
+      {...defaultProps}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+    {
+      useDnd: true,
+    },
+  );
+
+  const newPropsWithUpdatedColNames = {
+    ...defaultProps,
+    columns: [
+      {
+        column_name: 'column_a',
+        verbose_name: 'new col A name',
+      },
+      {
+        column_name: 'column_b',
+        verbose_name: 'new col B name',
+      },
+    ],
+  };
+
+  // rerender twice - first to update columns, second to update value
+  rerender(
+    <DndMetricSelect
+      {...newPropsWithUpdatedColNames}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+  );
+  rerender(
+    <DndMetricSelect
+      {...newPropsWithUpdatedColNames}
+      value={metricValues}
+      onChange={onChange}
+      multi
+    />,
+  );
+
+  expect(screen.getByText('metric_a')).toBeVisible();
+  expect(screen.getByText('Metric B')).toBeVisible();
+  expect(screen.getByText('SUM(new col A name)')).toBeVisible();
+  expect(screen.getByText('SUM(new col B name)')).toBeVisible();
+});

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) {