fix(Dashboard download): Download dashboard screenshot/PDF using Supe…

…rsetClient (apache#30212)

(cherry picked from commit d191e67)

superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx

@@ -84,7 +84,7 @@ describe('DownloadScreenshot component', () => {
     const props = defaultProps();
 
     fetchMock.post(
-      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot`,
+      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot/`,
       {
         status: 400,
         body: {},
@@ -105,19 +105,23 @@ describe('DownloadScreenshot component', () => {
   test('displays success message when API call succeeds', async () => {
     const props = defaultProps();
     fetchMock.post(
-      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot`,
+      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot/`,
       {
         status: 200,
         body: {
           image_url: 'mocked_image_url',
+          cache_key: 'mocked_cache_key',
         },
       },
     );
 
-    fetchMock.get('glob:*/mocked_image_url?download_format=pdf', {
-      status: 200,
-      body: {},
-    });
+    fetchMock.get(
+      `glob:*/api/v1/dashboard/${props.dashboardId}/screenshot/mocked_cache_key/?download_format=pdf`,
+      {
+        status: 200,
+        body: {},
+      },
+    );
 
     renderComponent();
 
@@ -130,14 +134,14 @@ describe('DownloadScreenshot component', () => {
     });
   });
 
-  test('throws error when no image URL is provided', async () => {
+  test('throws error when no image cache key is provided', async () => {
     const props = defaultProps();
     fetchMock.post(
-      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot`,
+      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot/`,
       {
         status: 200,
         body: {
-          image_url: '',
+          cache_key: '',
         },
       },
     );
@@ -156,24 +160,27 @@ describe('DownloadScreenshot component', () => {
 
   test('displays success message when image retrieval succeeds', async () => {
     const props = defaultProps();
-    const imageUrl = 'glob:*/mocked_image_url?download_format=pdf';
     fetchMock.post(
-      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot`,
+      `glob:*/api/v1/dashboard/${props.dashboardId}/cache_dashboard_screenshot/`,
       {
         status: 200,
         body: {
           image_url: 'mocked_image_url',
+          cache_key: 'mocked_cache_key',
         },
       },
     );
 
-    fetchMock.get(imageUrl, {
-      status: 200,
-      headers: {
-        'Content-Type': 'image/png',
+    fetchMock.get(
+      `glob:*/api/v1/dashboard/${props.dashboardId}/screenshot/mocked_cache_key/?download_format=pdf`,
+      {
+        status: 200,
+        headers: {
+          'Content-Type': 'application/pdf',
+        },
+        body: new Blob([], { type: 'application/pdf' }),
       },
-      body: new Blob([], { type: 'image/png' }),
-    });
+    );
 
     global.URL.createObjectURL = jest.fn(() => 'mockedObjectURL');
     global.URL.revokeObjectURL = jest.fn();
@@ -185,7 +192,11 @@ describe('DownloadScreenshot component', () => {
     userEvent.click(screen.getByRole('button', { name: 'Download' }));
 
     await waitFor(() => {
-      expect(fetchMock.calls(imageUrl).length).toBe(1);
+      expect(
+        fetchMock.calls(
+          `glob:*/api/v1/dashboard/${props.dashboardId}/screenshot/mocked_cache_key/?download_format=pdf`,
+        ).length,
+      ).toBe(1);
     });
 
     // Wait for the successful image retrieval message

superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx

@@ -17,7 +17,12 @@
  * under the License.
  */
 
-import { logging, t, SupersetClient } from '@superset-ui/core';
+import {
+  logging,
+  t,
+  SupersetClient,
+  SupersetApiError,
+} from '@superset-ui/core';
 import { Menu } from 'src/components/Menu';
 import {
   LOG_ACTIONS_DASHBOARD_DOWNLOAD_AS_IMAGE,
@@ -63,14 +68,13 @@ export default function DownloadScreenshot({
     let retries = 0;
 
     // this function checks if the image is ready
-    const checkImageReady = (imageUrl: string) =>
-      fetch(`${imageUrl}?download_format=${format}`)
-        .then(response => {
-          if (response.status === 404) {
-            throw new Error('Image not ready');
-          }
-          return response.blob();
-        })
+    const checkImageReady = (cacheKey: string) =>
+      SupersetClient.get({
+        endpoint: `/api/v1/dashboard/${dashboardId}/screenshot/${cacheKey}/?download_format=${format}`,
+        headers: { Accept: 'application/pdf, image/png' },
+        parseMethod: 'raw',
+      })
+        .then((response: Response) => response.blob())
         .then(blob => {
           const url = window.URL.createObjectURL(blob);
           const a = document.createElement('a');
@@ -80,11 +84,16 @@ export default function DownloadScreenshot({
           a.click();
           document.body.removeChild(a);
           window.URL.revokeObjectURL(url);
+        })
+        .catch(err => {
+          if ((err as SupersetApiError).status === 404) {
+            throw new Error('Image not ready');
+          }
         });
 
     // this is the functions that handles the retries
-    const fetchImageWithRetry = (imageUrl: string) => {
-      checkImageReady(imageUrl)
+    const fetchImageWithRetry = (cacheKey: string) => {
+      checkImageReady(cacheKey)
         .then(() => {
           addSuccessToast(t('The screenshot is now being downloaded.'));
         })
@@ -100,7 +109,7 @@ export default function DownloadScreenshot({
                 noDuplicate: true,
               },
             );
-            setTimeout(() => fetchImageWithRetry(imageUrl), RETRY_INTERVAL);
+            setTimeout(() => fetchImageWithRetry(cacheKey), RETRY_INTERVAL);
           } else {
             addDangerToast(
               t(
@@ -113,24 +122,24 @@ export default function DownloadScreenshot({
     };
 
     SupersetClient.post({
-      endpoint: `/api/v1/dashboard/${dashboardId}/cache_dashboard_screenshot`,
+      endpoint: `/api/v1/dashboard/${dashboardId}/cache_dashboard_screenshot/`,
       jsonPayload: {
         anchor,
         activeTabs,
         dataMask,
       },
     })
       .then(({ json }) => {
-        const imageUrl = json?.image_url;
-        if (!imageUrl) {
+        const cacheKey = json?.cache_key;
+        if (!cacheKey) {
           throw new Error('No image URL in response');
         }
         addInfoToast(
           t(
             'The screenshot is being generated. Please, do not leave the page.',
           ),
         );
-        fetchImageWithRetry(imageUrl);
+        fetchImageWithRetry(cacheKey);
       })
       .catch(error => {
         logging.error(error);