Fix incorrect exception handling in db query (#23582)

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

api/commands.py

@@ -9,7 +9,7 @@ import sqlalchemy as sa
 from flask import current_app
 from pydantic import TypeAdapter
 from sqlalchemy import select
-from werkzeug.exceptions import NotFound
+from sqlalchemy.exc import SQLAlchemyError
 
 from configs import dify_config
 from constants.languages import languages
@@ -181,8 +181,8 @@ def migrate_annotation_vector_database():
             )
             if not apps:
                 break
-        except NotFound:
+        except SQLAlchemyError:
-            break
+            raise
 
         page += 1
         for app in apps:
@@ -308,8 +308,8 @@ def migrate_knowledge_vector_database():
             )
 
             datasets = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False)
-        except NotFound:
+        except SQLAlchemyError:
-            break
+            raise
 
         page += 1
         for dataset in datasets:
@@ -561,8 +561,8 @@ def old_metadata_migration():
                 .order_by(DatasetDocument.created_at.desc())
             )
             documents = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False)
-        except NotFound:
+        except SQLAlchemyError:
-            break
+            raise
         if not documents:
             break
         for document in documents:

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:
+        except SQLAlchemyError:
-            break
+            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:
+        except SQLAlchemyError:
-            break
+            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:
+        except SQLAlchemyError:
-            break
+            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:
+        except SQLAlchemyError:
-            break
+            raise
         if datasets.items is None or len(datasets.items) == 0:
             break
         for dataset in datasets:

api/services/workflow_draft_variable_service.py

@@ -256,7 +256,7 @@ class WorkflowDraftVariableService:
     def _reset_node_var_or_sys_var(
         self, workflow: Workflow, variable: WorkflowDraftVariable
     ) -> WorkflowDraftVariable | None:
-        # If a variable does not allow updating, it makes no sence to resetting it.
+        # If a variable does not allow updating, it makes no sense to reset it.
         if not variable.editable:
             return variable
         # No execution record for this variable, delete the variable instead.
@@ -478,7 +478,7 @@ def _batch_upsert_draft_variable(
                 "node_execution_id": stmt.excluded.node_execution_id,
             },
         )
-    elif _UpsertPolicy.IGNORE:
+    elif policy == _UpsertPolicy.IGNORE:
         stmt = stmt.on_conflict_do_nothing(index_elements=WorkflowDraftVariable.unique_app_id_node_id_name())
     else:
         raise Exception("Invalid value for update policy.")