fix: Add Cascade Delete Function for Transactions and Builds Associated with Flows (#3848)

* refactor: Add cascade delete functionality for flows

This commit adds a new function `cascade_delete_flow` to the `utils.py` file in the `langflow.api` module. This function is responsible for deleting related records when a flow is deleted. It uses the `delete` method from SQLAlchemy to delete records from the `TransactionTable` and `VertexBuildTable` tables based on the flow ID. Finally, it deletes the flow record itself from the `Flow` table.

The function is wrapped in a try-except block to handle any exceptions that may occur during the deletion process. If an exception is raised, a `RuntimeError` is raised with an appropriate error message.

This refactor improves the code by encapsulating the cascade delete logic in a separate function, making it more modular and easier to maintain.

* refactor: Add cascade delete functionality for flows

* refactor: Add cascade delete functionality for flows and folders

* refactor: Remove unused delete_flow_by_id function

* refactor: Add cascade delete functionality for flows and folders

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

src/backend/base/langflow/api/utils.py

@@ -3,6 +3,9 @@ import warnings
 from typing import TYPE_CHECKING, Any
 
 from fastapi import HTTPException
+from langflow.services.database.models.transactions.model import TransactionTable
+from langflow.services.database.models.vertex_builds.model import VertexBuildTable
+from sqlalchemy import delete
 from sqlmodel import Session
 
 from langflow.graph.graph.base import Graph
@@ -241,3 +244,12 @@ def parse_value(value: Any, input_type: str) -> Any:
         return float(value) if value is not None else None
     else:
         return value
+
+
+async def cascade_delete_flow(session: Session, flow: Flow):
+    try:
+        session.exec(delete(TransactionTable).where(TransactionTable.flow_id == flow.id))  # type: ignore
+        session.exec(delete(VertexBuildTable).where(VertexBuildTable.flow_id == flow.id))  # type: ignore
+        session.exec(delete(Flow).where(Flow.id == flow.id))  # type: ignore
+    except Exception as e:
+        raise RuntimeError(f"Unable to cascade delete flow: ${flow.id}", e)

src/backend/base/langflow/api/v1/flows.py

@@ -12,12 +12,12 @@ from fastapi.responses import StreamingResponse
 from loguru import logger
 from sqlmodel import Session, and_, col, select
 
-from langflow.api.utils import remove_api_keys, validate_is_component
+from langflow.api.utils import cascade_delete_flow, remove_api_keys, validate_is_component
 from langflow.api.v1.schemas import FlowListCreate
 from langflow.initial_setup.setup import STARTER_FOLDER_NAME
 from langflow.services.auth.utils import get_current_active_user
 from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate
-from langflow.services.database.models.flow.utils import delete_flow_by_id, get_webhook_component_in_flow
+from langflow.services.database.models.flow.utils import get_webhook_component_in_flow
 from langflow.services.database.models.folder.constants import DEFAULT_FOLDER_NAME
 from langflow.services.database.models.folder.model import Folder
 from langflow.services.database.models.transactions.crud import get_transactions_by_flow_id
@@ -251,7 +251,7 @@ def update_flow(
 
 
 @router.delete("/{flow_id}", status_code=200)
-def delete_flow(
+async def delete_flow(
     *,
     session: Session = Depends(get_session),
     flow_id: UUID,
@@ -267,7 +267,7 @@ def delete_flow(
     )
     if not flow:
         raise HTTPException(status_code=404, detail="Flow not found")
-    delete_flow_by_id(str(flow_id), session)
+    await cascade_delete_flow(session, flow)
     session.commit()
     return {"message": "Flow deleted successfully"}
 

src/backend/base/langflow/api/v1/folders.py

@@ -1,3 +1,4 @@
+from langflow.api.utils import cascade_delete_flow
 import orjson
 from fastapi import APIRouter, Depends, File, HTTPException, Response, UploadFile, status
 from sqlalchemy import or_, update
@@ -171,22 +172,24 @@ def update_folder(
 
 
 @router.delete("/{folder_id}", status_code=204)
-def delete_folder(
+async def delete_folder(
     *,
     session: Session = Depends(get_session),
     folder_id: str,
     current_user: User = Depends(get_current_active_user),
 ):
     try:
+        flows = session.exec(select(Flow).where(Flow.folder_id == folder_id, Folder.user_id == current_user.id)).all()
+        if len(flows) > 0:
+            for flow in flows:
+                await cascade_delete_flow(session, flow)
+
         folder = session.exec(select(Folder).where(Folder.id == folder_id, Folder.user_id == current_user.id)).first()
         if not folder:
             raise HTTPException(status_code=404, detail="Folder not found")
         session.delete(folder)
         session.commit()
-        flows = session.exec(select(Flow).where(Flow.folder_id == folder_id, Folder.user_id == current_user.id)).all()
-        for flow in flows:
-            session.delete(flow)
-        session.commit()
+
         return Response(status_code=status.HTTP_204_NO_CONTENT)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

src/backend/base/langflow/services/database/models/flow/utils.py

@@ -3,13 +3,10 @@ from typing import Optional
 from fastapi import Depends
 from langflow.utils.version import get_version_info
 from sqlmodel import Session
-from sqlalchemy import delete
 
 from langflow.services.deps import get_session
 
 from .model import Flow
-from .. import TransactionTable, MessageTable
-from loguru import logger
 
 
 def get_flow_by_id(session: Session = Depends(get_session), flow_id: Optional[str] = None) -> Flow | None:
@@ -21,15 +18,6 @@ def get_flow_by_id(session: Session = Depends(get_session), flow_id: Optional[st
     return session.get(Flow, flow_id)
 
 
-def delete_flow_by_id(flow_id: str, session: Session) -> None:
-    """Delete flow by id."""
-    # Manually delete flow, transactions and messages because foreign key constraints might be disabled
-    session.exec(delete(Flow).where(Flow.id == flow_id))  # type: ignore
-    session.exec(delete(TransactionTable).where(TransactionTable.flow_id == flow_id))  #  type: ignore
-    session.exec(delete(MessageTable).where(MessageTable.flow_id == flow_id))  #  type: ignore
-    logger.info(f"Deleted flow {flow_id}")
-
-
 def get_webhook_component_in_flow(flow_data: dict):
     """Get webhook component in flow data."""