feat(streaming): Streaming CSV uploads for over 100k records for constant memory usage (#35478)

superset-frontend/src/SqlLab/components/ResultSet/index.tsx

@@ -83,6 +83,8 @@ import {
 } from 'src/logger/LogUtils';
 import { Icons } from '@superset-ui/core/components/Icons';
 import { findPermission } from 'src/utils/findPermission';
+import { StreamingExportModal } from 'src/components/StreamingExportModal';
+import { useStreamingExport } from 'src/components/StreamingExportModal/useStreamingExport';
 import { ensureAppRoot } from 'src/utils/pathUtils';
 import { useConfirmModal } from 'src/hooks/useConfirmModal';
 import ExploreCtasResultsButton from '../ExploreCtasResultsButton';
@@ -184,6 +186,10 @@ const ResultSet = ({
   defaultQueryLimit,
 }: ResultSetProps) => {
   const user = useSelector(({ user }: SqlLabRootState) => user, shallowEqual);
+  const streamingThreshold = useSelector(
+    (state: SqlLabRootState) =>
+      state.common?.conf?.CSV_STREAMING_ROW_THRESHOLD || 1000,
+  );
   const query = useSelector(
     ({ sqlLab: { queries } }: SqlLabRootState) =>
       pick(queries[queryId], [
@@ -224,12 +230,21 @@ const ResultSet = ({
   const [searchText, setSearchText] = useState('');
   const [cachedData, setCachedData] = useState<Record<string, unknown>[]>([]);
   const [showSaveDatasetModal, setShowSaveDatasetModal] = useState(false);
+  const [showStreamingModal, setShowStreamingModal] = useState(false);
 
   const history = useHistory();
   const dispatch = useDispatch();
   const logAction = useLogAction({ queryId, sqlEditorId: query.sqlEditorId });
   const { showConfirm, ConfirmModal } = useConfirmModal();
 
+  const { progress, startExport, resetExport, retryExport, cancelExport } =
+    useStreamingExport({
+      onComplete: () => {},
+      onError: error => {
+        addDangerToast(t('Export failed: %s', error));
+      },
+    });
+
   const reRunQueryIfSessionTimeoutErrorOnMount = useCallback(() => {
     if (
       query.errorMessage &&
@@ -302,6 +317,28 @@ const ResultSet = ({
   const getExportCsvUrl = (clientId: string) =>
     ensureAppRoot(`/api/v1/sqllab/export/${clientId}/`);
 
+  const handleCloseStreamingModal = () => {
+    cancelExport();
+    setShowStreamingModal(false);
+    resetExport();
+  };
+
+  const shouldUseStreamingExport = () => {
+    const { rows, queryLimit, limitingFactor } = query;
+    const limit = queryLimit || query.results?.query?.limit;
+    const rowsCount = Math.min(rows || 0, query.results?.data?.length || 0);
+
+    let actualRowCount = rowsCount;
+
+    if (limitingFactor === LimitingFactor.NotLimited && rows) {
+      actualRowCount = rows;
+    } else if (limit) {
+      actualRowCount = Math.max(actualRowCount, limit);
+    }
+
+    return actualRowCount >= streamingThreshold;
+  };
+
   const renderControls = () => {
     if (search || visualize || csv) {
       const { limitingFactor, queryLimit, results, rows } = query;
@@ -372,9 +409,27 @@ const ResultSet = ({
               <CopyStyledButton
                 buttonSize="small"
                 buttonStyle="secondary"
-                href={getExportCsvUrl(query.id)}
+                {...(!shouldUseStreamingExport() && {
+                  href: getExportCsvUrl(query.id),
+                })}
                 data-test="export-csv-button"
-                onClick={handleDownloadCsv}
+                onClick={e => {
+                  const useStreaming = shouldUseStreamingExport();
+
+                  if (useStreaming) {
+                    e.preventDefault();
+                    setShowStreamingModal(true);
+
+                    startExport({
+                      url: '/api/v1/sqllab/export_streaming/',
+                      payload: { client_id: query.id },
+                      exportType: 'csv',
+                      expectedRows: rows,
+                    });
+                  } else {
+                    handleDownloadCsv(e);
+                  }
+                }}
               >
                 <Icons.DownloadOutlined iconSize="m" /> {t('Download to CSV')}
               </CopyStyledButton>
@@ -723,43 +778,75 @@ const ResultSet = ({
               </AutoSizer>
             </div>
           </ResultContainer>
+          <StreamingExportModal
+            visible={showStreamingModal}
+            onCancel={handleCloseStreamingModal}
+            onRetry={retryExport}
+            progress={progress}
+          />
           {ConfirmModal}
         </>
       );
     }
     if (data && data.length === 0) {
-      return <Alert type="warning" message={t('The query returned no data')} />;
+      return (
+        <>
+          <Alert type="warning" message={t('The query returned no data')} />
+          <StreamingExportModal
+            visible={showStreamingModal}
+            onCancel={handleCloseStreamingModal}
+            onRetry={retryExport}
+            progress={progress}
+          />
+        </>
+      );
     }
   }
 
   if (query.cached || (query.state === QueryState.Success && !query.results)) {
     if (query.isDataPreview) {
       return (
-        <Button
-          buttonSize="small"
-          buttonStyle="primary"
-          onClick={() =>
-            dispatch(
-              reFetchQueryResults({
-                ...query,
-                isDataPreview: true,
-              }),
-            )
-          }
-        >
-          {t('Fetch data preview')}
-        </Button>
+        <>
+          <Button
+            buttonSize="small"
+            buttonStyle="primary"
+            onClick={() =>
+              dispatch(
+                reFetchQueryResults({
+                  ...query,
+                  isDataPreview: true,
+                }),
+              )
+            }
+          >
+            {t('Fetch data preview')}
+          </Button>
+          <StreamingExportModal
+            visible={showStreamingModal}
+            onCancel={handleCloseStreamingModal}
+            onRetry={retryExport}
+            progress={progress}
+          />
+        </>
       );
     }
     if (query.resultsKey) {
       return (
-        <Button
-          buttonSize="small"
-          buttonStyle="primary"
-          onClick={() => fetchResults(query)}
-        >
-          {t('Refetch results')}
-        </Button>
+        <>
+          <Button
+            buttonSize="small"
+            buttonStyle="primary"
+            onClick={() => fetchResults(query)}
+          >
+            {t('Refetch results')}
+          </Button>
+          <StreamingExportModal
+            visible={showStreamingModal}
+            onCancel={handleCloseStreamingModal}
+            onRetry={retryExport}
+            progress={progress}
+          />
+        </>
       );
     }
   }
@@ -774,15 +861,24 @@ const ResultSet = ({
   const progressMsg = query?.extra?.progress ?? null;
 
   return (
-    <ResultlessStyles>
-      <div>{!progressBar && <Loading position="normal" />}</div>
-      {/* show loading bar whenever progress bar is completed but needs time to render */}
-      <div>{query.progress === 100 && <Loading position="normal" />}</div>
-      <QueryStateLabel query={query} />
-      <div>{progressMsg && <Alert type="success" message={progressMsg} />}</div>
-      <div>{query.progress !== 100 && progressBar}</div>
-      {trackingUrl && <div>{trackingUrl}</div>}
-    </ResultlessStyles>
+    <>
+      <ResultlessStyles>
+        <div>{!progressBar && <Loading position="normal" />}</div>
+        {/* show loading bar whenever progress bar is completed but needs time to render */}
+        <div>{query.progress === 100 && <Loading position="normal" />}</div>
+        <QueryStateLabel query={query} />
+        <div>
+          {progressMsg && <Alert type="success" message={progressMsg} />}
+        </div>
+        <div>{query.progress !== 100 && progressBar}</div>
+        {trackingUrl && <div>{trackingUrl}</div>}
+      </ResultlessStyles>
+      <StreamingExportModal
+        visible={showStreamingModal}
+        onCancel={handleCloseStreamingModal}
+        progress={progress}
+      />
+    </>
   );
 };