diff --git a/api/tasks/document_indexing_update_task.py b/api/tasks/document_indexing_update_task.py index 59ffae8d5b8..18e2338e97a 100644 --- a/api/tasks/document_indexing_update_task.py +++ b/api/tasks/document_indexing_update_task.py @@ -9,9 +9,11 @@ from core.db.session_factory import session_factory from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from extensions.ext_storage import storage from libs.datetime_utils import naive_utc_now -from models.dataset import Dataset, Document, DocumentSegment +from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding from models.enums import IndexingStatus +from models.model import UploadFile from tasks.generate_summary_index_task import generate_summary_index_task logger = logging.getLogger(__name__) @@ -54,6 +56,7 @@ def document_indexing_update_task(dataset_id: str, document_id: str): session.commit() clean_success = False + index_processor = None try: index_processor = IndexProcessorFactory(index_type).init_index_processor() if index_node_ids: @@ -87,11 +90,80 @@ def document_indexing_update_task(dataset_id: str, document_id: str): document.processing_started_at = naive_utc_now() session.commit() - if clean_success: + if clean_success and index_processor is not None: + attachment_storage_keys: list[str] = [] + if dataset.is_multimodal: + segment_attachment_bindings = session.scalars( + select(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.tenant_id == dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == dataset.id, + SegmentAttachmentBinding.document_id == document_id, + ) + ).all() + attachment_ids = list( + dict.fromkeys(binding.attachment_id for binding in segment_attachment_bindings) + ) + + if segment_attachment_bindings: + session.execute( + delete(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.tenant_id == dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == dataset.id, + SegmentAttachmentBinding.document_id == document_id, + ) + ) + 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), + ) + ) + segment_delete_stmt = delete(DocumentSegment).where(DocumentSegment.document_id == document_id) session.execute(segment_delete_stmt) session.commit() + for storage_key in attachment_storage_keys: + try: + storage.delete(storage_key) + except Exception: + logger.exception( + "Failed to delete document attachment from storage during re-indexing, key: %s", + storage_key, + ) + indexing_runner = IndexingRunner() indexing_runner.run([document], session) session.commit() diff --git a/api/tests/unit_tests/tasks/test_document_indexing_update_task.py b/api/tests/unit_tests/tasks/test_document_indexing_update_task.py index 62a39e9ebb0..3756a02a194 100644 --- a/api/tests/unit_tests/tasks/test_document_indexing_update_task.py +++ b/api/tests/unit_tests/tasks/test_document_indexing_update_task.py @@ -4,6 +4,7 @@ from __future__ import annotations import uuid from collections.abc import Callable +from datetime import UTC, datetime from unittest.mock import MagicMock import pytest @@ -13,8 +14,10 @@ from sqlalchemy.orm import Session import tasks.document_indexing_update_task as task_module from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from extensions.storage.storage_type import StorageType +from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.model import UploadFile from tasks.document_indexing_update_task import document_indexing_update_task @@ -102,6 +105,39 @@ def _complete_indexing(documents: list[Document], _session: Session) -> None: document.indexing_status = IndexingStatus.COMPLETED +def _persist_attachment( + session: Session, + *, + dataset: Dataset, + document: Document, + segment: DocumentSegment, + key: str, +) -> tuple[UploadFile, SegmentAttachmentBinding]: + attachment = UploadFile( + tenant_id=dataset.tenant_id, + storage_type=StorageType.LOCAL, + key=key, + name="image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=document.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, + ) + session.add_all([attachment, binding]) + session.commit() + return attachment, binding + + def test_queues_summary_when_all_persisted_conditions_match( sqlite_session: Session, task_harness: tuple[MagicMock, MagicMock], @@ -274,3 +310,147 @@ def test_cleans_and_deletes_persisted_segments_with_real_session( assert isinstance(processor.clean.call_args.kwargs["session"], Session) assert sqlite_session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document.id)).all() == [] delay.assert_called_once_with(dataset.id, document.id, None) + + +def test_removes_orphaned_multimodal_attachments_during_reindex( + sqlite_session: Session, + task_harness: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, processor = task_harness + dataset, document = _persist_rows(sqlite_session, with_segment=True) + dataset.is_multimodal = True + segment = sqlite_session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document.id)) + assert segment is not None + attachment, binding = _persist_attachment( + sqlite_session, + dataset=dataset, + document=document, + segment=segment, + key="attachments/orphaned-image.png", + ) + attachment_id = attachment.id + binding_id = binding.id + storage_delete = MagicMock() + monkeypatch.setattr(task_module.storage, "delete", storage_delete) + + document_indexing_update_task(dataset.id, document.id) + + assert processor.clean.call_count == 2 + assert processor.clean.call_args_list[0].args[1] == ["node-1"] + assert processor.clean.call_args_list[1].kwargs["node_ids"] == [attachment_id] + assert processor.clean.call_args_list[1].kwargs["with_keywords"] is False + storage_delete.assert_called_once_with("attachments/orphaned-image.png") + sqlite_session.expire_all() + assert sqlite_session.get(SegmentAttachmentBinding, binding_id) is None + assert sqlite_session.get(UploadFile, attachment_id) is None + assert sqlite_session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document.id)).all() == [] + runner.run.assert_called_once() + + +def test_preserves_multimodal_attachment_referenced_by_another_document_during_reindex( + sqlite_session: Session, + task_harness: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, processor = task_harness + dataset, document = _persist_rows(sqlite_session, with_segment=True) + dataset.is_multimodal = True + segment = sqlite_session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document.id)) + assert segment is not None + attachment, binding = _persist_attachment( + sqlite_session, + dataset=dataset, + document=document, + segment=segment, + key="attachments/shared-image.png", + ) + + other_document = Document( + id=str(uuid.uuid4()), + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + position=2, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-2", + name="other-document.txt", + created_from=DocumentCreatedFrom.WEB, + created_by=document.created_by, + indexing_status=IndexingStatus.COMPLETED, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + ) + sqlite_session.add(other_document) + sqlite_session.flush() + other_segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=other_document.id, + position=1, + content="other segment", + word_count=2, + tokens=2, + created_by=document.created_by, + index_node_id="node-2", + ) + sqlite_session.add(other_segment) + sqlite_session.flush() + shared_binding = SegmentAttachmentBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=other_document.id, + segment_id=other_segment.id, + attachment_id=attachment.id, + ) + sqlite_session.add(shared_binding) + sqlite_session.commit() + attachment_id = attachment.id + binding_id = binding.id + shared_binding_id = shared_binding.id + storage_delete = MagicMock() + monkeypatch.setattr(task_module.storage, "delete", storage_delete) + + document_indexing_update_task(dataset.id, document.id) + + processor.clean.assert_called_once() + assert processor.clean.call_args.args[1] == ["node-1"] + 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 + runner.run.assert_called_once() + + +def test_keeps_database_cleanup_when_reindex_attachment_storage_delete_fails( + sqlite_session: Session, + task_harness: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + runner, _processor = task_harness + dataset, document = _persist_rows(sqlite_session, with_segment=True) + dataset.is_multimodal = True + segment = sqlite_session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document.id)) + assert segment is not None + attachment, binding = _persist_attachment( + sqlite_session, + dataset=dataset, + document=document, + segment=segment, + key="attachments/failing-image.png", + ) + attachment_id = attachment.id + binding_id = binding.id + storage_delete = MagicMock(side_effect=RuntimeError("storage unavailable")) + monkeypatch.setattr(task_module.storage, "delete", storage_delete) + + with caplog.at_level("ERROR", logger="tasks.document_indexing_update_task"): + document_indexing_update_task(dataset.id, document.id) + + storage_delete.assert_called_once_with("attachments/failing-image.png") + assert "Failed to delete document attachment from storage during re-indexing" 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 + assert sqlite_session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document.id)).all() == [] + runner.run.assert_called_once()