chore(file-download): enable downloading folder/file contents from fi…

…le directory (#346)

* chore(file-download): enable downloading folder/file contents from file directory

* chore(optimize-code): Updated router file and created a utils file to add helper functions code and use it from there and removed unwnated code from the components

* fix(update-file-donwlaoder-UI): Updated file download to be in list item ,updated dilaog for download and updated api path formation

setup.cfg

@@ -32,6 +32,7 @@ install_requires =
     setuptools<=68.2.2
     pygwalker<=0.4.8.1
     openai<=1.1.0
+    aiofiles<=24.1.0
     python-multipart<=0.0.9
     ruff<=0.6.7
 python_requires = >=3.9

zt_backend/models/api/request.py

@@ -171,3 +171,8 @@ class RenameItemRequest(BaseModel):
 class DeleteItemRequest(BaseModel):
     path: str
     name: str
+
+class DownloadRequest(BaseModel):
+    path: str
+    filename: str
+    isFolder: bool

zt_backend/router.py

@@ -1,7 +1,7 @@
-import subprocess
 import shutil
 from fastapi import (
     APIRouter,
+    Depends,
     WebSocket,
     WebSocketDisconnect,
     Query,
@@ -34,7 +34,8 @@
 from zt_backend.models.state.user_state import UserState
 from zt_backend.models.state.app_state import AppState
 from zt_backend.models.state.notebook_state import UploadState
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.requests import Request
 from pathlib import Path
 import logging
 import uuid
@@ -46,6 +47,12 @@
 import requests
 import re
 from zt_backend.utils.file_utils import upload_queue, process_upload
+import mimetypes
+from typing import Dict, Tuple, Optional
+import aiofiles
+import tempfile
+from zt_backend.utils.file_utils import *
+
 
 router = APIRouter()
 manager = ConnectionManager()
@@ -510,67 +517,15 @@ def share_notebook(shareRequest: request.ShareRequest):
             response_json = response.json()
             signed_url = response_json.get("uploadURL")
             if not signed_url:
-                raise HTTPException(
-                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                    detail="Failed to get a signed URL",
-                )
-
-            python_warning = response_json.get("pythonWarning", "")
-            zt_warning = response_json.get("ztWarning", "")
-            project_warning = response_json.get("projectWarning", "")
-            warning_message = ""
-            if python_warning:
-                warning_message += f"\n{python_warning}"
-            if zt_warning:
-                warning_message += f"\n{zt_warning}"
-            if project_warning:
-                warning_message += f"\n{project_warning}"
-            if warning_message:
-                warning_message += "\nSelect confirm if you would like to proceed"
-                upload_state.signed_url = signed_url
-                return {"warning": warning_message}
-
-            publish_files(project_name, signed_url)
-
-    except HTTPException as e:
-        raise e
-    except Exception as e:
-        logger.error("Error submitting share request: %s", str(e))
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail="Error submitting share request",
-        )
-
+                return {"Error": "Failed to get signed URL"}
 
-@router.post("/api/confirm_share")
-def confirm_share(shareRequest: request.ShareRequest):
-    if app_state.run_mode == "dev":
-        if upload_state.signed_url:
-            try:
-                publish_files(
-                    shareRequest.projectName.lower().strip(), upload_state.signed_url
-                )
-                upload_state.signed_url = None
-            except Exception as e:
-                raise HTTPException(
-                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                    detail="Error submitting share request",
-                )
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST, detail="No active signed URL"
+            output_filename = f"{project_name}"
+            project_source = os.path.normpath(os.getcwd())
+            logger.info(project_source)
+            shutil.make_archive(
+                base_name=output_filename, format="gztar", root_dir=project_source
             )
 
-
-def publish_files(project_name, signed_url):
-    try:
-        output_filename = Path(settings.zt_path) / f"{project_name}.tar.gz"
-        tar_base = str(Path(settings.zt_path) / project_name)
-
-        shutil.make_archive(
-            base_name=tar_base, format="gztar", root_dir=settings.zt_path
-        )
-
         with output_filename.open("rb") as file:
             upload_files = {"file": file}
             upload_response = requests.post(
@@ -732,37 +687,59 @@ def delete_item(delete_request: request.DeleteItemRequest):
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
+@router.get("/api/download")
+async def download_item(
+    request: Request,  # For headers
+    download_req: request.DownloadRequest = Depends()  # This is the key change
+):
+    """Stream download with range support for files and folders."""
+    try:
+        chunk_size: int = 8192
+        file_path = Path(download_req.path).resolve()
+
+        if not file_path.exists():
+            raise HTTPException(status_code=404, detail="Path not found")
+
+        if download_req.isFolder:
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as temp_zip:
+                temp_zip_path = temp_zip.name
+
+            await create_zip_file(file_path, temp_zip_path)
+            file_path = Path(temp_zip_path)
+            download_req.filename = f"{download_req.filename}.zip"
+
+        file_size = file_path.stat().st_size
+        start, end = await parse_range_header(
+            request.headers.get('range'),
+            file_size
+        )
 
-def get_file_type(name):
-    extension = name.split(".")[-1]
-    if extension in ["html", "js", "json", "md", "pdf", "png", "txt", "xls"]:
-        return extension
-    return None
+        response = StreamingResponse(
+            stream_file_range(str(file_path), start, end, chunk_size),
+            status_code=206 if request.headers.get('range') else 200,
+            media_type=get_mime_type(download_req.filename)
+        )
 
+        response.headers.update({
+            "Content-Range": f"bytes {start}-{end}/{file_size}",
+            "Content-Length": str(end - start + 1),
+            "Content-Disposition": f"attachment; filename={download_req.filename}",
+            "Accept-Ranges": "bytes"
+        })
 
-def list_dir(path):
-    items = []
-    for item in path.iterdir():
-        if item.is_dir():
-            items.append(
-                {
-                    "title": item.name,
-                    "file": "folder",
-                    "id": item.as_posix(),
-                    "children": [],
-                }
-            )
-        else:
-            file_type = get_file_type(item.name)
-            if file_type:
-                items.append(
-                    {"title": item.name, "file": file_type, "id": item.as_posix()}
-                )
-            else:
-                items.append(
-                    {"title": item.name, "file": "file", "id": item.as_posix()}
-                )
-    return items
+        if download_req.isFolder:
+            async def cleanup_temp_file():
+                yield
+                os.unlink(temp_zip_path)
+            response.background = cleanup_temp_file()
+
+        return response
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
 
 
 @router.get("/api/get_files")

zt_backend/utils/file_utils.py

@@ -1,9 +1,114 @@
+import zipfile
+import mimetypes
+import os
+from typing import Dict, Optional, Tuple
+from fastapi import HTTPException
+from pathlib import Path
+import aiofiles
 import asyncio
 import os
 from contextlib import asynccontextmanager
 from pathlib import Path
 from fastapi import HTTPException, BackgroundTasks
 
+
+def initialize_mime_types():
+    """Initialize and customize MIME types."""
+    mimetypes.init()
+
+    custom_mime_types: Dict[str, str] = {
+        '.md': 'text/markdown',
+        '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
+        '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+        '.7z': 'application/x-7z-compressed',
+        '.rar': 'application/vnd.rar',
+        '.webp': 'image/webp',
+        # Add more custom MIME types here as needed
+    }
+
+    for ext, mime_type in custom_mime_types.items():
+        mimetypes.add_type(mime_type, ext)
+
+def get_mime_type(filename: str) -> str:
+    """Determine the MIME type of a file based on its filename."""
+    initialize_mime_types()
+    return mimetypes.guess_type(filename)[0] or 'application/octet-stream'
+
+def validate_path(path: Path, filename: str):
+    """Validate the path and raise appropriate exceptions if invalid."""
+    if not path.exists():
+        raise HTTPException(status_code=404, detail=f"Path not found: {filename}")
+
+async def create_zip_file(folder_path: Path, temp_zip_path: str):
+    """Create a zip file from a folder."""
+    with zipfile.ZipFile(temp_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(folder_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, start=str(folder_path))
+                zipf.write(file_path, arcname)
+
+async def parse_range_header(range_header: Optional[str], file_size: int) -> Tuple[int, int]:
+    """Parse Range header and return start and end bytes."""
+    if not range_header:
+        return 0, file_size - 1
+
+    try:
+        range_str = range_header.replace('bytes=', '')
+        start_str, end_str = range_str.split('-')
+        start = int(start_str)
+        end = int(end_str) if end_str else file_size - 1
+        return start, min(end, file_size - 1)
+    except ValueError:
+        raise HTTPException(status_code=400, detail="Invalid range header")
+
+async def stream_file_range(file_path: str, start: int, end: int, chunk_size: int = 8192):
+    """Stream file content for the specified byte range."""
+    async with aiofiles.open(file_path, mode='rb') as file:
+        await file.seek(start)
+        bytes_remaining = end - start + 1
+
+        while bytes_remaining > 0:
+            chunk_size = min(chunk_size, bytes_remaining)
+            chunk = await file.read(chunk_size)
+            if not chunk:
+                break
+            yield chunk
+            bytes_remaining -= len(chunk)
+
+def get_file_type(name):
+    extension = name.split(".")[-1]
+    if extension in ["html", "js", "json", "md", "pdf", "png", "txt", "xls"]:
+        return extension
+    return None
+
+
+def list_dir(path):
+    items = []
+    for item in path.iterdir():
+        if item.is_dir():
+            items.append(
+                {
+                    "title": item.name,
+                    "file": "folder",
+                    "id": item.as_posix(),
+                    "children": [],
+                }
+            )
+        else:
+            file_type = get_file_type(item.name)
+            if file_type:
+                items.append(
+                    {"title": item.name, "file": file_type, "id": item.as_posix()}
+                )
+            else:
+                items.append(
+                    {"title": item.name, "file": "file", "id": item.as_posix()}
+                )
+    return items
+
+
 # Semaphore to control concurrent uploads
 upload_semaphore = asyncio.Semaphore(5)
 # Locks to ensure each file is written in isolation