Line data    Source code

       1              : import * as React from 'react';
       2              : import type { TopLevelSpec } from 'vega-lite';
       3              : 
       4              : import { QueryExecutionStatus, type QueryExecutionState } from '../../query/query_execution_state.js';
       5              : import { ButtonSize, ButtonVariant, IconButton } from '../../ui/foundations/button.js';
       6              : import { DownloadIcon } from '../../ui/foundations/symbol_icon.js';
       7              : import { arrowTableToRows } from '../../ui/vegalite/vegalite_rows.js';
       8              : import { renderVegaLitePng } from '../../ui/vegalite/vegalite_image.js';
       9              : import { useFileDownloader } from '../../platform/file/file_downloader_hook.js';
      10              : import { useLogger } from '../../platform/logger/logger_provider.js';
      11              : import { stringifyError } from '../../platform/logger/logger.js';
      12              : 
      13              : interface Props {
      14              :     query: QueryExecutionState;
      15              :     /// The Vega-Lite spec of the chart, the query result is its data
      16              :     spec: TopLevelSpec;
      17              :     dependencies?: {
      18              :         useFileDownloader: () => Pick<ReturnType<typeof useFileDownloader>, 'downloadBufferAsFile'>;
      19              :         useLogger: () => Pick<ReturnType<typeof useLogger>, 'error'>;
      20              :         renderVegaLitePng: typeof renderVegaLitePng;
      21              :     };
      22              : }
      23              : 
      24            1 : const DEFAULT_DEPENDENCIES: NonNullable<Props['dependencies']> = {
      25            1 :     useFileDownloader,
      26            1 :     useLogger,
      27            1 :     renderVegaLitePng,
      28            1 : };
      29              : 
      30              : /// Download the chart of a query result as PNG
      31            1 : export function ChartDownloadButton({ query, spec, dependencies = DEFAULT_DEPENDENCIES }: Props) {
      32            1 :     const fileDownloader = dependencies.useFileDownloader();
      33            1 :     const logger = dependencies.useLogger();
      34            1 :     const [downloading, setDownloading] = React.useState(false);
      35            1 :     const hasResult = query.status === QueryExecutionStatus.SUCCEEDED && query.resultTable != null;
      36              : 
      37            1 :     const download = async () => {
      38            1 :         if (!hasResult || downloading) return;
      39            1 :         setDownloading(true);
      40            1 :         try {
      41            1 :             const image = await dependencies.renderVegaLitePng(spec, arrowTableToRows(query.resultTable!));
      42            1 :             await fileDownloader.downloadBufferAsFile(image.png, `chart-${query.queryId}.png`);
      43            1 :         } catch (error) {
      44            0 :             logger.error('Failed to download chart', { error: stringifyError(error) }, 'chart_download');
      45            1 :         } finally {
      46            1 :             setDownloading(false);
      47            1 :         }
      48            1 :     };
      49              : 
      50            1 :     return (
      51            1 :         <IconButton
      52            1 :             variant={ButtonVariant.Invisible}
      53            1 :             size={ButtonSize.Small}
      54            1 :             aria-label="Download chart as PNG"
      55            1 :             data-testid="download-chart-png"
      56            1 :             disabled={!hasResult || downloading}
      57            1 :             onClick={download}
      58              :         >
      59            1 :             <DownloadIcon size={16} />
      60            1 :         </IconButton>
      61              :     );
      62            1 : }