From 9e8715cf4c0b8a216d9c195441f5a71ffc8cad78 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Tue, 19 Nov 2024 15:41:37 -0300 Subject: [PATCH] fix: enhance folder download endpoint with zip file support (#4706) * expanding download folder to zip file * run formatter * Update src/backend/base/langflow/api/v1/folders.py Co-authored-by: Gabriel Luiz Freitas Almeida * changing model to flowRead * run formatter --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- src/backend/base/langflow/api/v1/folders.py | 55 +++++++++++++++---- src/backend/base/langflow/api/v1/schemas.py | 2 +- .../components/sideBarFolderButtons/index.tsx | 31 +++++++---- .../folders/use-get-download-folders.ts | 22 ++++++-- 4 files changed, 80 insertions(+), 30 deletions(-) diff --git a/src/backend/base/langflow/api/v1/folders.py b/src/backend/base/langflow/api/v1/folders.py index 910cfa19b..4b9a6b279 100644 --- a/src/backend/base/langflow/api/v1/folders.py +++ b/src/backend/base/langflow/api/v1/folders.py @@ -1,16 +1,22 @@ +import io +import json +import zipfile +from datetime import datetime, timezone from typing import Annotated import orjson from fastapi import APIRouter, Depends, File, HTTPException, Response, UploadFile, status +from fastapi.encoders import jsonable_encoder +from fastapi.responses import StreamingResponse from fastapi_pagination import Params from fastapi_pagination.ext.sqlmodel import paginate from sqlalchemy import or_, update from sqlalchemy.orm import selectinload from sqlmodel import select -from langflow.api.utils import AsyncDbSession, CurrentActiveUser, cascade_delete_flow, custom_params +from langflow.api.utils import AsyncDbSession, CurrentActiveUser, cascade_delete_flow, custom_params, remove_api_keys from langflow.api.v1.flows import create_flows -from langflow.api.v1.schemas import FlowListCreate, FlowListReadWithFolderName +from langflow.api.v1.schemas import FlowListCreate from langflow.helpers.flow import generate_unique_flow_name from langflow.helpers.folders import generate_unique_folder_name from langflow.initial_setup.setup import STARTER_FOLDER_NAME @@ -248,28 +254,53 @@ async def delete_folder( raise HTTPException(status_code=500, detail=str(e)) from e -@router.get("/download/{folder_id}", response_model=FlowListReadWithFolderName, status_code=200) +@router.get("/download/{folder_id}", status_code=200) async def download_file( *, session: AsyncDbSession, folder_id: str, current_user: CurrentActiveUser, ): - """Download all flows from folder.""" + """Download all flows from folder as a zip file.""" try: - folder = ( - await session.exec(select(Folder).where(Folder.id == folder_id, Folder.user_id == current_user.id)) - ).first() + query = select(Folder).where(Folder.id == folder_id, Folder.user_id == current_user.id) + result = await session.exec(query) + folder = result.first() + + if not folder: + raise HTTPException(status_code=404, detail="Folder not found") + + flows_query = select(Flow).where(Flow.folder_id == folder_id) + flows_result = await session.exec(flows_query) + flows = [FlowRead.model_validate(flow, from_attributes=True) for flow in flows_result.all()] + + if not flows: + raise HTTPException(status_code=404, detail="No flows found in folder") + + flows_without_api_keys = [remove_api_keys(flow.model_dump()) for flow in flows] + zip_stream = io.BytesIO() + + with zipfile.ZipFile(zip_stream, "w") as zip_file: + for flow in flows_without_api_keys: + flow_json = json.dumps(jsonable_encoder(flow)) + zip_file.writestr(f"{flow['name']}.json", flow_json) + + zip_stream.seek(0) + + current_time = datetime.now(tz=timezone.utc).astimezone().strftime("%Y%m%d_%H%M%S") + filename = f"{current_time}_{folder.name}_flows.zip" + + return StreamingResponse( + zip_stream, + media_type="application/x-zip-compressed", + headers={"Content-Disposition": f"attachment; filename={filename}"}, + ) + except Exception as e: if "No result found" in str(e): raise HTTPException(status_code=404, detail="Folder not found") from e raise HTTPException(status_code=500, detail=str(e)) from e - if not folder: - raise HTTPException(status_code=404, detail="Folder not found") - - return folder - @router.post("/upload/", response_model=list[FlowRead], status_code=201) async def upload_file( diff --git a/src/backend/base/langflow/api/v1/schemas.py b/src/backend/base/langflow/api/v1/schemas.py index 9d9b21018..87b137bfd 100644 --- a/src/backend/base/langflow/api/v1/schemas.py +++ b/src/backend/base/langflow/api/v1/schemas.py @@ -162,7 +162,7 @@ class FlowListRead(BaseModel): class FlowListReadWithFolderName(BaseModel): flows: list[FlowRead] - name: str + folder_name: str description: str diff --git a/src/frontend/src/components/folderSidebarComponent/components/sideBarFolderButtons/index.tsx b/src/frontend/src/components/folderSidebarComponent/components/sideBarFolderButtons/index.tsx index f7031f8bc..c2aa99bcc 100644 --- a/src/frontend/src/components/folderSidebarComponent/components/sideBarFolderButtons/index.tsx +++ b/src/frontend/src/components/folderSidebarComponent/components/sideBarFolderButtons/index.tsx @@ -125,7 +125,7 @@ const SideBarFoldersButtonsComponent = ({ }); }; - const { mutate: mutateDownloadFolder } = useGetDownloadFolders(); + const { mutate: mutateDownloadFolder } = useGetDownloadFolders({}); const handleDownloadFolder = (id: string) => { mutateDownloadFolder( @@ -133,20 +133,29 @@ const SideBarFoldersButtonsComponent = ({ folderId: id, }, { - onSuccess: (data) => { - data.folder_name = data?.name || "folder"; - data.folder_description = data?.description || ""; - - const jsonString = `data:text/json;charset=utf-8,${encodeURIComponent( - JSON.stringify(data), - )}`; + onSuccess: (response) => { + // Create a blob from the response data + const blob = new Blob([response.data], { + type: "application/x-zip-compressed", + }); + const url = window.URL.createObjectURL(blob); const link = document.createElement("a"); - link.href = jsonString; - link.download = `${data?.name}.json`; + link.href = url; + // Get filename from header or use default + const filename = + response.headers?.["content-disposition"] + ?.split("filename=")[1] + ?.replace(/['"]/g, "") ?? "flows.zip"; + + link.setAttribute("download", filename); + document.body.appendChild(link); link.click(); - track("Folder Exported", { folderId: id! }); + link.remove(); + window.URL.revokeObjectURL(url); + + track("Folder Exported", { folderId: id }); }, onError: () => { setErrorData({ diff --git a/src/frontend/src/controllers/API/queries/folders/use-get-download-folders.ts b/src/frontend/src/controllers/API/queries/folders/use-get-download-folders.ts index 8d4dec403..36a2cc868 100644 --- a/src/frontend/src/controllers/API/queries/folders/use-get-download-folders.ts +++ b/src/frontend/src/controllers/API/queries/folders/use-get-download-folders.ts @@ -1,4 +1,5 @@ import { useMutationFunctionType } from "@/types/api"; +import { UseMutationResult } from "@tanstack/react-query"; import { api } from "../../api"; import { getURL } from "../../helpers/constants"; import { UseRequestProcessor } from "../../services/request-processor"; @@ -8,22 +9,31 @@ interface IGetDownloadFolders { } export const useGetDownloadFolders: useMutationFunctionType< - undefined, + any, // Changed to any since we're getting the full response IGetDownloadFolders > = (options?) => { const { mutate } = UseRequestProcessor(); const downloadFoldersFn = async ( - data: IGetDownloadFolders, - ): Promise => { - const res = await api.get(`${getURL("FOLDERS")}/download/${data.folderId}`); - return res.data; + payload: IGetDownloadFolders, + ): Promise => { + const response = await api.get( + `${getURL("FOLDERS")}/download/${payload.folderId}`, + { + responseType: "blob", + headers: { + Accept: "application/x-zip-compressed", + }, + }, + ); + return response; }; - const mutation = mutate( + const mutation: UseMutationResult = mutate( ["useGetDownloadFolders"], downloadFoldersFn, options, ); + return mutation; };