fix(dataset): clean multimodal attachments during reindex (#41999)

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()

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()