opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-08-15 04:59:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): bind remaining RAG resources to owners (#40503)

Browse Source
This commit is contained in:
WH-2099 2026-08-12 06:05:16 +00:00 committed by GitHub
parent d813edb945
commit af382be837
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
31 changed files with 957 additions and 322 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
RESOURCE_BOUNDARY_CHANGE_GUIDE.md Normal file
View File

@ -0,0 +1,12 @@
# Resource Boundary Change Guide
- Resolve the tenant-scoped parent at the request boundary, then pass the validated model, owner reference, and actor downstream.
- Put the complete owner tuple in the database query; do not load by a bare ID and check ownership afterward.
- Treat missing and foreign-owned resources alike as `404` before locks, rate limits, tasks, plugin calls, network calls, or writes.
- Reuse existing owner resolvers and trusted objects instead of adding parallel helpers or refetching the same resource.
- Pass tenant, actor, and session explicitly; authenticated code must not depend on ambient account or tenant fallbacks.
- Raise typed domain errors in services and translate them to HTTP errors in controllers; reserve `ValueError` for invalid values or state.
- Let RBAC own authorization when enabled, and run legacy dataset permission checks only when RBAC is disabled.
- Preserve successful HTTP responses and shared runtime contracts, especially Celery task names and argument shapes during rolling upgrades.
- Keep runtime validation and OpenAPI schemas aligned, then regenerate Markdown and TypeScript contracts after schema changes.
- Prove the boundary with a foreign-owner decoy and assert that rejected requests trigger no downstream side effects.

4
api/controllers/common/wraps.py
View File

@ -171,7 +171,9 @@ def _extract_resource_id(
pipeline_id = matched_args.get("pipeline_id")
if pipeline_id:
dataset = db.session.scalar(select(Dataset).where(Dataset.pipeline_id == str(pipeline_id)))
dataset = db.session.scalar(
select(Dataset).where(Dataset.pipeline_id == str(pipeline_id), Dataset.tenant_id == tenant_id)
)
if not dataset:
raise NotFound("Dataset not found for pipeline")
return str(dataset.id) # pyrefly: ignore[unnecessary-type-conversion]

75
api/controllers/console/datasets/rag_pipeline/rag_pipeline.py
View File

@ -3,10 +3,10 @@ from typing import Any
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from werkzeug.exceptions import NotFound
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
from configs import dify_config
from controllers.common.fields import SimpleDataResponse
from controllers.common.schema import (
JsonResponseWithStatus,
@ -16,11 +16,15 @@ from controllers.common.schema import (
)
from controllers.console import console_ns
from controllers.console.app.wraps import with_session
from controllers.console.datasets.wraps import get_rag_pipeline
from controllers.console.wraps import (
RBACPermission,
RBACResourceScope,
account_initialization_required,
enterprise_license_required,
knowledge_pipeline_publish_enabled,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
with_current_user,
@ -30,8 +34,11 @@ from fields.base import ResponseModel
from libs.helper import dump_response
from libs.login import login_required
from models.account import Account
from models.dataset import PipelineCustomizedTemplate
from models.dataset import Pipeline
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity
from services.errors.account import NoPermissionError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.rag_pipeline import RagPipelineService
logger: logging.Logger = logging.getLogger(__name__)
@ -125,15 +132,24 @@ class PipelineTemplateListApi(Resource):
class PipelineTemplateDetailApi(Resource):
@console_ns.doc(params=query_params_from_model(PipelineTemplateDetailQuery))
@console_ns.response(200, "Pipeline template", console_ns.models[PipelineTemplateDetailResponse.__name__])
@console_ns.response(404, "Pipeline template not found")
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@with_session
@with_current_tenant_id
@with_session(write=False)
@model_validate(PipelineTemplateDetailQuery)
def get(self, req_data: PipelineTemplateDetailQuery, session: Session, template_id: str) -> JsonResponseWithStatus:
def get(
self,
req_data: PipelineTemplateDetailQuery,
session: Session,
current_tenant_id: str,
template_id: str,
) -> JsonResponseWithStatus:
pipeline_template = RagPipelineService.get_pipeline_template_detail(
template_id,
current_tenant_id,
type=req_data.type,
session=session,
)
@ -181,38 +197,57 @@ class CustomizedPipelineTemplateApi(Resource):
@account_initialization_required
@enterprise_license_required
@console_ns.response(200, "Success", console_ns.models[SimpleDataResponse.__name__])
def post(self, template_id: str) -> JsonResponseWithStatus:
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
template = session.scalar(
select(PipelineCustomizedTemplate).where(PipelineCustomizedTemplate.id == template_id).limit(1)
@console_ns.response(404, "Customized pipeline template not found")
@with_current_tenant_id
@with_session(write=False)
def post(self, session: Session, current_tenant_id: str, template_id: str) -> JsonResponseWithStatus:
try:
yaml_content = RagPipelineService.get_customized_pipeline_template_yaml(
template_id, current_tenant_id, session=session
)
if not template:
raise ValueError("Customized pipeline template not found.")
except RagPipelineResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(SimpleDataResponse, {"data": template.yaml_content}), 200
return dump_response(SimpleDataResponse, {"data": yaml_content}), 200
@console_ns.route("/rag/pipelines/<string:pipeline_id>/customized/publish")
class PublishCustomizedPipelineTemplateApi(Resource):
@console_ns.expect(console_ns.models[CustomizedPipelineTemplatePayload.__name__])
@console_ns.response(204, "Pipeline template published")
@console_ns.response(404, "Pipeline, workflow, or dataset not found")
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@knowledge_pipeline_publish_enabled
@with_current_user
@with_current_tenant_id
@get_rag_pipeline
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_PIPELINE_RELEASE)
@model_validate(CustomizedPipelineTemplatePayload)
def post(
self,
req_data: CustomizedPipelineTemplatePayload,
current_tenant_id: str,
current_user: Account,
pipeline_id: str,
pipeline: Pipeline,
) -> tuple[str, int]:
rag_pipeline_service = RagPipelineService(db.session())
rag_pipeline_service.publish_customized_pipeline_template(
pipeline_id, req_data.model_dump(), current_user, current_tenant_id, session=db.session()
)
session = db.session()
dataset = pipeline.retrieve_dataset(session=session)
if dataset is None:
raise NotFound("Dataset not found")
if not dify_config.RBAC_ENABLED:
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, session)
except NoPermissionError as exc:
raise Forbidden(str(exc)) from exc
try:
RagPipelineService.publish_customized_pipeline_template(
pipeline, dataset, req_data.model_dump(), current_user, session=session
)
except RagPipelineResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return "", 204

29
api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py
View File

@ -10,6 +10,7 @@ from sqlalchemy.orm import Session, sessionmaker
from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
import services
from configs import dify_config
from controllers.common.controller_schemas import DefaultBlockConfigQuery, WorkflowListQuery, WorkflowUpdatePayload
from controllers.common.fields import SimpleResultResponse
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
@ -60,8 +61,10 @@ from models import Account
from models.dataset import Pipeline
from models.model import EndUser
from models.workflow import Workflow
from services.dataset_service import DatasetService
from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
from services.errors.llm import InvokeRateLimitError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService
from services.rag_pipeline.rag_pipeline import RagPipelineService
from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService
@ -1016,19 +1019,31 @@ class RagPipelineWorkflowLastRunApi(Resource):
@console_ns.route("/rag/pipelines/transform/datasets/<uuid:dataset_id>")
class RagPipelineTransformApi(Resource):
@console_ns.response(200, "Success", console_ns.models[RagPipelineOpaqueResponse.__name__])
@console_ns.response(404, "Dataset or pipeline not found")
@setup_required
@login_required
@account_initialization_required
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@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()
def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), current_tenant_id, session=session)
if dataset is None:
raise NotFound("Dataset not found.")
dataset_id_str = str(dataset_id)
rag_pipeline_transform_service = RagPipelineTransformService()
result = rag_pipeline_transform_service.transform_dataset(dataset_id_str, session)
return result
if not dify_config.RBAC_ENABLED:
if not (current_user.has_edit_permission or current_user.is_dataset_operator):
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as exc:
raise Forbidden(str(exc)) from exc
try:
return RagPipelineTransformService().transform_dataset(dataset, current_user.id, session)
except RagPipelineResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
@console_ns.route("/rag/pipelines/<uuid:pipeline_id>/workflows/draft/datasource/variables-inspect")

15
api/core/rag/retrieval/dataset_retrieval.py
View File

@ -651,16 +651,19 @@ class DatasetRetrieval:
self._record_usage(router_usage)
timer = None
if dataset_id:
# get retrieval model config
dataset_stmt = select(Dataset).where(Dataset.id == dataset_id)
selected_dataset = session.scalar(dataset_stmt)
allowed_dataset = next((dataset for dataset in available_datasets if dataset.id == dataset_id), None)
selected_dataset = (
session.scalar(select(Dataset).where(Dataset.id == allowed_dataset.id, Dataset.tenant_id == tenant_id))
if allowed_dataset
else None
)
if selected_dataset:
results = []
if selected_dataset.provider == "external":
external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval(
session=session,
tenant_id=selected_dataset.tenant_id,
dataset_id=dataset_id,
dataset_id=selected_dataset.id,
query=query,
external_retrieval_parameters=selected_dataset.retrieval_model,
metadata_condition=metadata_condition,
@ -674,7 +677,7 @@ class DatasetRetrieval:
if document.metadata is not None:
document.metadata["score"] = external_document.get("score")
document.metadata["title"] = external_document.get("title")
document.metadata["dataset_id"] = dataset_id
document.metadata["dataset_id"] = selected_dataset.id
document.metadata["dataset_name"] = selected_dataset.name
results.append(document)
else:
@ -724,7 +727,7 @@ class DatasetRetrieval:
weights=retrieval_model_config.get("weights", None),
document_ids_filter=document_ids_filter,
)
self._on_query(query, None, [dataset_id], app_id, user_from, user_id)
self._on_query(query, None, [selected_dataset.id], app_id, user_from, user_id)
if results:
thread = threading.Thread(

10
api/models/dataset.py
View File

@ -347,7 +347,11 @@ class Dataset(Base):
def get_doc_form(self, *, session: Session) -> str | None:
if self.chunk_structure:
return self.chunk_structure
return session.scalar(select(Document.doc_form).where(Document.dataset_id == self.id).limit(1))
return session.scalar(
select(Document.doc_form)
.where(Document.dataset_id == self.id, Document.tenant_id == self.tenant_id)
.limit(1)
)
@property
def retrieval_model_dict(self):
@ -1744,7 +1748,9 @@ class Pipeline(TypeBase):
)
def retrieve_dataset(self, session: Session | scoped_session):
return session.scalar(select(Dataset).where(Dataset.pipeline_id == self.id))
return session.scalar(
select(Dataset).where(Dataset.pipeline_id == self.id, Dataset.tenant_id == self.tenant_id)
)
class DocumentPipelineExecutionLog(TypeBase):

4
api/openapi/markdown/console-openapi.md
View File

@ -7891,6 +7891,7 @@ Update account-level Step-by-step Tour state
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [SimpleDataResponse](#simpledataresponse)<br> |
| 404 | Customized pipeline template not found | |
### [POST] /rag/pipeline/dataset
#### Request Body
@ -7939,6 +7940,7 @@ Update account-level Step-by-step Tour state
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Pipeline template | **application/json**: [PipelineTemplateDetailResponse](#pipelinetemplatedetailresponse)<br> |
| 404 | Pipeline template not found | |
### [GET] /rag/pipelines/datasource-plugins
#### Responses
@ -8014,6 +8016,7 @@ Update account-level Step-by-step Tour state
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [RagPipelineOpaqueResponse](#ragpipelineopaqueresponse)<br> |
| 404 | Dataset or pipeline not found | |
### [POST] /rag/pipelines/{pipeline_id}/customized/publish
#### Parameters
@ -8033,6 +8036,7 @@ Update account-level Step-by-step Tour state
| Code | Description |
| ---- | ----------- |
| 204 | Pipeline template published |
| 404 | Pipeline, workflow, or dataset not found |
### [GET] /rag/pipelines/{pipeline_id}/exports
#### Parameters

22
api/services/dataset_service.py
View File

@ -1384,7 +1384,11 @@ class DatasetService:
if dataset.maintainer != user.id:
user_permission = session.scalar(
select(DatasetPermission)
.where(DatasetPermission.dataset_id == dataset.id, DatasetPermission.account_id == user.id)
.where(
DatasetPermission.dataset_id == dataset.id,
DatasetPermission.account_id == user.id,
DatasetPermission.tenant_id == dataset.tenant_id,
)
.limit(1)
)
if not user_permission:
@ -1407,12 +1411,16 @@ class DatasetService:
raise NoPermissionError("You do not have permission to access this dataset.")
elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
if not any(
dp.dataset_id == dataset.id
for dp in session.scalars(
select(DatasetPermission).where(DatasetPermission.account_id == user.id)
).all()
):
user_permission = session.scalar(
select(DatasetPermission.id)
.where(
DatasetPermission.dataset_id == dataset.id,
DatasetPermission.account_id == user.id,
DatasetPermission.tenant_id == dataset.tenant_id,
)
.limit(1)
)
if user_permission is None:
raise NoPermissionError("You do not have permission to access this dataset.")
@staticmethod

2
api/services/errors/rag_pipeline.py Normal file
View File

@ -0,0 +1,2 @@
class RagPipelineResourceNotFoundError(Exception):
pass

6
api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py
View File

@ -30,8 +30,10 @@ class BuiltInPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
return result
@override
def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None:
del session
def get_pipeline_template_detail(
self, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None:
del current_tenant_id, session
result = self.fetch_pipeline_template_detail_from_builtin(template_id)
return result

17
api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py
View File

@ -49,8 +49,10 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
)
@override
def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None:
return self.fetch_pipeline_template_detail_from_db(template_id, session=session)
def get_pipeline_template_detail(
self, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None:
return self.fetch_pipeline_template_detail_from_db(template_id, current_tenant_id, session=session)
@override
def get_type(self) -> str:
@ -89,13 +91,20 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
return {"pipeline_templates": recommended_pipelines_results}
@classmethod
def fetch_pipeline_template_detail_from_db(cls, template_id: str, *, session: Session) -> dict[str, Any] | None:
def fetch_pipeline_template_detail_from_db(
cls, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None:
"""
Fetch pipeline template detail from db.
:param template_id: Template ID
:return:
"""
pipeline_template = session.get(PipelineCustomizedTemplate, template_id)
pipeline_template = session.scalar(
select(PipelineCustomizedTemplate).where(
PipelineCustomizedTemplate.id == template_id,
PipelineCustomizedTemplate.tenant_id == current_tenant_id,
)
)
if not pipeline_template:
return None

5
api/services/rag_pipeline/pipeline_template/database/database_retrieval.py
View File

@ -47,7 +47,10 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
return self.fetch_pipeline_templates_from_db(language, session=session)
@override
def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None:
def get_pipeline_template_detail(
self, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None:
del current_tenant_id
return self.fetch_pipeline_template_detail_from_db(template_id, session=session)
@override

4
api/services/rag_pipeline/pipeline_template/pipeline_template_base.py
View File

@ -10,6 +10,8 @@ class PipelineTemplateRetrievalBase(Protocol):
self, language: str, current_tenant_id: str | None = None, *, session: Session
) -> dict[str, Any]: ...
def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: ...
def get_pipeline_template_detail(
self, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None: ...
def get_type(self) -> str: ...

5
api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py
View File

@ -18,7 +18,10 @@ class RemotePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
"""
@override
def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None:
def get_pipeline_template_detail(
self, template_id: str, current_tenant_id: str, *, session: Session
) -> dict[str, Any] | None:
del current_tenant_id
try:
return self.fetch_pipeline_template_detail_from_dify_official(template_id)
except Exception as e:

99
api/services/rag_pipeline/rag_pipeline.py
View File

@ -82,6 +82,7 @@ from services.entities.knowledge_entities.rag_pipeline_entities import (
PipelineTemplateInfoEntity,
)
from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.pipeline_template.pipeline_template_factory import PipelineTemplateRetrievalFactory
from services.tools.builtin_tools_manage_service import BuiltinToolManageService
from services.workflow_draft_variable_service import DraftVariableSaver, DraftVarLoader
@ -145,7 +146,12 @@ class RagPipelineService:
@classmethod
def get_pipeline_template_detail(
cls, template_id: str, type: str = "built-in", *, session: Session
cls,
template_id: str,
current_tenant_id: str,
type: str = "built-in",
*,
session: Session,
) -> dict[str, Any] | None:
"""
Get pipeline template detail.
@ -158,7 +164,7 @@ class RagPipelineService:
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, session=session
template_id, current_tenant_id, session=session
)
if built_in_result is None:
logger.warning(
@ -171,10 +177,22 @@ class RagPipelineService:
mode = "customized"
retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
customized_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(
template_id, session=session
template_id, current_tenant_id, session=session
)
return customized_result
@staticmethod
def get_customized_pipeline_template_yaml(template_id: str, current_tenant_id: str, *, session: Session) -> str:
yaml_content = session.scalar(
select(PipelineCustomizedTemplate.yaml_content).where(
PipelineCustomizedTemplate.id == template_id,
PipelineCustomizedTemplate.tenant_id == current_tenant_id,
)
)
if yaml_content is None:
raise RagPipelineResourceNotFoundError("Customized pipeline template not found.")
return yaml_content
@classmethod
def update_customized_pipeline_template(
cls,
@ -1244,45 +1262,48 @@ class RagPipelineService:
)
return assemble_workflow_node_execution_traces(node_executions, self._node_execution_service_repo)
@classmethod
@staticmethod
def publish_customized_pipeline_template(
cls,
pipeline_id: str,
pipeline: Pipeline,
dataset: Dataset,
args: dict[str, Any],
current_user: Account | None = None,
current_tenant_id: str | None = None,
current_user: Account,
*,
session: Session,
):
"""
Publish customized pipeline template
"""
current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
pipeline = session.get(Pipeline, pipeline_id)
if not pipeline:
raise ValueError("Pipeline not found")
) -> None:
"""Publish a customized template from a caller-validated pipeline and dataset."""
if not pipeline.workflow_id:
raise ValueError("Pipeline workflow not found")
workflow = session.get(Workflow, pipeline.workflow_id)
raise RagPipelineResourceNotFoundError("Pipeline workflow not found")
workflow = session.scalar(
select(Workflow).where(
Workflow.id == pipeline.workflow_id,
Workflow.tenant_id == pipeline.tenant_id,
Workflow.app_id == pipeline.id,
)
)
if not workflow:
raise ValueError("Workflow not found")
dataset = pipeline.retrieve_dataset(session=session)
if not dataset:
raise ValueError("Dataset not found")
raise RagPipelineResourceNotFoundError("Workflow not found")
draft_workflow_id = session.scalar(
select(Workflow.id).where(
Workflow.tenant_id == pipeline.tenant_id,
Workflow.app_id == pipeline.id,
Workflow.version == Workflow.VERSION_DRAFT,
)
)
if not draft_workflow_id:
raise RagPipelineResourceNotFoundError("Draft workflow not found")
# 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)
template = session.scalar(
select(PipelineCustomizedTemplate)
.where(
PipelineCustomizedTemplate.name == args["name"],
PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id,
)
if template:
raise ValueError("Template name is already exists")
.limit(1)
)
if template:
raise ValueError("Template name is already exists")
max_position = session.scalar(
select(func.max(PipelineCustomizedTemplate.position)).where(
@ -1294,16 +1315,10 @@ class RagPipelineService:
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 {},
name=args["name"],
description=args["description"],
icon=args["icon_info"],
tenant_id=pipeline.tenant_id,
yaml_content=dsl,
install_count=0,

69
api/services/rag_pipeline/rag_pipeline_transform_service.py
View File

@ -6,7 +6,6 @@ from typing import Any
from uuid import uuid4
import yaml
from flask_login import current_user
from sqlalchemy import select
from sqlalchemy.orm import Session
@ -19,28 +18,34 @@ from core.rag.retrieval.retrieval_methods import RetrievalMethod
from factories import variable_factory
from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline
from models.enums import DatasetRuntimeMode, DataSourceType
from models.model import UploadFile
from models.workflow import Workflow, WorkflowType
from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration, RetrievalSetting
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.file_service import FileService
from services.plugin.plugin_migration import PluginMigration
logger = logging.getLogger(__name__)
class RagPipelineTransformService:
def transform_dataset(self, dataset_id: str, session: Session):
def transform_dataset(self, dataset: Dataset, account_id: str, session: Session):
"""Transform a vendor dataset within the caller-owned transaction.
Dataset and document state is read through ``session`` so uncommitted caller changes remain visible. The
transformation commits only after the pipeline and migrated document metadata have been persisted.
The caller must resolve and authorize ``dataset`` before entering this service. Plugin provisioning is an
intentionally non-transactional prerequisite; database changes are committed only after the pipeline and
migrated document metadata have been persisted.
"""
dataset = session.get(Dataset, dataset_id)
if not dataset:
raise ValueError("Dataset not found")
if dataset.pipeline_id and dataset.runtime_mode == DatasetRuntimeMode.RAG_PIPELINE:
pipeline = session.scalar(
select(Pipeline)
.where(Pipeline.id == dataset.pipeline_id, Pipeline.tenant_id == dataset.tenant_id)
.limit(1)
)
if pipeline is None:
raise RagPipelineResourceNotFoundError("Pipeline not found")
return {
"pipeline_id": dataset.pipeline_id,
"dataset_id": dataset_id,
"pipeline_id": pipeline.id,
"dataset_id": dataset.id,
"status": "success",
}
if dataset.provider != "vendor":
@ -49,11 +54,11 @@ class RagPipelineTransformService:
indexing_technique = dataset.indexing_technique
if not datasource_type and not indexing_technique:
return self._transform_to_empty_pipeline(dataset, session=session)
return self._transform_to_empty_pipeline(dataset, account_id=account_id, session=session)
doc_form = dataset.get_doc_form(session=session)
if not doc_form:
return self._transform_to_empty_pipeline(dataset, session=session)
return self._transform_to_empty_pipeline(dataset, account_id=account_id, 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
@ -74,8 +79,6 @@ class RagPipelineTransformService:
node = self._deal_file_extensions(node)
if node.get("data", {}).get("type") == "knowledge-index":
knowledge_configuration = KnowledgeConfiguration.model_validate(node.get("data", {}))
if dataset.tenant_id != current_user.current_tenant_id:
raise ValueError("Unauthorized")
node = self._deal_knowledge_index(
knowledge_configuration, dataset, indexing_technique, retrieval_model, node
)
@ -85,7 +88,12 @@ class RagPipelineTransformService:
workflow_data["graph"] = graph
pipeline_yaml["workflow"] = workflow_data
# create pipeline
pipeline = self._create_pipeline(pipeline_yaml, session=session)
pipeline = self._create_pipeline(
pipeline_yaml,
tenant_id=dataset.tenant_id,
account_id=account_id,
session=session,
)
# save chunk structure to dataset
if doc_form == IndexStructureType.PARENT_CHILD_INDEX:
@ -104,7 +112,7 @@ class RagPipelineTransformService:
session.commit()
return {
"pipeline_id": pipeline.id,
"dataset_id": dataset_id,
"dataset_id": dataset.id,
"status": "success",
}
@ -200,6 +208,8 @@ class RagPipelineTransformService:
self,
data: dict[str, Any],
*,
tenant_id: str,
account_id: str,
session: Session,
) -> Pipeline:
"""Create a new app or update an existing one."""
@ -223,11 +233,11 @@ class RagPipelineTransformService:
# Create new app
pipeline = Pipeline(
tenant_id=current_user.current_tenant_id,
tenant_id=tenant_id,
name=pipeline_data.get("name", ""),
description=pipeline_data.get("description", ""),
created_by=current_user.id,
updated_by=current_user.id,
created_by=account_id,
updated_by=account_id,
is_published=True,
is_public=True,
)
@ -243,7 +253,7 @@ class RagPipelineTransformService:
type=WorkflowType.RAG_PIPELINE,
version="draft",
graph=json.dumps(graph),
created_by=current_user.id,
created_by=account_id,
environment_variables=environment_variables,
conversation_variables=conversation_variables,
rag_pipeline_variables=rag_pipeline_variables_list,
@ -255,7 +265,7 @@ class RagPipelineTransformService:
type=WorkflowType.RAG_PIPELINE,
version=str(datetime.now(UTC).replace(tzinfo=None)),
graph=json.dumps(graph),
created_by=current_user.id,
created_by=account_id,
environment_variables=environment_variables,
conversation_variables=conversation_variables,
rag_pipeline_variables=rag_pipeline_variables_list,
@ -297,19 +307,19 @@ class RagPipelineTransformService:
logger.debug("Installing missing pipeline plugins %s", package_identifiers_to_install)
PluginService.install_from_marketplace_pkg(tenant_id, package_identifiers_to_install)
def _transform_to_empty_pipeline(self, dataset: Dataset, *, session: Session):
def _transform_to_empty_pipeline(self, dataset: Dataset, *, account_id: str, session: Session):
pipeline = Pipeline(
tenant_id=dataset.tenant_id,
name=dataset.name,
description=dataset.description,
created_by=current_user.id,
created_by=account_id,
)
session.add(pipeline)
session.flush()
dataset.pipeline_id = pipeline.id
dataset.runtime_mode = DatasetRuntimeMode.RAG_PIPELINE
dataset.updated_by = current_user.id
dataset.updated_by = account_id
dataset.updated_at = datetime.now(UTC).replace(tzinfo=None)
session.add(dataset)
session.commit()
@ -325,18 +335,23 @@ class RagPipelineTransformService:
jina_node_id = "1752491761974"
firecrawl_node_id = "1752565402678"
documents = session.scalars(select(Document).where(Document.dataset_id == dataset.id)).all()
documents = session.scalars(
select(Document).where(Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id)
).all()
for document in documents:
data_source_info_dict = document.data_source_info_dict
if not data_source_info_dict:
continue
if document.data_source_type == DataSourceType.UPLOAD_FILE:
document.data_source_type = DataSourceType.LOCAL_FILE
file_id = data_source_info_dict.get("upload_file_id")
if file_id:
file = session.get(UploadFile, file_id)
file_id = str(file_id)
file = FileService.get_upload_files_by_ids(dataset.tenant_id, [file_id], session=session).get(
file_id
)
if file:
document.data_source_type = DataSourceType.LOCAL_FILE
data_source_info = json.dumps(
{
"real_file_id": file_id,

2
api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py
View File

@ -36,7 +36,7 @@ def test_export_customized_pipeline_template_from_database(
db_session_with_containers.expire_all()
with flask_app_with_containers.test_request_context("/"):
response, status = method(api, template.id)
response, status = method(api, db_session_with_containers, template.tenant_id, template.id)
assert status == 200
assert response == {"data": "yaml-data"}

291
api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py
View File

@ -1,15 +1,14 @@
from __future__ import annotations
from collections.abc import Iterator
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import PropertyMock, patch
from inspect import getclosurevars, unwrap
from unittest.mock import ANY, PropertyMock, patch
import pytest
from flask import Flask
from sqlalchemy import Engine
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from werkzeug.exceptions import Forbidden, NotFound
from controllers.console import console_ns
from controllers.console.datasets.rag_pipeline import rag_pipeline as module
@ -22,10 +21,12 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline import (
PipelineTemplateListQuery,
PublishCustomizedPipelineTemplateApi,
)
from models.account import Account
from models.dataset import PipelineCustomizedTemplate
from models.account import Account, TenantAccountRole
from models.dataset import Pipeline, PipelineCustomizedTemplate
from models.engine import db
from services.entities.knowledge_entities.rag_pipeline_entities import PipelineTemplateInfoEntity
from services.errors.account import NoPermissionError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
def _template_item() -> dict[str, object]:
@ -65,6 +66,12 @@ def _account() -> Account:
return account
def _pipeline() -> Pipeline:
pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline")
pipeline.id = "pipeline-1"
return pipeline
@pytest.fixture
def database_app() -> Iterator[Flask]:
app = Flask(__name__)
@ -136,11 +143,13 @@ class TestPipelineTemplateDetailApi:
def test_get_serializes_template_detail(self, app: Flask, sqlite_engine: Engine) -> None:
api = PipelineTemplateDetailApi()
method = unwrap(api.get)
service_calls: list[tuple[str, str]] = []
service_calls: list[tuple[str, str, str]] = []
def get_pipeline_template_detail(template_id: str, type: str, *, session) -> dict[str, object]:
def get_pipeline_template_detail(
template_id: str, current_tenant_id: str, type: str, *, session
) -> dict[str, object]:
del session
service_calls.append((template_id, type))
service_calls.append((template_id, current_tenant_id, type))
return _template_detail()
with (
@ -152,18 +161,20 @@ class TestPipelineTemplateDetailApi:
side_effect=get_pipeline_template_detail,
),
):
response, status = method(api, PipelineTemplateDetailQuery(type="customized"), session, "template-1")
response, status = method(
api, PipelineTemplateDetailQuery(type="customized"), session, "tenant-1", "template-1"
)
assert status == 200
assert response == {**_template_detail(), "created_by": None}
assert service_calls == [("template-1", "customized")]
assert service_calls == [("template-1", "tenant-1", "customized")]
def test_get_raises_not_found_without_custom_response_body(self, app: Flask, sqlite_engine: Engine) -> None:
api = PipelineTemplateDetailApi()
method = unwrap(api.get)
def get_pipeline_template_detail(template_id: str, type: str, *, session) -> None:
del template_id, type, session
def get_pipeline_template_detail(template_id: str, current_tenant_id: str, type: str, *, session) -> None:
del template_id, current_tenant_id, type, session
with (
Session(sqlite_engine) as session,
@ -175,7 +186,7 @@ class TestPipelineTemplateDetailApi:
),
pytest.raises(NotFound),
):
method(api, PipelineTemplateDetailQuery(), session, "missing")
method(api, PipelineTemplateDetailQuery(), session, "tenant-1", "missing")
class TestCustomizedPipelineTemplateApi:
@ -286,9 +297,7 @@ class TestCustomizedPipelineTemplateApi:
assert deleted_templates == [("template-1", tenant_id)]
@pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True)
def test_post_exports_yaml_from_orm_template(
self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session
) -> None:
def test_post_exports_yaml_from_orm_template(self, app: Flask, sqlite_session: Session) -> None:
api = CustomizedPipelineTemplateApi()
method = unwrap(api.post)
template = PipelineCustomizedTemplate(
@ -306,115 +315,201 @@ class TestCustomizedPipelineTemplateApi:
template.id = "template-1"
sqlite_session.add(template)
sqlite_session.commit()
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine))
with app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"):
response, status = method(api, "template-1")
response, status = method(
api,
sqlite_session,
"00000000-0000-0000-0000-000000000001",
"template-1",
)
assert status == 200
assert response == {"data": "dsl: value"}
@pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True)
def test_post_raises_when_template_is_missing(
self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session
) -> None:
def test_post_returns_not_found_for_other_tenant(self, app: Flask, sqlite_session: Session) -> None:
api = CustomizedPipelineTemplateApi()
method = unwrap(api.post)
assert sqlite_session.get(PipelineCustomizedTemplate, "missing") is None
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine))
template = PipelineCustomizedTemplate(
tenant_id="00000000-0000-0000-0000-000000000002",
name="Other tenant template",
description="Description",
chunk_structure="general",
icon={},
position=1,
yaml_content="secret: value",
install_count=0,
language="en-US",
created_by="00000000-0000-0000-0000-000000000003",
)
template.id = "template-1"
sqlite_session.add(template)
sqlite_session.commit()
with app.test_request_context("/rag/pipeline/customized/templates/missing", method="POST"):
with pytest.raises(ValueError, match="Customized pipeline template not found"):
method(api, "missing")
with (
app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"),
pytest.raises(NotFound, match="Customized pipeline template not found"),
):
method(
api,
sqlite_session,
"00000000-0000-0000-0000-000000000001",
"template-1",
)
class TestPublishCustomizedPipelineTemplateApi:
def test_post_validates_payload_and_returns_empty_204(self, app: Flask) -> None:
def test_post_uses_pipeline_release_rbac_scene(self) -> None:
method = PublishCustomizedPipelineTemplateApi.post
while "scene" not in getclosurevars(method).nonlocals:
method = method.__wrapped__
assert getclosurevars(method).nonlocals["scene"] == module.RBACPermission.DATASET_PIPELINE_RELEASE
def test_post_validates_payload_and_returns_empty_204(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
payload = _payload()
account = _account()
tenant_id = "tenant-1"
service_calls: list[tuple[str, dict[str, object], Account, str]] = []
class Service:
def __init__(self, *args, **kwargs) -> None:
pass
def publish_customized_pipeline_template(
self,
pipeline_id: str,
data: dict[str, object],
current_user: Account,
current_tenant_id: str,
*,
session,
) -> None:
del session
service_calls.append((pipeline_id, data, current_user, current_tenant_id))
pipeline = _pipeline()
dataset = object()
with (
app.test_request_context("/rag/pipelines/pipeline-1/customized/publish", method="POST", json=payload),
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch.object(module, "RagPipelineService", Service),
patch.object(module.dify_config, "RBAC_ENABLED", True),
patch.object(Pipeline, "retrieve_dataset", return_value=dataset),
patch.object(module.DatasetService, "check_dataset_permission") as legacy_acl,
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
):
response, status = method(
api, CustomizedPipelineTemplatePayload.model_validate(payload), tenant_id, account, "pipeline-1"
)
response, status = method(api, CustomizedPipelineTemplatePayload.model_validate(payload), account, pipeline)
assert (response, status) == ("", 204)
assert service_calls == [("pipeline-1", payload, account, tenant_id)]
publish.assert_called_once_with(pipeline, dataset, payload, account, session=ANY)
legacy_acl.assert_not_called()
def test_post_allows_missing_icon_info_for_publish_service_fallback(self, app: Flask) -> None:
@pytest.mark.parametrize(
("payload", "expected_icon_info"),
[
(
{"name": "Published template", "description": "Description"},
{"icon": "", "icon_background": None, "icon_type": None, "icon_url": None},
),
({"name": "Published template", "description": "Description", "icon_info": {}}, {}),
],
)
def test_post_preserves_valid_icon_info(
self,
payload: dict[str, object],
expected_icon_info: dict[str, object | None],
) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
payload: dict[str, object] = {
"name": "Published template",
"description": "Description",
}
account = _account()
tenant_id = "tenant-1"
service_calls: list[tuple[str, dict[str, object], Account, str]] = []
class Service:
def __init__(self, *args, **kwargs) -> None:
pass
def publish_customized_pipeline_template(
self,
pipeline_id: str,
data: dict[str, object],
current_user: Account,
current_tenant_id: str,
*,
session,
) -> None:
del session
service_calls.append((pipeline_id, data, current_user, current_tenant_id))
pipeline = _pipeline()
dataset = object()
with (
app.test_request_context("/rag/pipelines/pipeline-1/customized/publish", method="POST", json=payload),
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch.object(module, "RagPipelineService", Service),
patch.object(module.dify_config, "RBAC_ENABLED", True),
patch.object(Pipeline, "retrieve_dataset", return_value=dataset),
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
):
response, status = method(
api, CustomizedPipelineTemplatePayload.model_validate(payload), tenant_id, account, "pipeline-1"
)
response, status = method(api, CustomizedPipelineTemplatePayload.model_validate(payload), account, pipeline)
assert (response, status) == ("", 204)
assert service_calls == [
(
"pipeline-1",
{
**payload,
"icon_info": {
"icon": "",
"icon_background": None,
"icon_type": None,
"icon_url": None,
},
},
account,
tenant_id,
)
]
publish.assert_called_once_with(pipeline, dataset, ANY, account, session=ANY)
assert publish.call_args.args[2]["icon_info"] == expected_icon_info
def test_post_translates_missing_owned_resource_to_not_found(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
payload = _payload()
with (
patch.object(module.dify_config, "RBAC_ENABLED", True),
patch.object(Pipeline, "retrieve_dataset", return_value=object()),
patch.object(
module.RagPipelineService,
"publish_customized_pipeline_template",
side_effect=RagPipelineResourceNotFoundError("Workflow not found"),
),
pytest.raises(NotFound, match="Workflow not found"),
):
method(api, CustomizedPipelineTemplatePayload.model_validate(payload), _account(), _pipeline())
def test_post_allows_legacy_dataset_operator_after_dataset_acl(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
account = _account()
account.role = TenantAccountRole.DATASET_OPERATOR
pipeline = _pipeline()
dataset = object()
payload = _payload()
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(Pipeline, "retrieve_dataset", return_value=dataset),
patch.object(module.DatasetService, "check_dataset_permission") as check_permission,
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
):
response = method(api, CustomizedPipelineTemplatePayload.model_validate(payload), account, pipeline)
assert response == ("", 204)
assert check_permission.call_args.args[:2] == (dataset, account)
publish.assert_called_once_with(pipeline, dataset, payload, account, session=ANY)
def test_post_rejects_legacy_non_editor_before_dataset_acl(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
account = _account()
account.role = TenantAccountRole.NORMAL
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(Pipeline, "retrieve_dataset", return_value=object()),
patch.object(module.DatasetService, "check_dataset_permission") as check_permission,
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
pytest.raises(Forbidden),
):
method(api, CustomizedPipelineTemplatePayload.model_validate(_payload()), account, _pipeline())
check_permission.assert_not_called()
publish.assert_not_called()
def test_post_rejects_legacy_dataset_acl_before_publish(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
account = _account()
account.role = TenantAccountRole.EDITOR
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(Pipeline, "retrieve_dataset", return_value=object()),
patch.object(
module.DatasetService,
"check_dataset_permission",
side_effect=NoPermissionError("Dataset is private"),
),
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
pytest.raises(Forbidden, match="Dataset is private"),
):
method(api, CustomizedPipelineTemplatePayload.model_validate(_payload()), account, _pipeline())
publish.assert_not_called()
def test_post_rejects_missing_legacy_dataset_before_publish(self) -> None:
api = PublishCustomizedPipelineTemplateApi()
method = unwrap(api.post)
account = _account()
account.role = TenantAccountRole.EDITOR
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(Pipeline, "retrieve_dataset", return_value=None),
patch.object(module.DatasetService, "check_dataset_permission") as check_permission,
patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish,
pytest.raises(NotFound, match="Dataset not found"),
):
method(api, CustomizedPipelineTemplatePayload.model_validate(_payload()), account, _pipeline())
check_permission.assert_not_called()
publish.assert_not_called()

139
api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py
View File

@ -13,7 +13,7 @@ import pytest
from flask import Flask
from sqlalchemy import Engine
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
from werkzeug.exceptions import Forbidden, NotFound
from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as module
from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import (
@ -25,18 +25,21 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import (
WorkflowUpdatePayload,
)
from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
from models.account import Account, TenantAccountRole
from models.dataset import Pipeline
from models.account import Account, Tenant, TenantAccountRole
from models.dataset import Dataset, Pipeline
from models.engine import db
from models.enums import PermissionEnum
from models.tools import WorkflowToolProvider
from models.workflow import Workflow, WorkflowType
from services.errors.llm import InvokeRateLimitError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.rag_pipeline import RagPipelineService
DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001"
DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002"
DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003"
DEFAULT_WORKFLOW_ID = "00000000-0000-0000-0000-000000000004"
DEFAULT_DATASET_ID = "44444444-4444-4444-4444-444444444444"
def _make_workflow(**overrides: object) -> Workflow:
@ -67,6 +70,9 @@ def _account() -> Account:
account = Account(name="Alice", email="alice@example.com")
account.id = DEFAULT_WORKFLOW_CREATED_BY
account.role = TenantAccountRole.EDITOR
tenant = Tenant(name="Tenant")
tenant.id = DEFAULT_WORKFLOW_TENANT_ID
account._current_tenant = tenant
return account
@ -76,6 +82,18 @@ def _pipeline() -> Pipeline:
return pipeline
def _dataset(*, tenant_id: str = DEFAULT_WORKFLOW_TENANT_ID, maintainer: str = DEFAULT_WORKFLOW_CREATED_BY) -> Dataset:
return Dataset(
id=DEFAULT_DATASET_ID,
tenant_id=tenant_id,
name="Dataset",
created_by=maintainer,
maintainer=maintainer,
permission=PermissionEnum.ONLY_ME,
provider="vendor",
)
def _persist_workflow(workflow: Workflow) -> None:
db.session.add(workflow)
db.session.commit()
@ -231,18 +249,119 @@ def test_rag_pipeline_recommended_plugins_serializes_known_envelope(database_app
assert response == recommended_plugins
def test_rag_pipeline_transform_rejects_read_only_member(app: Flask, sqlite_engine: Engine) -> None:
def test_rag_pipeline_transform_rejects_read_only_member(sqlite_engine: Engine) -> None:
account = _account()
account.role = TenantAccountRole.NORMAL
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
with (
Session(sqlite_engine) as session,
app.test_request_context("/"),
pytest.raises(Forbidden),
):
handler(api, session, account, UUID("44444444-4444-4444-4444-444444444444"))
with Session(sqlite_engine) as session:
session.add(_dataset())
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
pytest.raises(Forbidden),
):
handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID))
def test_rag_pipeline_transform_rejects_dataset_from_another_tenant_before_service_call(
sqlite_engine: Engine,
) -> None:
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
with Session(sqlite_engine) as session:
session.add(_dataset(tenant_id="00000000-0000-0000-0000-000000000099"))
with (
patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset,
pytest.raises(NotFound),
):
handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID))
transform_dataset.assert_not_called()
def test_rag_pipeline_transform_enforces_legacy_dataset_permission_before_service_call(
sqlite_engine: Engine,
) -> None:
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
with Session(sqlite_engine) as session:
session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099"))
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset,
pytest.raises(Forbidden),
):
handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID))
transform_dataset.assert_not_called()
def test_rag_pipeline_transform_passes_authorized_dataset_and_account_to_service(
sqlite_engine: Engine,
) -> None:
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
account = _account()
expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"}
with Session(sqlite_engine) as session:
dataset = _dataset()
session.add(dataset)
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(module.RagPipelineTransformService, "transform_dataset", return_value=expected) as transform,
):
response = handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID))
transform.assert_called_once_with(dataset, account.id, session)
assert response == expected
def test_rag_pipeline_transform_maps_missing_pipeline_to_not_found(sqlite_engine: Engine) -> None:
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
with Session(sqlite_engine) as session:
session.add(_dataset())
with (
patch.object(module.dify_config, "RBAC_ENABLED", False),
patch.object(
module.RagPipelineTransformService,
"transform_dataset",
side_effect=RagPipelineResourceNotFoundError("Pipeline not found"),
),
pytest.raises(NotFound, match="Pipeline not found"),
):
handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID))
def test_rag_pipeline_transform_skips_legacy_acl_when_rbac_is_enabled(sqlite_engine: Engine) -> None:
api = module.RagPipelineTransformApi()
handler = unwrap_all(api.post)
account = _account()
account.role = TenantAccountRole.NORMAL
expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"}
with Session(sqlite_engine) as session:
session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099"))
with (
patch.object(module.dify_config, "RBAC_ENABLED", True),
patch.object(module.RagPipelineTransformService, "transform_dataset", return_value=expected) as transform,
):
response = handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID))
assert response == expected
transform.assert_called_once()
@pytest.mark.parametrize(

32
api/tests/unit_tests/controllers/console/test_wraps.py
View File

@ -40,7 +40,7 @@ from libs.login import AccountWithTenant
from machinery.context import RequestContext
from models import Account
from models.account import AccountStatus, TenantAccountRole
from models.dataset import RateLimitLog
from models.dataset import Dataset, RateLimitLog
from services.entities.feature_entities import LicenseStatus
@ -396,6 +396,36 @@ class TestRbacPermissionRequired:
request.view_args = {"resource_id": "dataset-1"}
assert _extract_resource_id(RBACResourceScope.DATASET, "tenant-1") == "dataset-1"
def test_extract_resource_id_scopes_pipeline_resolution_to_the_calling_tenant(self, sqlite_session: Session):
app = Flask(__name__)
pipeline_id = "00000000-0000-0000-0000-000000000001"
current_tenant_id = "00000000-0000-0000-0000-000000000002"
foreign_dataset = Dataset(
id="00000000-0000-0000-0000-000000000003",
tenant_id="00000000-0000-0000-0000-000000000004",
name="Foreign decoy",
created_by="00000000-0000-0000-0000-000000000005",
pipeline_id=pipeline_id,
)
current_dataset = Dataset(
id="00000000-0000-0000-0000-000000000006",
tenant_id=current_tenant_id,
name="Current tenant dataset",
created_by="00000000-0000-0000-0000-000000000007",
pipeline_id=pipeline_id,
)
sqlite_session.add_all([foreign_dataset, current_dataset])
unscoped_dataset = sqlite_session.scalar(select(Dataset).where(Dataset.pipeline_id == pipeline_id))
assert unscoped_dataset is foreign_dataset
with (
app.test_request_context("/rag/pipelines/pipeline-1"),
patch("controllers.common.wraps.db", SimpleNamespace(session=sqlite_session)),
):
request.view_args = {"pipeline_id": pipeline_id}
assert _extract_resource_id(RBACResourceScope.DATASET, current_tenant_id) == current_dataset.id
def test_extract_resource_id_resolves_agent_to_its_authz_app(self):
app = Flask(__name__)

91
api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py
View File

@ -8,6 +8,7 @@ from uuid import uuid4
import pytest
from flask import Flask, current_app
from sqlalchemy.orm import Session
from core.app.app_config.entities import (
DatasetEntity,
@ -4961,6 +4962,7 @@ class TestSingleAndMultipleRetrieveCoverage:
assert len(result) == 1
assert result[0].provider == "external"
session.scalar.assert_called_once()
mock_end.assert_called_once()
assert retrieval.llm_usage.total_tokens == 2
@ -5037,6 +5039,95 @@ class TestSingleAndMultipleRetrieveCoverage:
)
assert results == []
def test_single_retrieve_rejects_dataset_outside_available_datasets(self, retrieval: DatasetRetrieval) -> None:
available_dataset = _dataset(id="ds-1", name="Available DS", description=None)
session = MagicMock()
session.scalar.return_value = _dataset(
id="ds-2",
name="Foreign DS",
provider="external",
tenant_id="tenant-2",
retrieval_model={},
)
with (
patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls,
patch(
"core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval",
return_value=[],
) as mock_external_retrieve,
patch.object(retrieval, "_on_query") as mock_on_query,
):
mock_router_cls.return_value.invoke.return_value = ("ds-2", LLMUsage.empty_usage())
results = retrieval.single_retrieve(
session,
app_id="app-1",
tenant_id="tenant-1",
user_id="user-1",
user_from="workflow",
query="python",
available_datasets=[available_dataset],
model_instance=Mock(),
model_config=Mock(),
planning_strategy=PlanningStrategy.REACT_ROUTER,
)
assert results == []
session.scalar.assert_not_called()
mock_external_retrieve.assert_not_called()
mock_on_query.assert_not_called()
def test_single_retrieve_rejects_allowlisted_dataset_owned_by_another_tenant(
self, retrieval: DatasetRetrieval, sqlite_session: Session
) -> None:
dataset_id = str(uuid4())
caller_tenant_id = str(uuid4())
foreign_dataset = Dataset(
id=dataset_id,
tenant_id=str(uuid4()),
name="Foreign DS",
provider="external",
indexing_technique="high_quality",
retrieval_model={},
created_by=str(uuid4()),
)
sqlite_session.add(foreign_dataset)
available_dataset = _dataset(
id=dataset_id,
tenant_id=caller_tenant_id,
name="Available DS",
description=None,
)
with (
patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls,
patch(
"core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval",
) as mock_external_retrieve,
patch(
"core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve",
) as mock_internal_retrieve,
patch.object(retrieval, "_on_query") as mock_on_query,
):
mock_router_cls.return_value.invoke.return_value = (dataset_id, LLMUsage.empty_usage())
results = retrieval.single_retrieve(
sqlite_session,
app_id="app-1",
tenant_id=caller_tenant_id,
user_id="user-1",
user_from="workflow",
query="python",
available_datasets=[available_dataset],
model_instance=Mock(),
model_config=Mock(),
planning_strategy=PlanningStrategy.REACT_ROUTER,
)
assert results == []
mock_internal_retrieve.assert_not_called()
mock_external_retrieve.assert_not_called()
mock_on_query.assert_not_called()
def test_single_retrieve_respects_metadata_filter_shortcuts(self, retrieval: DatasetRetrieval) -> None:
dataset = _dataset(
id="ds-1",

27
api/tests/unit_tests/models/test_dataset_models.py
View File

@ -18,6 +18,7 @@ from urllib.parse import parse_qs, urlparse
from uuid import uuid4
import pytest
from sqlalchemy.orm import Session
from core.rag.entities import ParentMode
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
@ -112,6 +113,32 @@ class TestDatasetModelValidation:
session.get.assert_called_once_with(Account, dataset.created_by)
assert session.scalar.call_count == 2
def test_get_doc_form_ignores_foreign_tenant_document(self, sqlite_session: Session) -> None:
dataset_id = str(uuid4())
tenant_id = str(uuid4())
created_by = str(uuid4())
dataset = Dataset(
id=dataset_id,
tenant_id=tenant_id,
name="Dataset",
data_source_type=DataSourceType.UPLOAD_FILE,
created_by=created_by,
)
foreign_document = Document(
tenant_id=str(uuid4()),
dataset_id=dataset_id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="foreign",
name="Foreign",
created_from=DocumentCreatedFrom.WEB,
created_by=created_by,
doc_form=IndexStructureType.PARENT_CHILD_INDEX,
)
sqlite_session.add_all([dataset, foreign_document])
assert dataset.get_doc_form(session=sqlite_session) is None
def test_get_dataset_keyword_table_uses_caller_session(self):
dataset = Dataset(
tenant_id=str(uuid4()),

4
api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py
View File

@ -45,7 +45,7 @@ def test_get_pipeline_template_detail(mocker: MockerFixture, sqlite_session: Ses
)
retrieval = BuiltInPipelineTemplateRetrieval()
detail = retrieval.get_pipeline_template_detail("tpl-1", session=sqlite_session)
detail = retrieval.get_pipeline_template_detail("tpl-1", "tenant-1", session=sqlite_session)
assert detail == {"id": "tpl-1", "name": "Template 1"}
assert not sqlite_session.in_transaction()
@ -79,7 +79,7 @@ def test_get_pipeline_template_detail_returns_none_for_unknown_id(
)
retrieval = BuiltInPipelineTemplateRetrieval()
result = retrieval.get_pipeline_template_detail("nonexistent-id", session=sqlite_session)
result = retrieval.get_pipeline_template_detail("nonexistent-id", "tenant-1", session=sqlite_session)
assert result is None
assert not sqlite_session.in_transaction()

8
api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py
View File

@ -81,7 +81,7 @@ def test_get_pipeline_template_detail_returns_detail(monkeypatch: pytest.MonkeyP
monkeypatch.setattr("models.dataset.db", SimpleNamespace(session=sqlite_session))
retrieval = CustomizedPipelineTemplateRetrieval()
detail = retrieval.get_pipeline_template_detail(TEMPLATE_ID, session=sqlite_session)
detail = retrieval.get_pipeline_template_detail(TEMPLATE_ID, TENANT_ID, session=sqlite_session)
assert detail == {
"id": TEMPLATE_ID,
@ -97,10 +97,12 @@ def test_get_pipeline_template_detail_returns_detail(monkeypatch: pytest.MonkeyP
@pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True)
def test_get_pipeline_template_detail_returns_none_when_not_found(sqlite_session: Session) -> None:
def test_get_pipeline_template_detail_rejects_other_tenant(sqlite_session: Session) -> None:
sqlite_session.add(_template(tenant_id=OTHER_TENANT_ID))
sqlite_session.commit()
retrieval = CustomizedPipelineTemplateRetrieval()
result = retrieval.get_pipeline_template_detail(TEMPLATE_ID, session=sqlite_session)
result = retrieval.get_pipeline_template_detail(TEMPLATE_ID, TENANT_ID, session=sqlite_session)
assert result is None
assert sqlite_session.in_transaction()

4
api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py
View File

@ -68,7 +68,7 @@ def test_get_pipeline_template_detail_returns_detail(sqlite_session: Session) ->
sqlite_session.commit()
retrieval = DatabasePipelineTemplateRetrieval()
detail = retrieval.get_pipeline_template_detail(TEMPLATE_ID, session=sqlite_session)
detail = retrieval.get_pipeline_template_detail(TEMPLATE_ID, "tenant-1", session=sqlite_session)
assert detail == {
"id": TEMPLATE_ID,
@ -86,7 +86,7 @@ def test_get_pipeline_template_detail_returns_detail(sqlite_session: Session) ->
def test_get_pipeline_template_detail_returns_none_when_not_found(sqlite_session: Session) -> None:
retrieval = DatabasePipelineTemplateRetrieval()
result = retrieval.get_pipeline_template_detail(TEMPLATE_ID, session=sqlite_session)
result = retrieval.get_pipeline_template_detail(TEMPLATE_ID, "tenant-1", session=sqlite_session)
assert result is None
assert sqlite_session.in_transaction()

6
api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py
View File

@ -9,8 +9,8 @@ class DummyRetrieval(PipelineTemplateRetrievalBase):
del session
return {"language": language}
def get_pipeline_template_detail(self, template_id: str, *, session) -> dict | None:
del session
def get_pipeline_template_detail(self, template_id: str, current_tenant_id: str, *, session) -> dict | None:
del current_tenant_id, session
return {"id": template_id}
def get_type(self) -> str:
@ -22,6 +22,6 @@ def test_pipeline_template_retrieval_base_concrete_implementation(sqlite_session
retrieval = DummyRetrieval()
assert retrieval.get_pipeline_templates("en-US", session=sqlite_session) == {"language": "en-US"}
assert retrieval.get_pipeline_template_detail("tpl-1", session=sqlite_session) == {"id": "tpl-1"}
assert retrieval.get_pipeline_template_detail("tpl-1", "tenant-1", session=sqlite_session) == {"id": "tpl-1"}
assert retrieval.get_type() == "dummy"
assert not sqlite_session.in_transaction()

2
api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py
View File

@ -46,7 +46,7 @@ def test_get_pipeline_template_detail_fallbacks_to_database_on_error(
)
retrieval = RemotePipelineTemplateRetrieval()
result = retrieval.get_pipeline_template_detail("tpl-1", session=sqlite_session)
result = retrieval.get_pipeline_template_detail("tpl-1", "tenant-1", session=sqlite_session)
assert result == {"id": "db-1"}
fetch_mock.assert_called_once_with("tpl-1")

101
api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py
View File

@ -30,6 +30,7 @@ from models.dataset import (
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
from models.workflow import Workflow
from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.rag_pipeline import RagPipelineService
from services.workflow_ref_service import WorkflowRef
@ -95,6 +96,10 @@ def _make_pipeline(
return pipeline
def _make_template_args(name: str = "New Template") -> dict[str, object]:
return {"name": name, "description": "Desc", "icon_info": {"icon": "star"}}
def _make_workflow(
*,
workflow_id: str = "wf-1",
@ -103,13 +108,14 @@ def _make_workflow(
graph: dict[str, object] | None = None,
features: dict[str, object] | None = None,
created_by: str = "u1",
version: str = Workflow.VERSION_DRAFT,
) -> Workflow:
workflow = Workflow(
id=workflow_id,
tenant_id=tenant_id,
app_id=app_id,
type="workflow",
version="draft",
version=version,
marked_name="",
marked_comment="",
graph=json.dumps(graph or {"nodes": []}),
@ -267,12 +273,12 @@ def test_get_pipeline_template_detail_uses_expected_mode(
factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory")
factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval
result = RagPipelineService.get_pipeline_template_detail("tpl-1", type=template_type, session=session)
result = RagPipelineService.get_pipeline_template_detail("tpl-1", "tenant-1", type=template_type, session=session)
assert result == {"id": "tpl-1"}
expected_mode = "remote" if template_type == "built-in" else "customized"
factory_mock.get_pipeline_template_factory.assert_called_with(expected_mode)
retrieval.get_pipeline_template_detail.assert_called_once_with("tpl-1", session=session)
retrieval.get_pipeline_template_detail.assert_called_once_with("tpl-1", "tenant-1", session=session)
def test_get_published_workflow_returns_none_when_pipeline_has_no_workflow_id(
@ -884,8 +890,9 @@ def test_publish_customized_pipeline_template_success(
account = _make_account(account_id="user-123")
args = {"name": "New Template", "description": "Desc", "icon_info": {"icon": "star"}, "tags": ["tag1"]}
rag_pipeline_service.service.publish_customized_pipeline_template("p1", args, account, "t1", session=session)
rag_pipeline_service.service.publish_customized_pipeline_template(
pipeline, dataset, _make_template_args(), account, session=session
)
mock_dsl_service.export_rag_pipeline_dsl.assert_called_once_with(pipeline=pipeline, include_secret=True)
templates = session.query(PipelineCustomizedTemplate).all()
@ -1847,24 +1854,15 @@ def test_run_datasource_node_preview_raises_for_unsupported_provider(
)
def test_publish_customized_pipeline_template_raises_for_missing_pipeline(
rag_pipeline_service: RagPipelineServiceTestContext,
) -> None:
with pytest.raises(ValueError, match="Pipeline not found"):
rag_pipeline_service.service.publish_customized_pipeline_template(
"p1", {}, _make_account(), "t1", session=rag_pipeline_service.session
)
def test_publish_customized_pipeline_template_raises_for_missing_workflow_id(
rag_pipeline_service: RagPipelineServiceTestContext,
) -> None:
pipeline = _make_pipeline(workflow_id=None)
_persist(rag_pipeline_service.session, pipeline)
with pytest.raises(ValueError, match="Pipeline workflow not found"):
with pytest.raises(RagPipelineResourceNotFoundError, match="Pipeline workflow not found"):
rag_pipeline_service.service.publish_customized_pipeline_template(
"p1", {"name": "template-name"}, _make_account(), "t1", session=rag_pipeline_service.session
pipeline, _make_dataset(), _make_template_args(), _make_account(), session=rag_pipeline_service.session
)
@ -2153,30 +2151,69 @@ def test_run_free_workflow_node_delegates_to_handle_result(
handle.assert_called_once()
def test_publish_customized_pipeline_template_raises_when_workflow_missing(
@pytest.mark.parametrize(("workflow_tenant_id", "workflow_app_id"), [("t2", "p1"), ("t1", "p2")])
def test_publish_customized_pipeline_template_rejects_unowned_workflow_before_export(
mocker: MockerFixture,
rag_pipeline_service: RagPipelineServiceTestContext,
workflow_tenant_id: str,
workflow_app_id: str,
) -> None:
pipeline = _make_pipeline(workflow_id="wf-1")
_persist(rag_pipeline_service.session, pipeline)
with pytest.raises(ValueError, match="Workflow not found"):
rag_pipeline_service.service.publish_customized_pipeline_template(
"p1", {}, _make_account(), "t1", session=rag_pipeline_service.session
)
def test_publish_customized_pipeline_template_raises_when_dataset_missing(
rag_pipeline_service: RagPipelineServiceTestContext,
) -> None:
pipeline = _make_pipeline(workflow_id="wf-1")
workflow = _make_workflow(workflow_id="wf-1")
workflow = _make_workflow(workflow_id="wf-1", tenant_id=workflow_tenant_id, app_id=workflow_app_id)
_persist(rag_pipeline_service.session, pipeline, workflow)
dsl_service = mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.RagPipelineDslService")
with pytest.raises(ValueError, match="Dataset not found"):
with pytest.raises(RagPipelineResourceNotFoundError, match="Workflow not found"):
rag_pipeline_service.service.publish_customized_pipeline_template(
"p1", {}, _make_account(), "t1", session=rag_pipeline_service.session
pipeline,
_make_dataset(),
_make_template_args(),
_make_account(),
session=rag_pipeline_service.session,
)
dsl_service.assert_not_called()
assert rag_pipeline_service.session.query(PipelineCustomizedTemplate).count() == 0
def test_pipeline_retrieve_dataset_rejects_unowned_dataset(
rag_pipeline_service: RagPipelineServiceTestContext,
) -> None:
pipeline = _make_pipeline(workflow_id="wf-1")
other_tenant_dataset = _make_dataset(tenant_id="t2")
_persist(rag_pipeline_service.session, pipeline, other_tenant_dataset)
assert pipeline.retrieve_dataset(session=rag_pipeline_service.session) is None
@pytest.mark.parametrize(
("draft_tenant_id", "draft_app_id"),
[(None, None), ("t2", "p1"), ("t1", "p2")],
)
def test_publish_customized_pipeline_template_rejects_missing_or_unowned_draft_before_side_effects(
mocker: MockerFixture,
rag_pipeline_service: RagPipelineServiceTestContext,
draft_tenant_id: str | None,
draft_app_id: str | None,
) -> None:
session = rag_pipeline_service.session
pipeline = _make_pipeline(workflow_id="wf-published")
published_workflow = _make_workflow(workflow_id="wf-published", version="published")
dataset = _make_dataset()
resources = [pipeline, published_workflow, dataset]
if draft_tenant_id and draft_app_id:
resources.append(_make_workflow(workflow_id="wf-draft", tenant_id=draft_tenant_id, app_id=draft_app_id))
_persist(session, *resources)
dsl_service = mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.RagPipelineDslService")
with pytest.raises(RagPipelineResourceNotFoundError, match="Draft workflow not found"):
rag_pipeline_service.service.publish_customized_pipeline_template(
pipeline, dataset, _make_template_args(), _make_account(), session=session
)
dsl_service.assert_not_called()
assert session.query(PipelineCustomizedTemplate).count() == 0
def test_get_recommended_plugins_skips_manifest_when_missing(
mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext

151
api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py
View File

@ -13,6 +13,7 @@ from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipe
from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom
from models.model import UploadFile
from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.rag_pipeline.rag_pipeline_transform_service import RagPipelineTransformService
@ -46,6 +47,24 @@ def _document(**overrides: object) -> Document:
return Document(**values)
def _upload_file(*, file_id: str = "file-1", tenant_id: str = "tenant-1") -> UploadFile:
upload_file = UploadFile(
tenant_id=tenant_id,
storage_type=StorageType.LOCAL,
key="files/f.txt",
name="f.txt",
size=10,
extension="txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by="user-1",
created_at=datetime.now(UTC).replace(tzinfo=None),
used=False,
)
upload_file.id = file_id
return upload_file
@pytest.mark.parametrize(
("doc_form", "datasource_type", "indexing_technique"),
[
@ -128,20 +147,14 @@ def test_deal_dependencies_installs_missing_marketplace_plugins(mocker: MockerFi
@pytest.mark.parametrize("sqlite_session", [(Dataset, Pipeline)], indirect=True)
def test_transform_to_empty_pipeline_updates_dataset_and_commits(
mocker: MockerFixture, sqlite_session: Session
) -> None:
def test_transform_to_empty_pipeline_updates_dataset_and_commits(sqlite_session: Session) -> None:
service = RagPipelineTransformService()
mocker.patch(
"services.rag_pipeline.rag_pipeline_transform_service.current_user",
SimpleNamespace(id="user-1"),
)
dataset = _dataset()
sqlite_session.add(dataset)
sqlite_session.commit()
result = service._transform_to_empty_pipeline(dataset, session=sqlite_session)
result = service._transform_to_empty_pipeline(dataset, account_id="user-1", session=sqlite_session)
pipeline = sqlite_session.get(Pipeline, result["pipeline_id"])
assert pipeline is not None
@ -155,33 +168,54 @@ def test_transform_to_empty_pipeline_updates_dataset_and_commits(
# --- transform_dataset ---
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
def test_transform_dataset_returns_early_when_pipeline_exists(sqlite_session: Session) -> None:
service = RagPipelineTransformService()
dataset = _dataset(id="d1", pipeline_id="p1", runtime_mode="rag_pipeline")
sqlite_session.add(dataset)
sqlite_session.commit()
pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline", description="")
pipeline.id = "p1"
sqlite_session.add_all([dataset, pipeline])
result = service.transform_dataset("d1", sqlite_session)
result = service.transform_dataset(dataset, "user-1", sqlite_session)
assert result == {"pipeline_id": "p1", "dataset_id": "d1", "status": "success"}
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
def test_transform_dataset_raises_for_dataset_not_found(sqlite_session: Session) -> None:
@pytest.mark.parametrize("pipeline_tenant_id", [None, "tenant-2"])
def test_transform_dataset_rejects_missing_or_foreign_pipeline_before_side_effects(
mocker: MockerFixture,
sqlite_session: Session,
pipeline_tenant_id: str | None,
) -> None:
service = RagPipelineTransformService()
with pytest.raises(ValueError, match="Dataset not found"):
service.transform_dataset("d1", sqlite_session)
dataset = _dataset(id="d1", pipeline_id="p1", runtime_mode="rag_pipeline")
sqlite_session.add(dataset)
if pipeline_tenant_id is not None:
pipeline = Pipeline(tenant_id=pipeline_tenant_id, name="Pipeline", description="")
pipeline.id = "p1"
sqlite_session.add(pipeline)
install_plugins = mocker.patch(
"services.rag_pipeline.rag_pipeline_transform_service.PluginService.install_from_marketplace_pkg"
)
create_pipeline = mocker.patch.object(service, "_create_pipeline")
transform_empty = mocker.patch.object(service, "_transform_to_empty_pipeline")
with pytest.raises(RagPipelineResourceNotFoundError, match="Pipeline not found"):
service.transform_dataset(dataset, "user-1", sqlite_session)
install_plugins.assert_not_called()
create_pipeline.assert_not_called()
transform_empty.assert_not_called()
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
def test_transform_dataset_raises_for_external_dataset(sqlite_session: Session) -> None:
service = RagPipelineTransformService()
sqlite_session.add(_dataset(id="d1", provider="external"))
dataset = _dataset(id="d1", provider="external")
sqlite_session.add(dataset)
sqlite_session.commit()
with pytest.raises(ValueError, match="External dataset is not supported"):
service.transform_dataset("d1", sqlite_session)
service.transform_dataset(dataset, "user-1", sqlite_session)
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
@ -189,13 +223,14 @@ def test_transform_dataset_calls_empty_pipeline_when_no_datasource(
mocker: MockerFixture, sqlite_session: Session
) -> None:
service = RagPipelineTransformService()
sqlite_session.add(_dataset(id="d1", data_source_type=None, indexing_technique=None))
dataset = _dataset(id="d1", data_source_type=None, indexing_technique=None)
sqlite_session.add(dataset)
sqlite_session.commit()
empty_result = {"pipeline_id": "p-empty", "dataset_id": "d1", "status": "success"}
mocker.patch.object(service, "_transform_to_empty_pipeline", return_value=empty_result)
result = service.transform_dataset("d1", sqlite_session)
result = service.transform_dataset(dataset, "user-1", sqlite_session)
assert result == empty_result
@ -205,15 +240,14 @@ def test_transform_dataset_calls_empty_pipeline_when_no_doc_form(
mocker: MockerFixture, sqlite_session: Session
) -> None:
service = RagPipelineTransformService()
sqlite_session.add(
_dataset(id="d1", data_source_type="upload_file", indexing_technique="high_quality", chunk_structure=None)
)
dataset = _dataset(id="d1", data_source_type="upload_file", indexing_technique="high_quality", chunk_structure=None)
sqlite_session.add(dataset)
sqlite_session.commit()
empty_result = {"pipeline_id": "p-empty", "dataset_id": "d1", "status": "success"}
mocker.patch.object(service, "_transform_to_empty_pipeline", return_value=empty_result)
result = service.transform_dataset("d1", sqlite_session)
result = service.transform_dataset(dataset, "user-1", sqlite_session)
assert result == empty_result
@ -354,18 +388,19 @@ def test_transform_dataset_full_flow(mocker: MockerFixture, sqlite_session: Sess
mocker.patch.object(service, "_deal_dependencies")
mocker.patch.object(service, "_deal_document_data")
# Mock current_user to have the same tenant_id as dataset
mock_current_user = SimpleNamespace(current_tenant_id="t1")
mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.current_user", mock_current_user)
pipeline = SimpleNamespace(id="p-new")
mocker.patch.object(service, "_create_pipeline", return_value=pipeline)
create_pipeline = mocker.patch.object(service, "_create_pipeline", return_value=pipeline)
result = service.transform_dataset("d1", sqlite_session)
result = service.transform_dataset(dataset, "user-1", sqlite_session)
assert result["pipeline_id"] == "p-new"
assert dataset.runtime_mode == "rag_pipeline"
assert dataset.chunk_structure == "text_model"
assert create_pipeline.call_args.kwargs == {
"tenant_id": "t1",
"account_id": "user-1",
"session": sqlite_session,
}
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
@ -393,7 +428,7 @@ def test_transform_dataset_raises_for_unsupported_doc_form_after_pipeline_create
mocker.patch.object(service, "_create_pipeline", return_value=SimpleNamespace(id="p-new"))
with pytest.raises(ValueError, match="Unsupported doc form"):
service.transform_dataset("d1", sqlite_session)
service.transform_dataset(dataset, "user-1", sqlite_session)
@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
@ -420,7 +455,7 @@ def test_transform_dataset_raises_when_transform_yaml_missing_workflow(
mocker.patch.object(service, "_deal_dependencies")
with pytest.raises(ValueError, match="Missing workflow data for rag pipeline"):
service.transform_dataset("d1", sqlite_session)
service.transform_dataset(dataset, "user-1", sqlite_session)
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
@ -428,7 +463,12 @@ def test_create_pipeline_raises_when_workflow_data_missing(sqlite_session: Sessi
service = RagPipelineTransformService()
with pytest.raises(ValueError, match="Missing workflow data for rag pipeline"):
service._create_pipeline({"rag_pipeline": {"name": "N"}}, session=sqlite_session)
service._create_pipeline(
{"rag_pipeline": {"name": "N"}},
tenant_id="tenant-1",
account_id="user-1",
session=sqlite_session,
)
@pytest.mark.parametrize("sqlite_session", [(Document, DocumentPipelineExecutionLog, UploadFile)], indirect=True)
@ -442,21 +482,7 @@ def test_deal_document_data_upload_file_with_existing_file(sqlite_session: Sessi
data_source_info='{"upload_file_id":"file-1"}',
name="Doc",
)
upload_file = UploadFile(
tenant_id="tenant-1",
storage_type=StorageType.LOCAL,
key="files/f.txt",
name="f.txt",
size=10,
extension="txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by="user-1",
created_at=datetime.now(UTC).replace(tzinfo=None),
used=False,
)
upload_file.id = "file-1"
sqlite_session.add_all([document, upload_file])
sqlite_session.add_all([document, _upload_file()])
sqlite_session.commit()
service._deal_document_data(dataset, sqlite_session)
@ -469,6 +495,35 @@ def test_deal_document_data_upload_file_with_existing_file(sqlite_session: Sessi
assert log.document_id == document.id
@pytest.mark.parametrize(
("document_tenant_id", "upload_file_tenant_id"),
[("tenant-2", "tenant-1"), ("tenant-1", "tenant-2")],
)
@pytest.mark.parametrize("sqlite_session", [(Document, DocumentPipelineExecutionLog, UploadFile)], indirect=True)
def test_deal_document_data_scopes_documents_and_upload_files_to_dataset_tenant(
sqlite_session: Session,
document_tenant_id: str,
upload_file_tenant_id: str,
) -> None:
service = RagPipelineTransformService()
dataset = _dataset(id="d1", tenant_id="tenant-1", pipeline_id="p1")
document = _document(
id="doc-1",
tenant_id=document_tenant_id,
dataset_id="d1",
data_source_type="upload_file",
data_source_info='{"upload_file_id":"file-1"}',
)
sqlite_session.add_all([document, _upload_file(tenant_id=upload_file_tenant_id)])
sqlite_session.commit()
service._deal_document_data(dataset, sqlite_session)
sqlite_session.flush()
assert document.data_source_type == DataSourceType.UPLOAD_FILE
assert sqlite_session.scalar(select(DocumentPipelineExecutionLog)) is None
def _make_service():
return RagPipelineTransformService.__new__(RagPipelineTransformService)

27
api/tests/unit_tests/services/test_dataset_service_dataset.py
View File

@ -1,5 +1,9 @@
"""Unit tests for DatasetService and dataset-related collaborators."""
from sqlalchemy.orm import Session
from models.dataset import DatasetPermission
from .dataset_service_test_helpers import (
DatasetNameDuplicateError,
DatasetPermissionEnum,
@ -45,6 +49,29 @@ class TestDatasetServiceValidation:
with pytest.raises(ValueError, match="doc_form is different"):
DatasetService.check_doc_form(dataset, "text_model", session=session)
@pytest.mark.parametrize("operator_check", [False, True])
def test_dataset_permission_checks_ignore_foreign_tenant_binding(
self, sqlite_session: Session, operator_check: bool
) -> None:
dataset = DatasetServiceUnitDataFactory.create_dataset_mock(
dataset_id="dataset-1",
tenant_id="tenant-1",
permission=DatasetPermissionEnum.PARTIAL_TEAM,
maintainer="owner-1",
)
user = DatasetServiceUnitDataFactory.create_user_mock(
user_id="user-1",
tenant_id="tenant-1",
role=TenantAccountRole.NORMAL,
)
sqlite_session.add(DatasetPermission(dataset_id=dataset.id, account_id=user.id, tenant_id="tenant-2"))
with pytest.raises(NoPermissionError):
if operator_check:
DatasetService.check_dataset_operator_permission(user, dataset, session=sqlite_session)
else:
DatasetService.check_dataset_permission(dataset, user, sqlite_session)
def test_check_dataset_model_setting_skips_non_high_quality_datasets(self):
dataset = DatasetServiceUnitDataFactory.create_dataset_mock(indexing_technique="economy")

16
packages/contracts/generated/api/console/rag/types.gen.ts
View File

@ -585,6 +585,10 @@ export type PostRagPipelineCustomizedTemplatesByTemplateIdData = {
url: '/rag/pipeline/customized/templates/{template_id}'
}
export type PostRagPipelineCustomizedTemplatesByTemplateIdErrors = {
404: unknown
}
export type PostRagPipelineCustomizedTemplatesByTemplateIdResponses = {
200: SimpleDataResponse
}
@ -648,6 +652,10 @@ export type GetRagPipelineTemplatesByTemplateIdData = {
url: '/rag/pipeline/templates/{template_id}'
}
export type GetRagPipelineTemplatesByTemplateIdErrors = {
404: unknown
}
export type GetRagPipelineTemplatesByTemplateIdResponses = {
200: PipelineTemplateDetailResponse
}
@ -755,6 +763,10 @@ export type PostRagPipelinesTransformDatasetsByDatasetIdData = {
url: '/rag/pipelines/transform/datasets/{dataset_id}'
}
export type PostRagPipelinesTransformDatasetsByDatasetIdErrors = {
404: unknown
}
export type PostRagPipelinesTransformDatasetsByDatasetIdResponses = {
200: RagPipelineOpaqueResponse
}
@ -771,6 +783,10 @@ export type PostRagPipelinesByPipelineIdCustomizedPublishData = {
url: '/rag/pipelines/{pipeline_id}/customized/publish'
}
export type PostRagPipelinesByPipelineIdCustomizedPublishErrors = {
404: unknown
}
export type PostRagPipelinesByPipelineIdCustomizedPublishResponses = {
204: void
}