deal db session in celery worker (#17549)

api/tasks/enable_segment_to_index_task.py

@@ -4,7 +4,6 @@ import time
 
 import click
 from celery import shared_task  # type: ignore
-from werkzeug.exceptions import NotFound
 
 from core.rag.index_processor.constant.index_type import IndexType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
@@ -27,10 +26,14 @@ def enable_segment_to_index_task(segment_id: str):
 
     segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_id).first()
     if not segment:
-        raise NotFound("Segment not found")
+        logging.info(click.style("Segment not found: {}".format(segment_id), fg="red"))
+        db.session.close()
+        return
 
     if segment.status != "completed":
-        raise NotFound("Segment is not completed, enable action is not allowed.")
+        logging.info(click.style("Segment is not completed, enable is not allowed: {}".format(segment_id), fg="red"))
+        db.session.close()
+        return
 
     indexing_cache_key = "segment_{}_indexing".format(segment.id)
 
@@ -94,3 +97,4 @@ def enable_segment_to_index_task(segment_id: str):
         db.session.commit()
     finally:
         redis_client.delete(indexing_cache_key)
+        db.session.close()