From 05404c765b24897eeeb76eec2e1aeb596f79ad15 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Tue, 10 Mar 2026 16:44:17 -0300 Subject: [PATCH] fix(api): Handle Windows ChromaDB file locks when deleting Knowledge Bases (#12132) Co-authored-by: Claude Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../base/langflow/api/utils/kb_helpers.py | 113 ++++- .../base/langflow/api/v1/knowledge_bases.py | 31 +- .../base/langflow/utils/kb_constants.py | 4 + .../tests/unit/test_kb_storage_deletion.py | 392 ++++++++++++++++++ .../tests/unit/test_knowledge_bases_api.py | 12 +- 5 files changed, 521 insertions(+), 31 deletions(-) create mode 100644 src/backend/tests/unit/test_kb_storage_deletion.py diff --git a/src/backend/base/langflow/api/utils/kb_helpers.py b/src/backend/base/langflow/api/utils/kb_helpers.py index 466125bed3..3d5481255b 100644 --- a/src/backend/base/langflow/api/utils/kb_helpers.py +++ b/src/backend/base/langflow/api/utils/kb_helpers.py @@ -2,6 +2,8 @@ import asyncio import contextlib import gc import json +import shutil +import time import uuid from datetime import datetime, timezone from functools import lru_cache @@ -25,8 +27,10 @@ from langflow.services.database.models.jobs.model import JobStatus from langflow.services.deps import get_settings_service from langflow.services.jobs.service import JobService from langflow.utils.kb_constants import ( + DELETE_BACKOFF_SECONDS, EXPONENTIAL_BACKOFF_MULTIPLIER, INGESTION_BATCH_SIZE, + MAX_DELETE_RETRIES, MAX_RETRY_ATTEMPTS, ) @@ -81,8 +85,31 @@ class KBStorageHelper: ) @staticmethod - def teardown_storage(kb_path: Path, kb_name: str) -> None: - """Explicitly flush and invalidate Chroma clients before directory deletion.""" + def release_chroma_resources(kb_path: Path) -> None: + """Release ChromaDB resources by clearing the registry entry and forcing GC.""" + path_key = str(kb_path) + try: + if path_key in SharedSystemClient._identifier_to_system: # noqa: SLF001 + del SharedSystemClient._identifier_to_system[path_key] # noqa: SLF001 + except KeyError: + pass + gc.collect() + + @staticmethod + def delete_storage(kb_path: Path, kb_name: str) -> bool: + """Teardown ChromaDB connections and delete KB directory with retry logic. + + Handles ChromaDB SQLite file locks that can prevent deletion, particularly + on Windows where mandatory file locks block deletion of open files. + Uses retry with exponential backoff and rename-as-fallback strategy. + + Returns: + True if deletion succeeded (or path already gone), False otherwise. + """ + if not kb_path.exists(): + return True + + # Teardown ChromaDB collection to release handles try: has_data = any((kb_path / m).exists() for m in ["chroma", "chroma.sqlite3", "index"]) if has_data: @@ -91,9 +118,70 @@ class KBStorageHelper: with contextlib.suppress(Exception): chroma.delete_collection() chroma = None - gc.collect() + client = None except (OSError, ValueError, TypeError, chromadb.errors.ChromaError) as e: - logger.debug(f"Storage teardown failed for {kb_path.name} (ignoring): {e}") + logger.debug("Collection teardown failed for %s: %s", kb_path.name, e) + + gc.collect() + + for attempt in range(MAX_DELETE_RETRIES): + try: + if attempt > 0: + time.sleep(DELETE_BACKOFF_SECONDS * (2**attempt)) + + _remove_sqlite_lock_files(kb_path) + _truncate_sqlite_files(kb_path) + gc.collect() + + shutil.rmtree(kb_path, ignore_errors=False) + + if not kb_path.exists(): + logger.info("Deleted knowledge base %s on attempt %d", kb_name, attempt + 1) + return True + + except OSError as e: + if attempt < MAX_DELETE_RETRIES - 1: + logger.debug("KB deletion attempt %d failed for %s: %s", attempt + 1, kb_name, e) + else: + logger.warning( + "KB deletion failed for %s after %d attempts: %s", + kb_name, + MAX_DELETE_RETRIES, + e, + ) + + # Last resort: rename for deferred cleanup + if kb_path.exists(): + try: + deferred = kb_path.with_name(f".deleted_{kb_name}_{int(time.time())}") + kb_path.rename(deferred) + except OSError as e: + logger.warning("Deferred rename failed for %s: %s", kb_name, e) + else: + logger.info("Renamed %s for deferred cleanup", kb_name) + return True + + return False + + +def _remove_sqlite_lock_files(kb_path: Path) -> None: + """Remove SQLite auxiliary files (WAL, SHM, journal) that hold locks.""" + for pattern in ["*.sqlite3-wal", "*.sqlite3-shm", "*.sqlite3-journal"]: + for lock_file in kb_path.glob(pattern): + try: + lock_file.unlink() + except OSError as e: + logger.debug("Could not remove lock file %s: %s", lock_file.name, e) + + +def _truncate_sqlite_files(kb_path: Path) -> None: + """Truncate SQLite database files to release locks.""" + for sqlite_file in kb_path.glob("*.sqlite3"): + try: + with sqlite_file.open("r+b") as f: + f.truncate(0) + except OSError as e: + logger.debug("Could not truncate %s: %s", sqlite_file.name, e) class KBAnalysisHelper: @@ -154,11 +242,15 @@ class KBAnalysisHelper: @staticmethod def update_text_metrics(kb_path: Path, metadata: dict, chroma: Chroma | None = None) -> None: """Update text metrics (chunks, words, characters) for a knowledge base.""" + created_locally = chroma is None + client = None try: - if chroma is None: + if created_locally: client = KBStorageHelper.get_fresh_chroma_client(kb_path) chroma = Chroma(client=client, collection_name=kb_path.name) + if chroma is None: + return collection = chroma._collection # noqa: SLF001 metadata["chunks"] = collection.count() @@ -190,6 +282,11 @@ class KBAnalysisHelper: ) except (OSError, ValueError, TypeError, json.JSONDecodeError, chromadb.errors.ChromaError) as e: logger.debug(f"Metrics update failed for {kb_path.name}: {e}") + finally: + if created_locally: + client = None + chroma = None + KBStorageHelper.release_chroma_resources(kb_path) @staticmethod def _detect_embedding_provider(kb_path: Path) -> str: @@ -417,8 +514,9 @@ class KBIngestionHelper: await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name) raise finally: + client = None chroma = None - gc.collect() + KBStorageHelper.release_chroma_resources(kb_path) @staticmethod async def cleanup_chroma_chunks_by_job( @@ -438,8 +536,9 @@ class KBIngestionHelper: except (OSError, ValueError, TypeError, chromadb.errors.ChromaError) as cleanup_error: await logger.aerror(f"Failed to clean up chunks for job {job_id}: {cleanup_error}") finally: + client = None chroma = None - gc.collect() + KBStorageHelper.release_chroma_resources(kb_path) @staticmethod async def _is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool: diff --git a/src/backend/base/langflow/api/v1/knowledge_bases.py b/src/backend/base/langflow/api/v1/knowledge_bases.py index ec07311ee9..6f11838159 100644 --- a/src/backend/base/langflow/api/v1/knowledge_bases.py +++ b/src/backend/base/langflow/api/v1/knowledge_bases.py @@ -1,7 +1,5 @@ import asyncio -import gc import json -import shutil import uuid from datetime import datetime, timezone from http import HTTPStatus @@ -78,11 +76,11 @@ async def create_knowledge_base( try: client = KBStorageHelper.get_fresh_chroma_client(kb_path) client.create_collection(name=kb_name) - # Explicitly delete reference to help release handle - client = None - gc.collect() except (OSError, ValueError, chromadb.errors.ChromaError) as e: logger.warning("Initial Chroma setup for %s failed: %s", kb_name, e) + finally: + client = None + KBStorageHelper.release_chroma_resources(kb_path) # Serialize column_config for persistence column_config_dicts = None @@ -130,7 +128,7 @@ async def create_knowledge_base( except Exception as e: # Clean up if something went wrong if kb_path.exists(): - shutil.rmtree(kb_path) + KBStorageHelper.delete_storage(kb_path, kb_name) await logger.aerror("Error creating knowledge base: %s", e) raise HTTPException(status_code=500, detail="Internal error creating knowledge base") from e @@ -593,8 +591,9 @@ async def get_knowledge_base_chunks( await logger.aerror("Error getting chunks for '%s': %s", kb_name, e) raise HTTPException(status_code=500, detail="Error getting chunks.") from e finally: + client = None chroma = None - gc.collect() + KBStorageHelper.release_chroma_resources(kb_path) @router.delete("/{kb_name}", status_code=HTTPStatus.OK) @@ -603,11 +602,11 @@ async def delete_knowledge_base(kb_name: str, current_user: CurrentActiveUser) - try: kb_path = _resolve_kb_path(kb_name, current_user) - # Explicitly teardown KB storage to flush Chroma handles before directory deletion - KBStorageHelper.teardown_storage(kb_path, kb_name) - - # Delete the entire knowledge base directory - shutil.rmtree(kb_path) + if not KBStorageHelper.delete_storage(kb_path, kb_name): + raise HTTPException( + status_code=500, + detail=f"Failed to delete knowledge base '{kb_name}'. The database may be in use.", + ) except HTTPException: raise @@ -636,12 +635,8 @@ async def delete_knowledge_bases_bulk(request: BulkDeleteRequest, current_user: continue try: - # Explicitly teardown KB storage to flush Chroma handles before directory deletion - KBStorageHelper.teardown_storage(kb_path, kb_name) - - # Delete the entire knowledge base directory - shutil.rmtree(kb_path) - deleted_count += 1 + if KBStorageHelper.delete_storage(kb_path, kb_name): + deleted_count += 1 except (OSError, PermissionError) as e: await logger.aexception("Error deleting knowledge base '%s': %s", kb_name, e) # Continue with other deletions even if one fails diff --git a/src/backend/base/langflow/utils/kb_constants.py b/src/backend/base/langflow/utils/kb_constants.py index 1eaee8cb89..b05b876b21 100644 --- a/src/backend/base/langflow/utils/kb_constants.py +++ b/src/backend/base/langflow/utils/kb_constants.py @@ -3,3 +3,7 @@ INGESTION_BATCH_SIZE = 200 EXPONENTIAL_BACKOFF_MULTIPLIER = 2 MIN_KB_NAME_LENGTH = 3 CHUNK_PREVIEW_MULTIPLIER = 3 + +# KB deletion retry constants +MAX_DELETE_RETRIES = 5 +DELETE_BACKOFF_SECONDS = 0.5 diff --git a/src/backend/tests/unit/test_kb_storage_deletion.py b/src/backend/tests/unit/test_kb_storage_deletion.py new file mode 100644 index 0000000000..ba6345e4ea --- /dev/null +++ b/src/backend/tests/unit/test_kb_storage_deletion.py @@ -0,0 +1,392 @@ +"""Tests for unified Knowledge Base deletion and resource cleanup. + +Covers the delete_storage method, release_chroma_resources, private helpers +(_remove_sqlite_lock_files, _truncate_sqlite_files), and the delete endpoints. +""" + +import json +import shutil +import uuid +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def kb_dir(tmp_path): + """Create a fake KB directory with SQLite files simulating ChromaDB.""" + kb = tmp_path / "test_kb" + kb.mkdir() + (kb / "chroma.sqlite3").write_bytes(b"fake-sqlite-data") + (kb / "chroma.sqlite3-wal").write_bytes(b"wal-data") + (kb / "chroma.sqlite3-shm").write_bytes(b"shm-data") + (kb / "embedding_metadata.json").write_text(json.dumps({"id": str(uuid.uuid4())})) + return kb + + +@pytest.fixture +def empty_kb_dir(tmp_path): + """Create a KB directory with no ChromaDB data files.""" + kb = tmp_path / "empty_kb" + kb.mkdir() + return kb + + +# =========================================================================== +# Unit tests: _remove_sqlite_lock_files +# =========================================================================== + + +class TestRemoveSqliteLockFiles: + """Tests for _remove_sqlite_lock_files — removes WAL, SHM, journal files.""" + + def test_should_remove_all_lock_files(self, kb_dir): + from langflow.api.utils.kb_helpers import _remove_sqlite_lock_files + + (kb_dir / "chroma.sqlite3-journal").write_bytes(b"journal") + assert (kb_dir / "chroma.sqlite3-wal").exists() + assert (kb_dir / "chroma.sqlite3-shm").exists() + + _remove_sqlite_lock_files(kb_dir) + + assert not (kb_dir / "chroma.sqlite3-wal").exists() + assert not (kb_dir / "chroma.sqlite3-shm").exists() + assert not (kb_dir / "chroma.sqlite3-journal").exists() + assert (kb_dir / "chroma.sqlite3").exists() + + def test_should_not_raise_when_no_lock_files(self, empty_kb_dir): + from langflow.api.utils.kb_helpers import _remove_sqlite_lock_files + + _remove_sqlite_lock_files(empty_kb_dir) + + def test_should_handle_permission_error_gracefully(self, kb_dir): + from langflow.api.utils.kb_helpers import _remove_sqlite_lock_files + + with patch.object(Path, "unlink", side_effect=OSError("Permission denied")): + _remove_sqlite_lock_files(kb_dir) + + +# =========================================================================== +# Unit tests: _truncate_sqlite_files +# =========================================================================== + + +class TestTruncateSqliteFiles: + """Tests for _truncate_sqlite_files — truncates .sqlite3 files.""" + + def test_should_truncate_sqlite_files_to_zero(self, kb_dir): + from langflow.api.utils.kb_helpers import _truncate_sqlite_files + + assert (kb_dir / "chroma.sqlite3").stat().st_size > 0 + + _truncate_sqlite_files(kb_dir) + + assert (kb_dir / "chroma.sqlite3").stat().st_size == 0 + + def test_should_not_raise_when_no_sqlite_files(self, empty_kb_dir): + from langflow.api.utils.kb_helpers import _truncate_sqlite_files + + _truncate_sqlite_files(empty_kb_dir) + + def test_should_handle_locked_file_gracefully(self, kb_dir): + from langflow.api.utils.kb_helpers import _truncate_sqlite_files + + with patch("builtins.open", side_effect=OSError("File is locked")): + _truncate_sqlite_files(kb_dir) + + +# =========================================================================== +# Unit tests: KBStorageHelper.release_chroma_resources +# =========================================================================== + + +class TestReleaseChromaResources: + """Tests for release_chroma_resources — clears registry and forces GC.""" + + def test_should_clear_registry_entry_for_path(self, kb_dir): + from chromadb.api.shared_system_client import SharedSystemClient + from langflow.api.utils.kb_helpers import KBStorageHelper + + path_key = str(kb_dir) + SharedSystemClient._identifier_to_system[path_key] = MagicMock() + + KBStorageHelper.release_chroma_resources(kb_dir) + + assert path_key not in SharedSystemClient._identifier_to_system + + def test_should_not_raise_when_path_not_in_registry(self, tmp_path): + from langflow.api.utils.kb_helpers import KBStorageHelper + + KBStorageHelper.release_chroma_resources(tmp_path / "nonexistent") + + @patch("langflow.api.utils.kb_helpers.gc.collect") + def test_should_call_gc_collect(self, mock_gc, tmp_path): + from langflow.api.utils.kb_helpers import KBStorageHelper + + KBStorageHelper.release_chroma_resources(tmp_path) + + mock_gc.assert_called_once() + + +# =========================================================================== +# Unit tests: KBStorageHelper.delete_storage +# =========================================================================== + + +class TestDeleteStorage: + """Tests for KBStorageHelper.delete_storage — unified deletion with retry.""" + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_return_true_when_path_does_not_exist(self, tmp_path): + from langflow.api.utils.kb_helpers import KBStorageHelper + + non_existent = tmp_path / "does_not_exist" + result = KBStorageHelper.delete_storage(non_existent, "ghost_kb") + + assert result is True + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_delete_directory_on_first_attempt(self, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is True + assert not kb_dir.exists() + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_retry_and_succeed_on_second_attempt(self, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + original_rmtree = shutil.rmtree + call_count = 0 + + def rmtree_fails_once(path, *, ignore_errors=False): + nonlocal call_count + call_count += 1 + if call_count == 1: + msg = "[WinError 32] File in use" + raise OSError(msg) + original_rmtree(path, ignore_errors=ignore_errors) + + with patch("langflow.api.utils.kb_helpers.shutil.rmtree", side_effect=rmtree_fails_once): + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is True + assert call_count == 2 + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_rename_as_fallback_when_all_retries_fail(self, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + with patch( + "langflow.api.utils.kb_helpers.shutil.rmtree", + side_effect=OSError("[WinError 32] File in use"), + ): + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is True + assert not kb_dir.exists() + renamed_dirs = [p for p in kb_dir.parent.iterdir() if p.name.startswith(".deleted_test_kb_")] + assert len(renamed_dirs) == 1 + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_return_false_when_all_strategies_fail(self, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + with ( + patch( + "langflow.api.utils.kb_helpers.shutil.rmtree", + side_effect=OSError("[WinError 32] File in use"), + ), + patch.object(Path, "rename", side_effect=OSError("Cannot rename")), + ): + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is False + assert kb_dir.exists() + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep") + def test_should_use_exponential_backoff_on_retries(self, mock_sleep, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + real_rmtree = shutil.rmtree + call_count = 0 + + def rmtree_fails_three_times(path, *, ignore_errors=False): + nonlocal call_count + call_count += 1 + if call_count <= 3: + msg = "[WinError 32] File in use" + raise OSError(msg) + real_rmtree(path, ignore_errors=ignore_errors) + + with patch("langflow.api.utils.kb_helpers.shutil.rmtree", side_effect=rmtree_fails_three_times): + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is True + sleep_values = [call.args[0] for call in mock_sleep.call_args_list] + assert sleep_values == [1.0, 2.0, 4.0] + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client") + @patch("langflow.api.utils.kb_helpers.Chroma") + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_teardown_collection_before_deletion(self, mock_chroma_cls, mock_client_cls, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + mock_chroma = MagicMock() + mock_chroma_cls.return_value = mock_chroma + mock_client_cls.return_value = MagicMock() + + KBStorageHelper.delete_storage(kb_dir, "test_kb") + + mock_chroma.delete_collection.assert_called_once() + assert not kb_dir.exists() + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client") + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_not_fail_when_teardown_raises(self, mock_client_cls, kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + mock_client_cls.side_effect = OSError("Cannot open database") + + result = KBStorageHelper.delete_storage(kb_dir, "test_kb") + + assert result is True + assert not kb_dir.exists() + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + def test_should_skip_teardown_when_no_chroma_data(self, empty_kb_dir): + from langflow.api.utils.kb_helpers import KBStorageHelper + + result = KBStorageHelper.delete_storage(empty_kb_dir, "empty_kb") + + assert result is True + assert not empty_kb_dir.exists() + + +# =========================================================================== +# Integration tests: delete endpoints using unified delete_storage +# =========================================================================== + + +class TestDeleteEndpoint: + """Tests that delete endpoint uses KBStorageHelper.delete_storage.""" + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") + async def test_should_delete_kb_successfully(self, mock_root, client, logged_in_headers, tmp_path): + mock_root.return_value = tmp_path + (tmp_path / "activeuser" / "My_KB").mkdir(parents=True) + + response = await client.delete("api/v1/knowledge_bases/My_KB", headers=logged_in_headers) + + assert response.status_code == 200 + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.delete_storage", return_value=False) + @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") + async def test_should_return_500_when_deletion_fails( + self, mock_root, mock_delete, client, logged_in_headers, tmp_path + ): + mock_root.return_value = tmp_path + (tmp_path / "activeuser" / "My_KB").mkdir(parents=True) + + response = await client.delete("api/v1/knowledge_bases/My_KB", headers=logged_in_headers) + + assert response.status_code == 500 + assert "may be in use" in response.json()["detail"] + mock_delete.assert_called_once() + + async def test_should_return_404_when_kb_not_found(self, client, logged_in_headers): + response = await client.delete("api/v1/knowledge_bases/NonExistent_KB", headers=logged_in_headers) + + assert response.status_code == 404 + + +class TestBulkDeleteEndpoint: + """Tests that bulk delete endpoint uses KBStorageHelper.delete_storage.""" + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.Chroma", new=MagicMock()) + @patch("langflow.api.utils.kb_helpers.time.sleep", new=MagicMock()) + @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") + async def test_should_delete_multiple_kbs(self, mock_root, client, logged_in_headers, tmp_path): + mock_root.return_value = tmp_path + kb_user_path = tmp_path / "activeuser" + kb_user_path.mkdir(parents=True) + (kb_user_path / "KB1").mkdir() + (kb_user_path / "KB2").mkdir() + + response = await client.request( + "DELETE", + "api/v1/knowledge_bases", + headers=logged_in_headers, + json={"kb_names": ["KB1", "KB2"]}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["deleted_count"] == 2 + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.delete_storage") + @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") + async def test_should_handle_partial_failure(self, mock_root, mock_delete, client, logged_in_headers, tmp_path): + mock_root.return_value = tmp_path + kb_user_path = tmp_path / "activeuser" + kb_user_path.mkdir(parents=True) + (kb_user_path / "KB1").mkdir() + (kb_user_path / "KB2").mkdir() + + mock_delete.side_effect = [True, False] + + response = await client.request( + "DELETE", + "api/v1/knowledge_bases", + headers=logged_in_headers, + json={"kb_names": ["KB1", "KB2"]}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["deleted_count"] == 1 + + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.delete_storage", new=MagicMock(return_value=True)) + @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") + async def test_should_report_not_found_kbs(self, mock_root, client, logged_in_headers, tmp_path): + mock_root.return_value = tmp_path + kb_user_path = tmp_path / "activeuser" + kb_user_path.mkdir(parents=True) + (kb_user_path / "KB1").mkdir() + + response = await client.request( + "DELETE", + "api/v1/knowledge_bases", + headers=logged_in_headers, + json={"kb_names": ["KB1", "Ghost"]}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["deleted_count"] == 1 + assert "Ghost" in data["not_found"] diff --git a/src/backend/tests/unit/test_knowledge_bases_api.py b/src/backend/tests/unit/test_knowledge_bases_api.py index 07065d4a41..34f02e1c70 100644 --- a/src/backend/tests/unit/test_knowledge_bases_api.py +++ b/src/backend/tests/unit/test_knowledge_bases_api.py @@ -259,22 +259,22 @@ class TestKnowledgeBaseAPI: assert data["chunks"] == 5 assert data["name"] == "Detail KB" - @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.teardown_storage") + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.delete_storage", return_value=True) @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") async def test_delete_knowledge_base( - self, mock_root, mock_teardown, client: AsyncClient, logged_in_headers, tmp_path + self, mock_root, mock_delete, client: AsyncClient, logged_in_headers, tmp_path ): mock_root.return_value = tmp_path (tmp_path / "activeuser" / "To_Delete").mkdir(parents=True, exist_ok=True) response = await client.delete("api/v1/knowledge_bases/To_Delete", headers=logged_in_headers) assert response.status_code == 200 - mock_teardown.assert_called_once() + mock_delete.assert_called_once() - @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.teardown_storage") + @patch("langflow.api.utils.kb_helpers.KBStorageHelper.delete_storage", return_value=True) @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") async def test_bulk_delete_knowledge_bases( - self, mock_root, mock_teardown, client: AsyncClient, logged_in_headers, tmp_path + self, mock_root, mock_delete, client: AsyncClient, logged_in_headers, tmp_path ): mock_root.return_value = tmp_path kb_user_path = tmp_path / "activeuser" @@ -292,7 +292,7 @@ class TestKnowledgeBaseAPI: data = response.json() assert data["deleted_count"] == 2 assert "NonExistent" in data["not_found"] - assert mock_teardown.called + assert mock_delete.called @patch("langflow.api.v1.knowledge_bases.KBStorageHelper.get_root_path") @patch("langflow.api.v1.knowledge_bases.KBAnalysisHelper.get_metadata")