refactor(api): use sessionmaker in datasource provider service (#34811)

api/services/datasource_provider_service.py

@@ -5,7 +5,7 @@ from typing import Any
 
 from graphon.model_runtime.entities.provider_entities import FormType
 from sqlalchemy import func, select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from configs import dify_config
 from constants import HIDDEN_VALUE, UNKNOWN_VALUE
@@ -53,13 +53,12 @@ class DatasourceProviderService:
         """
         remove oauth custom client params
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             session.query(DatasourceOauthTenantParamConfig).filter_by(
                 tenant_id=tenant_id,
                 provider=datasource_provider_id.provider_name,
                 plugin_id=datasource_provider_id.plugin_id,
             ).delete()
-            session.commit()
 
     def decrypt_datasource_provider_credentials(
         self,
@@ -109,7 +108,7 @@ class DatasourceProviderService:
         """
         get credential by id
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             if credential_id:
                 datasource_provider = (
                     session.query(DatasourceProvider).filter_by(tenant_id=tenant_id, id=credential_id).first()
@@ -156,7 +155,6 @@ class DatasourceProviderService:
                     datasource_provider=datasource_provider,
                 )
                 datasource_provider.expires_at = refreshed_credentials.expires_at
-                session.commit()
 
             return self.decrypt_datasource_provider_credentials(
                 tenant_id=tenant_id,
@@ -174,7 +172,7 @@ class DatasourceProviderService:
         """
         get all datasource credentials by provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             datasource_providers = (
                 session.query(DatasourceProvider)
                 .filter_by(tenant_id=tenant_id, provider=provider, plugin_id=plugin_id)
@@ -224,7 +222,6 @@ class DatasourceProviderService:
                     provider=provider,
                 )
                 real_credentials_list.append(real_credentials)
-            session.commit()
 
             return real_credentials_list
 
@@ -234,7 +231,7 @@ class DatasourceProviderService:
         """
         update datasource provider name
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             target_provider = (
                 session.query(DatasourceProvider)
                 .filter_by(
@@ -266,7 +263,6 @@ class DatasourceProviderService:
                 raise ValueError("Authorization name is already exists")
 
             target_provider.name = name
-            session.commit()
         return
 
     def set_default_datasource_provider(
@@ -275,7 +271,7 @@ class DatasourceProviderService:
         """
         set default datasource provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             # get provider
             target_provider = (
                 session.query(DatasourceProvider)
@@ -300,7 +296,6 @@ class DatasourceProviderService:
 
             # set new default provider
             target_provider.is_default = True
-            session.commit()
         return {"result": "success"}
 
     def setup_oauth_custom_client_params(
@@ -315,7 +310,7 @@ class DatasourceProviderService:
         """
         if client_params is None and enabled is None:
             return
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             tenant_oauth_client_params = (
                 session.query(DatasourceOauthTenantParamConfig)
                 .filter_by(
@@ -349,7 +344,6 @@ class DatasourceProviderService:
 
             if enabled is not None:
                 tenant_oauth_client_params.enabled = enabled
-            session.commit()
 
     def is_system_oauth_params_exist(self, datasource_provider_id: DatasourceProviderID) -> bool:
         """
@@ -488,7 +482,7 @@ class DatasourceProviderService:
         """
         update datasource oauth provider
         """
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{CredentialType.OAUTH2.value}"
             with redis_client.lock(lock, timeout=20):
                 target_provider = (
@@ -535,7 +529,6 @@ class DatasourceProviderService:
                 target_provider.expires_at = expire_at
                 target_provider.encrypted_credentials = credentials
                 target_provider.avatar_url = avatar_url or target_provider.avatar_url
-                session.commit()
 
     def add_datasource_oauth_provider(
         self,
@@ -550,7 +543,7 @@ class DatasourceProviderService:
         add datasource oauth provider
         """
         credential_type = CredentialType.OAUTH2
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{credential_type.value}"
             with redis_client.lock(lock, timeout=60):
                 db_provider_name = name
@@ -604,7 +597,6 @@ class DatasourceProviderService:
                     expires_at=expire_at,
                 )
                 session.add(datasource_provider)
-                session.commit()
 
     def add_datasource_api_key_provider(
         self,
@@ -623,7 +615,7 @@ class DatasourceProviderService:
         provider_name = provider_id.provider_name
         plugin_id = provider_id.plugin_id
 
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{CredentialType.API_KEY}"
             with redis_client.lock(lock, timeout=20):
                 db_provider_name = name or self.generate_next_datasource_provider_name(
@@ -670,7 +662,6 @@ class DatasourceProviderService:
                     encrypted_credentials=credentials,
                 )
                 session.add(datasource_provider)
-                session.commit()
 
     def extract_secret_variables(self, tenant_id: str, provider_id: str, credential_type: CredentialType) -> list[str]:
         """
@@ -926,7 +917,7 @@ class DatasourceProviderService:
         update datasource credentials.
         """
 
-        with Session(db.engine) as session:
+        with sessionmaker(bind=db.engine).begin() as session:
             datasource_provider = (
                 session.query(DatasourceProvider)
                 .filter_by(tenant_id=tenant_id, id=auth_id, provider=provider, plugin_id=plugin_id)
@@ -980,7 +971,6 @@ class DatasourceProviderService:
                         encrypted_credentials[key] = value
 
                 datasource_provider.encrypted_credentials = encrypted_credentials
-            session.commit()
 
     def remove_datasource_credentials(self, tenant_id: str, auth_id: str, provider: str, plugin_id: str) -> None:
         """