From 109f51bbbda8f481959e44262b6320c682338dd2 Mon Sep 17 00:00:00 2001
From: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
Date: Wed, 26 Apr 2023 19:36:21 +0200
Subject: [PATCH] fix: Loading state when cols for drill by are loading
 (#23830)

---
 .../Chart/DrillBy/DrillByMenuItems.test.tsx   |  3 ++-
 .../Chart/DrillBy/DrillByMenuItems.tsx        | 19 +++++++++++++++----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.test.tsx b/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.test.tsx
index 7d411e050..485e19652 100644
--- a/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.test.tsx
+++ b/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.test.tsx
@@ -140,7 +140,8 @@ test('render disabled menu item for supported chart, no columns', async () => {
   fetchMock.get(DATASET_ENDPOINT, { result: { columns: [] } });
   renderMenu({});
   await waitFor(() => fetchMock.called(DATASET_ENDPOINT));
-  await expectDrillByDisabled('No dimensions available for drill by');
+  await expectDrillByEnabled();
+  screen.getByText('No columns found');
 });
 
 test('render menu item with submenu without searchbox', async () => {
diff --git a/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.tsx b/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.tsx
index a19960279..3354241ef 100644
--- a/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.tsx
+++ b/superset-frontend/src/components/Chart/DrillBy/DrillByMenuItems.tsx
@@ -40,6 +40,7 @@ import {
 import Icons from 'src/components/Icons';
 import { Input } from 'src/components/Input';
 import { useToasts } from 'src/components/MessageToasts/withToasts';
+import Loading from 'src/components/Loading';
 import {
   cachedSupersetGet,
   supersetGetCache,
@@ -79,6 +80,7 @@ export const DrillByMenuItems = ({
 }: DrillByMenuItemsProps) => {
   const theme = useTheme();
   const { addDangerToast } = useToasts();
+  const [isLoadingColumns, setIsLoadingColumns] = useState(true);
   const [searchInput, setSearchInput] = useState('');
   const [dataset, setDataset] = useState<Dataset>();
   const [columns, setColumns] = useState<Column[]>([]);
@@ -145,6 +147,9 @@ export const DrillByMenuItems = ({
         .catch(() => {
           supersetGetCache.delete(`/api/v1/dataset/${datasetId}`);
           addDangerToast(t('Failed to load dimensions for drill by'));
+        })
+        .finally(() => {
+          setIsLoadingColumns(false);
         });
     }
   }, [
@@ -192,11 +197,9 @@ export const DrillByMenuItems = ({
     tooltip = t('Drill by is not yet supported for this chart type');
   } else if (!hasDrillBy) {
     tooltip = t('Drill by is not available for this data point');
-  } else if (columns.length === 0) {
-    tooltip = t('No dimensions available for drill by');
   }
 
-  if (!handlesDimensionContextMenu || !hasDrillBy || columns.length === 0) {
+  if (!handlesDimensionContextMenu || !hasDrillBy) {
     return (
       <Menu.Item key="drill-by-disabled" disabled {...rest}>
         <div>
@@ -241,7 +244,15 @@ export const DrillByMenuItems = ({
               `}
             />
           )}
-          {filteredColumns.length ? (
+          {isLoadingColumns ? (
+            <div
+              css={css`
+                padding: ${theme.gridUnit * 3}px 0;
+              `}
+            >
+              <Loading position="inline-centered" />
+            </div>
+          ) : filteredColumns.length ? (
             <div
               css={css`
                 max-height: ${MAX_SUBMENU_HEIGHT}px;