diff --git a/api/tasks/delete_segment_from_index_task.py b/api/tasks/delete_segment_from_index_task.py index 6fc310138e8..df29e9e98af 100644 --- a/api/tasks/delete_segment_from_index_task.py +++ b/api/tasks/delete_segment_from_index_task.py @@ -7,6 +7,7 @@ from sqlalchemy import delete, select from core.db.session_factory import session_factory from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from extensions.ext_storage import storage from models.dataset import Dataset, Document, SegmentAttachmentBinding from models.model import UploadFile @@ -63,25 +64,73 @@ def delete_segment_from_index_task( if dataset.is_multimodal: # delete segment attachment binding segment_attachment_bindings = session.scalars( - select(SegmentAttachmentBinding).where(SegmentAttachmentBinding.segment_id.in_(segment_ids)) + select(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.tenant_id == dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == dataset.id, + SegmentAttachmentBinding.document_id == document_id, + SegmentAttachmentBinding.segment_id.in_(segment_ids), + ) ).all() if segment_attachment_bindings: - attachment_ids = [binding.attachment_id for binding in segment_attachment_bindings] - index_processor.clean( - session=session, dataset=dataset, node_ids=attachment_ids, with_keywords=False - ) segment_attachment_bind_ids = [i.id for i in segment_attachment_bindings] + attachment_ids = list( + dict.fromkeys(binding.attachment_id for binding in segment_attachment_bindings) + ) for i in range(0, len(segment_attachment_bind_ids), 1000): segment_attachment_bind_delete_stmt = delete(SegmentAttachmentBinding).where( - SegmentAttachmentBinding.id.in_(segment_attachment_bind_ids[i : i + 1000]) + SegmentAttachmentBinding.tenant_id == dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == dataset.id, + SegmentAttachmentBinding.document_id == document_id, + SegmentAttachmentBinding.id.in_(segment_attachment_bind_ids[i : i + 1000]), ) session.execute(segment_attachment_bind_delete_stmt) - # delete upload file - session.execute(delete(UploadFile).where(UploadFile.id.in_(attachment_ids))) + session.flush() + remaining_attachment_ids = set( + session.scalars( + select(SegmentAttachmentBinding.attachment_id).where( + SegmentAttachmentBinding.attachment_id.in_(attachment_ids) + ) + ).all() + ) + orphan_attachment_ids = [ + attachment_id + for attachment_id in attachment_ids + if attachment_id not in remaining_attachment_ids + ] + + if orphan_attachment_ids: + attachment_storage_keys = list( + dict.fromkeys( + session.scalars( + select(UploadFile.key).where( + UploadFile.tenant_id == dataset.tenant_id, + UploadFile.id.in_(orphan_attachment_ids), + ) + ).all() + ) + ) + index_processor.clean( + session=session, dataset=dataset, node_ids=orphan_attachment_ids, with_keywords=False + ) + session.execute( + delete(UploadFile).where( + UploadFile.tenant_id == dataset.tenant_id, + UploadFile.id.in_(orphan_attachment_ids), + ) + ) + else: + attachment_storage_keys = [] + session.commit() + for storage_key in attachment_storage_keys: + try: + storage.delete(storage_key) + except Exception: + logger.exception("Failed to delete segment attachment from storage, key: %s", storage_key) + end_at = time.perf_counter() logger.info(click.style(f"Segment deleted from index latency: {end_at - start_at}", fg="green")) except Exception: diff --git a/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py b/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py index 8af7ac6a3b5..6a256656ad3 100644 --- a/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py +++ b/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py @@ -1,6 +1,7 @@ import uuid from collections.abc import Generator from contextlib import contextmanager +from datetime import UTC, datetime from unittest.mock import MagicMock, patch import pytest @@ -8,8 +9,10 @@ from sqlalchemy import event from sqlalchemy.orm import Session, sessionmaker from core.rag.index_processor.constant.index_type import IndexStructureType -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from extensions.storage.storage_type import StorageType +from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from models.model import UploadFile from tasks.delete_segment_from_index_task import delete_segment_from_index_task from tasks.disable_segment_from_index_task import disable_segment_from_index_task from tasks.disable_segments_from_index_task import disable_segments_from_index_task @@ -152,3 +155,167 @@ def test_delete_segment_commits_index_cleanup_without_attachments( delete_segment_from_index_task.run(["node-1"], dataset.id, document.id, [segment.id]) assert phase_events == ["clean", "commit"] + + +def test_delete_segment_removes_attachment_blobs_from_storage( + indexed_segment: tuple[Dataset, Document, DocumentSegment], + sqlite_session: Session, +) -> None: + dataset, document, segment = indexed_segment + dataset.is_multimodal = True + attachment = UploadFile( + tenant_id=dataset.tenant_id, + storage_type=StorageType.LOCAL, + key="attachments/segment-image.png", + name="segment-image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=segment.created_by, + created_at=datetime.now(UTC), + used=True, + ) + binding = SegmentAttachmentBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + attachment_id=attachment.id, + ) + sqlite_session.add_all([dataset, attachment, binding]) + sqlite_session.commit() + attachment_id = attachment.id + attachment_key = attachment.key + binding_id = binding.id + + with ( + patch("tasks.delete_segment_from_index_task.IndexProcessorFactory") as processor_factory, + patch("tasks.delete_segment_from_index_task.storage.delete") as storage_delete, + ): + delete_segment_from_index_task.run(["node-1"], dataset.id, document.id, [segment.id]) + + storage_delete.assert_called_once_with(attachment_key) + sqlite_session.expire_all() + assert sqlite_session.get(SegmentAttachmentBinding, binding_id) is None + assert sqlite_session.get(UploadFile, attachment_id) is None + + +def test_delete_segment_preserves_attachment_shared_by_another_segment( + indexed_segment: tuple[Dataset, Document, DocumentSegment], + sqlite_session: Session, +) -> None: + dataset, document, segment = indexed_segment + dataset.is_multimodal = True + other_segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=2, + content="other content", + word_count=1, + tokens=1, + created_by=segment.created_by, + index_node_id="node-2", + index_node_hash="hash-2", + disabled_by=segment.disabled_by, + status=SegmentStatus.COMPLETED, + ) + attachment = UploadFile( + tenant_id=dataset.tenant_id, + storage_type=StorageType.LOCAL, + key="attachments/shared-segment-image.png", + name="shared-segment-image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=segment.created_by, + created_at=datetime.now(UTC), + used=True, + ) + binding = SegmentAttachmentBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + attachment_id=attachment.id, + ) + shared_binding = SegmentAttachmentBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=other_segment.id, + attachment_id=attachment.id, + ) + sqlite_session.add_all([dataset, other_segment, attachment, binding, shared_binding]) + sqlite_session.commit() + attachment_id = attachment.id + binding_id = binding.id + shared_binding_id = shared_binding.id + + with ( + patch("tasks.delete_segment_from_index_task.IndexProcessorFactory") as processor_factory, + patch("tasks.delete_segment_from_index_task.storage.delete") as storage_delete, + ): + delete_segment_from_index_task.run(["node-1"], dataset.id, document.id, [segment.id]) + + processor = processor_factory.return_value.init_index_processor.return_value + assert processor.clean.call_count == 1 + assert processor.clean.call_args.args[1] == ["node-1"] + assert processor.clean.call_args.kwargs["with_keywords"] is True + storage_delete.assert_not_called() + sqlite_session.expire_all() + assert sqlite_session.get(SegmentAttachmentBinding, binding_id) is None + assert sqlite_session.get(SegmentAttachmentBinding, shared_binding_id) is not None + assert sqlite_session.get(UploadFile, attachment_id) is not None + + +def test_delete_segment_keeps_database_cleanup_when_attachment_storage_delete_fails( + indexed_segment: tuple[Dataset, Document, DocumentSegment], + sqlite_session: Session, + caplog: pytest.LogCaptureFixture, +) -> None: + dataset, document, segment = indexed_segment + dataset.is_multimodal = True + attachment = UploadFile( + tenant_id=dataset.tenant_id, + storage_type=StorageType.LOCAL, + key="attachments/failing-segment-image.png", + name="failing-segment-image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=segment.created_by, + created_at=datetime.now(UTC), + used=True, + ) + binding = SegmentAttachmentBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + attachment_id=attachment.id, + ) + sqlite_session.add_all([dataset, attachment, binding]) + sqlite_session.commit() + attachment_id = attachment.id + attachment_key = attachment.key + binding_id = binding.id + + with ( + patch("tasks.delete_segment_from_index_task.IndexProcessorFactory") as processor_factory, + patch( + "tasks.delete_segment_from_index_task.storage.delete", + side_effect=RuntimeError("storage unavailable"), + ) as storage_delete, + caplog.at_level("ERROR", logger="tasks.delete_segment_from_index_task"), + ): + delete_segment_from_index_task.run(["node-1"], dataset.id, document.id, [segment.id]) + + storage_delete.assert_called_once_with(attachment_key) + assert "Failed to delete segment attachment from storage" in caplog.text + sqlite_session.expire_all() + assert sqlite_session.get(SegmentAttachmentBinding, binding_id) is None + assert sqlite_session.get(UploadFile, attachment_id) is None