deal db session in celery worker (#17549)

This commit is contained in:
Jyong
2025-04-07 20:31:26 +08:00
committed by GitHub
parent da7f8ad936
commit c8145ce581
26 changed files with 104 additions and 25 deletions

View File

@@ -4,7 +4,6 @@ import time
import click
from celery import shared_task # type: ignore
from werkzeug.exceptions import NotFound
from core.indexing_runner import DocumentIsPausedError, IndexingRunner
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
@@ -27,7 +26,9 @@ def document_indexing_update_task(dataset_id: str, document_id: str):
document = db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()
if not document:
raise NotFound("Document not found")
logging.info(click.style("Document not found: {}".format(document_id), fg="red"))
db.session.close()
return
document.indexing_status = "parsing"
document.processing_started_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
@@ -73,3 +74,5 @@ def document_indexing_update_task(dataset_id: str, document_id: str):
logging.info(click.style(str(ex), fg="yellow"))
except Exception:
pass
finally:
db.session.close()