diff --git a/api/tasks/batch_clean_document_task.py b/api/tasks/batch_clean_document_task.py index 11cf4b9835c..040b9b4f349 100644 --- a/api/tasks/batch_clean_document_task.py +++ b/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() diff --git a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py b/api/tests/unit_tests/tasks/test_batch_clean_document_task.py index 59fab89ab9c..8f80c5a7ec2 100644 --- a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py +++ b/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)