admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

perf: optimize save_document_with_dataset_id (#29550)

Browse Source
This commit is contained in:
wangxiaolei
2025-12-12 13:14:45 +08:00
committed by GitHub
parent 12e39365fa
commit ac40309850
1 changed files with 70 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
api/services/dataset_service.py
View File

@@ -1636,6 +1636,20 @@ class DocumentService:
return [], "" return [], ""
db.session.add(dataset_process_rule) db.session.add(dataset_process_rule)
db.session.flush() db.session.flush()
else:
# Fallback when no process_rule provided in knowledge_config:
# 1) reuse dataset.latest_process_rule if present
# 2) otherwise create an automatic rule
dataset_process_rule = getattr(dataset, "latest_process_rule", None)
if not dataset_process_rule:
dataset_process_rule = DatasetProcessRule(
dataset_id=dataset.id,
mode="automatic",
rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
created_by=account.id,
)
db.session.add(dataset_process_rule)
db.session.flush()
lock_name = f"add_document_lock_dataset_id_{dataset.id}" lock_name = f"add_document_lock_dataset_id_{dataset.id}"
try: try:
with redis_client.lock(lock_name, timeout=600): with redis_client.lock(lock_name, timeout=600):
@@ -1647,65 +1661,67 @@ class DocumentService:
if not knowledge_config.data_source.info_list.file_info_list: if not knowledge_config.data_source.info_list.file_info_list:
raise ValueError("File source info is required") raise ValueError("File source info is required")
upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids
for file_id in upload_file_list: files = (
file = ( db.session.query(UploadFile)
db.session.query(UploadFile) .where(
.where(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id) UploadFile.tenant_id == dataset.tenant_id,
.first() UploadFile.id.in_(upload_file_list),
) )
.all()
)
if len(files) != len(set(upload_file_list)):
raise FileNotExistsError("One or more files not found.")
# raise error if file not found file_names = [file.name for file in files]
if not file: db_documents = (
raise FileNotExistsError() db.session.query(Document)
.where(
file_name = file.name Document.dataset_id == dataset.id,
Document.tenant_id == current_user.current_tenant_id,
Document.data_source_type == "upload_file",
Document.enabled == True,
Document.name.in_(file_names),
)
.all()
)
documents_map = {document.name: document for document in db_documents}
for file in files:
data_source_info: dict[str, str | bool] = { data_source_info: dict[str, str | bool] = {
"upload_file_id": file_id, "upload_file_id": file.id,
} }
# check duplicate document = documents_map.get(file.name)
if knowledge_config.duplicate: if knowledge_config.duplicate and document:
document = ( document.dataset_process_rule_id = dataset_process_rule.id
db.session.query(Document) document.updated_at = naive_utc_now()
.filter_by( document.created_from = created_from
dataset_id=dataset.id, document.doc_form = knowledge_config.doc_form
tenant_id=current_user.current_tenant_id, document.doc_language = knowledge_config.doc_language
data_source_type="upload_file", document.data_source_info = json.dumps(data_source_info)
enabled=True, document.batch = batch
name=file_name, document.indexing_status = "waiting"
) db.session.add(document)
.first() documents.append(document)
duplicate_document_ids.append(document.id)
continue
else:
document = DocumentService.build_document(
dataset,
dataset_process_rule.id,
knowledge_config.data_source.info_list.data_source_type,
knowledge_config.doc_form,
knowledge_config.doc_language,
data_source_info,
created_from,
position,
account,
file.name,
batch,
) )
if document: db.session.add(document)
document.dataset_process_rule_id = dataset_process_rule.id db.session.flush()
document.updated_at = naive_utc_now() document_ids.append(document.id)
document.created_from = created_from documents.append(document)
document.doc_form = knowledge_config.doc_form position += 1
document.doc_language = knowledge_config.doc_language
document.data_source_info = json.dumps(data_source_info)
document.batch = batch
document.indexing_status = "waiting"
db.session.add(document)
documents.append(document)
duplicate_document_ids.append(document.id)
continue
document = DocumentService.build_document(
dataset,
dataset_process_rule.id,
knowledge_config.data_source.info_list.data_source_type,
knowledge_config.doc_form,
knowledge_config.doc_language,
data_source_info,
created_from,
position,
account,
file_name,
batch,
)
db.session.add(document)
db.session.flush()
document_ids.append(document.id)
documents.append(document)
position += 1
elif knowledge_config.data_source.info_list.data_source_type == "notion_import": elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore
if not notion_info_list: if not notion_info_list: