Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : import * as React from 'react';
       3              : import * as styles from './query_result_download_button.module.css';
       4              : 
       5              : import { QueryExecutionStatus, type QueryExecutionState } from '../../query/query_execution_state.js';
       6              : import { ButtonSize, ButtonVariant, IconButton } from '../../ui/foundations/button.js';
       7              : import * as ActionList from '../../ui/foundations/action_list.js';
       8              : import { AnchorAlignment, AnchorSide } from '../../ui/foundations/anchored_position.js';
       9              : import { AnchoredOverlay } from '../../ui/foundations/anchored_overlay.js';
      10              : import { DownloadIcon } from '../../ui/foundations/symbol_icon.js';
      11              : import { useComputeDatabase } from '../../compute/compute_connection_provider.js';
      12              : import { exportArrowTableAsParquet, supportsParquetExport } from '../../platform/hyperdb/hyperdb_parquet_export.js';
      13              : import { useFileDownloader } from '../../platform/file/file_downloader_hook.js';
      14              : import { useLogger } from '../../platform/logger/logger_provider.js';
      15              : import { stringifyError } from '../../platform/logger/logger.js';
      16              : 
      17            1 : const enum DownloadFormat {
      18            1 :     Arrow = 'arrow',
      19            1 :     Parquet = 'parquet',
      20            1 : }
      21              : 
      22              : interface Props {
      23              :     query: QueryExecutionState;
      24              :     dependencies?: {
      25              :         useFileDownloader: () => Pick<ReturnType<typeof useFileDownloader>, 'downloadBufferAsFile'>;
      26              :         useLogger: () => Pick<ReturnType<typeof useLogger>, 'error'>;
      27              :         useComputeDatabase: typeof useComputeDatabase;
      28              :         exportArrowTableAsParquet: typeof exportArrowTableAsParquet;
      29              :     };
      30              : }
      31              : 
      32            1 : const DEFAULT_DEPENDENCIES: NonNullable<Props['dependencies']> = {
      33            1 :     useFileDownloader,
      34            1 :     useLogger,
      35            1 :     useComputeDatabase,
      36            1 :     exportArrowTableAsParquet,
      37            1 : };
      38              : 
      39            1 : export function QueryResultDownloadButton({ query, dependencies = DEFAULT_DEPENDENCIES }: Props) {
      40            1 :     const fileDownloader = dependencies.useFileDownloader();
      41            1 :     const logger = dependencies.useLogger();
      42            1 :     const computeDatabase = dependencies.useComputeDatabase();
      43            1 :     const parquetDatabase = supportsParquetExport(computeDatabase) ? computeDatabase : null;
      44            1 :     const [open, setOpen] = React.useState(false);
      45            1 :     const [downloading, setDownloading] = React.useState(false);
      46            1 :     const triggerRef = React.useRef<HTMLButtonElement>(null);
      47            1 :     const hasResult = query.status === QueryExecutionStatus.SUCCEEDED && query.resultTable != null;
      48              : 
      49            1 :     const downloadResult = async (format: DownloadFormat) => {
      50            1 :         setOpen(false);
      51            1 :         if (!hasResult || downloading) return;
      52            1 :         setDownloading(true);
      53            1 :         try {
      54            1 :             switch (format) {
      55            1 :                 case DownloadFormat.Arrow: {
      56            1 :                     const bytes = arrow.tableToIPC(query.resultTable!, 'file');
      57            1 :                     await fileDownloader.downloadBufferAsFile(bytes, `query-result-${query.queryId}.arrow`);
      58            1 :                     break;
      59            1 :                 }
      60            1 :                 case DownloadFormat.Parquet: {
      61            1 :                     if (parquetDatabase == null) throw new Error('compute database cannot export Parquet');
      62            1 :                     const bytes = await dependencies.exportArrowTableAsParquet(parquetDatabase, query.resultTable!);
      63            1 :                     await fileDownloader.downloadBufferAsFile(bytes, `query-result-${query.queryId}.parquet`);
      64            1 :                     break;
      65            1 :                 }
      66            1 :             }
      67            1 :         } catch (error) {
      68            1 :             logger.error('Failed to download query result', { format, error: stringifyError(error) }, 'query_result_download');
      69            1 :         } finally {
      70            1 :             setDownloading(false);
      71            1 :         }
      72            1 :     };
      73              : 
      74            1 :     return (
      75            1 :         <AnchoredOverlay
      76            1 :             open={open}
      77            1 :             onOpen={() => setOpen(true)}
      78            1 :             onClose={() => setOpen(false)}
      79            1 :             side={AnchorSide.OutsideBottom}
      80            1 :             align={AnchorAlignment.End}
      81            1 :             anchorOffset={4}
      82            1 :             anchorRef={triggerRef}
      83            1 :             returnFocusRef={triggerRef}
      84            1 :             renderAnchor={(anchorProps) => (
      85            1 :                 <IconButton
      86            1 :                     {...anchorProps}
      87            1 :                     ref={triggerRef}
      88            1 :                     variant={ButtonVariant.Invisible}
      89            1 :                     size={ButtonSize.Small}
      90            1 :                     aria-label="Download query result"
      91            1 :                     data-testid="download-query-result"
      92            1 :                     disabled={!hasResult || downloading}
      93              :                 >
      94            1 :                     <DownloadIcon size={16} />
      95            1 :                 </IconButton>
      96              :             )}
      97              :         >
      98            1 :             <div className={styles.menu} role="dialog" aria-label="Download query result">
      99            1 :                 <ActionList.List aria-label="Download formats">
     100            1 :                     <ActionList.ListItem data-testid="download-format-arrow" onClick={() => downloadResult(DownloadFormat.Arrow)}>
     101            1 :                         <ActionList.ItemText>Download as Arrow</ActionList.ItemText>
     102            1 :                     </ActionList.ListItem>
     103            1 :                     <ActionList.ListItem
     104            1 :                         data-testid="download-format-parquet"
     105            1 :                         disabled={parquetDatabase == null}
     106            1 :                         onClick={() => downloadResult(DownloadFormat.Parquet)}
     107              :                     >
     108            1 :                         <ActionList.ItemText>Download as Parquet</ActionList.ItemText>
     109            1 :                     </ActionList.ListItem>
     110            1 :                 </ActionList.List>
     111            1 :             </div>
     112            1 :         </AnchoredOverlay>
     113              :     );
     114            1 : }