feat: Download data button

package.json

@@ -20,6 +20,7 @@
     "@react-buddy/ide-toolbox": "^2.4.0",
     "@react-buddy/palette-mui": "^5.0.1",
     "@tanstack/react-query": "^5",
+    "@zip.js/zip.js": "^2.7.52",
     "apache-arrow": "^17.0.0",
     "array-move": "^4.0.0",
     "arrow-js-ffi": "^0.4.2",
@@ -28,6 +29,7 @@
     "dayjs": "^1.11.12",
     "dotenv": "^16.4.5",
     "immer": "^10.1.1",
+    "native-file-system-adapter": "^3.0.1",
     "parquet-wasm": "0.6.1",
     "react": "^18.3.1",
     "react-code-blocks": "^0.1.6",

src/Components/AuthFile.tsx

@@ -2,13 +2,49 @@
 import React, { useEffect, useState } from 'react'
 import { useCurrentUser } from './CurrentUserContext'
 import { useQuery } from '@tanstack/react-query'
-import axios from 'axios'
+import axios, {AxiosResponse} from 'axios'
 import { Link } from 'react-router-dom'
 import CircularProgress from '@mui/material/CircularProgress'
 import Button from '@mui/material/Button'
 import { ICONS } from '../constants'
 import SafeTooltip from './SafeTooltip'
 
+const clean_filename = (filename: string) => {
+    return filename.replace(/\.parquet.*$/, '.parquet')
+}
+
+export async function fetchAuthFile({
+    url,
+    headers,
+}: {
+    url: string
+    headers: Record<string, unknown>
+}): Promise<{filename: string, content: AxiosResponse<Blob>}> {
+    let filename: string = 'file'
+    const response = await axios.get(url, {
+        headers,
+        responseType: 'blob',
+    })
+    const redirect_url = response.headers['galv-storage-redirect-url']
+    if (redirect_url) {
+        filename = redirect_url.split('/').pop() ?? filename
+    } else {
+        const disposition = response.headers['content-disposition']
+        if (disposition) {
+            filename =
+                disposition.split('filename=')[1].split('"')[0] ?? filename
+        } else {
+            filename = url.split('/').pop() ?? filename
+        }
+    }
+    return {
+        filename: clean_filename(filename),
+        content: redirect_url
+            ? await axios.get(redirect_url, { responseType: 'blob' })
+            : response,
+    }
+}
+
 export default function AuthFile({ url }: { url: string }) {
     const [dataUrl, setDataUrl] = useState('')
     const [filename, setFilename] = useState('file')
@@ -20,32 +56,13 @@ export default function AuthFile({ url }: { url: string }) {
     const query = useQuery({
         queryKey: [url],
         queryFn: async () => {
-            const response = await axios.get(url, {
-                headers,
-                responseType: 'blob',
-            })
-            const redirect_url = response.headers['galv-storage-redirect-url']
-            if (redirect_url) {
-                setFilename(redirect_url.split('/').pop() ?? 'file')
-            } else {
-                const disposition = response.headers['content-disposition']
-                if (disposition) {
-                    setFilename(disposition.split('filename=')[1].split('"')[0])
-                } else {
-                    setFilename(url.split('/').pop() ?? 'file')
-                }
-            }
-            return redirect_url
-                ? axios.get(redirect_url, { responseType: 'blob' })
-                : response
+            const { filename, content } = await fetchAuthFile({ url, headers })
+            setFilename(filename)
+            return content
         },
         enabled: downloading,
     })
 
-    const clean_filename = (filename: string) => {
-        return filename.replace(/\.parquet.*$/, '.parquet')
-    }
-
     useEffect(() => {
         if (query.data) {
             setDataUrl(URL.createObjectURL(query.data.data))

src/Components/Dev.tsx

@@ -1,8 +1,20 @@
 import React from 'react'
+import DownloadDataset from './DownloadDataset'
 
 export function Dev() {
     if (!import.meta.env.DEV) {
         return <></>
     }
-    return <></>
+    return (
+        <>
+            <DownloadDataset
+                file={{
+                    id: '72eb4cd4-65f2-4128-b822-694207cb5454',
+                    parquet_partitions: [
+                        'http://localhost:8080/parquet_partitions/5bc4ac7b-661e-485a-9eb0-87a112159ddf/',
+                    ],
+                }}
+            />
+        </>
+    )
 }

src/Components/DownloadDataset.tsx

@@ -0,0 +1,158 @@
+import React, { useState } from 'react'
+import { showSaveFilePicker } from 'native-file-system-adapter'
+import { BlobReader, BlobWriter, ZipWriter } from '@zip.js/zip.js'
+import { useCurrentUser } from './CurrentUserContext'
+import { Configuration, ObservedFile, ParquetPartitionsApi } from '@galv/galv'
+import { get_url_components, has } from './misc'
+import { fetchAuthFile } from './AuthFile'
+import CircularProgress from '@mui/material/CircularProgress'
+import { MdDownload } from 'react-icons/md'
+import Button, { ButtonProps } from '@mui/material/Button'
+import IconButton, { IconButtonProps } from '@mui/material/IconButton'
+
+// Don't use useQueries or useFetchResource here because we have a complex async flow
+
+const pathSplitterRegEx = /[/\\]/
+
+type ZipBlobOptions = {
+    file: ObservedFile
+    api_config: Configuration
+    headers: Record<string, unknown>
+    // if inDirectory is true use the file's name as the directory name
+    in_directory?: string | boolean
+}
+const zipBlobs = async ({
+    file,
+    api_config,
+    headers,
+    in_directory = true,
+}: ZipBlobOptions) => {
+    const dname = in_directory === true ? getFileName(file) : in_directory || ''
+    const dir_name = dname && !/\/$/.test(dname) ? `${dname}/` : dname
+    const partitions = await Promise.all(
+        file.parquet_partitions.map((partition_url) => {
+            // First, look up the ParquetPartitions for the file and get their file URLs
+            const components = get_url_components(partition_url)
+            if (!components?.resourceId) {
+                return Promise.resolve(undefined)
+            }
+            return new ParquetPartitionsApi(api_config)
+                .parquetPartitionsRetrieve({ id: components.resourceId })
+                .then((response) => {
+                    // Second, fetch the ParquetPartition file via getAuthFile
+                    if (
+                        !has(response.data, 'parquet_file') ||
+                        response.data.parquet_file === null
+                    ) {
+                        return undefined
+                    }
+                    return fetchAuthFile({
+                        url: response.data.parquet_file,
+                        headers,
+                    })
+                })
+        }),
+    )
+    const zipWriter = new ZipWriter(new BlobWriter('application/zip'))
+    await Promise.all(
+        partitions
+            .filter((p) => p !== undefined)
+            .map((p) =>
+                zipWriter.add(
+                    `${dir_name}${p.filename}`,
+                    new BlobReader(p.content.data),
+                ),
+            ),
+    )
+    return await zipWriter.close()
+}
+
+const getFileName = ({ name, path, id }: ObservedFile) => {
+    // Return name, or basename without extension, or id
+    if (name) {
+        return name
+    }
+    if (path) {
+        const basename = path.split(pathSplitterRegEx).pop()
+        if (basename) {
+            return basename.split('.')[0]
+        }
+    }
+    return id
+}
+
+const downloadZip = async (zipBlob: Blob, filename: string) => {
+    // Use showSaveFilePicker to prompt user for a file save location
+    const handle = await showSaveFilePicker({
+        suggestedName: `${filename}.zip`,
+        types: [
+            {
+                description: 'ZIP Archive',
+                accept: { 'application/zip': ['.zip'] },
+            },
+        ],
+    })
+
+    const writableStream = await handle.createWritable()
+    await writableStream.write(zipBlob)
+    await writableStream.close()
+}
+
+type DownloadDatasetProps<UseIconButton> = UseIconButton extends true
+    ? { file: ObservedFile; iconButton: UseIconButton } & {
+          buttonProps: Partial<Omit<IconButtonProps, 'onClick' | 'children'>>
+      }
+    : { file: ObservedFile; iconButton: UseIconButton } & {
+          buttonProps: Partial<Omit<ButtonProps, 'onClick' | 'children'>>
+      }
+
+export default function DownloadDataset<UseIconButton>({
+    file,
+    iconButton,
+    ...buttonProps
+}: DownloadDatasetProps<UseIconButton>) {
+    const { user, api_config } = useCurrentUser()
+    const [loading, setLoading] = useState(false)
+    const headers = {
+        authorization: `Bearer ${user?.token}`,
+        'Galv-Storage-No-Redirect': true,
+    }
+
+    const downloadZippedBlobs = async () => {
+        setLoading(true)
+        setTimeout(() => {
+            zipBlobs({ file, api_config, headers })
+                .then((zipBlob) => downloadZip(zipBlob, getFileName(file)))
+                .finally(() => setLoading(false))
+        })
+    }
+
+    if (!user) {
+        return <Button disabled={true}>Log in to download files</Button>
+    }
+
+    if (iconButton)
+        return (
+            <IconButton
+                disabled={loading}
+                onClick={downloadZippedBlobs}
+                {...buttonProps}
+            >
+                {loading ? <CircularProgress /> : <MdDownload />}
+            </IconButton>
+        )
+
+    return (
+        <Button
+            startIcon={loading ? <CircularProgress /> : <MdDownload />}
+            variant="contained"
+            color="primary"
+            onClick={downloadZippedBlobs}
+            disabled={loading}
+            sx={{ mt: 2 }}
+            {...buttonProps}
+        >
+            Download Dataset
+        </Button>
+    )
+}