feat(explore): adhoc column expressions [ID-3] (#17379)

* add support for adhoc columns to api and sqla model

* fix some types

* fix duplicates in column names

* fix more lint

* fix schema and dedup

* clean up some logic

* first pass at fixing viz.py

* Add frontend support for adhoc columns

* Add title edit

* Fix showing custom title

* Use column name as default value in sql editor

* fix: Adds a loading message when needed in the Select component (#16531)

* fix(tests): make parquet select deterministic with order by (#16570)

* bump emotion to help with cache clobbering (#16559)

* fix: Support Jinja template functions in global async queries (#16412)

* Support Jinja template functions in async queries

* Pylint

* Add tests for async tasks

* Remove redundant has_request_context check

* fix: impersonate user label/tooltip (#16573)

* docs: update for small typos (#16568)

* feat: Add Aurora Data API engine spec (#16535)

* feat: Add Aurora Data API engine spec

* Fix lint

* refactor: sql_json view endpoint: encapsulate ctas parameters (#16548)

* refactor sql_json view endpoint: encapsulate ctas parameters

* fix failed tests

* fix failed tests and ci issues

* refactor sql_json view endpoint: separate concern into ad hod method (#16595)

* feat: Experimental cross-filter plugins (#16594)

* fix:fix get permission function

* feat: add cross filter chart in charts gallery under FF

* chore(deps): bump superset-ui to 0.18.2 (#16601)

* update type guard references

* fix imports

* update series_columns schema

* Add changes that got lost in rebase

* Use current columns name or expression as sql editor init value

* add integration test and do minor fixes

* Bump superset-ui

* fix linting issue

* bump superset-ui to 0.18.22

* resolve merge conflict

* lint

* fix select filter infinite loop

* bump superset-ui to 0.18.23

* Fix auto setting column popover title

* Enable adhoc columns only if UX_BETA enabled

* put back removed test

* Move popover height and width to constants

* Refactor big ternary expression

Co-authored-by: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
Co-authored-by: Michael S. Molina <70410625+michael-s-molina@users.noreply.github.com>
Co-authored-by: Elizabeth Thompson <eschutho@gmail.com>
Co-authored-by: Rob DiCiuccio <rob.diciuccio@gmail.com>
Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>
Co-authored-by: joeADSP <75027008+joeADSP@users.noreply.github.com>
Co-authored-by: ofekisr <35701650+ofekisr@users.noreply.github.com>
Co-authored-by: simcha90 <56388545+simcha90@users.noreply.github.com>

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

@@ -17,15 +17,27 @@
  * under the License.
  */
 /* eslint-disable camelcase */
-import React, { useCallback, useMemo, useState } from 'react';
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
+import { AdhocColumn, t, styled, css } from '@superset-ui/core';
+import {
+  ColumnMeta,
+  isAdhocColumn,
+  isSavedExpression,
+} from '@superset-ui/chart-controls';
 import Tabs from 'src/components/Tabs';
 import Button from 'src/components/Button';
 import { Select } from 'src/components';
-import { t, styled } from '@superset-ui/core';
 
 import { Form, FormItem } from 'src/components/Form';
+import { SQLEditor } from 'src/components/AsyncAceEditor';
 import { StyledColumnOption } from 'src/explore/components/optionRenderers';
-import { ColumnMeta } from '@superset-ui/chart-controls';
+import { POPOVER_INITIAL_HEIGHT } from 'src/explore/constants';
 
 const StyledSelect = styled(Select)`
   .metric-option {
@@ -41,29 +53,58 @@ const StyledSelect = styled(Select)`
 
 interface ColumnSelectPopoverProps {
   columns: ColumnMeta[];
-  editedColumn?: ColumnMeta;
-  onChange: (column: ColumnMeta) => void;
+  editedColumn?: ColumnMeta | AdhocColumn;
+  onChange: (column: ColumnMeta | AdhocColumn) => void;
   onClose: () => void;
+  setLabel: (title: string) => void;
+  getCurrentTab: (tab: string) => void;
+  label: string;
+  isAdhocColumnsEnabled: boolean;
 }
 
+const getInitialColumnValues = (
+  editedColumn?: ColumnMeta | AdhocColumn,
+): [AdhocColumn?, ColumnMeta?, ColumnMeta?] => {
+  if (!editedColumn) {
+    return [undefined, undefined, undefined];
+  }
+  if (isAdhocColumn(editedColumn)) {
+    return [editedColumn, undefined, undefined];
+  }
+  if (isSavedExpression(editedColumn)) {
+    return [undefined, editedColumn, undefined];
+  }
+  return [undefined, undefined, editedColumn];
+};
+
 const ColumnSelectPopover = ({
   columns,
   editedColumn,
   onChange,
   onClose,
+  setLabel,
+  getCurrentTab,
+  label,
+  isAdhocColumnsEnabled,
 }: ColumnSelectPopoverProps) => {
+  const [initialLabel] = useState(label);
   const [
+    initialAdhocColumn,
     initialCalculatedColumn,
     initialSimpleColumn,
-  ] = editedColumn?.expression
-    ? [editedColumn, undefined]
-    : [undefined, editedColumn];
-  const [selectedCalculatedColumn, setSelectedCalculatedColumn] = useState(
-    initialCalculatedColumn,
-  );
-  const [selectedSimpleColumn, setSelectedSimpleColumn] = useState(
-    initialSimpleColumn,
+  ] = getInitialColumnValues(editedColumn);
+
+  const [adhocColumn, setAdhocColumn] = useState<AdhocColumn | undefined>(
+    initialAdhocColumn,
   );
+  const [selectedCalculatedColumn, setSelectedCalculatedColumn] = useState<
+    ColumnMeta | undefined
+  >(initialCalculatedColumn);
+  const [selectedSimpleColumn, setSelectedSimpleColumn] = useState<
+    ColumnMeta | undefined
+  >(initialSimpleColumn);
+
+  const sqlEditorRef = useRef(null);
 
   const [calculatedColumns, simpleColumns] = useMemo(
     () =>
@@ -81,6 +122,15 @@ const ColumnSelectPopover = ({
     [columns],
   );
 
+  const onSqlExpressionChange = useCallback(
+    sqlExpression => {
+      setAdhocColumn({ label, sqlExpression } as AdhocColumn);
+      setSelectedSimpleColumn(undefined);
+      setSelectedCalculatedColumn(undefined);
+    },
+    [label],
+  );
+
   const onCalculatedColumnChange = useCallback(
     selectedColumnName => {
       const selectedColumn = calculatedColumns.find(
@@ -88,8 +138,12 @@ const ColumnSelectPopover = ({
       );
       setSelectedCalculatedColumn(selectedColumn);
       setSelectedSimpleColumn(undefined);
+      setAdhocColumn(undefined);
+      setLabel(
+        selectedColumn?.verbose_name || selectedColumn?.column_name || '',
+      );
     },
-    [calculatedColumns],
+    [calculatedColumns, setLabel],
   );
 
   const onSimpleColumnChange = useCallback(
@@ -99,33 +153,79 @@ const ColumnSelectPopover = ({
       );
       setSelectedCalculatedColumn(undefined);
       setSelectedSimpleColumn(selectedColumn);
+      setAdhocColumn(undefined);
+      setLabel(
+        selectedColumn?.verbose_name || selectedColumn?.column_name || '',
+      );
     },
-    [simpleColumns],
+    [setLabel, simpleColumns],
   );
 
-  const defaultActiveTabKey =
-    initialSimpleColumn || calculatedColumns.length === 0 ? 'simple' : 'saved';
+  const defaultActiveTabKey = initialAdhocColumn
+    ? 'sqlExpression'
+    : initialSimpleColumn || calculatedColumns.length === 0
+    ? 'simple'
+    : 'saved';
+
+  useEffect(() => {
+    getCurrentTab(defaultActiveTabKey);
+  }, [defaultActiveTabKey, getCurrentTab]);
 
   const onSave = useCallback(() => {
-    const selectedColumn = selectedCalculatedColumn || selectedSimpleColumn;
+    if (adhocColumn && adhocColumn.label !== label) {
+      adhocColumn.label = label;
+    }
+    const selectedColumn =
+      adhocColumn || selectedCalculatedColumn || selectedSimpleColumn;
     if (!selectedColumn) {
       return;
     }
     onChange(selectedColumn);
     onClose();
-  }, [onChange, onClose, selectedCalculatedColumn, selectedSimpleColumn]);
+  }, [
+    adhocColumn,
+    label,
+    onChange,
+    onClose,
+    selectedCalculatedColumn,
+    selectedSimpleColumn,
+  ]);
 
   const onResetStateAndClose = useCallback(() => {
     setSelectedCalculatedColumn(initialCalculatedColumn);
     setSelectedSimpleColumn(initialSimpleColumn);
+    setAdhocColumn(initialAdhocColumn);
     onClose();
-  }, [initialCalculatedColumn, initialSimpleColumn, onClose]);
+  }, [
+    initialAdhocColumn,
+    initialCalculatedColumn,
+    initialSimpleColumn,
+    onClose,
+  ]);
 
-  const stateIsValid = selectedCalculatedColumn || selectedSimpleColumn;
+  const onTabChange = useCallback(
+    tab => {
+      getCurrentTab(tab);
+      // @ts-ignore
+      sqlEditorRef.current?.editor.focus();
+    },
+    [getCurrentTab],
+  );
+
+  const onSqlEditorFocus = useCallback(() => {
+    // @ts-ignore
+    sqlEditorRef.current?.editor.resize();
+  }, []);
+
+  const stateIsValid =
+    adhocColumn || selectedCalculatedColumn || selectedSimpleColumn;
   const hasUnsavedChanges =
+    initialLabel !== label ||
     selectedCalculatedColumn?.column_name !==
       initialCalculatedColumn?.column_name ||
-    selectedSimpleColumn?.column_name !== initialSimpleColumn?.column_name;
+    selectedSimpleColumn?.column_name !== initialSimpleColumn?.column_name ||
+    adhocColumn?.sqlExpression !== initialAdhocColumn?.sqlExpression;
+
   const savedExpressionsLabel = t('Saved expressions');
   const simpleColumnsLabel = t('Column');
 
@@ -134,8 +234,12 @@ const ColumnSelectPopover = ({
       <Tabs
         id="adhoc-metric-edit-tabs"
         defaultActiveKey={defaultActiveTabKey}
+        onChange={onTabChange}
         className="adhoc-metric-edit-tabs"
         allowOverflow
+        css={css`
+          height: ${POPOVER_INITIAL_HEIGHT}px;
+        `}
       >
         <Tabs.TabPane key="saved" tab={t('Saved')}>
           <FormItem label={savedExpressionsLabel}>
@@ -178,6 +282,28 @@ const ColumnSelectPopover = ({
             />
           </FormItem>
         </Tabs.TabPane>
+        {isAdhocColumnsEnabled && (
+          <Tabs.TabPane key="sqlExpression" tab={t('Custom SQL')}>
+            <SQLEditor
+              value={
+                adhocColumn?.sqlExpression ||
+                selectedSimpleColumn?.column_name ||
+                selectedCalculatedColumn?.expression
+              }
+              onFocus={onSqlEditorFocus}
+              showLoadingForImport
+              onChange={onSqlExpressionChange}
+              width="100%"
+              height={`${POPOVER_INITIAL_HEIGHT - 80}px`}
+              showGutter={false}
+              editorProps={{ $blockScrolling: true }}
+              enableLiveAutocompletion
+              className="filter-sql-editor"
+              wrapEnabled
+              ref={sqlEditorRef}
+            />
+          </Tabs.TabPane>
+        )}
       </Tabs>
       <div>
         <Button buttonSize="small" onClick={onResetStateAndClose} cta>

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

@@ -16,16 +16,27 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useCallback, useMemo, useState } from 'react';
-import { ColumnMeta } from '@superset-ui/chart-controls';
+import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import {
+  AdhocColumn,
+  FeatureFlag,
+  isFeatureEnabled,
+  t,
+} from '@superset-ui/core';
+import {
+  ColumnMeta,
+  isAdhocColumn,
+  isColumnMeta,
+} from '@superset-ui/chart-controls';
 import Popover from 'src/components/Popover';
 import { ExplorePopoverContent } from 'src/explore/components/ExploreContentPopover';
 import ColumnSelectPopover from './ColumnSelectPopover';
+import { DndColumnSelectPopoverTitle } from './DndColumnSelectPopoverTitle';
 
 interface ColumnSelectPopoverTriggerProps {
   columns: ColumnMeta[];
-  editedColumn?: ColumnMeta;
-  onColumnEdit: (editedColumn: ColumnMeta) => void;
+  editedColumn?: ColumnMeta | AdhocColumn;
+  onColumnEdit: (editedColumn: ColumnMeta | AdhocColumn) => void;
   isControlledComponent?: boolean;
   visible?: boolean;
   togglePopover?: (visible: boolean) => void;
@@ -33,6 +44,11 @@ interface ColumnSelectPopoverTriggerProps {
   children: React.ReactNode;
 }
 
+const defaultPopoverLabel = t('My column');
+const editableTitleTab = 'sqlExpression';
+
+const isAdhocColumnsEnabled = isFeatureEnabled(FeatureFlag.UX_BETA);
+
 const ColumnSelectPopoverTrigger = ({
   columns,
   editedColumn,
@@ -41,7 +57,21 @@ const ColumnSelectPopoverTrigger = ({
   children,
   ...props
 }: ColumnSelectPopoverTriggerProps) => {
+  const [popoverLabel, setPopoverLabel] = useState(defaultPopoverLabel);
   const [popoverVisible, setPopoverVisible] = useState(false);
+  const [isTitleEditDisabled, setIsTitleEditDisabled] = useState(true);
+  const [hasCustomLabel, setHasCustomLabel] = useState(false);
+
+  let initialPopoverLabel = defaultPopoverLabel;
+  if (editedColumn && isColumnMeta(editedColumn)) {
+    initialPopoverLabel = editedColumn.verbose_name || editedColumn.column_name;
+  } else if (editedColumn && isAdhocColumn(editedColumn)) {
+    initialPopoverLabel = editedColumn.label || defaultPopoverLabel;
+  }
+
+  useEffect(() => {
+    setPopoverLabel(initialPopoverLabel);
+  }, [initialPopoverLabel, popoverVisible]);
 
   const togglePopover = useCallback((visible: boolean) => {
     setPopoverVisible(visible);
@@ -67,6 +97,10 @@ const ColumnSelectPopoverTrigger = ({
         handleClosePopover: closePopover,
       };
 
+  const getCurrentTab = useCallback((tab: string) => {
+    setIsTitleEditDisabled(tab !== editableTitleTab);
+  }, []);
+
   const overlayContent = useMemo(
     () => (
       <ExplorePopoverContent>
@@ -75,10 +109,38 @@ const ColumnSelectPopoverTrigger = ({
           columns={columns}
           onClose={handleClosePopover}
           onChange={onColumnEdit}
+          label={popoverLabel}
+          setLabel={setPopoverLabel}
+          getCurrentTab={getCurrentTab}
+          isAdhocColumnsEnabled={isAdhocColumnsEnabled}
         />
       </ExplorePopoverContent>
     ),
-    [columns, editedColumn, handleClosePopover, onColumnEdit],
+    [
+      columns,
+      editedColumn,
+      getCurrentTab,
+      handleClosePopover,
+      onColumnEdit,
+      popoverLabel,
+    ],
+  );
+
+  const onLabelChange = useCallback((e: any) => {
+    setPopoverLabel(e.target.value);
+    setHasCustomLabel(true);
+  }, []);
+
+  const popoverTitle = useMemo(
+    () => (
+      <DndColumnSelectPopoverTitle
+        title={popoverLabel}
+        onChange={onLabelChange}
+        isEditDisabled={isTitleEditDisabled}
+        hasCustomLabel={hasCustomLabel}
+      />
+    ),
+    [hasCustomLabel, isTitleEditDisabled, onLabelChange, popoverLabel],
   );
 
   return (
@@ -89,6 +151,7 @@ const ColumnSelectPopoverTrigger = ({
       defaultVisible={visible}
       visible={visible}
       onVisibleChange={handleTogglePopover}
+      title={isAdhocColumnsEnabled && popoverTitle}
       destroyTooltipOnHide
     >
       {children}

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

@@ -17,8 +17,14 @@
  * under the License.
  */
 import React, { useCallback, useMemo, useState } from 'react';
-import { FeatureFlag, isFeatureEnabled, tn } from '@superset-ui/core';
-import { ColumnMeta } from '@superset-ui/chart-controls';
+import {
+  AdhocColumn,
+  FeatureFlag,
+  isFeatureEnabled,
+  tn,
+  QueryFormColumn,
+} from '@superset-ui/core';
+import { ColumnMeta, isColumnMeta } from '@superset-ui/chart-controls';
 import { isEmpty } from 'lodash';
 import DndSelectLabel from 'src/explore/components/controls/DndColumnSelectControl/DndSelectLabel';
 import OptionWrapper from 'src/explore/components/controls/DndColumnSelectControl/OptionWrapper';
@@ -29,7 +35,7 @@ import { useComponentDidUpdate } from 'src/common/hooks/useComponentDidUpdate';
 import ColumnSelectPopoverTrigger from './ColumnSelectPopoverTrigger';
 import { DndControlProps } from './types';
 
-export type DndColumnSelectProps = DndControlProps<string> & {
+export type DndColumnSelectProps = DndControlProps<QueryFormColumn> & {
   options: Record<string, ColumnMeta>;
 };
 
@@ -123,7 +129,8 @@ export function DndColumnSelect(props: DndColumnSelectProps) {
       Object.values(options).filter(
         col =>
           !optionSelector.values
-            .map(val => val.column_name)
+            .filter(isColumnMeta)
+            .map((val: ColumnMeta) => val.column_name)
             .includes(col.column_name),
       ),
     [optionSelector.values, options],
@@ -136,7 +143,11 @@ export function DndColumnSelect(props: DndColumnSelectProps) {
           <ColumnSelectPopoverTrigger
             columns={popoverOptions}
             onColumnEdit={newColumn => {
-              optionSelector.replace(idx, newColumn.column_name);
+              if (isColumnMeta(newColumn)) {
+                optionSelector.replace(idx, newColumn.column_name);
+              } else {
+                optionSelector.replace(idx, newColumn as AdhocColumn);
+              }
               onChange(optionSelector.getValues());
             }}
             editedColumn={column}
@@ -177,8 +188,12 @@ export function DndColumnSelect(props: DndColumnSelectProps) {
   );
 
   const addNewColumnWithPopover = useCallback(
-    (newColumn: ColumnMeta) => {
-      optionSelector.add(newColumn.column_name);
+    (newColumn: ColumnMeta | AdhocColumn) => {
+      if (isColumnMeta(newColumn)) {
+        optionSelector.add(newColumn.column_name);
+      } else {
+        optionSelector.add(newColumn as AdhocColumn);
+      }
       onChange(optionSelector.getValues());
     },
     [onChange, optionSelector],

superset-frontend/src/explore/components/controls/DndColumnSelectControl/DndColumnSelectPopoverTitle.jsx

@@ -0,0 +1,96 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React, { useCallback, useState } from 'react';
+import { t } from '@superset-ui/core';
+import { Input } from 'src/common/components';
+import { Tooltip } from 'src/components/Tooltip';
+
+export const DndColumnSelectPopoverTitle = ({
+  title,
+  onChange,
+  isEditDisabled,
+  hasCustomLabel,
+}) => {
+  const [isHovered, setIsHovered] = useState(false);
+  const [isEditMode, setIsEditMode] = useState(false);
+
+  const onMouseOver = useCallback(() => {
+    setIsHovered(true);
+  }, []);
+
+  const onMouseOut = useCallback(() => {
+    setIsHovered(false);
+  }, []);
+
+  const onClick = useCallback(() => {
+    setIsEditMode(true);
+  }, []);
+
+  const onBlur = useCallback(() => {
+    setIsEditMode(false);
+  }, []);
+
+  const onInputBlur = useCallback(
+    e => {
+      if (e.target.value === '') {
+        onChange(e);
+      }
+      onBlur();
+    },
+    [onBlur, onChange],
+  );
+
+  const defaultLabel = t('My column');
+
+  if (isEditDisabled) {
+    return <span>{title || defaultLabel}</span>;
+  }
+
+  return isEditMode ? (
+    <Input
+      className="metric-edit-popover-label-input"
+      type="text"
+      placeholder={title}
+      value={hasCustomLabel ? title : ''}
+      autoFocus
+      onChange={onChange}
+      onBlur={onInputBlur}
+    />
+  ) : (
+    <Tooltip placement="top" title={t('Click to edit label')}>
+      <span
+        className="AdhocMetricEditPopoverTitle inline-editable"
+        data-test="AdhocMetricEditTitle#trigger"
+        onMouseOver={onMouseOver}
+        onMouseOut={onMouseOut}
+        onClick={onClick}
+        onBlur={onBlur}
+        role="button"
+        tabIndex={0}
+      >
+        {title || defaultLabel}
+        &nbsp;
+        <i
+          className="fa fa-pencil"
+          style={{ color: isHovered ? 'black' : 'grey' }}
+        />
+      </span>
+    </Tooltip>
+  );
+};

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

@@ -31,7 +31,7 @@ import {
 import { Tooltip } from 'src/components/Tooltip';
 import { StyledColumnOption } from 'src/explore/components/optionRenderers';
 import { styled } from '@superset-ui/core';
-import { ColumnMeta } from '@superset-ui/chart-controls';
+import { ColumnMeta, isAdhocColumn } from '@superset-ui/chart-controls';
 import Option from './Option';
 
 export const OptionLabel = styled.div`
@@ -135,14 +135,20 @@ export default function OptionWrapper(
     );
   };
 
-  const ColumnOption = () => (
-    <StyledColumnOption
-      column={column as ColumnMeta}
-      labelRef={labelRef}
-      showTooltip={!!shouldShowTooltip}
-      showType
-    />
-  );
+  const ColumnOption = () => {
+    const transformedCol =
+      column && isAdhocColumn(column)
+        ? { verbose_name: column.label, expression: column.sqlExpression }
+        : column;
+    return (
+      <StyledColumnOption
+        column={transformedCol as ColumnMeta}
+        labelRef={labelRef}
+        showTooltip={!!shouldShowTooltip}
+        showType
+      />
+    );
+  };
 
   const Label = () => {
     if (label) {

superset-frontend/src/explore/components/controls/DndColumnSelectControl/types.ts

@@ -17,7 +17,7 @@
  * under the License.
  */
 import { ReactNode } from 'react';
-import { JsonValue } from '@superset-ui/core';
+import { AdhocColumn, JsonValue } from '@superset-ui/core';
 import { ControlComponentProps } from 'src/explore/components/Control';
 import { ColumnMeta } from '@superset-ui/chart-controls';
 
@@ -26,7 +26,7 @@ export interface OptionProps {
   index: number;
   label?: string;
   tooltipTitle?: string;
-  column?: ColumnMeta;
+  column?: ColumnMeta | AdhocColumn;
   clickClose: (index: number) => void;
   withCaret?: boolean;
   isExtra?: boolean;

superset-frontend/src/explore/components/controls/DndColumnSelectControl/utils/optionSelector.ts

@@ -16,11 +16,25 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { ColumnMeta } from '@superset-ui/chart-controls';
-import { ensureIsArray } from '@superset-ui/core';
+import { ColumnMeta, isColumnMeta } from '@superset-ui/chart-controls';
+import {
+  AdhocColumn,
+  ensureIsArray,
+  QueryFormColumn,
+  isPhysicalColumn,
+} from '@superset-ui/core';
+
+const getColumnNameOrAdhocColumn = (
+  column: ColumnMeta | AdhocColumn,
+): QueryFormColumn => {
+  if (isColumnMeta(column)) {
+    return column.column_name;
+  }
+  return column as AdhocColumn;
+};
 
 export class OptionSelector {
-  values: ColumnMeta[];
+  values: (ColumnMeta | AdhocColumn)[];
 
   options: Record<string, ColumnMeta>;
 
@@ -29,23 +43,28 @@ export class OptionSelector {
   constructor(
     options: Record<string, ColumnMeta>,
     multi: boolean,
-    initialValues?: string[] | string | null,
+    initialValues?: QueryFormColumn[] | QueryFormColumn | null,
   ) {
     this.options = options;
     this.multi = multi;
     this.values = ensureIsArray(initialValues)
       .map(value => {
-        if (value && value in options) {
+        if (value && isPhysicalColumn(value) && value in options) {
           return options[value];
         }
+        if (!isPhysicalColumn(value)) {
+          return value;
+        }
         return null;
       })
       .filter(Boolean) as ColumnMeta[];
   }
 
-  add(value: string) {
-    if (value in this.options) {
+  add(value: QueryFormColumn) {
+    if (isPhysicalColumn(value) && value in this.options) {
       this.values.push(this.options[value]);
+    } else if (!isPhysicalColumn(value)) {
+      this.values.push(value as AdhocColumn);
     }
   }
 
@@ -53,9 +72,9 @@ export class OptionSelector {
     this.values.splice(idx, 1);
   }
 
-  replace(idx: number, value: string) {
+  replace(idx: number, value: QueryFormColumn) {
     if (this.values[idx]) {
-      this.values[idx] = this.options[value];
+      this.values[idx] = isPhysicalColumn(value) ? this.options[value] : value;
     }
   }
 
@@ -63,14 +82,27 @@ export class OptionSelector {
     [this.values[a], this.values[b]] = [this.values[b], this.values[a]];
   }
 
-  has(value: string): boolean {
-    return ensureIsArray(this.getValues()).includes(value);
+  has(value: QueryFormColumn): boolean {
+    return this.values.some(col => {
+      if (isPhysicalColumn(value)) {
+        return (
+          (col as ColumnMeta).column_name === value ||
+          (col as AdhocColumn).label === value
+        );
+      }
+      return (
+        (col as ColumnMeta).column_name === value.label ||
+        (col as AdhocColumn).label === value.label
+      );
+    });
   }
 
-  getValues(): string[] | string | undefined {
+  getValues(): QueryFormColumn[] | QueryFormColumn | undefined {
     if (!this.multi) {
-      return this.values.length > 0 ? this.values[0].column_name : undefined;
+      return this.values.length > 0
+        ? getColumnNameOrAdhocColumn(this.values[0])
+        : undefined;
     }
-    return this.values.map(option => option.column_name);
+    return this.values.map(getColumnNameOrAdhocColumn);
   }
 }