From e7a20051887cb1f86c2c736ab3651c13986b0869 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Mon, 13 Jan 2025 19:42:53 -0300 Subject: [PATCH] fix: Improve `update_flow` data consistency, refine error handling, and add folder-moving tests (#5516) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐛 (flows.py): Fix issue where flow data was not being properly updated in the database during flow update 📝 (flows.py): Improve error handling and rollback database session in case of exceptions during flow update 📝 (flows.py): Refactor code to handle unique constraint errors and provide more informative error messages 📝 (utils.py): Refactor get_webhook_component_in_flow function to handle cases where flow_data may not have 'nodes' attribute * ✨ (sideBarFolderButtons/index.tsx): add unique id attribute to sidebar folder buttons for improved accessibility and testing ✨ (general-bugs-move-flow-from-folder.spec.ts): add test to ensure user can move flow from one folder to another in the frontend application * 🐛 (flows.py): remove unnecessary session rollback to prevent potential data inconsistency ♻️ (service.py): refactor with_session method to handle session commit and rollback more effectively * style: adjust line breaks for readability * style: reorder imports * fix: ruff error try300 * [autofix.ci] apply automated fixes * fix: mypy error module has no attribute "timeout" * 🐛 (flows.py): remove unnecessary error handling code and improve exception handling for better error propagation and clarity * [autofix.ci] apply automated fixes * Update src/backend/base/langflow/services/database/service.py Co-authored-by: Gabriel Luiz Freitas Almeida * use model dump besides overwrite value * [autofix.ci] apply automated fixes * 📝 (chat.py): improve code readability by refactoring session handling and adding comments for clarity 🔧 (chat.py): refactor code to create a fresh session for database operations and improve session management in build_flow function * [autofix.ci] apply automated fixes * refactor: remove unused session parameter from build_flow function in chat.py --------- Co-authored-by: italojohnny Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida --- src/backend/base/langflow/api/v1/chat.py | 22 ++++---- src/backend/base/langflow/api/v1/flows.py | 25 ++++----- .../services/database/models/flow/utils.py | 7 +-- .../langflow/services/database/service.py | 9 +++- .../components/sideBarFolderButtons/index.tsx | 1 + ...general-bugs-move-flow-from-folder.spec.ts | 54 +++++++++++++++++++ 6 files changed, 93 insertions(+), 25 deletions(-) create mode 100644 src/frontend/tests/extended/regression/general-bugs-move-flow-from-folder.spec.ts diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py index f4dc156d1..ad8f1cd7e 100644 --- a/src/backend/base/langflow/api/v1/chat.py +++ b/src/backend/base/langflow/api/v1/chat.py @@ -152,7 +152,6 @@ async def build_flow( start_component_id: str | None = None, log_builds: bool | None = True, current_user: CurrentActiveUser, - session: DbSession, ): chat_service = get_chat_service() telemetry_service = get_telemetry_service() @@ -164,15 +163,20 @@ async def build_flow( components_count = None try: flow_id_str = str(flow_id) - if not data: - graph = await build_graph_from_db(flow_id=flow_id, session=session, chat_service=chat_service) - else: - async with session_scope() as new_session: - result = await new_session.exec(select(Flow.name).where(Flow.id == flow_id)) + # Create a fresh session for database operations + async with session_scope() as fresh_session: + if not data: + graph = await build_graph_from_db(flow_id=flow_id, session=fresh_session, chat_service=chat_service) + else: + result = await fresh_session.exec(select(Flow.name).where(Flow.id == flow_id)) flow_name = result.first() - graph = await build_graph_from_data( - flow_id=flow_id_str, payload=data.model_dump(), user_id=str(current_user.id), flow_name=flow_name - ) + graph = await build_graph_from_data( + flow_id=flow_id_str, + payload=data.model_dump(), + user_id=str(current_user.id), + flow_name=flow_name, + ) + graph.validate_stream() if stop_component_id or start_component_id: try: diff --git a/src/backend/base/langflow/api/v1/flows.py b/src/backend/base/langflow/api/v1/flows.py index 6ce4d6d2b..b4a7150d3 100644 --- a/src/backend/base/langflow/api/v1/flows.py +++ b/src/backend/base/langflow/api/v1/flows.py @@ -271,18 +271,18 @@ async def update_flow( user_id=current_user.id, settings_service=settings_service, ) - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) from e - if not db_flow: - raise HTTPException(status_code=404, detail="Flow not found") + if not db_flow: + raise HTTPException(status_code=404, detail="Flow not found") + + update_data = flow.model_dump(exclude_unset=True, exclude_none=True) - try: - flow_data = flow.model_dump(exclude_unset=True) if settings_service.settings.remove_api_keys: - flow_data = remove_api_keys(flow_data) - for key, value in flow_data.items(): + update_data = remove_api_keys(update_data) + + for key, value in update_data.items(): setattr(db_flow, key, value) + webhook_component = get_webhook_component_in_flow(db_flow.data) db_flow.webhook = webhook_component is not None db_flow.updated_at = datetime.now(timezone.utc) @@ -291,13 +291,12 @@ async def update_flow( default_folder = (await session.exec(select(Folder).where(Folder.name == DEFAULT_FOLDER_NAME))).first() if default_folder: db_flow.folder_id = default_folder.id + session.add(db_flow) await session.commit() await session.refresh(db_flow) + except Exception as e: - # If it is a validation error, return the error message - if hasattr(e, "errors"): - raise HTTPException(status_code=400, detail=str(e)) from e if "UNIQUE constraint failed" in str(e): # Get the name of the column that failed columns = str(e).split("UNIQUE constraint failed: ")[1].split(".")[1].split("\n")[0] @@ -305,10 +304,12 @@ async def update_flow( # or UNIQUE constraint failed: flow.name # if the column has id in it, we want the other column column = columns.split(",")[1] if "id" in columns.split(",")[0] else columns.split(",")[0] - raise HTTPException( status_code=400, detail=f"{column.capitalize().replace('_', ' ')} must be unique" ) from e + + if hasattr(e, "status_code"): + raise HTTPException(status_code=e.status_code, detail=str(e)) from e raise HTTPException(status_code=500, detail=str(e)) from e return db_flow diff --git a/src/backend/base/langflow/services/database/models/flow/utils.py b/src/backend/base/langflow/services/database/models/flow/utils.py index 7023b19a5..32e639de0 100644 --- a/src/backend/base/langflow/services/database/models/flow/utils.py +++ b/src/backend/base/langflow/services/database/models/flow/utils.py @@ -5,9 +5,10 @@ from .model import Flow def get_webhook_component_in_flow(flow_data: dict): """Get webhook component in flow data.""" - for node in flow_data.get("nodes", []): - if "Webhook" in node.get("id"): - return node + if hasattr(flow_data, "nodes"): + for node in flow_data.get("nodes", []): + if "Webhook" in node.get("id"): + return node return None diff --git a/src/backend/base/langflow/services/database/service.py b/src/backend/base/langflow/services/database/service.py index fcc242b7c..d2874a520 100644 --- a/src/backend/base/langflow/services/database/service.py +++ b/src/backend/base/langflow/services/database/service.py @@ -136,7 +136,14 @@ class DatabaseService(Service): @asynccontextmanager async def with_session(self): async with AsyncSession(self.engine, expire_on_commit=False) as session: - yield session + try: + yield session + if session.is_active: + await session.commit() + except Exception: + logger.error("An error occurred during the session scope.") + await session.rollback() + raise async def assign_orphaned_flows_to_superuser(self) -> None: """Assign orphaned flows to the default superuser when auto login is enabled.""" diff --git a/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx b/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx index c72575ff0..9a5a82900 100644 --- a/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx +++ b/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx @@ -373,6 +373,7 @@ const SideBarFoldersButtonsComponent = ({ onDrop={(e) => onDrop(e, item.id!)} key={item.id} data-testid={`sidebar-nav-${item.name}`} + id={`sidebar-nav-${item.name}`} isActive={checkPathName(item.id!)} onClick={() => handleChangeFolder!(item.id!)} className={cn( diff --git a/src/frontend/tests/extended/regression/general-bugs-move-flow-from-folder.spec.ts b/src/frontend/tests/extended/regression/general-bugs-move-flow-from-folder.spec.ts new file mode 100644 index 000000000..d5efab06c --- /dev/null +++ b/src/frontend/tests/extended/regression/general-bugs-move-flow-from-folder.spec.ts @@ -0,0 +1,54 @@ +import { expect, test } from "@playwright/test"; +import { awaitBootstrapTest } from "../../utils/await-bootstrap-test"; + +test("user must be able to move flow from folder", async ({ page }) => { + const randomName = Math.random().toString(36).substring(2, 15); + + await awaitBootstrapTest(page); + + await page.getByTestId("side_nav_options_all-templates").click(); + await page.getByRole("heading", { name: "Basic Prompting" }).click(); + + await page.waitForSelector('[data-testid="flow_name"]', { + timeout: 3000, + }); + + await page.getByTestId("flow_name").click(); + await page.getByText("Flow Settings").first().click(); + await page.getByPlaceholder("Flow name").fill(randomName); + + await page.getByTestId("save-flow-settings").click(); + + await page.getByText("Changes saved successfully").isVisible(); + + await page.getByTestId("icon-ChevronLeft").click(); + await page.waitForSelector('[data-testid="add-folder-button"]', { + timeout: 3000, + }); + + await page.getByTestId("add-folder-button").click(); + + //wait for the folder to be created and changed to the new folder + await page.waitForTimeout(1000); + + await page.getByTestId("sidebar-nav-My Projects").click(); + + await page.getByText(randomName).hover(); + + await page + .getByTestId("list-card") + .first() + .dragTo(page.locator('//*[@id="sidebar-nav-New Folder"]')); + + //wait for the drag and drop to be completed + await page.waitForTimeout(1000); + + await page.getByTestId("sidebar-nav-New Folder").click(); + + await page.waitForSelector('[data-testid="list-card"]', { + timeout: 3000, + }); + + const flowNameCount = await page.getByText(randomName).count(); + expect(flowNameCount).toBeGreaterThan(0); +});