refactor(api): use sessionmaker in builtin tools manage service (#34812)

api/services/tools/builtin_tools_manage_service.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import Any
 
 from sqlalchemy import exists, select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from configs import dify_config
 from constants import HIDDEN_VALUE, UNKNOWN_VALUE
@@ -46,13 +46,12 @@ class BuiltinToolManageService:
         delete custom oauth client params
         """
         tool_provider = ToolProviderID(provider)
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             session.query(ToolOAuthTenantClient).filter_by(
                 tenant_id=tenant_id,
                 provider=tool_provider.provider_name,
                 plugin_id=tool_provider.plugin_id,
             ).delete()
-            session.commit()
         return {"result": "success"}
 
     @staticmethod
@@ -150,7 +149,7 @@ class BuiltinToolManageService:
         """
         update builtin tool provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             # get if the provider exists
             db_provider = (
                 session.query(BuiltinToolProvider)
@@ -203,9 +202,7 @@ class BuiltinToolManageService:
 
                     db_provider.name = name
 
-                session.commit()
             except Exception as e:
-                session.rollback()
                 raise ValueError(str(e))
         return {"result": "success"}
 
@@ -222,7 +219,7 @@ class BuiltinToolManageService:
         """
         add builtin tool provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             try:
                 lock = f"builtin_tool_provider_create_lock:{tenant_id}_{provider}"
                 with redis_client.lock(lock, timeout=20):
@@ -281,9 +278,7 @@ class BuiltinToolManageService:
                     )
 
                     session.add(db_provider)
-                    session.commit()
             except Exception as e:
-                session.rollback()
                 raise ValueError(str(e))
 
         return {"result": "success"}
@@ -379,7 +374,7 @@ class BuiltinToolManageService:
         """
         delete tool provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             db_provider = (
                 session.query(BuiltinToolProvider)
                 .where(
@@ -393,7 +388,6 @@ class BuiltinToolManageService:
                 raise ValueError(f"you have not added provider {provider}")
 
             session.delete(db_provider)
-            session.commit()
 
             # delete cache
             provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
@@ -409,7 +403,7 @@ class BuiltinToolManageService:
         """
         set default provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             # get provider
             target_provider = session.query(BuiltinToolProvider).filter_by(id=id, tenant_id=tenant_id).first()
             if target_provider is None:
@@ -422,7 +416,6 @@ class BuiltinToolManageService:
 
             # set new default provider
             target_provider.is_default = True
-            session.commit()
 
         return {"result": "success"}
 
@@ -654,7 +647,7 @@ class BuiltinToolManageService:
         if not isinstance(provider_controller, (BuiltinToolProviderController, PluginToolProviderController)):
             raise ValueError(f"Provider {provider} is not a builtin or plugin provider")
 
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             custom_client_params = (
                 session.query(ToolOAuthTenantClient)
                 .filter_by(
@@ -690,7 +683,6 @@ class BuiltinToolManageService:
             if enable_oauth_custom_client is not None:
                 custom_client_params.enabled = enable_oauth_custom_client
 
-            session.commit()
         return {"result": "success"}
 
     @staticmethod