feat: Improves the Drill By feature (#29242)

Co-authored-by: JUST.in DO IT <justin.park@airbnb.com>

superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts

@@ -23,6 +23,8 @@ import {
   ComponentType,
 } from 'react';
 import type { Editor } from 'brace';
+import { BaseFormData } from '../query';
+import { JsonResponse } from '../connection';
 
 /**
  * A function which returns text (or marked-up text)
@@ -30,6 +32,14 @@ import type { Editor } from 'brace';
  */
 type ReturningDisplayable<P = void> = (props: P) => string | ReactElement;
 
+/**
+ * A function which returns the drill by options for a given dataset and chart's formData.
+ */
+export type LoadDrillByOptions = (
+  datasetId: number,
+  formData: BaseFormData,
+) => Promise<JsonResponse>;
+
 /**
  * This type defines all available extensions of Superset's default UI.
  * Namespace the keys here to follow the form of 'some_domain.functionality.item'.
@@ -193,6 +203,7 @@ export interface CustomAutocomplete extends AutocompleteItem {
 
 export type Extensions = Partial<{
   'alertsreports.header.icon': ComponentType;
+  'load.drillby.options': LoadDrillByOptions;
   'embedded.documentation.configuration_details': ComponentType<ConfigDetailsProps>;
   'embedded.documentation.description': ReturningDisplayable;
   'embedded.documentation.url': string;

superset-frontend/src/components/Chart/ChartContextMenu/ChartContextMenu.tsx

@@ -18,6 +18,7 @@
  */
 import {
   forwardRef,
+  Key,
   ReactNode,
   RefObject,
   useCallback,
@@ -104,6 +105,7 @@ const ChartContextMenu = (
   const crossFiltersEnabled = useSelector<RootState, boolean>(
     ({ dashboardInfo }) => dashboardInfo.crossFiltersEnabled,
   );
+  const [openKeys, setOpenKeys] = useState<Key[]>([]);
 
   const isDisplayed = (item: ContextMenuItem) =>
     displayedItems === ContextMenuItem.All ||
@@ -254,6 +256,8 @@ const ChartContextMenu = (
         formData={formData}
         contextMenuY={clientY}
         submenuIndex={submenuIndex}
+        open={openKeys.includes('drill-by-submenu')}
+        key="drill-by-submenu"
         {...(additionalConfig?.drillBy || {})}
       />,
     );
@@ -288,7 +292,13 @@ const ChartContextMenu = (
   return ReactDOM.createPortal(
     <Dropdown
       overlay={
-        <Menu className="chart-context-menu" data-test="chart-context-menu">
+        <Menu
+          className="chart-context-menu"
+          data-test="chart-context-menu"
+          onOpenChange={openKeys => {
+            setOpenKeys(openKeys);
+          }}
+        >
           {menuItems.length ? (
             menuItems
           ) : (

superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.test.tsx

@@ -31,11 +31,17 @@ import { DrillByMenuItems, DrillByMenuItemsProps } from './DrillByMenuItems';
 
 /* eslint jest/expect-expect: ["warn", { "assertFunctionNames": ["expect*"] }] */
 
-const DATASET_ENDPOINT = 'glob:*/api/v1/dataset/7';
+const DATASET_ENDPOINT = 'glob:*/api/v1/dataset/7*';
 const CHART_DATA_ENDPOINT = 'glob:*/api/v1/chart/data*';
 const FORM_DATA_KEY_ENDPOINT = 'glob:*/api/v1/explore/form_data';
 const { form_data: defaultFormData } = chartQueries[sliceId];
 
+jest.mock('lodash/debounce', () => (fn: Function & { debounce: Function }) => {
+  // eslint-disable-next-line no-param-reassign
+  fn.debounce = jest.fn();
+  return fn;
+});
+
 const defaultColumns = [
   { column_name: 'col1', groupby: true },
   { column_name: 'col2', groupby: true },
@@ -68,6 +74,7 @@ const renderMenu = ({
       <DrillByMenuItems
         formData={formData ?? defaultFormData}
         drillByConfig={drillByConfig}
+        open
         {...rest}
       />
     </Menu>,

superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.tsx

@@ -18,11 +18,12 @@
  */
 
 import {
-  ChangeEvent,
+  CSSProperties,
   ReactNode,
   useCallback,
   useEffect,
   useMemo,
+  useRef,
   useState,
 } from 'react';
 import { Menu } from 'src/components/Menu';
@@ -31,12 +32,20 @@ import {
   Behavior,
   Column,
   ContextMenuFilters,
+  FAST_DEBOUNCE,
+  JsonResponse,
   css,
   ensureIsArray,
   getChartMetadataRegistry,
+  getExtensionsRegistry,
+  logging,
   t,
   useTheme,
 } from '@superset-ui/core';
+import rison from 'rison';
+import { debounce } from 'lodash';
+import { FixedSizeList as List } from 'react-window';
+import { AntdInput } from 'src/components';
 import Icons from 'src/components/Icons';
 import { Input } from 'src/components/Input';
 import { useToasts } from 'src/components/MessageToasts/withToasts';
@@ -52,7 +61,7 @@ import { getSubmenuYOffset } from '../utils';
 import { MenuItemWithTruncation } from '../MenuItemWithTruncation';
 import { Dataset } from '../types';
 
-const MAX_SUBMENU_HEIGHT = 200;
+const SUBMENU_HEIGHT = 200;
 const SHOW_COLUMNS_SEARCH_THRESHOLD = 10;
 const SEARCH_INPUT_HEIGHT = 48;
 
@@ -65,8 +74,28 @@ export interface DrillByMenuItemsProps {
   onClick?: (event: MouseEvent) => void;
   openNewModal?: boolean;
   excludedColumns?: Column[];
+  open: boolean;
 }
 
+const loadDrillByOptions = getExtensionsRegistry().get('load.drillby.options');
+
+const queryString = rison.encode({
+  columns: [
+    'table_name',
+    'owners.first_name',
+    'owners.last_name',
+    'created_by.first_name',
+    'created_by.last_name',
+    'created_on_humanized',
+    'changed_by.first_name',
+    'changed_by.last_name',
+    'changed_on_humanized',
+    'columns.column_name',
+    'columns.verbose_name',
+    'columns.groupby',
+  ],
+});
+
 export const DrillByMenuItems = ({
   drillByConfig,
   formData,
@@ -76,6 +105,7 @@ export const DrillByMenuItems = ({
   onClick = () => {},
   excludedColumns,
   openNewModal = true,
+  open,
   ...rest
 }: DrillByMenuItemsProps) => {
   const theme = useTheme();
@@ -86,6 +116,9 @@ export const DrillByMenuItems = ({
   const [columns, setColumns] = useState<Column[]>([]);
   const [showModal, setShowModal] = useState(false);
   const [currentColumn, setCurrentColumn] = useState();
+  const ref = useRef<AntdInput>(null);
+  const showSearch =
+    loadDrillByOptions || columns.length > SHOW_COLUMNS_SEARCH_THRESHOLD;
   const handleSelection = useCallback(
     (event, column) => {
       onClick(event);
@@ -102,10 +135,14 @@ export const DrillByMenuItems = ({
   }, []);
 
   useEffect(() => {
-    // Input is displayed only when columns.length > SHOW_COLUMNS_SEARCH_THRESHOLD
+    if (open) {
-    // Reset search input in case Input gets removed
+      ref.current?.input.focus();
-    setSearchInput('');
+    } else {
-  }, [columns.length]);
+      // Reset search input when menu is closed
+      ref.current?.setValue('');
+      setSearchInput('');
+    }
+  }, [open]);
 
   const hasDrillBy = drillByConfig?.groupbyFieldName;
 
@@ -119,51 +156,59 @@ export const DrillByMenuItems = ({
   const verboseMap = useVerboseMap(dataset);
 
   useEffect(() => {
+    async function loadOptions() {
+      const datasetId = Number(formData.datasource.split('__')[0]);
+      try {
+        setIsLoadingColumns(true);
+        let response: JsonResponse;
+        if (loadDrillByOptions) {
+          response = await loadDrillByOptions(datasetId, formData);
+        } else {
+          response = await cachedSupersetGet({
+            endpoint: `/api/v1/dataset/${datasetId}?q=${queryString}`,
+          });
+        }
+        const { json } = response;
+        const { result } = json;
+        setDataset(result);
+        setColumns(
+          ensureIsArray(result.columns)
+            .filter(column => column.groupby)
+            .filter(
+              column =>
+                !ensureIsArray(
+                  formData[drillByConfig?.groupbyFieldName ?? ''],
+                ).includes(column.column_name) &&
+                column.column_name !== formData.x_axis &&
+                ensureIsArray(excludedColumns)?.every(
+                  excludedCol => excludedCol.column_name !== column.column_name,
+                ),
+            ),
+        );
+      } catch (error) {
+        logging.error(error);
+        supersetGetCache.delete(`/api/v1/dataset/${datasetId}`);
+        addDangerToast(t('Failed to load dimensions for drill by'));
+      } finally {
+        setIsLoadingColumns(false);
+      }
+    }
     if (handlesDimensionContextMenu && hasDrillBy) {
-      const datasetId = formData.datasource.split('__')[0];
+      loadOptions();
-      cachedSupersetGet({
-        endpoint: `/api/v1/dataset/${datasetId}`,
-      })
-        .then(({ json: { result } }) => {
-          setDataset(result);
-          setColumns(
-            ensureIsArray(result.columns)
-              .filter(column => column.groupby)
-              .filter(
-                column =>
-                  !ensureIsArray(
-                    formData[drillByConfig.groupbyFieldName ?? ''],
-                  ).includes(column.column_name) &&
-                  column.column_name !== formData.x_axis &&
-                  ensureIsArray(excludedColumns)?.every(
-                    excludedCol =>
-                      excludedCol.column_name !== column.column_name,
-                  ),
-              ),
-          );
-        })
-        .catch(() => {
-          supersetGetCache.delete(`/api/v1/dataset/${datasetId}`);
-          addDangerToast(t('Failed to load dimensions for drill by'));
-        })
-        .finally(() => {
-          setIsLoadingColumns(false);
-        });
     }
   }, [
     addDangerToast,
+    drillByConfig?.groupbyFieldName,
     excludedColumns,
     formData,
-    drillByConfig?.groupbyFieldName,
     handlesDimensionContextMenu,
     hasDrillBy,
   ]);
 
-  const handleInput = useCallback((e: ChangeEvent<HTMLInputElement>) => {
+  const handleInput = debounce(
-    e.stopPropagation();
+    (value: string) => setSearchInput(value),
-    const input = e?.target?.value;
+    FAST_DEBOUNCE,
-    setSearchInput(input);
+  );
-  }, []);
 
   const filteredColumns = useMemo(
     () =>
@@ -181,12 +226,10 @@ export const DrillByMenuItems = ({
         contextMenuY,
         filteredColumns.length || 1,
         submenuIndex,
-        MAX_SUBMENU_HEIGHT,
+        SUBMENU_HEIGHT,
-        columns.length > SHOW_COLUMNS_SEARCH_THRESHOLD
+        showSearch ? SEARCH_INPUT_HEIGHT : 0,
-          ? SEARCH_INPUT_HEIGHT
-          : 0,
       ),
-    [contextMenuY, filteredColumns.length, submenuIndex, columns.length],
+    [contextMenuY, filteredColumns.length, submenuIndex, showSearch],
   );
 
   let tooltip: ReactNode;
@@ -208,27 +251,53 @@ export const DrillByMenuItems = ({
     );
   }
 
+  const Row = ({
+    index,
+    data,
+    style,
+  }: {
+    index: number;
+    data: { columns: Column[] };
+    style: CSSProperties;
+  }) => {
+    const { columns, ...rest } = data;
+    const column = columns[index];
+    return (
+      <MenuItemWithTruncation
+        key={`drill-by-item-${column.column_name}`}
+        tooltipText={column.verbose_name || column.column_name}
+        {...rest}
+        onClick={e => handleSelection(e, column)}
+        style={style}
+      >
+        {column.verbose_name || column.column_name}
+      </MenuItemWithTruncation>
+    );
+  };
+
   return (
     <>
       <Menu.SubMenu
         title={t('Drill by')}
-        key="drill-by-submenu"
         popupClassName="chart-context-submenu"
         popupOffset={[0, submenuYOffset]}
         {...rest}
       >
         <div data-test="drill-by-submenu">
-          {columns.length > SHOW_COLUMNS_SEARCH_THRESHOLD && (
+          {showSearch && (
             <Input
+              ref={ref}
               prefix={
                 <Icons.Search
                   iconSize="l"
                   iconColor={theme.colors.grayscale.light1}
                 />
               }
-              onChange={handleInput}
+              onChange={e => {
+                e.stopPropagation();
+                handleInput(e.target.value);
+              }}
               placeholder={t('Search columns')}
-              value={searchInput}
               onClick={e => {
                 // prevent closing menu when clicking on input
                 e.nativeEvent.stopImmediatePropagation();
@@ -251,23 +320,16 @@ export const DrillByMenuItems = ({
               <Loading position="inline-centered" />
             </div>
           ) : filteredColumns.length ? (
-            <div
+            <List
-              css={css`
+              width="100%"
-                max-height: ${MAX_SUBMENU_HEIGHT}px;
+              height={SUBMENU_HEIGHT}
-                overflow: auto;
+              itemSize={35}
-              `}
+              itemCount={filteredColumns.length}
+              itemData={{ columns: filteredColumns, ...rest }}
+              overscanCount={20}
             >
-              {filteredColumns.map(column => (
+              {Row}
-                <MenuItemWithTruncation
+            </List>
-                  key={`drill-by-item-${column.column_name}`}
-                  tooltipText={column.verbose_name || column.column_name}
-                  {...rest}
-                  onClick={e => handleSelection(e, column)}
-                >
-                  {column.verbose_name || column.column_name}
-                </MenuItemWithTruncation>
-              ))}
-            </div>
           ) : (
             <Menu.Item disabled key="no-drill-by-columns-found" {...rest}>
               {t('No columns found')}

superset-frontend/src/components/Chart/MenuItemWithTruncation.tsx

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import { ReactNode } from 'react';
+import { ReactNode, CSSProperties } from 'react';
 import { css, truncationCSS, useCSSTextTruncation } from '@superset-ui/core';
 import { Menu } from 'src/components/Menu';
 import { Tooltip } from 'src/components/Tooltip';
@@ -27,6 +27,7 @@ export type MenuItemWithTruncationProps = {
   tooltipText: ReactNode;
   children: ReactNode;
   onClick?: MenuProps['onClick'];
+  style?: CSSProperties;
 };
 
 export const MenuItemWithTruncation = ({