From 61265cc32761f6e98dc21cb27e240687e8af277e Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa <72977554+Cristhianzl@users.noreply.github.com> Date: Tue, 23 Jul 2024 11:41:55 -0300 Subject: [PATCH] feat: change download flows to download as .zip file (#2870) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 📝 (flows.py): Add endpoint to download multiple flows as a zip file 📝 (constants.ts): Add FLOWS constant for API endpoint 📝 (index.ts): Add use-post-download-multiple-flows query function 📝 (use-post-download-multiple-flows.ts): Implement function to download multiple flows as a zip file 📝 (index.tsx): Import usePostDownloadMultipleFlows in componentsComponent 📝 (index.tsx): Implement handleExport function to download multiple flows as a zip file 📝 (index.tsx): Update handleExport function to handle downloading and saving the zip file 📝 (headerComponent.tsx): Add shouldSelectAll state and setShouldSelectAll function to manage select all functionality * formatting file * 📝 (flows.py): Remove download_file endpoint and refactor download_multiple_file to support downloading multiple flows as a zip file 📝 (index.tsx): Refactor export functionality to handle exporting multiple flows as a zip file or a single flow as a JSON file 🔧 (use-post-download-multiple-flows.ts): Update API endpoint for downloading flows 🔧 (index.tsx): Remove unused imports and functions related to exporting flows 🔧 (reactflowUtils.ts): Remove unused downloadFlows function and related imports * [autofix.ci] apply automated fixes * 📝 (flows.py): remove unnecessary whitespace and improve code readability by aligning function parameters in download_multiple_file function --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- src/backend/base/langflow/api/v1/flows.py | 61 +++++++++++---- src/backend/tests/unit/test_database.py | 26 ++++--- .../src/controllers/API/helpers/constants.ts | 2 + src/frontend/src/controllers/API/index.ts | 13 ---- .../controllers/API/queries/flows/index.ts | 1 + .../flows/use-post-download-multiple-flows.ts | 43 +++++++++++ .../hooks/use-handle-export.tsx | 54 ------------- .../components/componentsComponent/index.tsx | 76 ++++++++++++++----- .../components/headerComponent/index.tsx | 6 +- src/frontend/src/utils/reactflowUtils.ts | 17 ----- 10 files changed, 170 insertions(+), 129 deletions(-) create mode 100644 src/frontend/src/controllers/API/queries/flows/index.ts create mode 100644 src/frontend/src/controllers/API/queries/flows/use-post-download-multiple-flows.ts delete mode 100644 src/frontend/src/pages/MainPage/components/componentsComponent/hooks/use-handle-export.tsx diff --git a/src/backend/base/langflow/api/v1/flows.py b/src/backend/base/langflow/api/v1/flows.py index b66201fad..d6f18a2c5 100644 --- a/src/backend/base/langflow/api/v1/flows.py +++ b/src/backend/base/langflow/api/v1/flows.py @@ -1,13 +1,17 @@ +import io +import json import re from datetime import datetime, timezone from typing import List from uuid import UUID +import zipfile +from fastapi.responses import StreamingResponse import orjson from fastapi import APIRouter, Depends, File, HTTPException, UploadFile from fastapi.encoders import jsonable_encoder from loguru import logger -from sqlmodel import Session, col, select +from sqlmodel import Session, and_, col, select from langflow.api.utils import remove_api_keys, validate_is_component from langflow.api.v1.schemas import FlowListCreate, FlowListRead @@ -314,18 +318,6 @@ async def upload_file( return response_list -@router.get("/download/", response_model=FlowListRead, status_code=200) -async def download_file( - *, - session: Session = Depends(get_session), - settings_service: "SettingsService" = Depends(get_settings_service), - current_user: User = Depends(get_current_active_user), -): - """Download all flows as a file.""" - flows = read_flows(current_user=current_user, session=session, settings_service=settings_service) - return FlowListRead(flows=flows) - - @router.delete("/") async def delete_multiple_flows( flow_ids: List[UUID], user: User = Depends(get_current_active_user), db: Session = Depends(get_session) @@ -350,3 +342,46 @@ async def delete_multiple_flows( except Exception as exc: logger.exception(exc) raise HTTPException(status_code=500, detail=str(exc)) from exc + + +@router.post("/download/", status_code=200) +async def download_multiple_file( + flow_ids: List[UUID], + user: User = Depends(get_current_active_user), + db: Session = Depends(get_session), +): + """Download all flows as a zip file.""" + flows = db.exec(select(Flow).where(and_(Flow.user_id == user.id, Flow.id.in_(flow_ids)))).all() # type: ignore + + if not flows: + raise HTTPException(status_code=404, detail="No flows found.") + + flows_without_api_keys = [remove_api_keys(flow.model_dump()) for flow in flows] + + if len(flows_without_api_keys) > 1: + # Create a byte stream to hold the ZIP file + zip_stream = io.BytesIO() + + # Create a ZIP file + with zipfile.ZipFile(zip_stream, "w") as zip_file: + for flow in flows_without_api_keys: + # Convert the flow object to JSON + flow_json = json.dumps(jsonable_encoder(flow)) + + # Write the JSON to the ZIP file + zip_file.writestr(f"{flow['name']}.json", flow_json) + + # Seek to the beginning of the byte stream + zip_stream.seek(0) + + # Generate the filename with the current datetime + current_time = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"{current_time}_langflow_flows.zip" + + return StreamingResponse( + zip_stream, + media_type="application/x-zip-compressed", + headers={"Content-Disposition": f"attachment; filename={filename}"}, + ) + else: + return FlowListRead(flows=flows_without_api_keys) diff --git a/src/backend/tests/unit/test_database.py b/src/backend/tests/unit/test_database.py index e09f9f62e..8cc288266 100644 --- a/src/backend/tests/unit/test_database.py +++ b/src/backend/tests/unit/test_database.py @@ -1,3 +1,4 @@ +import json from uuid import UUID, uuid4 import orjson @@ -197,26 +198,27 @@ def test_download_file( ) db_manager = get_db_service() with session_getter(db_manager) as session: + saved_flows = [] for flow in flow_list.flows: flow.user_id = active_user.id db_flow = Flow.model_validate(flow, from_attributes=True) session.add(db_flow) + saved_flows.append(db_flow) session.commit() - # Make request to endpoint - response = client.get("api/v1/flows/download/", headers=logged_in_headers) + # Make request to endpoint inside the session context + flow_ids = [str(db_flow.id) for db_flow in saved_flows] # Convert UUIDs to strings + flow_ids_json = json.dumps(flow_ids) + response = client.post( + "api/v1/flows/download/", + data=flow_ids_json, + headers={**logged_in_headers, "Content-Type": "application/json"}, + ) # Check response status code assert response.status_code == 200, response.json() # Check response data - response_data = response.json()["flows"] - starter_projects = load_starter_projects() - number_of_projects = len(starter_projects) + len(flow_list.flows) - assert len(response_data) == number_of_projects, response_data - assert response_data[0]["name"] == "Flow 1" - assert response_data[0]["description"] == "description" - assert response_data[0]["data"] == data - assert response_data[1]["name"] == "Flow 2" - assert response_data[1]["description"] == "description" - assert response_data[1]["data"] == data + # Since the endpoint now returns a zip file, we need to check the content type and the filename in the headers + assert response.headers["Content-Type"] == "application/x-zip-compressed" + assert "attachment; filename=" in response.headers["Content-Disposition"] def test_create_flow_with_invalid_data(client: TestClient, active_user, logged_in_headers): diff --git a/src/frontend/src/controllers/API/helpers/constants.ts b/src/frontend/src/controllers/API/helpers/constants.ts index 1c5146bb1..dae638ace 100644 --- a/src/frontend/src/controllers/API/helpers/constants.ts +++ b/src/frontend/src/controllers/API/helpers/constants.ts @@ -13,6 +13,8 @@ export const URLs = { AUTOLOGIN: "auto_login", REFRESH: "refresh", BUILD: `build`, + CUSTOM_COMPONENT: `custom_component`, + FLOWS: `flows`, FOLDERS: `folders`, } as const; diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts index cf1c0f43e..3b4b4adc9 100644 --- a/src/frontend/src/controllers/API/index.ts +++ b/src/frontend/src/controllers/API/index.ts @@ -197,19 +197,6 @@ export async function readFlowsFromDatabase() { } } -export async function downloadFlowsFromDatabase() { - try { - const response = await api.get(`${BASE_URL_API}flows/download/`); - if (response && response?.status !== 200) { - throw new Error(`HTTP error! status: ${response?.status}`); - } - return response?.data; - } catch (error) { - console.error(error); - throw error; - } -} - export async function uploadFlowsToDatabase(flows: FormData) { try { const response = await api.post(`${BASE_URL_API}flows/upload/`, flows); diff --git a/src/frontend/src/controllers/API/queries/flows/index.ts b/src/frontend/src/controllers/API/queries/flows/index.ts new file mode 100644 index 000000000..ddfa2f4bd --- /dev/null +++ b/src/frontend/src/controllers/API/queries/flows/index.ts @@ -0,0 +1 @@ +export * from "./use-post-download-multiple-flows"; diff --git a/src/frontend/src/controllers/API/queries/flows/use-post-download-multiple-flows.ts b/src/frontend/src/controllers/API/queries/flows/use-post-download-multiple-flows.ts new file mode 100644 index 000000000..1ef801124 --- /dev/null +++ b/src/frontend/src/controllers/API/queries/flows/use-post-download-multiple-flows.ts @@ -0,0 +1,43 @@ +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"; + +interface IPostDownloadMultipleFlows { + flow_ids: string[]; +} + +export const usePostDownloadMultipleFlows: useMutationFunctionType< + undefined, + IPostDownloadMultipleFlows +> = (options?) => { + const { mutate } = UseRequestProcessor(); + + const postDownloadMultipleFlowsFn = async ( + payload: IPostDownloadMultipleFlows, + ): Promise => { + const response = await api.post( + `${getURL("FLOWS")}/download/`, + payload.flow_ids, + { responseType: "blob" }, + ); + + return response.data; + }; + + const mutation: UseMutationResult< + IPostDownloadMultipleFlows, + any, + IPostDownloadMultipleFlows + > = mutate( + ["usePostDownloadMultipleFlows"], + async (payload: IPostDownloadMultipleFlows) => { + const res = await postDownloadMultipleFlowsFn(payload); + return res; + }, + options, + ); + + return mutation; +}; diff --git a/src/frontend/src/pages/MainPage/components/componentsComponent/hooks/use-handle-export.tsx b/src/frontend/src/pages/MainPage/components/componentsComponent/hooks/use-handle-export.tsx deleted file mode 100644 index 4b5e2f925..000000000 --- a/src/frontend/src/pages/MainPage/components/componentsComponent/hooks/use-handle-export.tsx +++ /dev/null @@ -1,54 +0,0 @@ -import { useCallback } from "react"; -import { FlowType } from "../../../../../types/flow"; - -const useExportFlows = ( - selectedFlowsComponentsCards: string[], - allFlows: Array, - downloadFlow: (flow: any, name: string, description: string) => void, - removeApiKeys: (flow: any) => any, - version: string, - setSuccessData: (data: { title: string }) => void, - setSelectedFlowsComponentsCards: ( - selectedFlowsComponentsCards: string[], - ) => void, - handleSelectAll: (select: boolean) => void, - cardTypes: string, -) => { - const handleExport = useCallback(() => { - selectedFlowsComponentsCards.forEach((selectedFlowId) => { - const selectedFlow = allFlows.find((flow) => flow.id === selectedFlowId); - if (selectedFlow) { - downloadFlow( - removeApiKeys({ - id: selectedFlow.id, - data: selectedFlow.data, - description: selectedFlow.description, - name: selectedFlow.name, - last_tested_version: version, - endpoint_name: selectedFlow.endpoint_name, - is_component: false, - }), - selectedFlow.name, - selectedFlow.description, - ); - } - }); - setSuccessData({ title: `${cardTypes} exported successfully` }); - setSelectedFlowsComponentsCards([]); - handleSelectAll(false); - }, [ - selectedFlowsComponentsCards, - allFlows, - downloadFlow, - removeApiKeys, - version, - setSuccessData, - setSelectedFlowsComponentsCards, - handleSelectAll, - cardTypes, - ]); - - return { handleExport }; -}; - -export default useExportFlows; diff --git a/src/frontend/src/pages/MainPage/components/componentsComponent/index.tsx b/src/frontend/src/pages/MainPage/components/componentsComponent/index.tsx index 16077b4a8..4a1816598 100644 --- a/src/frontend/src/pages/MainPage/components/componentsComponent/index.tsx +++ b/src/frontend/src/pages/MainPage/components/componentsComponent/index.tsx @@ -1,18 +1,15 @@ +import { usePostDownloadMultipleFlows } from "@/controllers/API/queries/flows"; import { useEffect, useMemo, useState } from "react"; import { FormProvider, useForm, useWatch } from "react-hook-form"; -import { Link, useLocation, useNavigate } from "react-router-dom"; +import { useLocation } from "react-router-dom"; import CardsWrapComponent from "../../../../components/cardsWrapComponent"; -import IconComponent from "../../../../components/genericIconComponent"; import PaginatorComponent from "../../../../components/paginatorComponent"; import { SkeletonCardComponent } from "../../../../components/skeletonCardComponent"; -import { Button } from "../../../../components/ui/button"; import DeleteConfirmationModal from "../../../../modals/deleteConfirmationModal"; import useAlertStore from "../../../../stores/alertStore"; -import { useDarkStore } from "../../../../stores/darkStore"; import useFlowsManagerStore from "../../../../stores/flowsManagerStore"; import { useFolderStore } from "../../../../stores/foldersStore"; import { FlowType } from "../../../../types/flow"; -import { downloadFlow, removeApiKeys } from "../../../../utils/reactflowUtils"; import useFileDrop from "../../hooks/use-on-file-drop"; import { getNameByType } from "../../utils/get-name-by-type"; import { sortFlows } from "../../utils/sort-flows"; @@ -23,7 +20,6 @@ import useDeleteMultipleFlows from "./hooks/use-delete-multiple"; import useDescriptionModal from "./hooks/use-description-modal"; import useFilteredFlows from "./hooks/use-filtered-flows"; import useDuplicateFlows from "./hooks/use-handle-duplicate"; -import useExportFlows from "./hooks/use-handle-export"; import useSelectAll from "./hooks/use-handle-select-all"; import useSelectOptionsChange from "./hooks/use-select-options-change"; import useSelectedFlows from "./hooks/use-selected-flows"; @@ -78,6 +74,8 @@ export default function ComponentsComponent({ const isLoadingFolders = useFolderStore((state) => state.isLoadingFolders); const setSelectedFolder = useFolderStore((state) => state.setSelectedFolder); + const [shouldSelectAll, setShouldSelectAll] = useState(true); + const cardTypes = useMemo(() => { if (window.location.pathname.includes("components")) { return "Components"; @@ -128,19 +126,59 @@ export default function ComponentsComponent({ cardTypes, ); - const version = useDarkStore((state) => state.version); + const { mutate: mutateDownloadMultipleFlows } = + usePostDownloadMultipleFlows(); - const { handleExport } = useExportFlows( - selectedFlowsComponentsCards, - allFlows, - downloadFlow, - removeApiKeys, - version, - setSuccessData, - setSelectedFlowsComponentsCards, - handleSelectAll, - cardTypes, - ); + const handleExport = () => { + mutateDownloadMultipleFlows( + { + flow_ids: selectedFlowsComponentsCards, + }, + { + onSuccess: (data) => { + const selectedFlow = allFlows.find( + (flow) => flow.id === selectedFlowsComponentsCards[0], + ); + + const blobType = + selectedFlowsComponentsCards.length > 1 + ? "application/zip" + : "application/json"; + + const fileNameSuffix = + selectedFlowsComponentsCards.length > 1 + ? "_langflow_flows.zip" + : `${selectedFlow!.name}.json`; + + const blob = new Blob([data], { type: blobType }); + + const link = document.createElement("a"); + link.href = window.URL.createObjectURL(blob); + + let current_time = new Date().toISOString().replace(/[:.]/g, ""); + + current_time = current_time + .replace(/-/g, "") + .replace(/T/g, "") + .replace(/Z/g, ""); + + link.download = + selectedFlowsComponentsCards.length > 1 + ? `${current_time}${fileNameSuffix}` + : `${fileNameSuffix}`; + + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + + setSuccessData({ title: `${cardTypes} exported successfully` }); + setSelectedFlowsComponentsCards([]); + handleSelectAll(false); + setShouldSelectAll(true); + }, + }, + ); + }; const { handleSelectOptionsChange } = useSelectOptionsChange( selectedFlowsComponentsCards, @@ -184,6 +222,8 @@ export default function ComponentsComponent({
{allFlows?.length > 0 && ( handleSelectOptionsChange("delete")} handleSelectAll={handleSelectAll} handleDuplicate={() => handleSelectOptionsChange("duplicate")} diff --git a/src/frontend/src/pages/MainPage/components/headerComponent/index.tsx b/src/frontend/src/pages/MainPage/components/headerComponent/index.tsx index cdd152229..4669323a6 100644 --- a/src/frontend/src/pages/MainPage/components/headerComponent/index.tsx +++ b/src/frontend/src/pages/MainPage/components/headerComponent/index.tsx @@ -11,6 +11,8 @@ type HeaderComponentProps = { handleDuplicate: () => void; handleExport: () => void; disableFunctions: boolean; + setShouldSelectAll: (select) => void; + shouldSelectAll: boolean; }; const HeaderComponent = ({ @@ -19,9 +21,9 @@ const HeaderComponent = ({ handleDuplicate, handleExport, disableFunctions, + setShouldSelectAll, + shouldSelectAll, }: HeaderComponentProps) => { - const [shouldSelectAll, setShouldSelectAll] = useState(true); - const handleClick = () => { handleSelectAll(shouldSelectAll); setShouldSelectAll((prevState) => !prevState); diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts index 22ee55786..5911b13cb 100644 --- a/src/frontend/src/utils/reactflowUtils.ts +++ b/src/frontend/src/utils/reactflowUtils.ts @@ -17,7 +17,6 @@ import { SUCCESS_BUILD, specialCharsRegex, } from "../constants/constants"; -import { downloadFlowsFromDatabase } from "../controllers/API"; import { DESCRIPTIONS } from "../flow_constants"; import { APIClassType, @@ -1423,22 +1422,6 @@ export function downloadFlow( link.click(); } -export function downloadFlows() { - downloadFlowsFromDatabase().then((flows) => { - const jsonString = `data:text/json;chatset=utf-8,${encodeURIComponent( - JSON.stringify(flows), - )}`; - - // create a link element and set its properties - const link = document.createElement("a"); - link.href = jsonString; - link.download = `flows.json`; - - // simulate a click on the link element to trigger the download - link.click(); - }); -} - export function getRandomElement(array: T[]): T { return array[Math.floor(Math.random() * array.length)]; }