refactor: manage rag pipeline sessions explicitly (#38274)

api/controllers/console/datasets/rag_pipeline/rag_pipeline.py

@@ -5,7 +5,7 @@ from flask import request
 from flask_restx import Resource
 from pydantic import BaseModel, Field
 from sqlalchemy import select
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import NotFound
 
 from controllers.common.fields import SimpleDataResponse
@@ -16,6 +16,7 @@ from controllers.common.schema import (
     register_schema_models,
 )
 from controllers.console import console_ns
+from controllers.console.app.wraps import with_session
 from controllers.console.wraps import (
     account_initialization_required,
     enterprise_license_required,
@@ -102,10 +103,13 @@ class PipelineTemplateListApi(Resource):
     @account_initialization_required
     @enterprise_license_required
     @with_current_tenant_id
-    def get(self, current_tenant_id: str) -> JsonResponseWithStatus:
+    @with_session
+    def get(self, session: Session, current_tenant_id: str) -> JsonResponseWithStatus:
         query = PipelineTemplateListQuery.model_validate(request.args.to_dict(flat=True))
         # get pipeline templates
-        pipeline_templates = RagPipelineService.get_pipeline_templates(query.type, query.language, current_tenant_id)
+        pipeline_templates = RagPipelineService.get_pipeline_templates(
+            session, query.type, query.language, current_tenant_id
+        )
         return dump_response(PipelineTemplateListResponse, pipeline_templates), 200
 
 
@@ -117,10 +121,11 @@ class PipelineTemplateDetailApi(Resource):
     @login_required
     @account_initialization_required
     @enterprise_license_required
-    def get(self, template_id: str) -> JsonResponseWithStatus:
+    @with_session
+    def get(self, session: Session, template_id: str) -> JsonResponseWithStatus:
         query = PipelineTemplateDetailQuery.model_validate(request.args.to_dict(flat=True))
         rag_pipeline_service = RagPipelineService()
-        pipeline_template = rag_pipeline_service.get_pipeline_template_detail(template_id, query.type)
+        pipeline_template = rag_pipeline_service.get_pipeline_template_detail(session, template_id, query.type)
         if pipeline_template is None:
             raise NotFound("Pipeline template not found from upstream service.")
         return dump_response(PipelineTemplateDetailResponse, pipeline_template), 200

api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py

@@ -6,7 +6,7 @@ from uuid import UUID
 from flask import abort, request
 from flask_restx import Resource
 from pydantic import BaseModel, Field, RootModel, ValidationError
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
 import services
@@ -26,6 +26,7 @@ from controllers.console.app.workflow import (
     WorkflowPaginationResponse,
     WorkflowResponse,
 )
+from controllers.console.app.wraps import with_session
 from controllers.console.datasets.wraps import get_rag_pipeline
 from controllers.console.wraps import (
     RBACPermission,
@@ -343,7 +344,8 @@ class DraftRagPipelineRunApi(Resource):
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_current_user
     @get_rag_pipeline
-    def post(self, current_user: Account, pipeline: Pipeline):
+    @with_session
+    def post(self, session: Session, current_user: Account, pipeline: Pipeline):
         """
         Run draft workflow
         """
@@ -352,6 +354,7 @@ class DraftRagPipelineRunApi(Resource):
 
         try:
             response = PipelineGenerateService.generate(
+                session=session,
                 pipeline=pipeline,
                 user=current_user,
                 args=args,
@@ -375,7 +378,8 @@ class PublishedRagPipelineRunApi(Resource):
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_current_user
     @get_rag_pipeline
-    def post(self, current_user: Account, pipeline: Pipeline):
+    @with_session
+    def post(self, session: Session, current_user: Account, pipeline: Pipeline):
         """
         Run published workflow
         """
@@ -385,6 +389,7 @@ class PublishedRagPipelineRunApi(Resource):
 
         try:
             response = PipelineGenerateService.generate(
+                session=session,
                 pipeline=pipeline,
                 user=current_user,
                 args=args,
@@ -1014,13 +1019,14 @@ class RagPipelineTransformApi(Resource):
     @login_required
     @account_initialization_required
     @with_current_user
-    def post(self, current_user: Account, dataset_id: UUID):
+    @with_session
+    def post(self, session: Session, current_user: Account, dataset_id: UUID):
         if not (current_user.has_edit_permission or current_user.is_dataset_operator):
             raise Forbidden()
 
         dataset_id_str = str(dataset_id)
         rag_pipeline_transform_service = RagPipelineTransformService()
-        result = rag_pipeline_transform_service.transform_dataset(dataset_id_str, db.session)
+        result = rag_pipeline_transform_service.transform_dataset(dataset_id_str, session)
         return result
 
 

api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py

@@ -5,6 +5,7 @@ from uuid import UUID
 from flask import request
 from pydantic import BaseModel, Field, RootModel
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -16,6 +17,7 @@ from controllers.common.schema import (
     register_schema_model,
     register_schema_models,
 )
+from controllers.console.app.wraps import with_session
 from controllers.service_api import service_api_ns
 from controllers.service_api.dataset.error import PipelineRunError
 from controllers.service_api.dataset.rag_pipeline.serializers import serialize_upload_file
@@ -264,7 +266,8 @@ class PipelineRunApi(DatasetApiResource):
         "Pipeline run successfully",
         service_api_ns.models[GeneratedAppResponse.__name__],
     )
-    def post(self, tenant_id: str, dataset_id: UUID):
+    @with_session
+    def post(self, session: Session, tenant_id: str, dataset_id: UUID):
         """Resource for running a rag pipeline."""
         dataset_id_str = str(dataset_id)
         # Verify dataset ownership
@@ -282,6 +285,7 @@ class PipelineRunApi(DatasetApiResource):
         pipeline: Pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str)
         try:
             response: dict[Any, Any] | Generator[str, Any, None] = PipelineGenerateService.generate(
+                session=session,
                 pipeline=pipeline,
                 user=current_user,
                 args=payload.model_dump(),

api/services/rag_pipeline/pipeline_generate_service.py

@@ -1,10 +1,11 @@
 from collections.abc import Mapping
 from typing import Any
 
+from sqlalchemy.orm import Session
+
 from configs import dify_config
 from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
 from core.app.entities.app_invoke_entities import InvokeFrom
-from extensions.ext_database import db
 from models.dataset import Document, Pipeline
 from models.enums import IndexingStatus
 from models.model import Account, App, EndUser
@@ -16,6 +17,7 @@ class PipelineGenerateService:
     @classmethod
     def generate(
         cls,
+        session: Session,
         pipeline: Pipeline,
         user: Account | EndUser,
         args: Mapping[str, Any],
@@ -35,7 +37,7 @@ class PipelineGenerateService:
             workflow = cls._get_workflow(pipeline, invoke_from)
             if original_document_id := args.get("original_document_id"):
                 # update document status to waiting
-                cls.update_document_status(original_document_id)
+                cls.update_document_status(original_document_id, session)
             return PipelineGenerator.convert_to_event_stream(
                 PipelineGenerator().generate(
                     pipeline=pipeline,
@@ -105,13 +107,12 @@ class PipelineGenerateService:
         return workflow
 
     @classmethod
-    def update_document_status(cls, document_id: str):
+    def update_document_status(cls, document_id: str, session: Session):
         """
         Update document status to waiting
         :param document_id: document id
         """
-        document = db.session.get(Document, document_id)
+        document = session.get(Document, document_id)
         if document:
             document.indexing_status = IndexingStatus.WAITING
-            db.session.add(document)
-            db.session.commit()
+            session.add(document)

api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Any, override
 
 from flask import current_app
+from sqlalchemy.orm import Session
 
 from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase
 from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType
@@ -21,13 +22,15 @@ class BuiltInPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
         return PipelineTemplateType.BUILTIN
 
     @override
-    def get_pipeline_templates(self, language: str, current_tenant_id: str | None = None) -> dict[str, Any]:
+    def get_pipeline_templates(
+        self, session: Session, language: str, current_tenant_id: str | None = None
+    ) -> dict[str, Any]:
         del current_tenant_id
         result = self.fetch_pipeline_templates_from_builtin(language)
         return result
 
     @override
-    def get_pipeline_template_detail(self, template_id: str) -> dict[str, Any] | None:
+    def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None:
         result = self.fetch_pipeline_template_detail_from_builtin(template_id)
         return result
 

api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py

@@ -2,8 +2,8 @@ from typing import Any, TypedDict, override
 
 import yaml
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 
-from extensions.ext_database import db
 from libs.login import resolve_tenant_id_fallback
 from models.dataset import PipelineCustomizedTemplate
 from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase
@@ -40,29 +40,38 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
     """
 
     @override
-    def get_pipeline_templates(self, language: str, current_tenant_id: str | None = None) -> dict[str, Any]:
+    def get_pipeline_templates(
+        self, session: Session, language: str, current_tenant_id: str | None = None
+    ) -> dict[str, Any]:
         current_tenant_id = resolve_tenant_id_fallback(current_tenant_id)
-        return self.fetch_pipeline_templates_from_customized(tenant_id=current_tenant_id, language=language)
+        return self.fetch_pipeline_templates_from_customized(
+            session=session, tenant_id=current_tenant_id, language=language
+        )
 
     @override
-    def get_pipeline_template_detail(self, template_id: str) -> dict[str, Any] | None:
-        return self.fetch_pipeline_template_detail_from_db(template_id)
+    def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None:
+        return self.fetch_pipeline_template_detail_from_db(session, template_id)
 
     @override
     def get_type(self) -> str:
         return PipelineTemplateType.CUSTOMIZED
 
     @classmethod
-    def fetch_pipeline_templates_from_customized(cls, tenant_id: str, language: str) -> dict[str, Any]:
+    def fetch_pipeline_templates_from_customized(
+        cls, session: Session, tenant_id: str, language: str
+    ) -> dict[str, Any]:
         """
         Fetch pipeline templates from db.
         :param tenant_id: tenant id
         :param language: language
         :return:
         """
-        pipeline_customized_templates = db.session.scalars(
+        pipeline_customized_templates = session.scalars(
             select(PipelineCustomizedTemplate)
-            .where(PipelineCustomizedTemplate.tenant_id == tenant_id, PipelineCustomizedTemplate.language == language)
+            .where(
+                PipelineCustomizedTemplate.tenant_id == tenant_id,
+                PipelineCustomizedTemplate.language == language,
+            )
             .order_by(PipelineCustomizedTemplate.position.asc(), PipelineCustomizedTemplate.created_at.desc())
         ).all()
         recommended_pipelines_results: list[CustomizedTemplateItemDict] = []
@@ -80,13 +89,13 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
         return {"pipeline_templates": recommended_pipelines_results}
 
     @classmethod
-    def fetch_pipeline_template_detail_from_db(cls, template_id: str) -> dict[str, Any] | None:
+    def fetch_pipeline_template_detail_from_db(cls, session: Session, template_id: str) -> dict[str, Any] | None:
         """
         Fetch pipeline template detail from db.
         :param template_id: Template ID
         :return:
         """
-        pipeline_template = db.session.get(PipelineCustomizedTemplate, template_id)
+        pipeline_template = session.get(PipelineCustomizedTemplate, template_id)
         if not pipeline_template:
             return None
 

api/services/rag_pipeline/pipeline_template/database/database_retrieval.py

@@ -2,8 +2,8 @@ from typing import Any, TypedDict, override
 
 import yaml
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 
-from extensions.ext_database import db
 from models.dataset import PipelineBuiltInTemplate
 from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase
 from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType
@@ -40,20 +40,22 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
     """
 
     @override
-    def get_pipeline_templates(self, language: str, current_tenant_id: str | None = None) -> dict[str, Any]:
+    def get_pipeline_templates(
+        self, session: Session, language: str, current_tenant_id: str | None = None
+    ) -> dict[str, Any]:
         del current_tenant_id
-        return self.fetch_pipeline_templates_from_db(language)
+        return self.fetch_pipeline_templates_from_db(session, language)
 
     @override
-    def get_pipeline_template_detail(self, template_id: str) -> dict[str, Any] | None:
-        return self.fetch_pipeline_template_detail_from_db(template_id)
+    def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None:
+        return self.fetch_pipeline_template_detail_from_db(session, template_id)
 
     @override
     def get_type(self) -> str:
         return PipelineTemplateType.DATABASE
 
     @classmethod
-    def fetch_pipeline_templates_from_db(cls, language: str) -> dict[str, Any]:
+    def fetch_pipeline_templates_from_db(cls, session: Session, language: str) -> dict[str, Any]:
         """
         Fetch pipeline templates from db.
         :param language: language
@@ -61,9 +63,7 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
         """
 
         pipeline_built_in_templates = list(
-            db.session.scalars(
-                select(PipelineBuiltInTemplate).where(PipelineBuiltInTemplate.language == language)
-            ).all()
+            session.scalars(select(PipelineBuiltInTemplate).where(PipelineBuiltInTemplate.language == language)).all()
         )
 
         recommended_pipelines_results: list[PipelineTemplateItemDict] = []
@@ -83,14 +83,14 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
         return {"pipeline_templates": recommended_pipelines_results}
 
     @classmethod
-    def fetch_pipeline_template_detail_from_db(cls, template_id: str) -> dict[str, Any] | None:
+    def fetch_pipeline_template_detail_from_db(cls, session: Session, template_id: str) -> dict[str, Any] | None:
         """
         Fetch pipeline template detail from db.
         :param pipeline_id: Pipeline ID
         :return:
         """
         # is in public recommended list
-        pipeline_template = db.session.get(PipelineBuiltInTemplate, template_id)
+        pipeline_template = session.get(PipelineBuiltInTemplate, template_id)
 
         if not pipeline_template:
             return None

api/services/rag_pipeline/pipeline_template/pipeline_template_base.py

@@ -1,11 +1,15 @@
 from typing import Any, Protocol
 
+from sqlalchemy.orm import Session
+
 
 class PipelineTemplateRetrievalBase(Protocol):
     """Interface for pipeline template retrieval."""
 
-    def get_pipeline_templates(self, language: str, current_tenant_id: str | None = None) -> dict[str, Any]: ...
+    def get_pipeline_templates(
+        self, session: Session, language: str, current_tenant_id: str | None = None
+    ) -> dict[str, Any]: ...
 
-    def get_pipeline_template_detail(self, template_id: str) -> dict[str, Any] | None: ...
+    def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: ...
 
     def get_type(self) -> str: ...

api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py

@@ -2,6 +2,7 @@ import logging
 from typing import Any, override
 
 import httpx
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from services.rag_pipeline.pipeline_template.database.database_retrieval import DatabasePipelineTemplateRetrieval
@@ -17,21 +18,23 @@ class RemotePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
     """
 
     @override
-    def get_pipeline_template_detail(self, template_id: str) -> dict[str, Any] | None:
+    def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None:
         try:
             return self.fetch_pipeline_template_detail_from_dify_official(template_id)
         except Exception as e:
             logger.warning("fetch recommended app detail from dify official failed: %r, switch to database.", e)
-            return DatabasePipelineTemplateRetrieval.fetch_pipeline_template_detail_from_db(template_id)
+            return DatabasePipelineTemplateRetrieval.fetch_pipeline_template_detail_from_db(session, template_id)
 
     @override
-    def get_pipeline_templates(self, language: str, current_tenant_id: str | None = None) -> dict[str, Any]:
+    def get_pipeline_templates(
+        self, session: Session, language: str, current_tenant_id: str | None = None
+    ) -> dict[str, Any]:
         del current_tenant_id
         try:
             return self.fetch_pipeline_templates_from_dify_official(language)
         except Exception as e:
             logger.warning("fetch pipeline templates from dify official failed: %r, switch to database.", e)
-            return DatabasePipelineTemplateRetrieval.fetch_pipeline_templates_from_db(language)
+            return DatabasePipelineTemplateRetrieval.fetch_pipeline_templates_from_db(session, language)
 
     @override
     def get_type(self) -> str:

api/services/rag_pipeline/rag_pipeline.py

@@ -27,6 +27,7 @@ from core.datasource.entities.datasource_entities import (
 from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
 from core.datasource.online_drive.online_drive_plugin import OnlineDriveDatasourcePlugin
 from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin
+from core.db.session_factory import session_factory
 from core.helper import marketplace
 from core.rag.entities import DatasourceCompletedEvent, DatasourceErrorEvent, DatasourceProcessingEvent
 from core.repositories.factory import DifyCoreRepositoryFactory, OrderConfig
@@ -98,7 +99,8 @@ class RagPipelineService:
     def __init__(self, session_maker: sessionmaker | None = None):
         """Initialize RagPipelineService with repository dependencies."""
         if session_maker is None:
-            session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
+            session_maker = session_factory.get_session_maker()
+        self._session_maker = session_maker
         self._node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
             session_maker
         )
@@ -107,6 +109,7 @@ class RagPipelineService:
     @classmethod
     def get_pipeline_templates(
         cls,
+        session: Session,
         type: str = "built-in",
         language: str = "en-US",
         current_tenant_id: str | None = None,
@@ -114,7 +117,7 @@ class RagPipelineService:
         if type == "built-in":
             mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE
             retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
-            result = retrieval_instance.get_pipeline_templates(language, current_tenant_id)
+            result = retrieval_instance.get_pipeline_templates(session, language, current_tenant_id)
             if not result.get("pipeline_templates") and language != "en-US":
                 template_retrieval = PipelineTemplateRetrievalFactory.get_built_in_pipeline_template_retrieval()
                 result = template_retrieval.fetch_pipeline_templates_from_builtin("en-US")
@@ -122,11 +125,13 @@ class RagPipelineService:
         else:
             mode = "customized"
             retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
-            result = retrieval_instance.get_pipeline_templates(language, current_tenant_id)
+            result = retrieval_instance.get_pipeline_templates(session, language, current_tenant_id)
             return result
 
     @classmethod
-    def get_pipeline_template_detail(cls, template_id: str, type: str = "built-in") -> dict[str, Any] | None:
+    def get_pipeline_template_detail(
+        cls, session: Session, template_id: str, type: str = "built-in"
+    ) -> dict[str, Any] | None:
         """
         Get pipeline template detail.
 
@@ -137,7 +142,9 @@ class RagPipelineService:
         if type == "built-in":
             mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE
             retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
-            built_in_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(template_id)
+            built_in_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(
+                session, template_id
+            )
             if built_in_result is None:
                 logger.warning(
                     "pipeline template retrieval returned empty result, template_id: %s, mode: %s",
@@ -148,7 +155,9 @@ class RagPipelineService:
         else:
             mode = "customized"
             retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
-            customized_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(template_id)
+            customized_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(
+                session, template_id
+            )
             return customized_result
 
     @classmethod
@@ -158,6 +167,7 @@ class RagPipelineService:
         template_info: PipelineTemplateInfoEntity,
         current_user: Account | None = None,
         current_tenant_id: str | None = None,
+        session: Session | None = None,
     ):
         """
         Update pipeline template.
@@ -165,7 +175,17 @@ class RagPipelineService:
         :param template_info: template info
         """
         current_user, current_tenant_id = resolve_account_fallback(current_user, current_tenant_id)
-        customized_template: PipelineCustomizedTemplate | None = db.session.scalar(
+        if session is None:
+            with session_factory.get_session_maker().begin() as new_session:
+                return cls.update_customized_pipeline_template(
+                    template_id,
+                    template_info,
+                    current_user,
+                    current_tenant_id,
+                    session=new_session,
+                )
+
+        customized_template: PipelineCustomizedTemplate | None = session.scalar(
             select(PipelineCustomizedTemplate)
             .where(
                 PipelineCustomizedTemplate.id == template_id,
@@ -178,7 +198,7 @@ class RagPipelineService:
         # check template name is exist
         template_name = template_info.name
         if template_name:
-            template = db.session.scalar(
+            template = session.scalar(
                 select(PipelineCustomizedTemplate)
                 .where(
                     PipelineCustomizedTemplate.name == template_name,
@@ -193,16 +213,22 @@ class RagPipelineService:
         customized_template.description = template_info.description
         customized_template.icon = template_info.icon_info.model_dump()
         customized_template.updated_by = current_user.id
-        db.session.commit()
         return customized_template
 
     @classmethod
-    def delete_customized_pipeline_template(cls, template_id: str, current_tenant_id: str | None = None):
+    def delete_customized_pipeline_template(
+        cls, template_id: str, current_tenant_id: str | None = None, session: Session | None = None
+    ):
         """
         Delete customized pipeline template.
         """
         current_tenant_id = resolve_tenant_id_fallback(current_tenant_id)
-        customized_template: PipelineCustomizedTemplate | None = db.session.scalar(
+        if session is None:
+            with session_factory.get_session_maker().begin() as new_session:
+                cls.delete_customized_pipeline_template(template_id, current_tenant_id, session=new_session)
+                return
+
+        customized_template: PipelineCustomizedTemplate | None = session.scalar(
             select(PipelineCustomizedTemplate)
             .where(
                 PipelineCustomizedTemplate.id == template_id,
@@ -212,23 +238,23 @@ class RagPipelineService:
         )
         if not customized_template:
             raise ValueError("Customized pipeline template not found.")
-        db.session.delete(customized_template)
-        db.session.commit()
+        session.delete(customized_template)
 
     def get_draft_workflow(self, pipeline: Pipeline) -> Workflow | None:
         """
         Get draft workflow
         """
         # fetch draft workflow by rag pipeline
-        workflow = db.session.scalar(
-            select(Workflow)
-            .where(
-                Workflow.tenant_id == pipeline.tenant_id,
-                Workflow.app_id == pipeline.id,
-                Workflow.version == "draft",
+        with self._session_maker() as session:
+            workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == pipeline.tenant_id,
+                    Workflow.app_id == pipeline.id,
+                    Workflow.version == "draft",
+                )
+                .limit(1)
             )
-            .limit(1)
-        )
 
         # return draft workflow
         return workflow
@@ -242,29 +268,31 @@ class RagPipelineService:
             return None
 
         # fetch published workflow by workflow_id
-        workflow = db.session.scalar(
-            select(Workflow)
-            .where(
-                Workflow.tenant_id == pipeline.tenant_id,
-                Workflow.app_id == pipeline.id,
-                Workflow.id == pipeline.workflow_id,
+        with self._session_maker() as session:
+            workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == pipeline.tenant_id,
+                    Workflow.app_id == pipeline.id,
+                    Workflow.id == pipeline.workflow_id,
+                )
+                .limit(1)
             )
-            .limit(1)
-        )
 
         return workflow
 
     def get_published_workflow_by_id(self, pipeline: Pipeline, workflow_id: str) -> Workflow | None:
         """Fetch a published workflow snapshot by ID for restore operations."""
-        workflow = db.session.scalar(
-            select(Workflow)
-            .where(
-                Workflow.tenant_id == pipeline.tenant_id,
-                Workflow.app_id == pipeline.id,
-                Workflow.id == workflow_id,
+        with self._session_maker() as session:
+            workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == pipeline.tenant_id,
+                    Workflow.app_id == pipeline.id,
+                    Workflow.id == workflow_id,
+                )
+                .limit(1)
             )
-            .limit(1)
-        )
         if workflow and workflow.version == Workflow.VERSION_DRAFT:
             raise IsDraftWorkflowError("source workflow must be published")
         return workflow
@@ -322,39 +350,51 @@ class RagPipelineService:
         Sync draft workflow
         :raises WorkflowHashNotEqualError
         """
-        # fetch draft workflow by app_model
-        workflow = self.get_draft_workflow(pipeline=pipeline)
+        with self._session_maker.begin() as session:
+            managed_pipeline = session.get(Pipeline, pipeline.id)
+            if not managed_pipeline:
+                raise ValueError("Pipeline not found")
 
-        if workflow and workflow.unique_hash != unique_hash:
-            raise WorkflowHashNotEqualError()
-
-        # create draft workflow if not found
-        if not workflow:
-            workflow = Workflow(
-                tenant_id=pipeline.tenant_id,
-                app_id=pipeline.id,
-                features="{}",
-                type=WorkflowType.RAG_PIPELINE.value,
-                version="draft",
-                graph=json.dumps(graph),
-                created_by=account.id,
-                environment_variables=environment_variables,
-                conversation_variables=conversation_variables,
-                rag_pipeline_variables=rag_pipeline_variables,
+            # fetch draft workflow by app_model
+            workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == managed_pipeline.tenant_id,
+                    Workflow.app_id == managed_pipeline.id,
+                    Workflow.version == "draft",
+                )
+                .limit(1)
             )
-            db.session.add(workflow)
-            db.session.flush()
-            pipeline.workflow_id = workflow.id
-        # update draft workflow if found
-        else:
-            workflow.graph = json.dumps(graph)
-            workflow.updated_by = account.id
-            workflow.updated_at = datetime.now(UTC).replace(tzinfo=None)
-            workflow.environment_variables = environment_variables
-            workflow.conversation_variables = conversation_variables
-            workflow.rag_pipeline_variables = rag_pipeline_variables
-        # commit db session changes
-        db.session.commit()
+
+            if workflow and workflow.unique_hash != unique_hash:
+                raise WorkflowHashNotEqualError()
+
+            # create draft workflow if not found
+            if not workflow:
+                workflow = Workflow(
+                    tenant_id=managed_pipeline.tenant_id,
+                    app_id=managed_pipeline.id,
+                    features="{}",
+                    type=WorkflowType.RAG_PIPELINE.value,
+                    version="draft",
+                    graph=json.dumps(graph),
+                    created_by=account.id,
+                    environment_variables=environment_variables,
+                    conversation_variables=conversation_variables,
+                    rag_pipeline_variables=rag_pipeline_variables,
+                )
+                session.add(workflow)
+                session.flush()
+                managed_pipeline.workflow_id = workflow.id
+                pipeline.workflow_id = workflow.id
+            # update draft workflow if found
+            else:
+                workflow.graph = json.dumps(graph)
+                workflow.updated_by = account.id
+                workflow.updated_at = datetime.now(UTC).replace(tzinfo=None)
+                workflow.environment_variables = environment_variables
+                workflow.conversation_variables = conversation_variables
+                workflow.rag_pipeline_variables = rag_pipeline_variables
 
         # trigger  workflow events TODO
         # app_draft_workflow_was_synced.send(pipeline, synced_draft_workflow=workflow)
@@ -375,26 +415,48 @@ class RagPipelineService:
         the pipeline-specific flush/link step that wires a newly created draft
         back onto ``pipeline.workflow_id``.
         """
-        source_workflow = self.get_published_workflow_by_id(pipeline=pipeline, workflow_id=workflow_id)
-        if not source_workflow:
-            raise WorkflowNotFoundError("Workflow not found.")
+        with self._session_maker.begin() as session:
+            managed_pipeline = session.get(Pipeline, pipeline.id)
+            if not managed_pipeline:
+                raise ValueError("Pipeline not found")
 
-        draft_workflow = self.get_draft_workflow(pipeline=pipeline)
-        draft_workflow, is_new_draft = apply_published_workflow_snapshot_to_draft(
-            tenant_id=pipeline.tenant_id,
-            app_id=pipeline.id,
-            source_workflow=source_workflow,
-            draft_workflow=draft_workflow,
-            account=account,
-            updated_at_factory=lambda: datetime.now(UTC).replace(tzinfo=None),
-        )
+            source_workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == managed_pipeline.tenant_id,
+                    Workflow.app_id == managed_pipeline.id,
+                    Workflow.id == workflow_id,
+                )
+                .limit(1)
+            )
+            if source_workflow and source_workflow.version == Workflow.VERSION_DRAFT:
+                raise IsDraftWorkflowError("source workflow must be published")
+            if not source_workflow:
+                raise WorkflowNotFoundError("Workflow not found.")
 
-        if is_new_draft:
-            db.session.add(draft_workflow)
-            db.session.flush()
-            pipeline.workflow_id = draft_workflow.id
+            draft_workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == managed_pipeline.tenant_id,
+                    Workflow.app_id == managed_pipeline.id,
+                    Workflow.version == Workflow.VERSION_DRAFT,
+                )
+                .limit(1)
+            )
+            draft_workflow, is_new_draft = apply_published_workflow_snapshot_to_draft(
+                tenant_id=managed_pipeline.tenant_id,
+                app_id=managed_pipeline.id,
+                source_workflow=source_workflow,
+                draft_workflow=draft_workflow,
+                account=account,
+                updated_at_factory=lambda: datetime.now(UTC).replace(tzinfo=None),
+            )
 
-        db.session.commit()
+            if is_new_draft:
+                session.add(draft_workflow)
+                session.flush()
+                managed_pipeline.workflow_id = draft_workflow.id
+                pipeline.workflow_id = draft_workflow.id
 
         return draft_workflow
 
@@ -571,7 +633,7 @@ class RagPipelineService:
             workflow_node_execution.id
         )
 
-        with sessionmaker(bind=db.engine).begin() as session:
+        with self._session_maker.begin() as session:
             draft_var_saver = DraftVariableSaver(
                 session=session,
                 app_id=pipeline.id,
@@ -988,23 +1050,22 @@ class RagPipelineService:
                     dataset_id = get_system_segment(variable_pool, SystemVariableKey.DATASET_ID)
                     pipeline_id = get_system_segment(variable_pool, SystemVariableKey.APP_ID)
                     if document_id and dataset_id and pipeline_id:
-                        document = db.session.scalar(
-                            select(Document)
-                            .join(Dataset, Dataset.id == Document.dataset_id)
-                            .where(
-                                Document.id == document_id.value,
-                                Document.tenant_id == tenant_id,
-                                Document.dataset_id == dataset_id.value,
-                                Dataset.tenant_id == tenant_id,
-                                Dataset.pipeline_id == pipeline_id.value,
+                        with self._session_maker.begin() as session:
+                            document = session.scalar(
+                                select(Document)
+                                .join(Dataset, Dataset.id == Document.dataset_id)
+                                .where(
+                                    Document.id == document_id.value,
+                                    Document.tenant_id == tenant_id,
+                                    Document.dataset_id == dataset_id.value,
+                                    Dataset.tenant_id == tenant_id,
+                                    Dataset.pipeline_id == pipeline_id.value,
+                                )
+                                .limit(1)
                             )
-                            .limit(1)
-                        )
-                        if document:
-                            document.indexing_status = IndexingStatus.ERROR
-                            document.error = error
-                            db.session.add(document)
-                            db.session.commit()
+                            if document:
+                                document.indexing_status = IndexingStatus.ERROR
+                                document.error = error
 
         return workflow_node_execution
 
@@ -1220,82 +1281,81 @@ class RagPipelineService:
         Publish customized pipeline template
         """
         current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
-        pipeline = db.session.get(Pipeline, pipeline_id)
-        if not pipeline:
-            raise ValueError("Pipeline not found")
-        if not pipeline.workflow_id:
-            raise ValueError("Pipeline workflow not found")
-        workflow = db.session.get(Workflow, pipeline.workflow_id)
-        if not workflow:
-            raise ValueError("Workflow not found")
-        with sessionmaker(db.engine).begin() as session:
+        with session_factory.get_session_maker().begin() as session:
+            pipeline = session.get(Pipeline, pipeline_id)
+            if not pipeline:
+                raise ValueError("Pipeline not found")
+            if not pipeline.workflow_id:
+                raise ValueError("Pipeline workflow not found")
+            workflow = session.get(Workflow, pipeline.workflow_id)
+            if not workflow:
+                raise ValueError("Workflow not found")
             dataset = pipeline.retrieve_dataset(session=session)
             if not dataset:
                 raise ValueError("Dataset not found")
 
-        # check template name is exist
-        template_name = args.get("name")
-        if template_name:
-            template = db.session.scalar(
-                select(PipelineCustomizedTemplate)
-                .where(
-                    PipelineCustomizedTemplate.name == template_name,
-                    PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id,
+            # check template name is exist
+            template_name = args.get("name")
+            if template_name:
+                template = session.scalar(
+                    select(PipelineCustomizedTemplate)
+                    .where(
+                        PipelineCustomizedTemplate.name == template_name,
+                        PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id,
+                    )
+                    .limit(1)
+                )
+                if template:
+                    raise ValueError("Template name is already exists")
+
+            max_position = session.scalar(
+                select(func.max(PipelineCustomizedTemplate.position)).where(
+                    PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id
                 )
-                .limit(1)
             )
-            if template:
-                raise ValueError("Template name is already exists")
 
-        max_position = db.session.scalar(
-            select(func.max(PipelineCustomizedTemplate.position)).where(
-                PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id
-            )
-        )
+            from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService
 
-        from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService
-
-        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:
-            args["icon_info"] = {}
-        if args.get("description") is None:
-            raise ValueError("Description is required")
-        if args.get("name") is None:
-            raise ValueError("Name is required")
-        pipeline_customized_template = PipelineCustomizedTemplate(
-            name=args.get("name") or "",
-            description=args.get("description") or "",
-            icon=args.get("icon_info") or {},
-            tenant_id=pipeline.tenant_id,
-            yaml_content=dsl,
-            install_count=0,
-            position=max_position + 1 if max_position else 1,
-            chunk_structure=dataset.chunk_structure,
-            language="en-US",
-            created_by=current_user.id,
-        )
-        db.session.add(pipeline_customized_template)
-        db.session.commit()
+            if args.get("icon_info") is None:
+                args["icon_info"] = {}
+            if args.get("description") is None:
+                raise ValueError("Description is required")
+            if args.get("name") is None:
+                raise ValueError("Name is required")
+            pipeline_customized_template = PipelineCustomizedTemplate(
+                name=args.get("name") or "",
+                description=args.get("description") or "",
+                icon=args.get("icon_info") or {},
+                tenant_id=pipeline.tenant_id,
+                yaml_content=dsl,
+                install_count=0,
+                position=max_position + 1 if max_position else 1,
+                chunk_structure=dataset.chunk_structure,
+                language="en-US",
+                created_by=current_user.id,
+            )
+            session.add(pipeline_customized_template)
 
     def is_workflow_exist(self, pipeline: Pipeline) -> bool:
-        return (
-            db.session.scalar(
-                select(func.count(Workflow.id)).where(
-                    Workflow.tenant_id == pipeline.tenant_id,
-                    Workflow.app_id == pipeline.id,
-                    Workflow.version == Workflow.VERSION_DRAFT,
+        with self._session_maker() as session:
+            return (
+                session.scalar(
+                    select(func.count(Workflow.id)).where(
+                        Workflow.tenant_id == pipeline.tenant_id,
+                        Workflow.app_id == pipeline.id,
+                        Workflow.version == Workflow.VERSION_DRAFT,
+                    )
                 )
-            )
-            or 0
-        ) > 0
+                or 0
+            ) > 0
 
     def get_node_last_run(
         self, pipeline: Pipeline, workflow: Workflow, node_id: str
     ) -> WorkflowNodeExecutionModel | None:
         node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
-            sessionmaker(db.engine)
+            self._session_maker
         )
 
         node_exec = node_execution_service_repo.get_node_last_execution(
@@ -1371,7 +1431,7 @@ class RagPipelineService:
         # Convert node_execution to WorkflowNodeExecution after save
         workflow_node_execution_db_model = repository._to_db_model(workflow_node_execution)  # type: ignore
 
-        with sessionmaker(bind=db.engine).begin() as session:
+        with self._session_maker.begin() as session:
             draft_var_saver = DraftVariableSaver(
                 session=session,
                 app_id=pipeline.id,
@@ -1405,7 +1465,10 @@ class RagPipelineService:
         if type and type != "all":
             stmt = stmt.where(PipelineRecommendedPlugin.type == type)
 
-        pipeline_recommended_plugins = db.session.scalars(stmt.order_by(PipelineRecommendedPlugin.position.asc())).all()
+        with self._session_maker() as session:
+            pipeline_recommended_plugins = session.scalars(
+                stmt.order_by(PipelineRecommendedPlugin.position.asc())
+            ).all()
 
         if not pipeline_recommended_plugins:
             return {
@@ -1444,139 +1507,173 @@ class RagPipelineService:
         """
         Retry error document
         """
-        document_pipeline_execution_log = db.session.scalar(
-            select(DocumentPipelineExecutionLog).where(DocumentPipelineExecutionLog.document_id == document.id).limit(1)
-        )
-        if not document_pipeline_execution_log:
-            raise ValueError("Document pipeline execution log not found")
-        pipeline = db.session.get(Pipeline, document_pipeline_execution_log.pipeline_id)
-        if not pipeline:
-            raise ValueError("Pipeline not found")
-        # convert to app config
-        workflow = self.get_published_workflow(pipeline)
-        if not workflow:
-            raise ValueError("Workflow not found")
-        PipelineGenerator().generate(
-            pipeline=pipeline,
-            workflow=workflow,
-            user=user,
-            args={
-                "inputs": document_pipeline_execution_log.input_data,
-                "start_node_id": document_pipeline_execution_log.datasource_node_id,
-                "datasource_type": document_pipeline_execution_log.datasource_type,
-                "datasource_info_list": [json.loads(document_pipeline_execution_log.datasource_info)],
-                "original_document_id": document.id,
-            },
-            invoke_from=InvokeFrom.PUBLISHED_PIPELINE,
-            streaming=False,
-            call_depth=0,
-            workflow_thread_pool_id=None,
-            is_retry=True,
-        )
+        with self._session_maker() as session:
+            document_pipeline_execution_log = session.scalar(
+                select(DocumentPipelineExecutionLog)
+                .where(DocumentPipelineExecutionLog.document_id == document.id)
+                .limit(1)
+            )
+            if not document_pipeline_execution_log:
+                raise ValueError("Document pipeline execution log not found")
+            pipeline = session.get(Pipeline, document_pipeline_execution_log.pipeline_id)
+            if not pipeline:
+                raise ValueError("Pipeline not found")
+            # convert to app config
+            workflow = session.scalar(
+                select(Workflow)
+                .where(
+                    Workflow.tenant_id == pipeline.tenant_id,
+                    Workflow.app_id == pipeline.id,
+                    Workflow.id == pipeline.workflow_id,
+                )
+                .limit(1)
+            )
+            if not workflow:
+                raise ValueError("Workflow not found")
+            PipelineGenerator().generate(
+                pipeline=pipeline,
+                workflow=workflow,
+                user=user,
+                args={
+                    "inputs": document_pipeline_execution_log.input_data,
+                    "start_node_id": document_pipeline_execution_log.datasource_node_id,
+                    "datasource_type": document_pipeline_execution_log.datasource_type,
+                    "datasource_info_list": [json.loads(document_pipeline_execution_log.datasource_info)],
+                    "original_document_id": document.id,
+                },
+                invoke_from=InvokeFrom.PUBLISHED_PIPELINE,
+                streaming=False,
+                call_depth=0,
+                workflow_thread_pool_id=None,
+                is_retry=True,
+            )
 
     def get_datasource_plugins(self, tenant_id: str, dataset_id: str, is_published: bool) -> list[dict]:
         """
         Get datasource plugins
         """
-        dataset: Dataset | None = db.session.scalar(
-            select(Dataset)
-            .where(
-                Dataset.id == dataset_id,
-                Dataset.tenant_id == tenant_id,
-            )
-            .limit(1)
-        )
-        if not dataset:
-            raise ValueError("Dataset not found")
-        pipeline: Pipeline | None = db.session.scalar(
-            select(Pipeline)
-            .where(
-                Pipeline.id == dataset.pipeline_id,
-                Pipeline.tenant_id == tenant_id,
-            )
-            .limit(1)
-        )
-        if not pipeline:
-            raise ValueError("Pipeline not found")
-
-        workflow: Workflow | None = None
-        if is_published:
-            workflow = self.get_published_workflow(pipeline=pipeline)
-        else:
-            workflow = self.get_draft_workflow(pipeline=pipeline)
-        if not pipeline or not workflow:
-            raise ValueError("Pipeline or workflow not found")
-
-        datasource_nodes = workflow.graph_dict.get("nodes", [])
-        datasource_plugins = []
-        for datasource_node in datasource_nodes:
-            if datasource_node.get("data", {}).get("type") == "datasource":
-                datasource_node_data = datasource_node["data"]
-                if not datasource_node_data:
-                    continue
-
-                variables = workflow.rag_pipeline_variables
-                if variables:
-                    variables_map = {item["variable"]: item for item in variables}
-                else:
-                    variables_map = {}
-
-                datasource_parameters = datasource_node_data.get("datasource_parameters", {})
-                user_input_variables_keys = []
-                user_input_variables = []
-
-                for _, value in datasource_parameters.items():
-                    if value.get("value") and isinstance(value.get("value"), str):
-                        pattern = r"\{\{#([a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z0-9_][a-zA-Z0-9_]{0,29}){1,10})#\}\}"
-                        match = re.match(pattern, value["value"])
-                        if match:
-                            full_path = match.group(1)
-                            last_part = full_path.split(".")[-1]
-                            user_input_variables_keys.append(last_part)
-                    elif value.get("value") and isinstance(value.get("value"), list):
-                        last_part = value.get("value")[-1]
-                        user_input_variables_keys.append(last_part)
-                for key, value in variables_map.items():
-                    if key in user_input_variables_keys:
-                        user_input_variables.append(value)
-
-                # get credentials
-                datasource_provider_service: DatasourceProviderService = DatasourceProviderService()
-                credentials: list[dict[Any, Any]] = datasource_provider_service.list_datasource_credentials(
-                    tenant_id=tenant_id,
-                    provider=datasource_node_data.get("provider_name"),
-                    plugin_id=datasource_node_data.get("plugin_id"),
+        with self._session_maker() as session:
+            dataset: Dataset | None = session.scalar(
+                select(Dataset)
+                .where(
+                    Dataset.id == dataset_id,
+                    Dataset.tenant_id == tenant_id,
                 )
-                credential_info_list: list[Any] = []
-                for credential in credentials:
-                    credential_info_list.append(
+                .limit(1)
+            )
+            if not dataset:
+                raise ValueError("Dataset not found")
+            pipeline: Pipeline | None = session.scalar(
+                select(Pipeline)
+                .where(
+                    Pipeline.id == dataset.pipeline_id,
+                    Pipeline.tenant_id == tenant_id,
+                )
+                .limit(1)
+            )
+            if not pipeline:
+                raise ValueError("Pipeline not found")
+
+            if is_published:
+                workflow = session.scalar(
+                    select(Workflow)
+                    .where(
+                        Workflow.tenant_id == pipeline.tenant_id,
+                        Workflow.app_id == pipeline.id,
+                        Workflow.id == pipeline.workflow_id,
+                    )
+                    .limit(1)
+                )
+            else:
+                workflow = session.scalar(
+                    select(Workflow)
+                    .where(
+                        Workflow.tenant_id == pipeline.tenant_id,
+                        Workflow.app_id == pipeline.id,
+                        Workflow.version == Workflow.VERSION_DRAFT,
+                    )
+                    .limit(1)
+                )
+            if not pipeline or not workflow:
+                raise ValueError("Pipeline or workflow not found")
+
+            datasource_nodes = workflow.graph_dict.get("nodes", [])
+            datasource_plugins = []
+            for datasource_node in datasource_nodes:
+                if datasource_node.get("data", {}).get("type") == "datasource":
+                    datasource_node_data = datasource_node["data"]
+                    if not datasource_node_data:
+                        continue
+
+                    variables = workflow.rag_pipeline_variables
+                    if variables:
+                        variables_map = {item["variable"]: item for item in variables}
+                    else:
+                        variables_map = {}
+
+                    datasource_parameters = datasource_node_data.get("datasource_parameters", {})
+                    user_input_variables_keys = []
+                    user_input_variables = []
+
+                    for _, value in datasource_parameters.items():
+                        if value.get("value") and isinstance(value.get("value"), str):
+                            pattern = (
+                                r"\{\{#([a-zA-Z0-9_]{1,50}"
+                                r"(?:\.[a-zA-Z0-9_][a-zA-Z0-9_]{0,29}){1,10})#\}\}"
+                            )
+                            match = re.match(pattern, value["value"])
+                            if match:
+                                full_path = match.group(1)
+                                last_part = full_path.split(".")[-1]
+                                user_input_variables_keys.append(last_part)
+                        elif value.get("value") and isinstance(value.get("value"), list):
+                            last_part = value.get("value")[-1]
+                            user_input_variables_keys.append(last_part)
+                    for key, value in variables_map.items():
+                        if key in user_input_variables_keys:
+                            user_input_variables.append(value)
+
+                    # get credentials
+                    datasource_provider_service: DatasourceProviderService = DatasourceProviderService()
+                    credentials: list[dict[Any, Any]] = datasource_provider_service.list_datasource_credentials(
+                        tenant_id=tenant_id,
+                        provider=datasource_node_data.get("provider_name"),
+                        plugin_id=datasource_node_data.get("plugin_id"),
+                    )
+                    credential_info_list: list[Any] = []
+                    for credential in credentials:
+                        credential_info_list.append(
+                            {
+                                "id": credential.get("id"),
+                                "name": credential.get("name"),
+                                "type": credential.get("type"),
+                                "is_default": credential.get("is_default"),
+                            }
+                        )
+
+                    datasource_plugins.append(
                         {
-                            "id": credential.get("id"),
-                            "name": credential.get("name"),
-                            "type": credential.get("type"),
-                            "is_default": credential.get("is_default"),
+                            "node_id": datasource_node.get("id"),
+                            "plugin_id": datasource_node_data.get("plugin_id"),
+                            "provider_name": datasource_node_data.get("provider_name"),
+                            "datasource_type": datasource_node_data.get("provider_type"),
+                            "title": datasource_node_data.get("title"),
+                            "user_input_variables": user_input_variables,
+                            "credentials": credential_info_list,
                         }
                     )
 
-                datasource_plugins.append(
-                    {
-                        "node_id": datasource_node.get("id"),
-                        "plugin_id": datasource_node_data.get("plugin_id"),
-                        "provider_name": datasource_node_data.get("provider_name"),
-                        "datasource_type": datasource_node_data.get("provider_type"),
-                        "title": datasource_node_data.get("title"),
-                        "user_input_variables": user_input_variables,
-                        "credentials": credential_info_list,
-                    }
-                )
+            return datasource_plugins
 
-        return datasource_plugins
-
-    def get_pipeline(self, tenant_id: str, dataset_id: str) -> Pipeline:
+    def get_pipeline(self, tenant_id: str, dataset_id: str, session: Session | None = None) -> Pipeline:
         """
         Get pipeline
         """
-        dataset: Dataset | None = db.session.scalar(
+        if session is None:
+            with self._session_maker() as new_session:
+                return self.get_pipeline(tenant_id, dataset_id, session=new_session)
+
+        dataset: Dataset | None = session.scalar(
             select(Dataset)
             .where(
                 Dataset.id == dataset_id,
@@ -1586,7 +1683,7 @@ class RagPipelineService:
         )
         if not dataset:
             raise ValueError("Dataset not found")
-        pipeline: Pipeline | None = db.session.scalar(
+        pipeline: Pipeline | None = session.scalar(
             select(Pipeline)
             .where(
                 Pipeline.id == dataset.pipeline_id,

api/services/rag_pipeline/rag_pipeline_transform_service.py

@@ -8,7 +8,7 @@ from uuid import uuid4
 import yaml
 from flask_login import current_user
 from sqlalchemy import select
-from sqlalchemy.orm import scoped_session
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from constants import DOCUMENT_EXTENSIONS
@@ -16,7 +16,6 @@ from core.plugin.impl.plugin import PluginInstaller
 from core.plugin.plugin_service import PluginService
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
-from extensions.ext_database import db
 from factories import variable_factory
 from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline
 from models.enums import DatasetRuntimeMode, DataSourceType
@@ -29,7 +28,7 @@ logger = logging.getLogger(__name__)
 
 
 class RagPipelineTransformService:
-    def transform_dataset(self, dataset_id: str, session: scoped_session):
+    def transform_dataset(self, dataset_id: str, session: Session):
         dataset = session.get(Dataset, dataset_id)
         if not dataset:
             raise ValueError("Dataset not found")
@@ -45,11 +44,11 @@ class RagPipelineTransformService:
         indexing_technique = dataset.indexing_technique
 
         if not datasource_type and not indexing_technique:
-            return self._transform_to_empty_pipeline(dataset)
+            return self._transform_to_empty_pipeline(dataset, session=session)
 
         doc_form = dataset.doc_form
         if not doc_form:
-            return self._transform_to_empty_pipeline(dataset)
+            return self._transform_to_empty_pipeline(dataset, session=session)
         retrieval_model = RetrievalSetting.model_validate(dataset.retrieval_model) if dataset.retrieval_model else None
         pipeline_yaml = self._get_transform_yaml(doc_form, datasource_type, indexing_technique)
         # deal dependencies
@@ -81,7 +80,7 @@ class RagPipelineTransformService:
             workflow_data["graph"] = graph
             pipeline_yaml["workflow"] = workflow_data
         # create pipeline
-        pipeline = self._create_pipeline(pipeline_yaml)
+        pipeline = self._create_pipeline(pipeline_yaml, session=session)
 
         # save chunk structure to dataset
         if doc_form == IndexStructureType.PARENT_CHILD_INDEX:
@@ -97,7 +96,7 @@ class RagPipelineTransformService:
         # deal document data
         self._deal_document_data(dataset, session)
 
-        session.commit()
+        session.flush()
         return {
             "pipeline_id": pipeline.id,
             "dataset_id": dataset_id,
@@ -195,6 +194,7 @@ class RagPipelineTransformService:
     def _create_pipeline(
         self,
         data: dict[str, Any],
+        session: Session,
     ) -> Pipeline:
         """Create a new app or update an existing one."""
         pipeline_data = data.get("rag_pipeline", {})
@@ -227,8 +227,8 @@ class RagPipelineTransformService:
         )
         pipeline.id = str(uuid4())
 
-        db.session.add(pipeline)
-        db.session.flush()
+        session.add(pipeline)
+        session.flush()
         # create draft workflow
         draft_workflow = Workflow(
             tenant_id=pipeline.tenant_id,
@@ -254,11 +254,11 @@ class RagPipelineTransformService:
             conversation_variables=conversation_variables,
             rag_pipeline_variables=rag_pipeline_variables_list,
         )
-        db.session.add(draft_workflow)
-        db.session.add(published_workflow)
-        db.session.flush()
+        session.add(draft_workflow)
+        session.add(published_workflow)
+        session.flush()
         pipeline.workflow_id = published_workflow.id
-        db.session.add(pipeline)
+        session.add(pipeline)
         return pipeline
 
     def _deal_dependencies(self, pipeline_yaml: dict[str, Any], tenant_id: str):
@@ -289,29 +289,29 @@ class RagPipelineTransformService:
             logger.debug("Installing missing pipeline plugins %s", need_install_plugin_unique_identifiers)
             PluginService.install_from_marketplace_pkg(tenant_id, need_install_plugin_unique_identifiers)
 
-    def _transform_to_empty_pipeline(self, dataset: Dataset):
+    def _transform_to_empty_pipeline(self, dataset: Dataset, session: Session):
         pipeline = Pipeline(
             tenant_id=dataset.tenant_id,
             name=dataset.name,
             description=dataset.description,
             created_by=current_user.id,
         )
-        db.session.add(pipeline)
-        db.session.flush()
+        session.add(pipeline)
+        session.flush()
 
         dataset.pipeline_id = pipeline.id
         dataset.runtime_mode = DatasetRuntimeMode.RAG_PIPELINE
         dataset.updated_by = current_user.id
         dataset.updated_at = datetime.now(UTC).replace(tzinfo=None)
-        db.session.add(dataset)
-        db.session.commit()
+        session.add(dataset)
+        session.flush()
         return {
             "pipeline_id": pipeline.id,
             "dataset_id": dataset.id,
             "status": "success",
         }
 
-    def _deal_document_data(self, dataset: Dataset, session: scoped_session):
+    def _deal_document_data(self, dataset: Dataset, session: Session):
         file_node_id = "1752479895761"
         notion_node_id = "1752489759475"
         jina_node_id = "1752491761974"

api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py

@@ -54,7 +54,7 @@ class TestPipelineTemplateListApi:
                 return_value=templates,
             ),
         ):
-            response, status = method(api, str(uuid4()))
+            response, status = method(api, MagicMock(), str(uuid4()))
 
         assert status == 200
         assert response == {
@@ -100,7 +100,7 @@ class TestPipelineTemplateDetailApi:
                 return_value=service,
             ),
         ):
-            response, status = method(api, "tpl-1")
+            response, status = method(api, MagicMock(), "tpl-1")
 
         assert status == 200
         assert response == {**template, "created_by": None}
@@ -120,7 +120,7 @@ class TestPipelineTemplateDetailApi:
             ),
         ):
             with pytest.raises(NotFound):
-                method(api, "non-existent-id")
+                method(api, MagicMock(), "non-existent-id")
 
     def test_get_returns_404_for_customized_type_not_found(self, app: Flask) -> None:
         api = PipelineTemplateDetailApi()
@@ -137,7 +137,7 @@ class TestPipelineTemplateDetailApi:
             ),
         ):
             with pytest.raises(NotFound):
-                method(api, "non-existent-id")
+                method(api, MagicMock(), "non-existent-id")
 
 
 class TestCustomizedPipelineTemplateApi:

api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py

@@ -470,7 +470,7 @@ class TestPipelineRunApis:
                 return_value={"ok": True},
             ),
         ):
-            assert method(api, user, pipeline) == {"ok": True}
+            assert method(api, MagicMock(), user, pipeline) == {"ok": True}
 
     def test_draft_run_rate_limit(self, app: Flask) -> None:
         api = DraftRagPipelineRunApi()
@@ -498,7 +498,7 @@ class TestPipelineRunApis:
             ),
         ):
             with pytest.raises(InvokeRateLimitHttpError):
-                method(api, user, pipeline)
+                method(api, MagicMock(), user, pipeline)
 
 
 class TestDraftNodeRun:
@@ -600,7 +600,7 @@ class TestMiscApis:
             app.test_request_context("/"),
         ):
             with pytest.raises(Forbidden):
-                method(api, user, "ds1")
+                method(api, MagicMock(spec=Session), user, "ds1")
 
     def test_recommended_plugins(self, app: Flask) -> None:
         api = RagPipelineRecommendedPluginApi()
@@ -655,7 +655,7 @@ class TestPublishedRagPipelineRunApi:
                 return_value={"ok": True},
             ),
         ):
-            result = method(api, user, pipeline)
+            result = method(api, MagicMock(), user, pipeline)
             assert result == {"ok": True}
 
     def test_published_run_rate_limit(self, app: Flask) -> None:
@@ -681,7 +681,7 @@ class TestPublishedRagPipelineRunApi:
             ),
         ):
             with pytest.raises(InvokeRateLimitHttpError):
-                method(api, user, pipeline)
+                method(api, MagicMock(), user, pipeline)
 
 
 class TestDefaultBlockConfigApi:

api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py

@@ -101,8 +101,8 @@ class TestRagPipelineServiceGetPipeline:
 
         service = self._make_service(flask_app_with_containers)
 
-        with pytest.raises(ValueError, match="(Dataset not found|Pipeline not found)"):
-            service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id)
+        with pytest.raises(ValueError, match="Pipeline not found"):
+            service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id, session=db_session_with_containers)
 
     def test_get_pipeline_returns_pipeline_when_found(
         self, db_session_with_containers: Session, flask_app_with_containers: Flask
@@ -117,7 +117,7 @@ class TestRagPipelineServiceGetPipeline:
 
         service = self._make_service(flask_app_with_containers)
 
-        result = service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id)
+        result = service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id, session=db_session_with_containers)
 
         assert result.id == pipeline.id
 
@@ -165,7 +165,9 @@ class TestUpdateCustomizedPipelineTemplate:
             description="Updated description",
             icon_info=IconInfo(icon="🔥"),
         )
-        result = RagPipelineService.update_customized_pipeline_template(template.id, info, account, tenant_id)
+        result = RagPipelineService.update_customized_pipeline_template(
+            template.id, info, account, tenant_id, session=db_session_with_containers
+        )
 
         assert result.name == "Updated Name"
         assert result.description == "Updated description"
@@ -203,7 +205,9 @@ class TestUpdateCustomizedPipelineTemplate:
             icon_info=IconInfo(icon="📄"),
         )
         with pytest.raises(ValueError, match="Template name is already exists"):
-            RagPipelineService.update_customized_pipeline_template(template1.id, info, account, tenant_id)
+            RagPipelineService.update_customized_pipeline_template(
+                template1.id, info, account, tenant_id, session=db_session_with_containers
+            )
 
 
 class TestDeleteCustomizedPipelineTemplate:
@@ -241,14 +245,14 @@ class TestDeleteCustomizedPipelineTemplate:
         template_id = template.id
         db_session_with_containers.flush()
 
-        RagPipelineService.delete_customized_pipeline_template(template_id, tenant_id)
+        RagPipelineService.delete_customized_pipeline_template(
+            template_id, tenant_id, session=db_session_with_containers
+        )
 
         # Verify the record is deleted within the same context
         from sqlalchemy import select
 
-        from extensions.ext_database import db as ext_db
-
-        remaining = ext_db.session.scalar(
+        remaining = db_session_with_containers.scalar(
             select(PipelineCustomizedTemplate).where(PipelineCustomizedTemplate.id == template_id)
         )
         assert remaining is None

api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from inspect import unwrap
-from unittest.mock import PropertyMock, patch
+from unittest.mock import Mock, PropertyMock, patch
 
 import pytest
 from flask import Flask
@@ -64,7 +64,9 @@ class TestPipelineTemplateListApi:
         tenant_id = "tenant-1"
         service_calls: list[tuple[str, str, str]] = []
 
-        def get_pipeline_templates(template_type: str, language: str, current_tenant_id: str) -> dict[str, object]:
+        def get_pipeline_templates(
+            session: Mock, template_type: str, language: str, current_tenant_id: str
+        ) -> dict[str, object]:
             service_calls.append((template_type, language, current_tenant_id))
             return {"pipeline_templates": [_template_item()]}
 
@@ -72,7 +74,7 @@ class TestPipelineTemplateListApi:
             app.test_request_context("/rag/pipeline/templates"),
             patch.object(module.RagPipelineService, "get_pipeline_templates", side_effect=get_pipeline_templates),
         ):
-            response, status = method(api, tenant_id)
+            response, status = method(api, Mock(), tenant_id)
 
         assert status == 200
         assert service_calls == [("built-in", "en-US", tenant_id)]
@@ -92,7 +94,9 @@ class TestPipelineTemplateListApi:
         tenant_id = "tenant-1"
         service_calls: list[tuple[str, str, str]] = []
 
-        def get_pipeline_templates(template_type: str, language: str, current_tenant_id: str) -> dict[str, object]:
+        def get_pipeline_templates(
+            session: Mock, template_type: str, language: str, current_tenant_id: str
+        ) -> dict[str, object]:
             service_calls.append((template_type, language, current_tenant_id))
             return {"pipeline_templates": []}
 
@@ -100,7 +104,7 @@ class TestPipelineTemplateListApi:
             app.test_request_context("/rag/pipeline/templates?type=customized&language=ja-JP"),
             patch.object(module.RagPipelineService, "get_pipeline_templates", side_effect=get_pipeline_templates),
         ):
-            response, status = method(api, tenant_id)
+            response, status = method(api, Mock(), tenant_id)
 
         assert status == 200
         assert response == {"pipeline_templates": []}
@@ -114,7 +118,9 @@ class TestPipelineTemplateDetailApi:
         service_calls: list[tuple[str, str]] = []
 
         class Service:
-            def get_pipeline_template_detail(self, template_id: str, template_type: str) -> dict[str, object]:
+            def get_pipeline_template_detail(
+                self, session: Mock, template_id: str, template_type: str
+            ) -> dict[str, object]:
                 service_calls.append((template_id, template_type))
                 return _template_detail()
 
@@ -122,7 +128,7 @@ class TestPipelineTemplateDetailApi:
             app.test_request_context("/rag/pipeline/templates/template-1?type=customized"),
             patch.object(module, "RagPipelineService", Service),
         ):
-            response, status = method(api, "template-1")
+            response, status = method(api, Mock(), "template-1")
 
         assert status == 200
         assert response == {**_template_detail(), "created_by": None}
@@ -133,7 +139,7 @@ class TestPipelineTemplateDetailApi:
         method = unwrap(api.get)
 
         class Service:
-            def get_pipeline_template_detail(self, template_id: str, template_type: str) -> None:
+            def get_pipeline_template_detail(self, session: Mock, template_id: str, template_type: str) -> None:
                 return None
 
         with (
@@ -141,7 +147,7 @@ class TestPipelineTemplateDetailApi:
             patch.object(module, "RagPipelineService", Service),
         ):
             with pytest.raises(NotFound):
-                method(api, "missing")
+                method(api, Mock(), "missing")
 
 
 class TestCustomizedPipelineTemplateApi:

api/tests/unit_tests/services/rag_pipeline/pipeline_template/conftest.py

@@ -0,0 +1,18 @@
+from collections.abc import Callable
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+
+@pytest.fixture
+def patch_session_factory(mocker: MockerFixture) -> Callable[[str, Mock], Mock]:
+    def _patch(module_path: str, session_mock: Mock) -> Mock:
+        session_context = mocker.MagicMock()
+        session_context.__enter__.return_value = session_mock
+        session_context.__exit__.return_value = None
+        session_maker = mocker.Mock(return_value=session_context)
+        mocker.patch(f"{module_path}.session_factory.get_session_maker", return_value=session_maker)
+        return session_maker
+
+    return _patch

api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py

@@ -23,7 +23,7 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None:
     )
     retrieval = BuiltInPipelineTemplateRetrieval()
 
-    templates = retrieval.get_pipeline_templates("en-US")
+    templates = retrieval.get_pipeline_templates(mocker.Mock(), "en-US")
 
     assert templates == {"pipeline_templates": [{"id": "tpl-1"}]}
 
@@ -40,7 +40,7 @@ def test_get_pipeline_template_detail(mocker: MockerFixture) -> None:
     )
     retrieval = BuiltInPipelineTemplateRetrieval()
 
-    detail = retrieval.get_pipeline_template_detail("tpl-1")
+    detail = retrieval.get_pipeline_template_detail(mocker.Mock(), "tpl-1")
 
     assert detail == {"id": "tpl-1", "name": "Template 1"}
 
@@ -53,7 +53,7 @@ def test_get_pipeline_templates_missing_language_returns_empty_dict(mocker: Mock
     )
     retrieval = BuiltInPipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_templates("fr-FR")
+    result = retrieval.get_pipeline_templates(mocker.Mock(), "fr-FR")
 
     assert result == {}
 
@@ -66,7 +66,7 @@ def test_get_pipeline_template_detail_returns_none_for_unknown_id(mocker: Mocker
     )
     retrieval = BuiltInPipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_template_detail("nonexistent-id")
+    result = retrieval.get_pipeline_template_detail(mocker.Mock(), "nonexistent-id")
 
     assert result is None
 

api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py

@@ -19,13 +19,9 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None:
     scalars_mock.all.return_value = [customized_template]
     session_mock = mocker.Mock()
     session_mock.scalars.return_value = scalars_mock
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.customized.customized_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = CustomizedPipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_templates("en-US", "tenant-id")
+    result = retrieval.get_pipeline_templates(session_mock, "en-US", "tenant-id")
 
     assert retrieval.get_type() == PipelineTemplateType.CUSTOMIZED
     assert result == {
@@ -53,13 +49,9 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N
         yaml_content="workflow:\n  graph:\n    edges: []",
         created_user_name="creator",
     )
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.customized.customized_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = CustomizedPipelineTemplateRetrieval()
 
-    detail = retrieval.get_pipeline_template_detail("tpl-1")
+    detail = retrieval.get_pipeline_template_detail(session_mock, "tpl-1")
 
     assert detail == {
         "id": "tpl-1",
@@ -76,12 +68,8 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N
 def test_get_pipeline_template_detail_returns_none_when_not_found(mocker: MockerFixture) -> None:
     session_mock = mocker.Mock()
     session_mock.get.return_value = None
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.customized.customized_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = CustomizedPipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_template_detail("missing")
+    result = retrieval.get_pipeline_template_detail(session_mock, "missing")
 
     assert result is None

api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py

@@ -21,13 +21,9 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None:
     scalars_mock.all.return_value = [built_in_template]
     session_mock = mocker.Mock()
     session_mock.scalars.return_value = scalars_mock
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.database.database_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = DatabasePipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_templates("en-US")
+    result = retrieval.get_pipeline_templates(session_mock, "en-US")
 
     assert retrieval.get_type() == PipelineTemplateType.DATABASE
     assert result == {
@@ -56,13 +52,9 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N
         chunk_structure="general",
         yaml_content="workflow:\n  graph:\n    nodes: []",
     )
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.database.database_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = DatabasePipelineTemplateRetrieval()
 
-    detail = retrieval.get_pipeline_template_detail("tpl-1")
+    detail = retrieval.get_pipeline_template_detail(session_mock, "tpl-1")
 
     assert detail == {
         "id": "tpl-1",
@@ -78,12 +70,8 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N
 def test_get_pipeline_template_detail_returns_none_when_not_found(mocker: MockerFixture) -> None:
     session_mock = mocker.Mock()
     session_mock.get.return_value = None
-    mocker.patch(
-        "services.rag_pipeline.pipeline_template.database.database_retrieval.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     retrieval = DatabasePipelineTemplateRetrieval()
 
-    result = retrieval.get_pipeline_template_detail("missing")
+    result = retrieval.get_pipeline_template_detail(session_mock, "missing")
 
     assert result is None

api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py

@@ -1,11 +1,15 @@
+from unittest.mock import Mock
+
 from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase
 
 
 class DummyRetrieval(PipelineTemplateRetrievalBase):
-    def get_pipeline_templates(self, language: str) -> dict:
+    def get_pipeline_templates(self, session: Mock, language: str, current_tenant_id: str | None = None) -> dict:
+        del session, current_tenant_id
         return {"language": language}
 
-    def get_pipeline_template_detail(self, template_id: str) -> dict | None:
+    def get_pipeline_template_detail(self, session: Mock, template_id: str) -> dict | None:
+        del session
         return {"id": template_id}
 
     def get_type(self) -> str:
@@ -14,7 +18,8 @@ class DummyRetrieval(PipelineTemplateRetrievalBase):
 
 def test_pipeline_template_retrieval_base_concrete_implementation() -> None:
     retrieval = DummyRetrieval()
+    session = Mock()
 
-    assert retrieval.get_pipeline_templates("en-US") == {"language": "en-US"}
-    assert retrieval.get_pipeline_template_detail("tpl-1") == {"id": "tpl-1"}
+    assert retrieval.get_pipeline_templates(session, "en-US") == {"language": "en-US"}
+    assert retrieval.get_pipeline_template_detail(session, "tpl-1") == {"id": "tpl-1"}
     assert retrieval.get_type() == "dummy"

api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py

@@ -18,13 +18,14 @@ def test_get_pipeline_templates_fallbacks_to_database_on_error(mocker: MockerFix
         return_value={"pipeline_templates": [{"id": "db-1"}]},
     )
     retrieval = RemotePipelineTemplateRetrieval()
+    session = mocker.Mock()
 
-    result = retrieval.get_pipeline_templates("en-US")
+    result = retrieval.get_pipeline_templates(session, "en-US")
 
     assert retrieval.get_type() == PipelineTemplateType.REMOTE
     assert result == {"pipeline_templates": [{"id": "db-1"}]}
     fetch_mock.assert_called_once_with("en-US")
-    fallback_mock.assert_called_once_with("en-US")
+    fallback_mock.assert_called_once_with(session, "en-US")
 
 
 def test_get_pipeline_template_detail_fallbacks_to_database_on_error(mocker: MockerFixture) -> None:
@@ -39,12 +40,13 @@ def test_get_pipeline_template_detail_fallbacks_to_database_on_error(mocker: Moc
         return_value={"id": "db-1"},
     )
     retrieval = RemotePipelineTemplateRetrieval()
+    session = mocker.Mock()
 
-    result = retrieval.get_pipeline_template_detail("tpl-1")
+    result = retrieval.get_pipeline_template_detail(session, "tpl-1")
 
     assert result == {"id": "db-1"}
     fetch_mock.assert_called_once_with("tpl-1")
-    fallback_mock.assert_called_once_with("tpl-1")
+    fallback_mock.assert_called_once_with(session, "tpl-1")
 
 
 def test_fetch_pipeline_templates_from_dify_official(mocker: MockerFixture) -> None:

api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py

@@ -1,15 +1,30 @@
+from collections.abc import Iterator
 from types import SimpleNamespace
 from typing import cast
+from uuid import uuid4
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy import create_engine, func, select
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.app.entities.app_invoke_entities import InvokeFrom
-from models.dataset import Pipeline
+from models.dataset import Document, Pipeline
+from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
 from models.model import Account, App, EndUser
 from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService
 
 
+@pytest.fixture
+def document_session() -> Iterator[Session]:
+    engine = create_engine("sqlite:///:memory:")
+    Document.__table__.create(engine)
+    session_factory = sessionmaker(bind=engine, expire_on_commit=False)
+    with session_factory() as session:
+        yield session
+    engine.dispose()
+
+
 def test_get_max_active_requests_uses_smallest_non_zero_limit(mocker: MockerFixture) -> None:
     mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_DEFAULT_ACTIVE_REQUESTS", 5)
     mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_MAX_ACTIVE_REQUESTS", 3)
@@ -63,6 +78,7 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
 
     mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1"))
     update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status")
+    session = mocker.Mock()
 
     generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator")
     generator_instance = generator_cls.return_value
@@ -70,6 +86,7 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
     generator_cls.convert_to_event_stream.return_value = "stream-events"
 
     result = PipelineGenerateService.generate(
+        session=session,
         pipeline=pipeline,
         user=user,
         args=args,
@@ -78,42 +95,40 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
     )
 
     assert result == "stream-events"
-    update_status_mock.assert_called_once_with("doc-1")
+    update_status_mock.assert_called_once_with("doc-1", session)
 
 
-def test_update_document_status_updates_existing_document(mocker: MockerFixture) -> None:
-    document = SimpleNamespace(indexing_status="completed")
-
-    session_mock = mocker.Mock()
-    session_mock.get.return_value = document
-    add_mock = session_mock.add
-    commit_mock = session_mock.commit
-    mocker.patch(
-        "services.rag_pipeline.pipeline_generate_service.db",
-        new=SimpleNamespace(session=session_mock),
+def test_update_document_status_updates_existing_document(document_session: Session) -> None:
+    session = document_session
+    document_id = str(uuid4())
+    document = Document(
+        id=document_id,
+        tenant_id=str(uuid4()),
+        dataset_id=str(uuid4()),
+        position=1,
+        data_source_type=DataSourceType.UPLOAD_FILE,
+        batch="batch-1",
+        name="Doc",
+        created_from=DocumentCreatedFrom.WEB,
+        created_by=str(uuid4()),
+        indexing_status=IndexingStatus.COMPLETED,
     )
+    session.add(document)
+    session.commit()
 
-    PipelineGenerateService.update_document_status("doc-1")
+    PipelineGenerateService.update_document_status(document_id, session)
 
-    assert document.indexing_status == "waiting"
-    add_mock.assert_called_once_with(document)
-    commit_mock.assert_called_once()
+    updated_document = session.get(Document, document_id)
+    assert updated_document is not None
+    assert updated_document.indexing_status == IndexingStatus.WAITING
 
 
-def test_update_document_status_skips_when_document_missing(mocker: MockerFixture) -> None:
-    session_mock = mocker.Mock()
-    session_mock.get.return_value = None
-    add_mock = session_mock.add
-    commit_mock = session_mock.commit
-    mocker.patch(
-        "services.rag_pipeline.pipeline_generate_service.db",
-        new=SimpleNamespace(session=session_mock),
-    )
+def test_update_document_status_skips_when_document_missing(document_session: Session) -> None:
+    session = document_session
 
-    PipelineGenerateService.update_document_status("missing")
+    PipelineGenerateService.update_document_status(str(uuid4()), session)
 
-    add_mock.assert_not_called()
-    commit_mock.assert_not_called()
+    assert session.scalar(select(func.count()).select_from(Document)) == 0
 
 
 # --- generate_single_iteration ---

api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py

@@ -152,7 +152,7 @@ def test_upload_invoke_entities_returns_file_id(mocker: MockerFixture, proxy) ->
     upload_file = SimpleNamespace(id="uploaded-file-1")
     file_service_cls = mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService")
     file_service_cls.return_value.upload_text.return_value = upload_file
-    mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.db", mocker.Mock(engine="fake-engine"))
+    mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.db", SimpleNamespace(engine="fake-engine"))
 
     result = proxy._upload_invoke_entities()
 

api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py

@@ -1,16 +1,31 @@
 import logging
+from collections.abc import Iterator
 from datetime import UTC, datetime
 from types import SimpleNamespace
 from typing import cast
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy import create_engine, select
+from sqlalchemy.orm import Session, sessionmaker
 
-from models.dataset import Dataset
+from models.dataset import Dataset, Pipeline
+from models.enums import DatasetRuntimeMode
 from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration
 from services.rag_pipeline.rag_pipeline_transform_service import RagPipelineTransformService
 
 
+@pytest.fixture
+def pipeline_session() -> Iterator[Session]:
+    engine = create_engine("sqlite:///:memory:")
+    Dataset.__table__.create(engine)
+    Pipeline.__table__.create(engine)
+    session_factory = sessionmaker(bind=engine, expire_on_commit=False)
+    with session_factory() as session:
+        yield session
+    engine.dispose()
+
+
 @pytest.mark.parametrize(
     ("doc_form", "datasource_type", "indexing_technique"),
     [
@@ -92,51 +107,37 @@ def test_deal_dependencies_installs_missing_marketplace_plugins(mocker: MockerFi
     install_mock.assert_called_once_with("tenant-1", ["missing-plugin:1.0.0"])
 
 
-def test_transform_to_empty_pipeline_updates_dataset_and_commits(mocker: MockerFixture) -> None:
+def test_transform_to_empty_pipeline_updates_dataset_and_flushes(
+    mocker: MockerFixture, pipeline_session: Session
+) -> None:
     service = RagPipelineTransformService()
     mocker.patch(
         "services.rag_pipeline.rag_pipeline_transform_service.current_user",
         SimpleNamespace(id="user-1"),
     )
 
-    class FakePipeline:
-        def __init__(self, **kwargs):
-            self.id = "pipeline-1"
-            self.tenant_id = kwargs["tenant_id"]
-            self.name = kwargs["name"]
-            self.description = kwargs["description"]
-            self.created_by = kwargs["created_by"]
-
-    mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.Pipeline", FakePipeline)
-    session_mock = mocker.Mock()
-    add_mock = session_mock.add
-    flush_mock = session_mock.flush
-    commit_mock = session_mock.commit
-    mocker.patch(
-        "services.rag_pipeline.rag_pipeline_transform_service.db",
-        new=SimpleNamespace(session=session_mock),
-    )
-
-    dataset = SimpleNamespace(
-        id="dataset-1",
+    session = pipeline_session
+    dataset = Dataset(
         tenant_id="tenant-1",
         name="Dataset",
         description="desc",
-        pipeline_id=None,
-        runtime_mode="general",
-        updated_by=None,
-        updated_at=None,
+        created_by="user-1",
     )
+    session.add(dataset)
+    session.commit()
+    flush_spy = mocker.spy(session, "flush")
+    commit_spy = mocker.spy(session, "commit")
 
-    result = service._transform_to_empty_pipeline(cast(Dataset, dataset))
+    result = service._transform_to_empty_pipeline(dataset, session)
 
-    assert result == {"pipeline_id": "pipeline-1", "dataset_id": "dataset-1", "status": "success"}
-    assert dataset.pipeline_id == "pipeline-1"
-    assert dataset.runtime_mode == "rag_pipeline"
+    assert flush_spy.call_count == 2
+    commit_spy.assert_not_called()
+    pipeline = session.scalar(select(Pipeline).where(Pipeline.id == dataset.pipeline_id))
+    assert pipeline is not None
+    assert result == {"pipeline_id": pipeline.id, "dataset_id": dataset.id, "status": "success"}
+    assert dataset.pipeline_id == pipeline.id
+    assert dataset.runtime_mode == DatasetRuntimeMode.RAG_PIPELINE
     assert dataset.updated_by == "user-1"
-    add_mock.assert_called()
-    flush_mock.assert_called_once()
-    commit_mock.assert_called_once()
 
 
 # --- transform_dataset ---
@@ -372,7 +373,6 @@ def test_transform_dataset_full_flow(mocker: MockerFixture) -> None:
 
     mocker.patch.object(service, "_deal_dependencies")
     mocker.patch.object(service, "_deal_document_data")
-    session_mock.commit = mocker.Mock()
 
     # Mock current_user to have the same tenant_id as dataset
     mock_current_user = SimpleNamespace(current_tenant_id="t1")
@@ -386,6 +386,8 @@ def test_transform_dataset_full_flow(mocker: MockerFixture) -> None:
     assert result["pipeline_id"] == "p-new"
     assert dataset.runtime_mode == "rag_pipeline"
     assert dataset.chunk_structure == "text_model"
+    session_mock.flush.assert_called_once_with()
+    session_mock.commit.assert_not_called()
 
 
 def test_transform_dataset_raises_for_unsupported_doc_form_after_pipeline_create(mocker: MockerFixture) -> None:
@@ -405,10 +407,6 @@ def test_transform_dataset_raises_for_unsupported_doc_form_after_pipeline_create
     )
     session_mock = mocker.Mock()
     session_mock.get.return_value = dataset
-    mocker.patch(
-        "services.rag_pipeline.rag_pipeline_transform_service.db",
-        new=SimpleNamespace(session=session_mock),
-    )
     mocker.patch.object(service, "_get_transform_yaml", return_value={"workflow": {"graph": {"nodes": []}}})
     mocker.patch.object(service, "_deal_dependencies")
     mocker.patch.object(service, "_create_pipeline", return_value=SimpleNamespace(id="p-new"))
@@ -441,11 +439,11 @@ def test_transform_dataset_raises_when_transform_yaml_missing_workflow(mocker: M
         service.transform_dataset("d1", session_mock)
 
 
-def test_create_pipeline_raises_when_workflow_data_missing() -> None:
+def test_create_pipeline_raises_when_workflow_data_missing(pipeline_session: Session) -> None:
     service = RagPipelineTransformService()
 
     with pytest.raises(ValueError, match="Missing workflow data for rag pipeline"):
-        service._create_pipeline({"rag_pipeline": {"name": "N"}})
+        service._create_pipeline({"rag_pipeline": {"name": "N"}}, pipeline_session)
 
 
 def test_deal_document_data_upload_file_with_existing_file(mocker: MockerFixture) -> None: