Repository: langgenius/dify
File: api/tasks/batch_clean_document_task.py (Python)

Code before:
                except Exception:
                    failed_batches += 1
                    logger.exception(
                        "Failed to delete DocumentSegment batch %d-%d for dataset_id: %s, document_ids: %s",
                        i,
                        i + len(batch),
                        dataset_id,
                        document_ids,
                    )
            if failed_batches > 0:
                total_batches = (len(segment_ids) + BATCH_SIZE - 1) // BATCH_SIZE
                logger.warning(
                    "DocumentSegment deletion: %d/%d batches failed, document_ids: %s",
                    failed_batches,
                    total_batches,
                    document_ids,
                )
            else:
                logger.debug(
                    "Deleted %d DocumentSegment records for dataset_id: %s, document_ids: %s",
                    deleted_segments,
                    dataset_id,
                    document_ids,
                )

        # ============ Step 6: Delete document-associated files (separate short transaction) ============
        if file_ids:
            try:
                with session_factory.create_session() as session:
                    stmt = delete(UploadFile).where(UploadFile.id.in_(file_ids))
                    result = session.execute(stmt)
                    session.commit()
                    logger.debug("Deleted %d document UploadFile records for file_ids: %s", result.rowcount, file_ids)
            except Exception:
                logger.exception(
                    "Failed to delete document UploadFile records for dataset_id: %s, file_ids: %s",
                    dataset_id,
                    file_ids,
                )

        # ============ Step 7: Delete storage files (I/O operations, no DB transaction) ============
        storage_delete_failures = 0
        for storage_key in storage_keys_to_delete:
            try:
                storage.delete(storage_key)
            except Exception:
                storage_delete_failures += 1
                logger.exception("Failed to delete file from storage, key: %s", storage_key)
        if storage_delete_failures > 0:
            logger.warning(
                "Storage file deletion completed with %d failures out of %d total files for dataset_id: %s",

Review comment (suggestion):
I would recommend putting the session management in `with` statement, like the following:

```
with session_factory.create_session() as session, session.begin():
```

This improves readability and prevent from forgetting calls to `session.commit`, especially for large with blocks.

Code after:
                        session.execute(segment_delete_stmt)
                        session.commit()
                except Exception:
                    failed_batches += 1
                    logger.exception(
                        "Failed to delete DocumentSegment batch %d-%d for dataset_id: %s, document_ids: %s",
                        i,
                        i + len(batch),
                        dataset_id,
                        document_ids,
                    )
            if failed_batches > 0:
                logger.warning(
                    "DocumentSegment deletion: %d/%d batches failed, document_ids: %s",
                    failed_batches,
                    total_batches,
                    document_ids,
                )

        # ============ Step 6: Delete document-associated files (separate short transaction) ============
        if file_ids:
            try:
                with session_factory.create_session() as session:
                    stmt = delete(UploadFile).where(UploadFile.id.in_(file_ids))
                    session.execute(stmt)
                    session.commit()
            except Exception:
                logger.exception(
                    "Failed to delete document UploadFile records for dataset_id: %s, file_ids: %s",
                    dataset_id,
                    file_ids,
                )

        # ============ Step 7: Delete storage files (I/O operations, no DB transaction) ============
        storage_delete_failures = 0
        for storage_key in storage_keys_to_delete:
            try:
                storage.delete(storage_key)
            except Exception:
                storage_delete_failures += 1
                logger.exception("Failed to delete file from storage, key: %s", storage_key)
        if storage_delete_failures > 0:
            logger.warning(
                "Storage file deletion completed with %d failures out of %d total files for dataset_id: %s",
                storage_delete_failures,
                len(storage_keys_to_delete),
                dataset_id,
            )

        end_at = time.perf_counter()
        logger.info(