Fix incorrect exception handling in db query (#23582)

Signed-off-by: Yongtao Huang <yongtaoh2022@gmail.com>

api/schedule/clean_embedding_cache_task.py

@@ -3,7 +3,7 @@ import time
 
 import click
 from sqlalchemy import text
-from werkzeug.exceptions import NotFound
+from sqlalchemy.exc import SQLAlchemyError
 
 import app
 from configs import dify_config
@@ -27,8 +27,8 @@ def clean_embedding_cache_task():
                 .all()
             )
             embedding_ids = [embedding_id[0] for embedding_id in embedding_ids]
-        except NotFound:
-            break
+        except SQLAlchemyError:
+            raise
         if embedding_ids:
             for embedding_id in embedding_ids:
                 db.session.execute(

api/schedule/clean_messages.py

@@ -3,7 +3,7 @@ import logging
 import time
 
 import click
-from werkzeug.exceptions import NotFound
+from sqlalchemy.exc import SQLAlchemyError
 
 import app
 from configs import dify_config
@@ -42,8 +42,8 @@ def clean_messages():
                 .all()
             )
 
-        except NotFound:
-            break
+        except SQLAlchemyError:
+            raise
         if not messages:
             break
         for message in messages:

api/schedule/clean_unused_datasets_task.py

@@ -3,7 +3,7 @@ import time
 
 import click
 from sqlalchemy import func, select
-from werkzeug.exceptions import NotFound
+from sqlalchemy.exc import SQLAlchemyError
 
 import app
 from configs import dify_config
@@ -65,8 +65,8 @@ def clean_unused_datasets_task():
 
             datasets = db.paginate(stmt, page=1, per_page=50)
 
-        except NotFound:
-            break
+        except SQLAlchemyError:
+            raise
         if datasets.items is None or len(datasets.items) == 0:
             break
         for dataset in datasets:
@@ -146,8 +146,8 @@ def clean_unused_datasets_task():
             )
             datasets = db.paginate(stmt, page=1, per_page=50)
 
-        except NotFound:
-            break
+        except SQLAlchemyError:
+            raise
         if datasets.items is None or len(datasets.items) == 0:
             break
         for dataset in datasets: