From 339d491dfc81d28d0caf175f215e58c1383066f0 Mon Sep 17 00:00:00 2001
From: "JUST.in DO IT" <justin.park@airbnb.com>
Date: Mon, 2 Dec 2024 11:15:25 -0800
Subject: [PATCH] feat(sqllab): Popup notification when download data can
 exceed row count (#31187)

Co-authored-by: Michael S. Molina <70410625+michael-s-molina@users.noreply.github.com>
---
 .../components/ResultSet/ResultSet.test.tsx   | 40 ++++++++++++++++++-
 .../src/SqlLab/components/ResultSet/index.tsx | 20 +++++++++-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/ResultSet/ResultSet.test.tsx b/superset-frontend/src/SqlLab/components/ResultSet/ResultSet.test.tsx
index bce0b95bb8cc2..31412afd22c42 100644
--- a/superset-frontend/src/SqlLab/components/ResultSet/ResultSet.test.tsx
+++ b/superset-frontend/src/SqlLab/components/ResultSet/ResultSet.test.tsx
@@ -16,7 +16,13 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { render, screen, waitFor } from 'spec/helpers/testing-library';
+import {
+  render,
+  screen,
+  waitFor,
+  fireEvent,
+  within,
+} from 'spec/helpers/testing-library';
 import configureStore from 'redux-mock-store';
 import { Store } from 'redux';
 import thunk from 'redux-thunk';
@@ -492,6 +498,38 @@ describe('ResultSet', () => {
     expect(queryByTestId('export-csv-button')).toBeInTheDocument();
   });
 
+  test('should display a popup message when the CSV content is limited to the dropdown limit', async () => {
+    const queryLimit = 2;
+    const { getByTestId, findByRole } = setup(
+      mockedProps,
+      mockStore({
+        ...initialState,
+        user: {
+          ...user,
+          roles: {
+            sql_lab: [['can_export_csv', 'SQLLab']],
+          },
+        },
+        sqlLab: {
+          ...initialState.sqlLab,
+          queries: {
+            [queries[0].id]: {
+              ...queries[0],
+              limitingFactor: 'DROPDOWN',
+              queryLimit,
+            },
+          },
+        },
+      }),
+    );
+    const downloadButton = getByTestId('export-csv-button');
+    fireEvent.click(downloadButton);
+    const warningModal = await findByRole('dialog');
+    expect(
+      within(warningModal).getByText(`Download is on the way`),
+    ).toBeInTheDocument();
+  });
+
   test('should not allow download as CSV when user does not have permission to export data', async () => {
     const { queryByTestId } = setup(
       mockedProps,
diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx
index 7d0df189a2f21..f6d73a4072caa 100644
--- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx
+++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx
@@ -64,6 +64,7 @@ import CopyToClipboard from 'src/components/CopyToClipboard';
 import { addDangerToast } from 'src/components/MessageToasts/actions';
 import { prepareCopyToClipboardTabularData } from 'src/utils/common';
 import { getItem, LocalStorageKeys } from 'src/utils/localStorageHelpers';
+import Modal from 'src/components/Modal';
 import {
   addQueryEditor,
   clearQueryResults,
@@ -296,6 +297,9 @@ const ResultSet = ({
 
   const renderControls = () => {
     if (search || visualize || csv) {
+      const { results, queryLimit, limitingFactor, rows } = query;
+      const limit = queryLimit || results.query.limit;
+      const rowsCount = Math.min(rows || 0, results?.data?.length || 0);
       let { data } = query.results;
       if (cache && query.cached) {
         data = cachedData;
@@ -342,7 +346,21 @@ const ResultSet = ({
                 buttonSize="small"
                 href={getExportCsvUrl(query.id)}
                 data-test="export-csv-button"
-                onClick={() => logAction(LOG_ACTIONS_SQLLAB_DOWNLOAD_CSV, {})}
+                onClick={() => {
+                  logAction(LOG_ACTIONS_SQLLAB_DOWNLOAD_CSV, {});
+                  if (
+                    limitingFactor === LimitingFactor.Dropdown &&
+                    limit === rowsCount
+                  ) {
+                    Modal.warning({
+                      title: t('Download is on the way'),
+                      content: t(
+                        'Downloading %(rows)s rows based on the LIMIT configuration. If you want the entire result set, you need to adjust the LIMIT.',
+                        { rows: rowsCount.toLocaleString() },
+                      ),
+                    });
+                  }
+                }}
               >
                 <i className="fa fa-file-text-o" /> {t('Download to CSV')}
               </Button>