mirror of
https://github.com/langgenius/dify.git
synced 2026-09-09 05:41:00 +08:00
fix(dataset): delete segment attachment blobs from storage (#41400)
This commit is contained in:
CoralGarden52
GitHub
parent
12d23062d4
commit
2cdfe24c18
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user