fix(api): clean multimodal attachments in batch deletion (#41256)

api/tasks/batch_clean_document_task.py

@@ -11,7 +11,7 @@ from core.db.session_factory import session_factory
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from core.tools.utils.web_reader_tool import get_image_upload_file_ids
 from extensions.ext_storage import storage
-from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment
+from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding
 from models.model import UploadFile
 from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh
 
@@ -65,6 +65,17 @@ def batch_clean_document_task(
                     image_upload_file_ids = get_image_upload_file_ids(segment.content)
                     total_image_upload_file_ids.extend(image_upload_file_ids)
 
+                total_image_upload_file_ids.extend(
+                    session.scalars(
+                        select(SegmentAttachmentBinding.attachment_id).where(
+                            SegmentAttachmentBinding.tenant_id == segments[0].tenant_id,
+                            SegmentAttachmentBinding.dataset_id == dataset_id,
+                            SegmentAttachmentBinding.document_id.in_(document_ids),
+                            SegmentAttachmentBinding.segment_id.in_(segment_ids),
+                        )
+                    ).all()
+                )
+
             # Query storage keys for image files
             if total_image_upload_file_ids:
                 image_files = session.scalars(
@@ -161,6 +172,13 @@ def batch_clean_document_task(
                 batch = segment_ids[i : i + BATCH_SIZE]
                 try:
                     with session_factory.create_session() as session:
+                        binding_delete_stmt = delete(SegmentAttachmentBinding).where(
+                            SegmentAttachmentBinding.tenant_id == segments[0].tenant_id,
+                            SegmentAttachmentBinding.dataset_id == dataset_id,
+                            SegmentAttachmentBinding.document_id.in_(document_ids),
+                            SegmentAttachmentBinding.segment_id.in_(batch),
+                        )
+                        session.execute(binding_delete_stmt)
                         segment_delete_stmt = delete(DocumentSegment).where(DocumentSegment.id.in_(batch))
                         session.execute(segment_delete_stmt)
                         session.commit()

api/tests/unit_tests/tasks/test_batch_clean_document_task.py

@@ -1,12 +1,15 @@
 import uuid
+from datetime import UTC, datetime
 from unittest.mock import patch
 
 import pytest
 from sqlalchemy.orm import Session
 
 import tasks.batch_clean_document_task as task_module
-from models.dataset import Dataset, DocumentSegment
-from models.enums import DataSourceType
+from extensions.storage.storage_type import StorageType
+from models.dataset import Dataset, DocumentSegment, SegmentAttachmentBinding
+from models.enums import CreatorUserRole, DataSourceType
+from models.model import UploadFile
 from tasks.batch_clean_document_task import batch_clean_document_task
 
 
@@ -83,3 +86,51 @@ def test_failed_vector_cleanup_does_not_schedule_billing_refresh(cleanup_rows: t
         )
 
     schedule_refresh.assert_not_called()
+
+
+def test_cleans_segment_attachment_bindings_and_files(cleanup_rows: tuple[str, str, str], sqlite_session: Session):
+    dataset_id, document_id, tenant_id = cleanup_rows
+    segment = sqlite_session.query(DocumentSegment).filter_by(document_id=document_id).one()
+    attachment = UploadFile(
+        tenant_id=tenant_id,
+        storage_type=StorageType.LOCAL,
+        key="attachments/image.png",
+        name="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=tenant_id,
+        dataset_id=dataset_id,
+        document_id=document_id,
+        segment_id=segment.id,
+        attachment_id=attachment.id,
+    )
+    sqlite_session.add_all([attachment, binding])
+    sqlite_session.commit()
+    attachment_id = attachment.id
+    attachment_key = attachment.key
+    binding_id = binding.id
+
+    with (
+        patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]),
+        patch("tasks.batch_clean_document_task.IndexProcessorFactory"),
+        patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh"),
+        patch("tasks.batch_clean_document_task.storage.delete") as storage_delete,
+    ):
+        batch_clean_document_task(
+            document_ids=[document_id],
+            dataset_id=dataset_id,
+            doc_form="paragraph",
+            file_ids=[],
+        )
+
+    sqlite_session.expire_all()
+    assert sqlite_session.get(SegmentAttachmentBinding, binding_id) is None
+    assert sqlite_session.get(UploadFile, attachment_id) is None
+    storage_delete.assert_called_once_with(attachment_key)