opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-04-27 02:36:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: migrate session.query to select API in summary and remove document tasks (#34650)

Browse Source
This commit is contained in:
Renzo 2026-04-07 00:55:31 -05:00 committed by GitHub
parent 84d8940dbf
commit 173e818a62
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
api/tasks/generate_summary_index_task.py
View File

@ -5,6 +5,7 @@ import time
import click import click
from celery import shared_task from celery import shared_task
from sqlalchemy import select, update
from core.db.session_factory import session_factory from core.db.session_factory import session_factory
from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.constant.index_type import IndexTechniqueType
@ -39,12 +40,12 @@ def generate_summary_index_task(dataset_id: str, document_id: str, segment_ids:
try: try:
with session_factory.create_session() as session: with session_factory.create_session() as session:
dataset = session.query(Dataset).where(Dataset.id == dataset_id).first() dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1))
if not dataset: if not dataset:
logger.error(click.style(f"Dataset not found: {dataset_id}", fg="red")) logger.error(click.style(f"Dataset not found: {dataset_id}", fg="red"))
return return
document = session.query(DatasetDocument).where(DatasetDocument.id == document_id).first() document = session.scalar(select(DatasetDocument).where(DatasetDocument.id == document_id).limit(1))
if not document: if not document:
logger.error(click.style(f"Document not found: {document_id}", fg="red")) logger.error(click.style(f"Document not found: {document_id}", fg="red"))
return return
@ -108,13 +109,12 @@ def generate_summary_index_task(dataset_id: str, document_id: str, segment_ids:
if segment_ids: if segment_ids:
error_message = f"Summary generation failed: {str(e)}" error_message = f"Summary generation failed: {str(e)}"
with session_factory.create_session() as session: with session_factory.create_session() as session:
session.query(DocumentSegment).filter( session.execute(
DocumentSegment.id.in_(segment_ids), update(DocumentSegment)
DocumentSegment.dataset_id == dataset_id, .where(
).update( DocumentSegment.id.in_(segment_ids),
{ DocumentSegment.dataset_id == dataset_id,
DocumentSegment.error: error_message, )
}, .values(error=error_message)
synchronize_session=False,
) )
session.commit() session.commit()

20
api/tasks/remove_document_from_index_task.py
View File

@ -3,7 +3,7 @@ import time
import click import click
from celery import shared_task from celery import shared_task
from sqlalchemy import select from sqlalchemy import select, update
from core.db.session_factory import session_factory from core.db.session_factory import session_factory
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
@ -26,7 +26,7 @@ def remove_document_from_index_task(document_id: str):
start_at = time.perf_counter() start_at = time.perf_counter()
with session_factory.create_session() as session: with session_factory.create_session() as session:
document = session.query(Document).where(Document.id == document_id).first() document = session.scalar(select(Document).where(Document.id == document_id).limit(1))
if not document: if not document:
logger.info(click.style(f"Document not found: {document_id}", fg="red")) logger.info(click.style(f"Document not found: {document_id}", fg="red"))
return return
@ -68,13 +68,15 @@ def remove_document_from_index_task(document_id: str):
except Exception: except Exception:
logger.exception("clean dataset %s from index failed", dataset.id) logger.exception("clean dataset %s from index failed", dataset.id)
# update segment to disable # update segment to disable
session.query(DocumentSegment).where(DocumentSegment.document_id == document.id).update( session.execute(
{ update(DocumentSegment)
DocumentSegment.enabled: False, .where(DocumentSegment.document_id == document.id)
DocumentSegment.disabled_at: naive_utc_now(), .values(
DocumentSegment.disabled_by: document.disabled_by, enabled=False,
DocumentSegment.updated_at: naive_utc_now(), disabled_at=naive_utc_now(),
} disabled_by=document.disabled_by,
updated_at=naive_utc_now(),
)
) )
session.commit() session.commit()