feat: Support further drill by in the modal (#23615)

superset-frontend/packages/superset-ui-core/src/chart/components/SuperChart.tsx

@@ -86,6 +86,10 @@ export type Props = Omit<SuperChartCoreProps, 'chartProps'> &
      * If not defined, NoResultsComponent is used
      */
     noResults?: ReactNode;
+    /**
+     * Determines is the context menu related to the chart is open
+     */
+    inContextMenu?: boolean;
   };
 
 type PropsWithDefault = Props & Readonly<typeof defaultProps>;

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

@@ -29,6 +29,7 @@ import { useDispatch, useSelector } from 'react-redux';
 import {
   Behavior,
   ContextMenuFilters,
+  ensureIsArray,
   FeatureFlag,
   getChartMetadataRegistry,
   isFeatureEnabled,
@@ -39,21 +40,33 @@ import {
 import { RootState } from 'src/dashboard/types';
 import { findPermission } from 'src/utils/findPermission';
 import { Menu } from 'src/components/Menu';
-import { AntdDropdown as Dropdown } from 'src/components';
+import { AntdDropdown as Dropdown } from 'src/components/index';
-import { DrillDetailMenuItems } from './DrillDetail';
+import { updateDataMask } from 'src/dataMask/actions';
-import { getMenuAdjustedY } from './utils';
+import { DrillDetailMenuItems } from '../DrillDetail';
-import { updateDataMask } from '../../dataMask/actions';
+import { getMenuAdjustedY } from '../utils';
-import { MenuItemTooltip } from './DisabledMenuItemTooltip';
+import { MenuItemTooltip } from '../DisabledMenuItemTooltip';
-import { DrillByMenuItems } from './DrillBy/DrillByMenuItems';
+import { DrillByMenuItems } from '../DrillBy/DrillByMenuItems';
 
+export enum ContextMenuItem {
+  CrossFilter,
+  DrillToDetail,
+  DrillBy,
+  All,
+}
 export interface ChartContextMenuProps {
   id: number;
   formData: QueryFormData;
   onSelection: () => void;
   onClose: () => void;
+  additionalConfig?: {
+    crossFilter?: Record<string, any>;
+    drillToDetail?: Record<string, any>;
+    drillBy?: Record<string, any>;
+  };
+  displayedItems?: ContextMenuItem[] | ContextMenuItem;
 }
 
-export interface Ref {
+export interface ChartContextMenuRef {
   open: (
     clientX: number,
     clientY: number,
@@ -62,8 +75,15 @@ export interface Ref {
 }
 
 const ChartContextMenu = (
-  { id, formData, onSelection, onClose }: ChartContextMenuProps,
+  {
-  ref: RefObject<Ref>,
+    id,
+    formData,
+    onSelection,
+    onClose,
+    displayedItems = ContextMenuItem.All,
+    additionalConfig,
+  }: ChartContextMenuProps,
+  ref: RefObject<ChartContextMenuRef>,
 ) => {
   const theme = useTheme();
   const dispatch = useDispatch();
@@ -74,6 +94,10 @@ const ChartContextMenu = (
     ({ dashboardInfo }) => dashboardInfo.crossFiltersEnabled,
   );
 
+  const isDisplayed = (item: ContextMenuItem) =>
+    displayedItems === ContextMenuItem.All ||
+    ensureIsArray(displayedItems).includes(item);
+
   const [{ filters, clientX, clientY }, setState] = useState<{
     clientX: number;
     clientY: number;
@@ -83,13 +107,19 @@ const ChartContextMenu = (
   const menuItems = [];
 
   const showDrillToDetail =
-    isFeatureEnabled(FeatureFlag.DRILL_TO_DETAIL) && canExplore;
+    isFeatureEnabled(FeatureFlag.DRILL_TO_DETAIL) &&
+    canExplore &&
+    isDisplayed(ContextMenuItem.DrillToDetail);
 
-  const showDrillBy = isFeatureEnabled(FeatureFlag.DRILL_BY) && canExplore;
+  const showDrillBy =
+    isFeatureEnabled(FeatureFlag.DRILL_BY) &&
+    canExplore &&
+    isDisplayed(ContextMenuItem.DrillBy);
+
+  const showCrossFilters =
+    isFeatureEnabled(FeatureFlag.DASHBOARD_CROSS_FILTERS) &&
+    isDisplayed(ContextMenuItem.CrossFilter);
 
-  const showCrossFilters = isFeatureEnabled(
-    FeatureFlag.DASHBOARD_CROSS_FILTERS,
-  );
   const isCrossFilteringSupportedByChart = getChartMetadataRegistry()
     .get(formData.viz_type)
     ?.behaviors?.includes(Behavior.INTERACTIVE_CHART);
@@ -108,7 +138,7 @@ const ChartContextMenu = (
     itemsCount = 1; // "No actions" appears if no actions in menu
   }
 
-  if (isFeatureEnabled(FeatureFlag.DASHBOARD_CROSS_FILTERS)) {
+  if (showCrossFilters) {
     const isCrossFilterDisabled =
       !isCrossFilteringSupportedByChart ||
       !crossFiltersEnabled ||
@@ -190,6 +220,7 @@ const ChartContextMenu = (
         contextMenuY={clientY}
         onSelection={onSelection}
         submenuIndex={showCrossFilters ? 2 : 1}
+        {...(additionalConfig?.drillToDetail || {})}
       />,
     );
   }
@@ -205,9 +236,11 @@ const ChartContextMenu = (
       <DrillByMenuItems
         filters={filters?.drillBy?.filters}
         groupbyFieldName={filters?.drillBy?.groupbyFieldName}
+        onSelection={onSelection}
         formData={formData}
         contextMenuY={clientY}
         submenuIndex={submenuIndex}
+        {...(additionalConfig?.drillBy || {})}
       />,
     );
   }
@@ -241,7 +274,7 @@ const ChartContextMenu = (
   return ReactDOM.createPortal(
     <Dropdown
       overlay={
-        <Menu>
+        <Menu className="chart-context-menu" data-test="chart-context-menu">
           {menuItems.length ? (
             menuItems
           ) : (

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

@@ -0,0 +1,86 @@
+/**
+ * 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 { FeatureFlag } from '@superset-ui/core';
+import { render, screen } from 'spec/helpers/testing-library';
+import { renderHook } from '@testing-library/react-hooks';
+import mockState from 'spec/fixtures/mockState';
+import { sliceId } from 'spec/fixtures/mockChartQueries';
+import { noOp } from 'src/utils/common';
+import { useContextMenu } from './useContextMenu';
+import { ContextMenuItem } from './ChartContextMenu';
+
+const CONTEXT_MENU_TEST_ID = 'chart-context-menu';
+
+// @ts-ignore
+global.featureFlags = {
+  [FeatureFlag.DASHBOARD_CROSS_FILTERS]: true,
+  [FeatureFlag.DRILL_TO_DETAIL]: true,
+  [FeatureFlag.DRILL_BY]: true,
+};
+
+const setup = ({
+  onSelection = noOp,
+  displayedItems = ContextMenuItem.All,
+  additionalConfig = {},
+}: {
+  onSelection?: () => void;
+  displayedItems?: ContextMenuItem | ContextMenuItem[];
+  additionalConfig?: Record<string, any>;
+} = {}) => {
+  const { result } = renderHook(() =>
+    useContextMenu(
+      sliceId,
+      { datasource: '1__table', viz_type: 'pie' },
+      onSelection,
+      displayedItems,
+      additionalConfig,
+    ),
+  );
+  render(result.current.contextMenu, {
+    useRedux: true,
+    initialState: {
+      ...mockState,
+      user: {
+        ...mockState.user,
+        roles: { Admin: [['can_explore', 'Superset']] },
+      },
+    },
+  });
+  return result;
+};
+
+test('Context menu renders', () => {
+  const result = setup();
+  expect(screen.queryByTestId(CONTEXT_MENU_TEST_ID)).not.toBeInTheDocument();
+  result.current.onContextMenu(0, 0, {});
+  expect(screen.getByTestId(CONTEXT_MENU_TEST_ID)).toBeInTheDocument();
+  expect(screen.getByText('Add cross-filter')).toBeInTheDocument();
+  expect(screen.getByText('Drill to detail')).toBeInTheDocument();
+  expect(screen.getByText('Drill by')).toBeInTheDocument();
+});
+
+test('Context menu contains all items only', () => {
+  const result = setup({
+    displayedItems: [ContextMenuItem.DrillToDetail, ContextMenuItem.DrillBy],
+  });
+  result.current.onContextMenu(0, 0, {});
+  expect(screen.queryByText('Add cross-filter')).not.toBeInTheDocument();
+  expect(screen.getByText('Drill to detail')).toBeInTheDocument();
+  expect(screen.getByText('Drill by')).toBeInTheDocument();
+});

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

@@ -0,0 +1,82 @@
+/**
+ * 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, useMemo, useRef, useState } from 'react';
+import { BaseFormData, ContextMenuFilters } from '@superset-ui/core';
+import ChartContextMenu, {
+  ChartContextMenuRef,
+  ContextMenuItem,
+} from './ChartContextMenu';
+
+export const useContextMenu = (
+  chartId: number,
+  formData: BaseFormData & { [key: string]: any },
+  onSelection?: (...args: any) => void,
+  displayedItems?: ContextMenuItem[] | ContextMenuItem,
+  additionalConfig?: {
+    crossFilter?: Record<string, any>;
+    drillToDetail?: Record<string, any>;
+    drillBy?: Record<string, any>;
+  },
+) => {
+  const contextMenuRef = useRef<ChartContextMenuRef>(null);
+  const [inContextMenu, setInContextMenu] = useState(false);
+  const onContextMenu = (
+    offsetX: number,
+    offsetY: number,
+    filters: ContextMenuFilters,
+  ) => {
+    contextMenuRef.current?.open(offsetX, offsetY, filters);
+    setInContextMenu(true);
+  };
+
+  const handleContextMenuSelected = useCallback(
+    (...args: any) => {
+      setInContextMenu(false);
+      onSelection?.(...args);
+    },
+    [onSelection],
+  );
+
+  const handleContextMenuClosed = useCallback(() => {
+    setInContextMenu(false);
+  }, []);
+
+  const contextMenu = useMemo(
+    () => (
+      <ChartContextMenu
+        ref={contextMenuRef}
+        id={chartId}
+        formData={formData}
+        onSelection={handleContextMenuSelected}
+        onClose={handleContextMenuClosed}
+        displayedItems={displayedItems}
+        additionalConfig={additionalConfig}
+      />
+    ),
+    [
+      additionalConfig,
+      chartId,
+      displayedItems,
+      formData,
+      handleContextMenuClosed,
+      handleContextMenuSelected,
+    ],
+  );
+  return { contextMenu, inContextMenu, onContextMenu };
+};

superset-frontend/src/components/Chart/ChartRenderer.jsx

@@ -31,7 +31,7 @@ import {
 import { Logger, LOG_ACTIONS_RENDER_CHART } from 'src/logger/LogUtils';
 import { EmptyStateBig, EmptyStateSmall } from 'src/components/EmptyState';
 import { ChartSource } from 'src/types/ChartSource';
-import ChartContextMenu from './ChartContextMenu';
+import ChartContextMenu from './ChartContextMenu/ChartContextMenu';
 
 const propTypes = {
   annotationData: PropTypes.object,

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

@@ -19,7 +19,7 @@
 import React from 'react';
 import { render, screen, waitFor } from 'spec/helpers/testing-library';
 import chartQueries, { sliceId } from 'spec/fixtures/mockChartQueries';
-import fetchMock from 'fetch-mock';
+import { noOp } from 'src/utils/common';
 import DrillByChart from './DrillByChart';
 
 const chart = chartQueries[sliceId];
@@ -28,6 +28,8 @@ const setup = (overrides: Record<string, any> = {}, result?: any) =>
   render(
     <DrillByChart
       formData={{ ...chart.form_data, ...overrides }}
+      onContextMenu={noOp}
+      inContextMenu={false}
       result={result}
     />,
     {
@@ -38,8 +40,6 @@ const setup = (overrides: Record<string, any> = {}, result?: any) =>
 const waitForRender = (overrides: Record<string, any> = {}) =>
   waitFor(() => setup(overrides));
 
-afterEach(fetchMock.restore);
-
 test('should render', async () => {
   const { container } = await waitForRender();
   expect(container).toBeInTheDocument();

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

@@ -16,21 +16,34 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React from 'react';
+import React, { useMemo } from 'react';
 import {
   BaseFormData,
-  Behavior,
   QueryData,
   SuperChart,
   css,
+  ContextMenuFilters,
 } from '@superset-ui/core';
 
 interface DrillByChartProps {
   formData: BaseFormData & { [key: string]: any };
   result: QueryData[];
+  onContextMenu: (
+    offsetX: number,
+    offsetY: number,
+    filters: ContextMenuFilters,
+  ) => void;
+  inContextMenu: boolean;
 }
 
-export default function DrillByChart({ formData, result }: DrillByChartProps) {
+export default function DrillByChart({
+  formData,
+  result,
+  onContextMenu,
+  inContextMenu,
+}: DrillByChartProps) {
+  const hooks = useMemo(() => ({ onContextMenu }), [onContextMenu]);
+
   return (
     <div
       css={css`
@@ -40,11 +53,12 @@ export default function DrillByChart({ formData, result }: DrillByChartProps) {
     >
       <SuperChart
         disableErrorBoundary
-        behaviors={[Behavior.INTERACTIVE_CHART]}
         chartType={formData.viz_type}
         enableNoResults
         formData={formData}
         queriesData={result}
+        hooks={hooks}
+        inContextMenu={inContextMenu}
         height="100%"
         width="100%"
       />

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

@@ -32,7 +32,9 @@ import { DrillByMenuItems, DrillByMenuItemsProps } from './DrillByMenuItems';
 
 /* eslint jest/expect-expect: ["warn", { "assertFunctionNames": ["expect*"] }] */
 
-const datasetEndpointMatcher = '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];
 
 const defaultColumns = [
@@ -60,6 +62,7 @@ const defaultFilters = [
 const renderMenu = ({
   formData = defaultFormData,
   filters = defaultFilters,
+  ...rest
 }: Partial<DrillByMenuItemsProps>) =>
   render(
     <Menu>
@@ -67,6 +70,7 @@ const renderMenu = ({
         formData={formData ?? defaultFormData}
         filters={filters}
         groupbyFieldName="groupby"
+        {...rest}
       />
     </Menu>,
     { useRouter: true, useRedux: true },
@@ -134,19 +138,19 @@ test('render disabled menu item for supported chart, no filters', async () => {
 });
 
 test('render disabled menu item for supported chart, no columns', async () => {
-  fetchMock.get(datasetEndpointMatcher, { result: { columns: [] } });
+  fetchMock.get(DATASET_ENDPOINT, { result: { columns: [] } });
   renderMenu({});
-  await waitFor(() => fetchMock.called(datasetEndpointMatcher));
+  await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
   await expectDrillByDisabled('No dimensions available for drill by');
 });
 
 test('render menu item with submenu without searchbox', async () => {
   const slicedColumns = defaultColumns.slice(0, 9);
-  fetchMock.get(datasetEndpointMatcher, {
+  fetchMock.get(DATASET_ENDPOINT, {
     result: { columns: slicedColumns },
   });
   renderMenu({});
-  await waitFor(() => fetchMock.called(datasetEndpointMatcher));
+  await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
   await expectDrillByEnabled();
   slicedColumns.forEach(column => {
     expect(screen.getByText(column.column_name)).toBeInTheDocument();
@@ -155,11 +159,11 @@ test('render menu item with submenu without searchbox', async () => {
 });
 
 test('render menu item with submenu and searchbox', async () => {
-  fetchMock.get(datasetEndpointMatcher, {
+  fetchMock.get(DATASET_ENDPOINT, {
     result: { columns: defaultColumns },
   });
   renderMenu({});
-  await waitFor(() => fetchMock.called(datasetEndpointMatcher));
+  await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
   await expectDrillByEnabled();
   defaultColumns.forEach(column => {
     expect(screen.getByText(column.column_name)).toBeInTheDocument();
@@ -184,3 +188,55 @@ test('render menu item with submenu and searchbox', async () => {
     expect(screen.getByText(colName)).toBeInTheDocument();
   });
 });
+
+test('Do not display excluded column in the menu', async () => {
+  fetchMock.get(DATASET_ENDPOINT, {
+    result: { columns: defaultColumns },
+  });
+
+  const excludedColNames = ['col3', 'col5'];
+  renderMenu({
+    excludedColumns: excludedColNames.map(colName => ({
+      column_name: colName,
+    })),
+  });
+
+  await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
+  await expectDrillByEnabled();
+
+  excludedColNames.forEach(colName => {
+    expect(screen.queryByText(colName)).not.toBeInTheDocument();
+  });
+
+  defaultColumns
+    .filter(column => !excludedColNames.includes(column.column_name))
+    .forEach(column => {
+      expect(screen.getByText(column.column_name)).toBeInTheDocument();
+    });
+});
+
+test('When menu item is clicked, call onSelection with clicked column and drill by filters', async () => {
+  fetchMock
+    .get(DATASET_ENDPOINT, {
+      result: { columns: defaultColumns },
+    })
+    .post(FORM_DATA_KEY_ENDPOINT, {})
+    .post(CHART_DATA_ENDPOINT, {});
+
+  const onSelectionMock = jest.fn();
+  renderMenu({
+    onSelection: onSelectionMock,
+  });
+
+  await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
+  await expectDrillByEnabled();
+
+  userEvent.click(screen.getByText('col1'));
+  expect(onSelectionMock).toHaveBeenCalledWith(
+    {
+      column_name: 'col1',
+      groupby: true,
+    },
+    defaultFilters,
+  );
+});

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

@@ -59,8 +59,10 @@ export interface DrillByMenuItemsProps {
   contextMenuY?: number;
   submenuIndex?: number;
   groupbyFieldName?: string;
-  onSelection?: () => void;
+  onSelection?: (...args: any) => void;
   onClick?: (event: MouseEvent) => void;
+  openNewModal?: boolean;
+  excludedColumns?: Column[];
 }
 
 export const DrillByMenuItems = ({
@@ -71,6 +73,8 @@ export const DrillByMenuItems = ({
   submenuIndex = 0,
   onSelection = () => {},
   onClick = () => {},
+  excludedColumns,
+  openNewModal = true,
   ...rest
 }: DrillByMenuItemsProps) => {
   const theme = useTheme();
@@ -80,14 +84,16 @@ export const DrillByMenuItems = ({
   const [showModal, setShowModal] = useState(false);
   const [currentColumn, setCurrentColumn] = useState();
 
-  const openModal = useCallback(
+  const handleSelection = useCallback(
     (event, column) => {
       onClick(event);
-      onSelection();
+      onSelection(column, filters);
       setCurrentColumn(column);
-      setShowModal(true);
+      if (openNewModal) {
+        setShowModal(true);
+      }
     },
-    [onClick, onSelection],
+    [filters, onClick, onSelection, openNewModal],
   );
   const closeModal = useCallback(() => {
     setShowModal(false);
@@ -142,12 +148,16 @@ export const DrillByMenuItems = ({
 
   const filteredColumns = useMemo(
     () =>
-      columns.filter(column =>
+      columns.filter(
-        (column.verbose_name || column.column_name)
+        column =>
-          .toLowerCase()
+          (column.verbose_name || column.column_name)
-          .includes(searchInput.toLowerCase()),
+            .toLowerCase()
+            .includes(searchInput.toLowerCase()) &&
+          !ensureIsArray(excludedColumns)?.some(
+            col => col.column_name === column.column_name,
+          ),
       ),
-    [columns, searchInput],
+    [columns, excludedColumns, searchInput],
   );
 
   const submenuYOffset = useMemo(
@@ -231,7 +241,7 @@ export const DrillByMenuItems = ({
                   key={`drill-by-item-${column.column_name}`}
                   tooltipText={column.verbose_name || column.column_name}
                   {...rest}
-                  onClick={e => openModal(e, column)}
+                  onClick={e => handleSelection(e, column)}
                 >
                   {column.verbose_name || column.column_name}
                 </MenuItemWithTruncation>

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

@@ -17,7 +17,13 @@
  * under the License.
  */
 
-import React, { useContext, useEffect, useMemo, useState } from 'react';
+import React, {
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useState,
+} from 'react';
 import {
   BaseFormData,
   BinaryQueryObjectFilterClause,
@@ -34,7 +40,7 @@ import Modal from 'src/components/Modal';
 import Loading from 'src/components/Loading';
 import Button from 'src/components/Button';
 import { Radio } from 'src/components/Radio';
-import { DashboardLayout, RootState } from 'src/dashboard/types';
+import { RootState } from 'src/dashboard/types';
 import { DashboardPageIdContext } from 'src/dashboard/containers/DashboardPage';
 import { postFormData } from 'src/explore/exploreUtils/formData';
 import { noOp } from 'src/utils/common';
@@ -43,6 +49,8 @@ import { useDatasetMetadataBar } from 'src/features/datasets/metadataBar/useData
 import { SingleQueryResultPane } from 'src/explore/components/DataTablesPane/components/SingleQueryResultPane';
 import { Dataset, DrillByType } from '../types';
 import DrillByChart from './DrillByChart';
+import { ContextMenuItem } from '../ChartContextMenu/ChartContextMenu';
+import { useContextMenu } from '../ChartContextMenu/useContextMenu';
 import { getChartDataRequest } from '../chartAction';
 
 const DATA_SIZE = 15;
@@ -119,31 +127,35 @@ export default function DrillByModal({
     () => formData.datasource.split('__'),
     [formData.datasource],
   );
-  const dashboardLayout = useSelector<RootState, DashboardLayout>(
+
-    state => state.dashboardLayout.present,
+  const [currentColumn, setCurrentColumn] = useState(column);
-  );
+  const [currentFormData, setCurrentFormData] = useState(formData);
-  const chartLayoutItem = Object.values(dashboardLayout).find(
+  const [currentFilters, setCurrentFilters] = useState(filters);
-    layoutItem => layoutItem.meta?.chartId === formData.slice_id,
+  const [usedGroupbyColumns, setUsedGroupbyColumns] = useState([
-  );
+    ...ensureIsArray(formData[groupbyFieldName]).map(colName =>
-  const chartName =
+      dataset.columns?.find(col => col.column_name === colName),
-    chartLayoutItem?.meta.sliceNameOverride || chartLayoutItem?.meta.sliceName;
+    ),
+    column,
+  ]);
 
   const updatedFormData = useMemo(() => {
-    let updatedFormData = { ...formData };
+    let updatedFormData = { ...currentFormData };
-    if (column) {
+    if (currentColumn) {
       updatedFormData[groupbyFieldName] = Array.isArray(
-        formData[groupbyFieldName],
+        currentFormData[groupbyFieldName],
       )
-        ? [column.column_name]
+        ? [currentColumn.column_name]
-        : column.column_name;
+        : currentColumn.column_name;
     }
 
-    if (filters) {
+    if (currentFilters) {
-      const adhocFilters = filters.map(filter => simpleFilterToAdhoc(filter));
+      const adhocFilters = currentFilters.map(filter =>
+        simpleFilterToAdhoc(filter),
+      );
       updatedFormData = {
         ...updatedFormData,
         adhoc_filters: [
-          ...ensureIsArray(formData.adhoc_filters),
+          ...ensureIsArray(currentFormData.adhoc_filters),
           ...adhocFilters,
         ],
       };
@@ -152,7 +164,46 @@ export default function DrillByModal({
     delete updatedFormData.slice_name;
     delete updatedFormData.dashboards;
     return updatedFormData;
-  }, [column, filters, formData, groupbyFieldName]);
+  }, [currentColumn, currentFormData, currentFilters, groupbyFieldName]);
+
+  useEffect(() => {
+    setUsedGroupbyColumns(cols =>
+      cols.includes(currentColumn) ? cols : [...cols, currentColumn],
+    );
+  }, [currentColumn]);
+
+  const onSelection = useCallback(
+    (newColumn: Column, filters: BinaryQueryObjectFilterClause[]) => {
+      setCurrentColumn(newColumn);
+      setCurrentFormData(updatedFormData);
+      setCurrentFilters(filters);
+    },
+    [updatedFormData],
+  );
+
+  const additionalConfig = useMemo(
+    () => ({
+      drillBy: { excludedColumns: usedGroupbyColumns, openNewModal: false },
+    }),
+    [usedGroupbyColumns],
+  );
+
+  const { contextMenu, inContextMenu, onContextMenu } = useContextMenu(
+    0,
+    currentFormData,
+    onSelection,
+    ContextMenuItem.DrillBy,
+    additionalConfig,
+  );
+
+  const chartName = useSelector<RootState, string | undefined>(state => {
+    const chartLayoutItem = Object.values(state.dashboardLayout.present).find(
+      layoutItem => layoutItem.meta?.chartId === formData.slice_id,
+    );
+    return (
+      chartLayoutItem?.meta.sliceNameOverride || chartLayoutItem?.meta.sliceName
+    );
+  });
 
   useEffect(() => {
     if (updatedFormData) {
@@ -228,7 +279,12 @@ export default function DrillByModal({
         </div>
         {!chartDataResult && <Loading />}
         {drillByDisplayMode === DrillByType.Chart && chartDataResult && (
-          <DrillByChart formData={updatedFormData} result={chartDataResult} />
+          <DrillByChart
+            formData={updatedFormData}
+            result={chartDataResult}
+            onContextMenu={onContextMenu}
+            inContextMenu={inContextMenu}
+          />
         )}
         {drillByDisplayMode === DrillByType.Table && chartDataResult && (
           <div
@@ -248,6 +304,7 @@ export default function DrillByModal({
             />
           </div>
         )}
+        {contextMenu}
       </div>
     </Modal>
   );

superset-frontend/src/components/Chart/types.ts

@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+import { Column } from '@superset-ui/core';
 
 export enum DrillByType {
   Chart,
@@ -39,4 +40,5 @@ export type Dataset = {
     first_name: string;
     last_name: string;
   }[];
+  columns?: Column[];
 };

superset-frontend/src/dashboard/styles.ts

@@ -89,6 +89,9 @@ export const filterCardPopoverStyle = (theme: SupersetTheme) => css`
 `;
 
 export const chartContextMenuStyles = (theme: SupersetTheme) => css`
+  .ant-dropdown-menu.chart-context-menu {
+    min-width: ${theme.gridUnit * 43}px;
+  }
   .ant-dropdown-menu-submenu.chart-context-submenu {
     max-width: ${theme.gridUnit * 60}px;
     min-width: ${theme.gridUnit * 40}px;