fix retry document db session issue (#23046)

This commit is contained in:
Jyong 2025-07-28 15:56:49 +08:00 • committed by GitHub
commit 3248d728ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -24,11 +24,10 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
""" """
documents: list[Document] = [] documents: list[Document] = []
start_at = time.perf_counter() start_at = time.perf_counter()
try:
dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first() dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
if not dataset: if not dataset:
logging.info(click.style(f"Dataset not found: {dataset_id}", fg="red")) logging.info(click.style(f"Dataset not found: {dataset_id}", fg="red"))
db.session.close()
return return
tenant_id = dataset.tenant_id tenant_id = dataset.tenant_id
for document_id in document_ids: for document_id in document_ids:
@ -45,7 +44,9 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
) )
except Exception as e: except Exception as e:
document = ( document = (
db.session.query(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).first() db.session.query(Document)
.where(Document.id == document_id, Document.dataset_id == dataset_id)
.first()
) )
if document: if document:
document.indexing_status = "error" document.indexing_status = "error"
@ -54,7 +55,6 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
db.session.add(document) db.session.add(document)
db.session.commit() db.session.commit()
redis_client.delete(retry_indexing_cache_key) redis_client.delete(retry_indexing_cache_key)
db.session.close()
return return
logging.info(click.style(f"Start retry document: {document_id}", fg="green")) logging.info(click.style(f"Start retry document: {document_id}", fg="green"))
@ -63,7 +63,6 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
) )
if not document: if not document:
logging.info(click.style(f"Document not found: {document_id}", fg="yellow")) logging.info(click.style(f"Document not found: {document_id}", fg="yellow"))
db.session.close()
return return
try: try:
# clean old data # clean old data
@ -96,7 +95,12 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
logging.info(click.style(str(ex), fg="yellow")) logging.info(click.style(str(ex), fg="yellow"))
redis_client.delete(retry_indexing_cache_key) redis_client.delete(retry_indexing_cache_key)
logging.exception("retry_document_indexing_task failed, document_id: %s", document_id) logging.exception("retry_document_indexing_task failed, document_id: %s", document_id)
finally:
db.session.close()
end_at = time.perf_counter() end_at = time.perf_counter()
logging.info(click.style(f"Retry dataset: {dataset_id} latency: {end_at - start_at}", fg="green")) logging.info(click.style(f"Retry dataset: {dataset_id} latency: {end_at - start_at}", fg="green"))
except Exception as e:
logging.exception(
"retry_document_indexing_task failed, dataset_id: %s, document_ids: %s", dataset_id, document_ids
)
raise e
finally:
db.session.close()