mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-25 05:13:10 +08:00
fix(api): Handle Windows ChromaDB file locks when deleting Knowledge Bases (#12132)
Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Cristhian Zanforlin Lousa
committed by
GitHub
GitHub
parent
d7cb9a08b1
commit
05404c765b
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
392
src/backend/tests/unit/test_kb_storage_deletion.py
Normal file
392
src/backend/tests/unit/test_kb_storage_deletion.py
Normal file
@ -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"]
|
||||
@ -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")
|
||||
|
||||
Reference in New Issue
Block a user