From 40568fd1ff3d17858fbbe9aac3ad7107bc351081 Mon Sep 17 00:00:00 2001 From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com> Date: Tue, 10 Sep 2024 14:27:19 -0300 Subject: [PATCH] fix(Dashboard download): Download dashboard screenshot/PDF using SupersetClient (#30212) (cherry picked from commit d191e67e519cf3867a45f833adec34e8e6ad276c) --- .../DownloadScreenshot.test.tsx | 47 ++++++++++++------- .../DownloadMenuItems/DownloadScreenshot.tsx | 41 +++++++++------- 2 files changed, 54 insertions(+), 34 deletions(-) diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx index 21bd7022eb362..01d756403a463 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx +++ b/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 diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx index 361f9880fec12..2d182e31fc754 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx +++ b/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,7 +122,7 @@ 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, @@ -121,8 +130,8 @@ export default function DownloadScreenshot({ }, }) .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( @@ -130,7 +139,7 @@ export default function DownloadScreenshot({ 'The screenshot is being generated. Please, do not leave the page.', ), ); - fetchImageWithRetry(imageUrl); + fetchImageWithRetry(cacheKey); }) .catch(error => { logging.error(error);