From 66c28d653f3f9381479ab817e6dba8f6598d8b8d Mon Sep 17 00:00:00 2001
From: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
Date: Fri, 16 Jul 2021 14:20:13 +0200
Subject: [PATCH] fix(explore): DndColumnSelect sometimes not working with
 multi: false (#15731)

* fix(explore): DndColumnSelect not working with multi: false

* fix values not synchronized when dataset changes
---
 .../DndColumnSelect.tsx                       | 44 ++++++++++++++-----
 .../utils/optionSelector.ts                   | 21 ++++-----
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndColumnSelect.tsx b/superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndColumnSelect.tsx
index 10c8bd60d2..29cd738d6d 100644
--- a/superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndColumnSelect.tsx
+++ b/superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndColumnSelect.tsx
@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useState } from 'react';
+import React, { useEffect } from 'react';
 import { tn } from '@superset-ui/core';
 import { ColumnMeta } from '@superset-ui/chart-controls';
 import { isEmpty } from 'lodash';
@@ -29,19 +29,41 @@ import { DndItemType } from 'src/explore/components/DndItemType';
 import { StyledColumnOption } from 'src/explore/components/optionRenderers';
 
 export const DndColumnSelect = (props: LabelProps) => {
-  const { value, options, multi = true } = props;
-  const optionSelector = new OptionSelector(options, value);
-  const [values, setValues] = useState<ColumnMeta[]>(optionSelector.values);
+  const { value, options, multi = true, onChange } = props;
+  const optionSelector = new OptionSelector(options, multi, value);
+
+  // synchronize values in case of dataset changes
+  useEffect(() => {
+    const optionSelectorValues = optionSelector.getValues();
+    if (typeof value !== typeof optionSelectorValues) {
+      onChange(optionSelectorValues);
+    }
+    if (
+      typeof value === 'string' &&
+      typeof optionSelectorValues === 'string' &&
+      value !== optionSelectorValues
+    ) {
+      onChange(optionSelectorValues);
+    }
+    if (
+      Array.isArray(optionSelectorValues) &&
+      Array.isArray(value) &&
+      (optionSelectorValues.length !== value.length ||
+        optionSelectorValues.every((val, index) => val === value[index]))
+    ) {
+      onChange(optionSelectorValues);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [JSON.stringify(value), JSON.stringify(optionSelector.getValues())]);
 
   const onDrop = (item: DatasourcePanelDndItem) => {
     const column = item.value as ColumnMeta;
-    if (!optionSelector.isArray && !isEmpty(optionSelector.values)) {
+    if (!optionSelector.multi && !isEmpty(optionSelector.values)) {
       optionSelector.replace(0, column.column_name);
     } else {
       optionSelector.add(column.column_name);
     }
-    setValues(optionSelector.values);
-    props.onChange(optionSelector.getValues());
+    onChange(optionSelector.getValues());
   };
 
   const canDrop = (item: DatasourcePanelDndItem) =>
@@ -50,18 +72,16 @@ export const DndColumnSelect = (props: LabelProps) => {
 
   const onClickClose = (index: number) => {
     optionSelector.del(index);
-    setValues(optionSelector.values);
-    props.onChange(optionSelector.getValues());
+    onChange(optionSelector.getValues());
   };
 
   const onShiftOptions = (dragIndex: number, hoverIndex: number) => {
     optionSelector.swap(dragIndex, hoverIndex);
-    setValues(optionSelector.values);
-    props.onChange(optionSelector.getValues());
+    onChange(optionSelector.getValues());
   };
 
   const valuesRenderer = () =>
-    values.map((column, idx) => (
+    optionSelector.values.map((column, idx) => (
       <OptionWrapper
         key={idx}
         index={idx}
diff --git a/superset-frontend/src/explore/components/controls/DndColumnSelectControl/utils/optionSelector.ts b/superset-frontend/src/explore/components/controls/DndColumnSelectControl/utils/optionSelector.ts
index 37833d943c..532aa33fc6 100644
--- a/superset-frontend/src/explore/components/controls/DndColumnSelectControl/utils/optionSelector.ts
+++ b/superset-frontend/src/explore/components/controls/DndColumnSelectControl/utils/optionSelector.ts
@@ -17,28 +17,23 @@
  * under the License.
  */
 import { ColumnMeta } from '@superset-ui/chart-controls';
+import { ensureIsArray } from '@superset-ui/core';
 
 export class OptionSelector {
   values: ColumnMeta[];
 
   options: { string: ColumnMeta };
 
-  isArray: boolean;
+  multi: boolean;
 
   constructor(
     options: { string: ColumnMeta },
+    multi: boolean,
     initialValues?: string[] | string,
   ) {
     this.options = options;
-    let values: string[];
-    if (Array.isArray(initialValues)) {
-      values = initialValues;
-      this.isArray = true;
-    } else {
-      values = initialValues ? [initialValues] : [];
-      this.isArray = false;
-    }
-    this.values = values
+    this.multi = multi;
+    this.values = ensureIsArray(initialValues)
       .map(value => {
         if (value in options) {
           return options[value];
@@ -68,12 +63,12 @@ export class OptionSelector {
     [this.values[a], this.values[b]] = [this.values[b], this.values[a]];
   }
 
-  has(groupBy: string): boolean {
-    return !!this.getValues()?.includes(groupBy);
+  has(value: string): boolean {
+    return !!this.getValues()?.includes(value);
   }
 
   getValues(): string[] | string | undefined {
-    if (!this.isArray) {
+    if (!this.multi) {
       return this.values.length > 0 ? this.values[0].column_name : undefined;
     }
     return this.values.map(option => option.column_name);