refactor: manage rag pipeline sessions explicitly (#38274)

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"