dify/api/tasks/clean_dataset_task.py

import logging
import time
from collections.abc import Callable, Sequence

import click
from celery import shared_task
from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from core.db.session_factory import session_factory
from core.rag.index_processor.index_processor import IndexProcessorFactory
from core.tools.utils.web_reader_tool import get_image_upload_file_ids
from extensions.ext_storage import storage
from models import WorkflowType
from models.dataset import (
    AppDatasetJoin,
    Dataset,
    DatasetMetadata,
    DatasetMetadataBinding,
    DatasetProcessRule,
    DatasetQuery,
    Document,
    DocumentSegment,
    Pipeline,
    SegmentAttachmentBinding,
)
from models.model import UploadFile
from models.workflow import Workflow
from services.dataset_knowledge_fs_upgrade_file_lease import reserve_upgrade_file_cleanup
from services.dataset_storage_cleanup_service import cleanup_dataset_crawl_cache
from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh

logger = logging.getLogger(__name__)

# Keep committed batches on failure, but preserve dependent rows until a retry
# finishes their prerequisite cleanup.
_DELETE_BATCH_SIZE = 1000
_MAX_RETRIES = 3
_RETRY_DELAY_SECONDS = 30


@shared_task(queue="dataset", bind=True, max_retries=_MAX_RETRIES, default_retry_delay=_RETRY_DELAY_SECONDS)
def clean_dataset_task(
    self,
    dataset_id: str,
    tenant_id: str,
    indexing_technique: str | None,
    index_struct: str | None,
    collection_binding_id: str | None,
    doc_form: str,
    pipeline_id: str | None = None,
    *,
    index_cleanup_done: bool = False,
) -> None:
    """
    Clean dataset when dataset deleted.
    :param dataset_id: dataset id
    :param tenant_id: tenant id
    :param indexing_technique: indexing technique
    :param index_struct: index struct dict
    :param collection_binding_id: collection binding id
    :param doc_form: dataset form
    :param index_cleanup_done: preserve successful index cleanup across explicit retries

    Usage: clean_dataset_task.delay(dataset_id, tenant_id, indexing_technique, index_struct)
    """
    logger.info(click.style(f"Start clean dataset when dataset deleted: {dataset_id}", fg="green"))
    start_at = time.perf_counter()

    # Enhanced validation: Check if doc_form is None, empty string, or contains only whitespace.
    # Treating all invalid values as the default paragraph index type keeps the vector
    # database cleanup reachable for empty/broken datasets.
    if doc_form is None or (isinstance(doc_form, str) and not doc_form.strip()):
        from core.rag.index_processor.constant.index_type import IndexStructureType

        doc_form = IndexStructureType.PARAGRAPH_INDEX
        logger.info(
            click.style(f"Invalid doc_form detected, using default index type for cleanup: {doc_form}", fg="yellow")
        )

    dataset = Dataset(
        id=dataset_id,
        tenant_id=tenant_id,
        indexing_technique=indexing_technique,
        index_struct=index_struct,
        collection_binding_id=collection_binding_id,
    )

    try:
        cleanup_dataset_crawl_cache(tenant_id, dataset_id)

        with session_factory.create_session() as session:
            if not index_cleanup_done:
                # Index processors also delete relational child rows, so the whole step must succeed first.
                _run_step(session, dataset_id, "dataset index", lambda: _clean_index(session, dataset, doc_form))
                index_cleanup_done = True
                schedule_billing_vector_space_refresh(dataset.tenant_id)

            # Stop on failure so documents and their resources outlive any remaining segments.
            _run_step(session, dataset_id, "segments", lambda: _delete_segments(session, dataset_id))
            _run_step(session, dataset_id, "documents", lambda: _delete_documents(session, dataset_id))
            _run_step(session, dataset_id, "attachments", lambda: _delete_attachments(session, tenant_id, dataset_id))
            _run_step(
                session,
                dataset_id,
                "dataset metadata",
                lambda: _delete_dataset_scoped_rows(session, tenant_id, dataset_id, pipeline_id),
            )

    except Exception as exc:
        if self.request.retries >= _MAX_RETRIES:
            logger.exception("Dataset cleanup retry budget exhausted, dataset_id=%s", dataset_id)
            raise

        logger.warning(
            "Dataset cleanup failed, scheduling retry %d/%d, dataset_id=%s",
            self.request.retries + 1,
            _MAX_RETRIES,
            dataset_id,
            exc_info=True,
        )
        # The session is closed before retrying. Normalize arguments for positional callers too.
        raise self.retry(
            exc=exc,
            args=(),
            kwargs={
                "dataset_id": dataset_id,
                "tenant_id": tenant_id,
                "indexing_technique": indexing_technique,
                "index_struct": index_struct,
                "collection_binding_id": collection_binding_id,
                "doc_form": doc_form,
                "pipeline_id": pipeline_id,
                "index_cleanup_done": index_cleanup_done,
            },
            countdown=_RETRY_DELAY_SECONDS * (2**self.request.retries),
        ) from exc

    end_at = time.perf_counter()
    logger.info(
        click.style(
            f"Cleaned dataset when dataset deleted: {dataset_id} latency: {end_at - start_at}",
            fg="green",
        )
    )


def _run_step(session: Session, dataset_id: str, name: str, step: Callable[[], None]) -> None:
    """Roll back uncommitted work on failure and stop dependent steps until retry."""
    try:
        step()
    except Exception:
        logger.exception(click.style(f"Failed to clean {name} for dataset {dataset_id}", fg="red"))
        try:
            session.rollback()
        except Exception:
            logger.exception("Failed to rollback database session")
        raise


