Skip to content

Commit

Permalink
fix: Update temp flow cleanup to use async file operations
Browse files Browse the repository at this point in the history
- Add explicit type hint for storage_service
- Use async methods for file and directory existence checks
- Improve error handling for file deletion during flow cleanup
  • Loading branch information
ogabrielluiz committed Feb 14, 2025
1 parent de81c16 commit 160c25a
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions src/backend/base/langflow/services/task/temp_flow_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import asyncio
import contextlib
from typing import TYPE_CHECKING

from loguru import logger
from sqlmodel import col, delete, select
Expand All @@ -11,6 +12,9 @@
from langflow.services.database.models.vertex_builds.model import VertexBuildTable
from langflow.services.deps import get_settings_service, get_storage_service, session_scope

if TYPE_CHECKING:
from langflow.services.storage.service import StorageService


async def cleanup_orphaned_records() -> None:
"""Clean up all records that reference non-existent flows."""
Expand Down Expand Up @@ -39,7 +43,7 @@ async def cleanup_orphaned_records() -> None:
await session.exec(delete(table).where(col(table.flow_id).in_(orphaned_flow_ids)))

# Clean up any associated storage files
storage_service = get_storage_service()
storage_service: StorageService = get_storage_service()
for flow_id in orphaned_flow_ids:
try:
files = await storage_service.list_files(str(flow_id))
Expand All @@ -50,8 +54,8 @@ async def cleanup_orphaned_records() -> None:
logger.error(f"Failed to delete file {file} for flow {flow_id}: {exc!s}")
# Delete the flow directory after all files are deleted
flow_dir = storage_service.data_dir / str(flow_id)
if flow_dir.exists():
flow_dir.rmdir()
if await flow_dir.exists():
await flow_dir.rmdir()
except Exception as exc: # noqa: BLE001
logger.error(f"Failed to list files for flow {flow_id}: {exc!s}")

Expand Down

0 comments on commit 160c25a

Please sign in to comment.