superset-frontend/src/dashboard/hooks/useDownloadScreenshot.ts

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
import { useCallback, useEffect, useRef } from 'react';
import { useSelector } from 'react-redux';
import { useToasts } from 'src/components/MessageToasts/withToasts';
import { last } from 'lodash';
import contentDisposition from 'content-disposition';
import {
  logging,
  t,
  SupersetClient,
  SupersetApiError,
} from '@superset-ui/core';
import {
  LOG_ACTIONS_DASHBOARD_DOWNLOAD_AS_IMAGE,
  LOG_ACTIONS_DASHBOARD_DOWNLOAD_AS_PDF,
} from 'src/logger/LogUtils';
import { RootState } from 'src/dashboard/types';
import { getDashboardUrlParams } from 'src/utils/urlUtils';
import { DownloadScreenshotFormat } from '../components/menu/DownloadMenuItems/types';

const RETRY_INTERVAL = 3000;
const MAX_RETRIES = 30;

export const useDownloadScreenshot = (
  dashboardId: number,
  logEvent?: Function,
) => {
  const activeTabs = useSelector(
    (state: RootState) => state.dashboardState.activeTabs || undefined,
  );
  const anchor = useSelector(
    (state: RootState) =>
      last(state.dashboardState.directPathToChild) || undefined,
  );
  const dataMask = useSelector(
    (state: RootState) => state.dataMask || undefined,
  );

  const { addDangerToast, addSuccessToast, addInfoToast } = useToasts();

  const currentIntervalIds = useRef<NodeJS.Timeout[]>([]);

  const stopIntervals = useCallback(
    (message?: 'success' | 'failure') => {
      currentIntervalIds.current.forEach(clearInterval);

      if (message === 'failure') {
        addDangerToast(
          t('The screenshot could not be downloaded. Please, try again later.'),
        );
      }
      if (message === 'success') {
        addSuccessToast(t('The screenshot has been downloaded.'));
      }
    },
    [addDangerToast, addSuccessToast],
  );

  const downloadScreenshot = useCallback(
    (format: DownloadScreenshotFormat) => {
      let retries = 0;

      const toastIntervalId = setInterval(
        () =>
          addInfoToast(
            t(
              'The screenshot is being generated. Please, do not leave the page.',
            ),
            { noDuplicate: true },
          ),
        RETRY_INTERVAL,
      );

      currentIntervalIds.current = [
        ...(currentIntervalIds.current || []),
        toastIntervalId,
      ];

      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) => {
            const disposition = response.headers.get('Content-Disposition');
            let fileName = `screenshot.${format}`; // default filename

            if (disposition) {
              try {
                const parsed = contentDisposition.parse(disposition);
                if (parsed?.parameters?.filename) {
                  fileName = parsed.parameters.filename;
                }
              } catch (error) {
                console.warn(
                  'Failed to parse Content-Disposition header:',
                  error,
                );
              }
            }

            return response.blob().then(blob => ({ blob, fileName }));
          })
          .then(({ blob, fileName }) => {
            const url = window.URL.createObjectURL(blob);
            const a = document.createElement('a');
            a.href = url;
            a.download = fileName;
            document.body.appendChild(a);
            a.click();
            document.body.removeChild(a);
            window.URL.revokeObjectURL(url);
            stopIntervals('success');
          })
          .catch(err => {
            if ((err as SupersetApiError).status === 404) {
              throw new Error('Image not ready');
            }
          });

      const fetchImageWithRetry = (cacheKey: string) => {
        if (retries >= MAX_RETRIES) {
          stopIntervals('failure');
          logging.error('Max retries reached');
          return;
        }
        checkImageReady(cacheKey).catch(() => {
          retries += 1;
        });
      };

      SupersetClient.post({
        endpoint: `/api/v1/dashboard/${dashboardId}/cache_dashboard_screenshot/`,
        jsonPayload: {
          anchor,
          activeTabs,
          dataMask,
          urlParams: getDashboardUrlParams(),
        },
      })
        .then(({ json }) => {
          const cacheKey = json?.cache_key;
          if (!cacheKey) {
            throw new Error('No image URL in response');
          }
          const retryIntervalId = setInterval(() => {
            fetchImageWithRetry(cacheKey);
          }, RETRY_INTERVAL);
          currentIntervalIds.current.push(retryIntervalId);
          fetchImageWithRetry(cacheKey);
        })
        .catch(error => {
          logging.error(error);
          stopIntervals('failure');
        })
        .finally(() => {
          logEvent?.(
            format === DownloadScreenshotFormat.PNG
              ? LOG_ACTIONS_DASHBOARD_DOWNLOAD_AS_IMAGE
              : LOG_ACTIONS_DASHBOARD_DOWNLOAD_AS_PDF,
          );
        });
    },
    [
      dashboardId,
      anchor,
      activeTabs,
      dataMask,
      addInfoToast,
      stopIntervals,
      logEvent,
    ],
  );

  useEffect(
    () => () => {
      if (currentIntervalIds.current.length > 0) {
        stopIntervals();
      }
      currentIntervalIds.current = [];
    },
    [stopIntervals],
  );

  return downloadScreenshot;
};