def _clean_index(session: Session, dataset: Dataset, doc_form: str) -> None:
    """Clean the dataset index and commit any relational rows removed by its processor."""
    index_processor = IndexProcessorFactory(doc_form).init_index_processor()
    index_processor.clean(dataset, None, with_keywords=True, delete_child_chunks=True, session=session)
    session.commit()
    logger.info(click.style(f"Successfully cleaned vector database for dataset: {dataset.id}", fg="green"))


def _delete_segments(session: Session, dataset_id: str) -> None:
    """Delete document segments in batches, removing their inline image files first."""
    while True:
        segments = session.scalars(
            select(DocumentSegment).where(DocumentSegment.dataset_id == dataset_id).limit(_DELETE_BATCH_SIZE)
        ).all()
        if not segments:
            break
        image_upload_file_ids: list[str] = []
        for segment in segments:
            image_upload_file_ids.extend(get_image_upload_file_ids(segment.content))
        _delete_upload_files(session, image_upload_file_ids)
        session.execute(delete(DocumentSegment).where(DocumentSegment.id.in_([segment.id for segment in segments])))
        session.commit()


def _delete_documents(session: Session, dataset_id: str) -> None:
    """Delete documents in batches, removing their uploaded source files first."""
    while True:
        documents = session.scalars(
            select(Document).where(Document.dataset_id == dataset_id).limit(_DELETE_BATCH_SIZE)
        ).all()
        if not documents:
            break
        file_ids: list[str] = []
        for document in documents:
            if document.data_source_type == "upload_file" and document.data_source_info:
                data_source_info = document.data_source_info_dict
                if data_source_info and "upload_file_id" in data_source_info:
                    file_ids.append(data_source_info["upload_file_id"])
        # Active KnowledgeFS upgrades still read these sources; their leases record the cleanup instead.
        leased_file_ids = reserve_upgrade_file_cleanup(session, file_ids)
        if leased_file_ids:
            logger.info(
                "Keep %d source files while KnowledgeFS upgrade leases are active, dataset_id=%s",
                len(leased_file_ids),
                dataset_id,
            )
        _delete_upload_files(session, [file_id for file_id in file_ids if file_id not in leased_file_ids])
        session.execute(delete(Document).where(Document.id.in_([document.id for document in documents])))
        session.commit()


def _delete_upload_files(session: Session, file_ids: Sequence[str]) -> None:
    """Delete blobs before their lookup rows so failures retain storage keys for retry."""
    if not file_ids:
        return
    files = session.scalars(select(UploadFile).where(UploadFile.id.in_(file_ids))).all()
    for file in files:
        try:
            storage.delete(file.key)
        except FileNotFoundError:
            logger.debug("Upload file already absent from storage, upload_file_id: %s", file.id)
    session.execute(delete(UploadFile).where(UploadFile.id.in_(file_ids)))


def _delete_attachments(session: Session, tenant_id: str, dataset_id: str) -> None:
    """Delete segment attachments (upload files + bindings) for the dataset."""
    # Inline-image cleanup may already have deleted the file; its binding still needs cleanup.
    attachments_with_bindings = session.execute(
        select(SegmentAttachmentBinding, UploadFile)
        .outerjoin(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
        .where(
            SegmentAttachmentBinding.tenant_id == tenant_id,
            SegmentAttachmentBinding.dataset_id == dataset_id,
        )
    ).all()
    if not attachments_with_bindings:
        return
    attachment_ids = [
        attachment_file.id for _, attachment_file in attachments_with_bindings if attachment_file is not None
    ]
    binding_ids = [binding.id for binding, _ in attachments_with_bindings]
    for binding, attachment_file in attachments_with_bindings:
        if attachment_file is None:
            continue
        try:
            storage.delete(attachment_file.key)
        except Exception:
            logger.exception(
                "Delete attachment_file failed when storage deleted, attachment_file_id: %s",
                binding.attachment_id,
            )
    if attachment_ids:
        session.execute(delete(UploadFile).where(UploadFile.id.in_(attachment_ids)))
    session.execute(delete(SegmentAttachmentBinding).where(SegmentAttachmentBinding.id.in_(binding_ids)))
    session.commit()


def _delete_dataset_scoped_rows(session: Session, tenant_id: str, dataset_id: str, pipeline_id: str | None) -> None:
    """Delete the remaining small dataset-scoped tables (and the pipeline/workflow)."""
    session.execute(delete(DatasetProcessRule).where(DatasetProcessRule.dataset_id == dataset_id))
    session.execute(delete(DatasetQuery).where(DatasetQuery.dataset_id == dataset_id))
    session.execute(delete(AppDatasetJoin).where(AppDatasetJoin.dataset_id == dataset_id))
    # delete dataset metadata
    session.execute(delete(DatasetMetadata).where(DatasetMetadata.dataset_id == dataset_id))
    session.execute(delete(DatasetMetadataBinding).where(DatasetMetadataBinding.dataset_id == dataset_id))
    # delete pipeline and workflow
    if pipeline_id:
        session.execute(delete(Pipeline).where(Pipeline.id == pipeline_id))
        session.execute(
            delete(Workflow).where(
                Workflow.tenant_id == tenant_id,
                Workflow.app_id == pipeline_id,
                Workflow.type == WorkflowType.RAG_PIPELINE,
            )
        )
    session.commit()