deal db session in celery worker (#17549)

This commit is contained in:
Jyong 2025-04-07 20:31:26 +08:00 committed by GitHub
commit c8145ce581
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 104 additions and 25 deletions

View file

@ -27,7 +27,9 @@ def duplicate_document_indexing_task(dataset_id: str, document_ids: list):
dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()
if dataset is None:
raise ValueError("Dataset not found")
logging.info(click.style("Dataset not found: {}".format(dataset_id), fg="red"))
db.session.close()
return
# check document limit
features = FeatureService.get_features(dataset.tenant_id)
@ -57,6 +59,8 @@ def duplicate_document_indexing_task(dataset_id: str, document_ids: list):
db.session.add(document)
db.session.commit()
return
finally:
db.session.close()
for document_id in document_ids:
logging.info(click.style("Start process document: {}".format(document_id), fg="green"))
@ -96,3 +100,5 @@ def duplicate_document_indexing_task(dataset_id: str, document_ids: list):
logging.info(click.style(str(ex), fg="yellow"))
except Exception:
pass
finally:
db.session.close()