pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refa: dataset operations to simplify error handling (#8132)

Browse Source 
### What problem does this PR solve?

- Consolidate database operations within single try-except blocks in the
methods

### Type of change

- [x] Refactoring
This commit is contained in:
Liu An
2025-06-09 13:29:56 +08:00
committed by GitHub
parent 2337bbf6ca
commit 968ffc7ef3
1 changed files with 68 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
api/apps/sdk/dataset.py
View File

@ -124,48 +124,36 @@ def create(tenant_id):
try: try:
if KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value): if KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_operating_result(message=f"Dataset name '{req['name']}' already exists") return get_error_operating_result(message=f"Dataset name '{req['name']}' already exists")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
req["parser_config"] = get_parser_config(req["parser_id"], req["parser_config"]) req["parser_config"] = get_parser_config(req["parser_id"], req["parser_config"])
req["id"] = get_uuid() req["id"] = get_uuid()
req["tenant_id"] = tenant_id req["tenant_id"] = tenant_id
req["created_by"] = tenant_id req["created_by"] = tenant_id
try:
ok, t = TenantService.get_by_id(tenant_id) ok, t = TenantService.get_by_id(tenant_id)
if not ok: if not ok:
return get_error_permission_result(message="Tenant not found") return get_error_permission_result(message="Tenant not found")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if not req.get("embd_id"): if not req.get("embd_id"):
req["embd_id"] = t.embd_id req["embd_id"] = t.embd_id
else: else:
ok, err = verify_embedding_availability(req["embd_id"], tenant_id) ok, err = verify_embedding_availability(req["embd_id"], tenant_id)
if not ok: if not ok:
return err return err
try:
if not KnowledgebaseService.save(**req): if not KnowledgebaseService.save(**req):
return get_error_data_result(message="Create dataset error.(Database error)") return get_error_data_result(message="Create dataset error.(Database error)")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
try:
ok, k = KnowledgebaseService.get_by_id(req["id"]) ok, k = KnowledgebaseService.get_by_id(req["id"])
if not ok: if not ok:
return get_error_data_result(message="Dataset created failed") return get_error_data_result(message="Dataset created failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
except OperationalError as e: except OperationalError as e:
logging.exception(e) logging.exception(e)
return get_error_data_result(message="Database operation failed") return get_error_data_result(message="Database operation failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
@manager.route("/datasets", methods=["DELETE"]) # noqa: F821 @manager.route("/datasets", methods=["DELETE"]) # noqa: F821
@token_required @token_required
@ -211,34 +199,27 @@ def delete(tenant_id):
if err is not None: if err is not None:
return get_error_argument_result(err) return get_error_argument_result(err)
kb_id_instance_pairs = [] try:
if req["ids"] is None: kb_id_instance_pairs = []
try: if req["ids"] is None:
kbs = KnowledgebaseService.query(tenant_id=tenant_id) kbs = KnowledgebaseService.query(tenant_id=tenant_id)
for kb in kbs: for kb in kbs:
kb_id_instance_pairs.append((kb.id, kb)) kb_id_instance_pairs.append((kb.id, kb))
except OperationalError as e:
logging.exception(e) else:
return get_error_data_result(message="Database operation failed") error_kb_ids = []
else: for kb_id in req["ids"]:
error_kb_ids = []
for kb_id in req["ids"]:
try:
kb = KnowledgebaseService.get_or_none(id=kb_id, tenant_id=tenant_id) kb = KnowledgebaseService.get_or_none(id=kb_id, tenant_id=tenant_id)
if kb is None: if kb is None:
error_kb_ids.append(kb_id) error_kb_ids.append(kb_id)
continue continue
kb_id_instance_pairs.append((kb_id, kb)) kb_id_instance_pairs.append((kb_id, kb))
except OperationalError as e: if len(error_kb_ids) > 0:
logging.exception(e) return get_error_permission_result(message=f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(error_kb_ids)}'""")
return get_error_data_result(message="Database operation failed")
if len(error_kb_ids) > 0:
return get_error_permission_result(message=f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(error_kb_ids)}'""")
errors = [] errors = []
success_count = 0 success_count = 0
for kb_id, kb in kb_id_instance_pairs: for kb_id, kb in kb_id_instance_pairs:
try:
for doc in DocumentService.query(kb_id=kb_id): for doc in DocumentService.query(kb_id=kb_id):
if not DocumentService.remove_document(doc, tenant_id): if not DocumentService.remove_document(doc, tenant_id):
errors.append(f"Remove document '{doc.id}' error for dataset '{kb_id}'") errors.append(f"Remove document '{doc.id}' error for dataset '{kb_id}'")
@ -256,18 +237,18 @@ def delete(tenant_id):
errors.append(f"Delete dataset error for {kb_id}") errors.append(f"Delete dataset error for {kb_id}")
continue continue
success_count += 1 success_count += 1
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if not errors: if not errors:
return get_result() return get_result()
error_message = f"Successfully deleted {success_count} datasets, {len(errors)} failed. Details: {'; '.join(errors)[:128]}..." error_message = f"Successfully deleted {success_count} datasets, {len(errors)} failed. Details: {'; '.join(errors)[:128]}..."
if success_count == 0: if success_count == 0:
return get_error_data_result(message=error_message) return get_error_data_result(message=error_message)
return get_result(data={"success_count": success_count, "errors": errors[:5]}, message=error_message) return get_result(data={"success_count": success_count, "errors": errors[:5]}, message=error_message)
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
@manager.route("/datasets/<dataset_id>", methods=["PUT"]) # noqa: F821 @manager.route("/datasets/<dataset_id>", methods=["PUT"]) # noqa: F821
@ -349,53 +330,41 @@ def update(tenant_id, dataset_id):
kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id) kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id)
if kb is None: if kb is None:
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'") return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if req.get("parser_config"): if req.get("parser_config"):
req["parser_config"] = deep_merge(kb.parser_config, req["parser_config"]) req["parser_config"] = deep_merge(kb.parser_config, req["parser_config"])
if (chunk_method := req.get("parser_id")) and chunk_method != kb.parser_id: if (chunk_method := req.get("parser_id")) and chunk_method != kb.parser_id:
if not req.get("parser_config"): if not req.get("parser_config"):
req["parser_config"] = get_parser_config(chunk_method, None) req["parser_config"] = get_parser_config(chunk_method, None)
elif "parser_config" in req and not req["parser_config"]: elif "parser_config" in req and not req["parser_config"]:
del req["parser_config"] del req["parser_config"]
if "name" in req and req["name"].lower() != kb.name.lower(): if "name" in req and req["name"].lower() != kb.name.lower():
try:
exists = KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value) exists = KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value)
if exists: if exists:
return get_error_data_result(message=f"Dataset name '{req['name']}' already exists") return get_error_data_result(message=f"Dataset name '{req['name']}' already exists")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if "embd_id" in req: if "embd_id" in req:
if kb.chunk_num != 0 and req["embd_id"] != kb.embd_id: if kb.chunk_num != 0 and req["embd_id"] != kb.embd_id:
return get_error_data_result(message=f"When chunk_num ({kb.chunk_num}) > 0, embedding_model must remain {kb.embd_id}") return get_error_data_result(message=f"When chunk_num ({kb.chunk_num}) > 0, embedding_model must remain {kb.embd_id}")
ok, err = verify_embedding_availability(req["embd_id"], tenant_id) ok, err = verify_embedding_availability(req["embd_id"], tenant_id)
if not ok: if not ok:
return err return err
try:
if not KnowledgebaseService.update_by_id(kb.id, req): if not KnowledgebaseService.update_by_id(kb.id, req):
return get_error_data_result(message="Update dataset error.(Database error)") return get_error_data_result(message="Update dataset error.(Database error)")
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
try:
ok, k = KnowledgebaseService.get_by_id(kb.id) ok, k = KnowledgebaseService.get_by_id(kb.id)
if not ok: if not ok:
return get_error_data_result(message="Dataset created failed") return get_error_data_result(message="Dataset created failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
except OperationalError as e: except OperationalError as e:
logging.exception(e) logging.exception(e)
return get_error_data_result(message="Database operation failed") return get_error_data_result(message="Database operation failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
@manager.route("/datasets", methods=["GET"]) # noqa: F821 @manager.route("/datasets", methods=["GET"]) # noqa: F821
@token_required @token_required
@ -459,26 +428,19 @@ def list_datasets(tenant_id):
if err is not None: if err is not None:
return get_error_argument_result(err) return get_error_argument_result(err)
kb_id = request.args.get("id")
name = args.get("name")
if kb_id:
try:
kbs = KnowledgebaseService.get_kb_by_id(kb_id, tenant_id)
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if not kbs:
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{kb_id}'")
if name:
try:
kbs = KnowledgebaseService.get_kb_by_name(name, tenant_id)
except OperationalError as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
if not kbs:
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{name}'")
try: try:
kb_id = request.args.get("id")
name = args.get("name")
if kb_id:
kbs = KnowledgebaseService.get_kb_by_id(kb_id, tenant_id)
if not kbs:
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{kb_id}'")
if name:
kbs = KnowledgebaseService.get_kb_by_name(name, tenant_id)
if not kbs:
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{name}'")
tenants = TenantService.get_joined_tenants_by_user_id(tenant_id) tenants = TenantService.get_joined_tenants_by_user_id(tenant_id)
kbs = KnowledgebaseService.get_list( kbs = KnowledgebaseService.get_list(
[m["tenant_id"] for m in tenants], [m["tenant_id"] for m in tenants],
@ -490,11 +452,11 @@ def list_datasets(tenant_id):
kb_id, kb_id,
name, name,
) )
response_data_list = []
for kb in kbs:
response_data_list.append(remap_dictionary_keys(kb))
return get_result(data=response_data_list)
except OperationalError as e: except OperationalError as e:
logging.exception(e) logging.exception(e)
return get_error_data_result(message="Database operation failed") return get_error_data_result(message="Database operation failed")
response_data_list = []
for kb in kbs:
response_data_list.append(remap_dictionary_keys(kb))
return get_result(data=response_data_list)