diff --git a/api/apps/restful_apis/document_api.py b/api/apps/restful_apis/document_api.py index 3a3f3cd30..0d3782f43 100644 --- a/api/apps/restful_apis/document_api.py +++ b/api/apps/restful_apis/document_api.py @@ -1027,6 +1027,13 @@ async def delete_documents(tenant_id, dataset_id): if delete_all: doc_ids = [doc.id for doc in DocumentService.query(kb_id=dataset_id)] + dataset_doc_ids = {doc.id for doc in DocumentService.query(kb_id=dataset_id)} + invalid_ids = [doc_id for doc_id in doc_ids if doc_id not in dataset_doc_ids] + if invalid_ids: + return get_error_data_result( + message=f"These documents do not belong to dataset {dataset_id} or Document not found: {', '.join(invalid_ids)}" + ) + # make sure each id is unique unique_doc_ids, duplicate_messages = check_duplicate_ids(doc_ids, "document") if duplicate_messages: diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_delete_documents.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_delete_documents.py index 0f9881bb1..82fdb413e 100644 --- a/test/testcases/test_http_api/test_file_management_within_dataset/test_delete_documents.py +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_delete_documents.py @@ -132,7 +132,7 @@ class TestDocumentsDeletion: res = delete_documents(HttpApiAuth, dataset_id, {"ids": document_ids}) assert res["code"] == 102 - assert "Document not found" in res["message"] + assert "or Document not found" in res["message"] @pytest.mark.p2 def test_duplicate_deletion(self, HttpApiAuth, add_documents_func): @@ -145,6 +145,24 @@ class TestDocumentsDeletion: assert len(res["data"]["docs"]) == 3 assert res["data"]["total"] == 3 + @pytest.mark.p2 + def test_cross_dataset_deletion_is_blocked(self, HttpApiAuth, add_dataset, add_documents_func, tmp_path): + dataset_id, _document_ids = add_documents_func + other_dataset_id = add_dataset + other_document_id = bulk_upload_documents(HttpApiAuth, other_dataset_id, 1, tmp_path)[0] + + res = delete_documents(HttpApiAuth, dataset_id, {"ids": [other_document_id]}) + assert res["code"] == 102 + assert f"These documents do not belong to dataset {dataset_id}" in res["message"] + + res = list_documents(HttpApiAuth, dataset_id) + assert len(res["data"]["docs"]) == 3 + assert res["data"]["total"] == 3 + + res = list_documents(HttpApiAuth, other_dataset_id) + assert len(res["data"]["docs"]) == 1 + assert res["data"]["total"] == 1 + @pytest.mark.p3 def test_concurrent_deletion(HttpApiAuth, add_dataset, tmp_path): diff --git a/test/testcases/test_web_api/test_document_app/test_rm_documents.py b/test/testcases/test_web_api/test_document_app/test_rm_documents.py index 1b799352b..2e8cefdbb 100644 --- a/test/testcases/test_web_api/test_document_app/test_rm_documents.py +++ b/test/testcases/test_web_api/test_document_app/test_rm_documents.py @@ -78,7 +78,7 @@ class TestDocumentsDeletion: for doc_id in document_ids: res = delete_document(WebApiAuth, kb_id, {"ids": [doc_id]}) assert res["code"] == 102, res - assert res["message"] == "Document not found!", res + assert "Document not found" in res["message"], res @pytest.mark.p2 def test_delete_all(self, WebApiAuth, add_documents_func):