refactor: use sessionmaker in small services 2 (langgenius#34696)

Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: 3 people

api/services/account_service.py

@@ -9,7 +9,7 @@
 
 from pydantic import BaseModel, TypeAdapter
 from sqlalchemy import delete, func, select, update
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 
 class InvitationData(TypedDict):
@@ -1516,7 +1516,7 @@ def invite_new_member(
 
         check_workspace_member_invite_permission(tenant.id)
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(email, session=session)
 
         if not account:

api/services/async_workflow_service.py

@@ -11,7 +11,7 @@
 
 from celery.result import AsyncResult
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from enums.quota_type import QuotaType
 from extensions.ext_database import db
@@ -237,7 +237,7 @@ def get_trigger_log(
         Returns:
             Trigger log as dictionary or None if not found
         """
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             trigger_log_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
             trigger_log = trigger_log_repo.get_by_id(workflow_trigger_log_id, tenant_id)
 
@@ -263,7 +263,7 @@ def get_recent_logs(
         Returns:
             List of trigger logs as dictionaries
         """
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             trigger_log_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
             logs = trigger_log_repo.get_recent_logs(
                 tenant_id=tenant_id, app_id=app_id, hours=hours, limit=limit, offset=offset
@@ -286,7 +286,7 @@ def get_failed_logs_for_retry(
         Returns:
             List of failed trigger logs as dictionaries
         """
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             trigger_log_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
             logs = trigger_log_repo.get_failed_for_retry(
                 tenant_id=tenant_id, max_retry_count=max_retry_count, limit=limit

api/services/clear_free_plan_tenant_expired_logs.py

@@ -346,7 +346,7 @@ def process(cls, days: int, batch: int, tenant_ids: list[str]):
         started_at = datetime.datetime(2023, 4, 3, 8, 59, 24)
         current_time = started_at
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             total_tenant_count = session.query(Tenant.id).count()
 
         click.echo(click.style(f"Total tenant count: {total_tenant_count}", fg="white"))
@@ -398,7 +398,7 @@ def process_tenant(flask_app: Flask, tenant_id: str):
                 # Initial interval of 1 day, will be dynamically adjusted based on tenant count
                 interval = datetime.timedelta(days=1)
                 # Process tenants in this batch
-                with Session(db.engine) as session:
+                with sessionmaker(db.engine).begin() as session:
                     # Calculate tenant count in next batch with current interval
                     # Try different intervals until we find one with a reasonable tenant count
                     test_intervals = [

api/services/credit_pool_service.py

@@ -1,7 +1,7 @@
 import logging
 
 from sqlalchemy import select, update
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 
 from configs import dify_config
 from core.errors.error import QuotaExceededError
@@ -71,7 +71,7 @@ def check_and_deduct_credits(
         actual_credits = min(credits_required, pool.remaining_credits)
 
         try:
-            with Session(db.engine) as session:
+            with sessionmaker(db.engine).begin() as session:
                 stmt = (
                     update(TenantCreditPool)
                     .where(
@@ -81,7 +81,6 @@ def check_and_deduct_credits(
                     .values(quota_used=TenantCreditPool.quota_used + actual_credits)
                 )
                 session.execute(stmt)
-                session.commit()
         except Exception:
             logger.exception("Failed to deduct credits for tenant %s", tenant_id)
             raise QuotaExceededError("Failed to deduct credits")

api/services/dataset_service.py

@@ -15,7 +15,7 @@
 from graphon.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from redis.exceptions import LockNotOwnedError
 from sqlalchemy import delete, exists, func, select, update
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import Forbidden, NotFound
 
 from configs import dify_config
@@ -551,22 +551,22 @@ def _update_external_knowledge_binding(dataset_id, external_knowledge_id, extern
             external_knowledge_id: External knowledge identifier
             external_knowledge_api_id: External knowledge API identifier
         """
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             external_knowledge_binding = (
                 session.query(ExternalKnowledgeBindings).filter_by(dataset_id=dataset_id).first()
             )
 
             if not external_knowledge_binding:
                 raise ValueError("External knowledge binding not found.")
 
-        # Update binding if values have changed
-        if (
-            external_knowledge_binding.external_knowledge_id != external_knowledge_id
-            or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id
-        ):
-            external_knowledge_binding.external_knowledge_id = external_knowledge_id
-            external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id
-            db.session.add(external_knowledge_binding)
+            # Update binding if values have changed
+            if (
+                external_knowledge_binding.external_knowledge_id != external_knowledge_id
+                or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id
+            ):
+                external_knowledge_binding.external_knowledge_id = external_knowledge_id
+                external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id
+                session.add(external_knowledge_binding)
 
     @staticmethod
     def _update_internal_dataset(dataset, data, user):

api/services/oauth_server.py

@@ -2,7 +2,7 @@
 import uuid
 
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest
 
 from extensions.ext_database import db
@@ -29,7 +29,7 @@ class OAuthServerService:
     def get_oauth_provider_app(client_id: str) -> OAuthProviderApp | None:
         query = select(OAuthProviderApp).where(OAuthProviderApp.client_id == client_id)
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
             return session.execute(query).scalar_one_or_none()
 
     @staticmethod

api/services/rag_pipeline/rag_pipeline.py

@@ -1182,7 +1182,7 @@ def publish_customized_pipeline_template(cls, pipeline_id: str, args: dict):
         workflow = db.session.get(Workflow, pipeline.workflow_id)
         if not workflow:
             raise ValueError("Workflow not found")
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             dataset = pipeline.retrieve_dataset(session=session)
             if not dataset:
                 raise ValueError("Dataset not found")
@@ -1209,7 +1209,7 @@ def publish_customized_pipeline_template(cls, pipeline_id: str, args: dict):
 
         from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             rag_pipeline_dsl_service = RagPipelineDslService(session)
             dsl = rag_pipeline_dsl_service.export_rag_pipeline_dsl(pipeline=pipeline, include_secret=True)
         if args.get("icon_info") is None:

api/services/workflow_service.py

@@ -834,7 +834,7 @@ def run_draft_workflow_node(
         if workflow_node_execution is None:
             raise ValueError(f"WorkflowNodeExecution with id {node_execution.id} not found after saving")
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             outputs = workflow_node_execution.load_full_outputs(session, storage)
 
         with Session(bind=db.engine) as session, session.begin():

api/tests/unit_tests/services/test_account_service.py

@@ -1427,16 +1427,18 @@ def test_invite_new_member_new_account(self, mock_db_dependencies, mock_redis_de
         mock_tenant.name = "Test Workspace"
         mock_inviter = TestAccountAssociatedDataFactory.create_account_mock(account_id="inviter-123", name="Inviter")
 
-        # Mock database queries - need to mock the Session query
+        # Mock database queries - need to mock the sessionmaker query
         mock_session = MagicMock()
         mock_session.query.return_value.filter_by.return_value.first.return_value = None  # No existing account
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = None
 
             # Mock RegisterService.register
@@ -1485,12 +1487,14 @@ def test_invite_new_member_normalizes_new_account_email(
         mixed_email = "Invitee@Example.com"
 
         mock_session = MagicMock()
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = None
 
             mock_new_account = TestAccountAssociatedDataFactory.create_account_mock(
@@ -1541,16 +1545,18 @@ def test_invite_new_member_existing_account(
             account_id="existing-user-456", email="existing@example.com", status="pending"
         )
 
-        # Mock database queries - need to mock the Session query
+        # Mock database queries - need to mock the sessionmaker query
         mock_session = MagicMock()
         mock_session.query.return_value.filter_by.return_value.first.return_value = mock_existing_account
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = mock_existing_account
 
             # Mock scalar for TenantAccountJoin lookup - no existing member

api/tests/unit_tests/services/test_async_workflow_service.py

@@ -357,11 +357,12 @@ def test_should_return_trigger_log_dict_or_none(self, repo_result, expected):
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=fake_engine)),
-            patch.object(
-                async_workflow_service_module, "Session", return_value=mock_session_context
-            ) as mock_session_class,
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",
@@ -373,7 +374,7 @@ def test_should_return_trigger_log_dict_or_none(self, repo_result, expected):
 
         # Assert
         assert result == expected
-        mock_session_class.assert_called_once_with(fake_engine)
+        mock_sessionmaker.assert_called_once_with(fake_engine)
         mock_repo.get_by_id.assert_called_once_with("trigger-log-123", "tenant-123")
 
     def test_should_return_recent_logs_as_dict_list(self):
@@ -391,9 +392,12 @@ def test_should_return_recent_logs_as_dict_list(self):
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=MagicMock())),
-            patch.object(async_workflow_service_module, "Session", return_value=mock_session_context),
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",
@@ -432,9 +436,12 @@ def test_should_return_failed_logs_for_retry_as_dict_list(self):
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=MagicMock())),
-            patch.object(async_workflow_service_module, "Session", return_value=mock_session_context),
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",