fix(api): bind dataset operations to owners (#40149)

api/controllers/console/datasets/datasets.py

@@ -53,6 +53,7 @@ from models.enums import ApiTokenType, SegmentStatus
 from models.provider_ids import ModelProviderID
 from services.api_token_service import ApiTokenCache
 from services.app_service import AppService
+from services.dataset_ref_service import DatasetRefService
 from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
 from services.enterprise import rbac_service as enterprise_rbac_service
 from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings
@@ -67,6 +68,18 @@ def _has_dataset_list_permission(permission_keys: list[str]) -> bool:
     return any(permission_key in DATASET_LIST_PERMISSION_KEYS for permission_key in permission_keys)
 
 
+def _get_accessible_dataset(dataset_id: UUID, tenant_id: str, current_user: Account, session: Session) -> Dataset:
+    dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), tenant_id, session=session)
+    if dataset is None:
+        raise NotFound("Dataset not found.")
+    if not dify_config.RBAC_ENABLED:
+        try:
+            DatasetService.check_dataset_permission(dataset, current_user, session)
+        except services.errors.account.NoPermissionError as e:
+            raise Forbidden(str(e))
+    return dataset
+
+
 def _validate_indexing_technique(value: str | None) -> str | None:
     if value is None:
         return value
@@ -811,12 +824,13 @@ class DatasetUseCheckApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
     @with_session(write=False)
-    def get(self, session: Session, dataset_id: UUID):
-        dataset_id_str = str(dataset_id)
-
-        dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, session)
+    def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
+        dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session)
+        dataset_is_using = DatasetService.dataset_use_check(DatasetRefService.create_dataset_ref(dataset), session)
         return UsageCheckResponse(is_using=dataset_is_using).model_dump(mode="json"), 200
 
 
@@ -1027,13 +1041,15 @@ class DatasetIndexingStatusApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
+    @with_current_user
     @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
     @with_session(write=False)
-    def get(self, session: Session, current_tenant_id: str, dataset_id: UUID):
-        dataset_id_str = str(dataset_id)
+    def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
+        dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session)
+
         documents = session.scalars(
-            select(Document).where(Document.dataset_id == dataset_id_str, Document.tenant_id == current_tenant_id)
+            select(Document).where(Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id)
         ).all()
         documents_status = []
         for document in documents:
@@ -1041,6 +1057,8 @@ class DatasetIndexingStatusApi(Resource):
                 session.scalar(
                     select(func.count(DocumentSegment.id)).where(
                         DocumentSegment.completed_at.isnot(None),
+                        DocumentSegment.tenant_id == dataset.tenant_id,
+                        DocumentSegment.dataset_id == dataset.id,
                         DocumentSegment.document_id == str(document.id),
                         DocumentSegment.status != SegmentStatus.RE_SEGMENT,
                     )
@@ -1050,6 +1068,8 @@ class DatasetIndexingStatusApi(Resource):
             total_segments = (
                 session.scalar(
                     select(func.count(DocumentSegment.id)).where(
+                        DocumentSegment.tenant_id == dataset.tenant_id,
+                        DocumentSegment.dataset_id == dataset.id,
                         DocumentSegment.document_id == str(document.id),
                         DocumentSegment.status != SegmentStatus.RE_SEGMENT,
                     )
@@ -1177,12 +1197,16 @@ class DatasetEnableApiApi(Resource):
     @login_required
     @account_initialization_required
     @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def post(self, session: Session, dataset_id: UUID, status: str):
-        dataset_id_str = str(dataset_id)
+    def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, status: str):
+        dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session)
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
 
-        DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", session)
+        DatasetService.update_dataset_api_status(dataset, status == "enable", current_user, session)
 
         return SimpleResultResponse(result="success").model_dump(mode="json"), 200
 
@@ -1248,14 +1272,15 @@ class DatasetErrorDocs(Resource):
     @setup_required
     @login_required
     @account_initialization_required
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
     @with_session(write=False)
-    def get(self, session: Session, dataset_id: UUID):
-        dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        if dataset is None:
-            raise NotFound("Dataset not found.")
-        results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, session)
+    def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
+        dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session)
+        results = DocumentService.get_error_documents_by_dataset_ref(
+            DatasetRefService.create_dataset_ref(dataset), session
+        )
 
         return dump_response(ErrorDocsResponse, {"data": results, "total": len(results)}), 200
 
@@ -1307,12 +1332,13 @@ class DatasetAutoDisableLogApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
     @with_session(write=False)
-    def get(self, session: Session, dataset_id: UUID):
-        dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        if dataset is None:
-            raise NotFound("Dataset not found.")
-        auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, session)
+    def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
+        dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session)
+        auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(
+            DatasetRefService.create_dataset_ref(dataset), session
+        )
         return dump_response(AutoDisableLogsResponse, auto_disable_logs), 200

api/controllers/console/datasets/datasets_document.py

@@ -16,6 +16,7 @@ from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
+from configs import dify_config
 from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload
 from controllers.common.fields import SimpleResultMessageResponse, SimpleResultResponse, UrlResponse
 from controllers.common.schema import register_response_schema_models, register_schema_models
@@ -78,6 +79,7 @@ from ..datasets.error import (
 )
 from ..wraps import (
     account_initialization_required,
+    check_knowledge_rate_limit,
     cloud_edition_billing_rate_limit_check,
     cloud_edition_billing_resource_check,
     setup_required,
@@ -294,23 +296,23 @@ class DocumentResource(Resource):
     def get_document(
         self, session: Session, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str
     ) -> Document:
-        dataset = DatasetService.get_dataset(dataset_id, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id, current_tenant_id, session=session)
         if not dataset:
             raise NotFound("Dataset not found.")
 
-        try:
-            DatasetService.check_dataset_permission(dataset, current_user, session)
-        except services.errors.account.NoPermissionError as e:
-            raise Forbidden(str(e))
+        if not dify_config.RBAC_ENABLED:
+            try:
+                DatasetService.check_dataset_permission(dataset, current_user, session)
+            except services.errors.account.NoPermissionError as e:
+                raise Forbidden(str(e))
 
-        document = DocumentService.get_document(dataset_id, document_id, session=session)
+        dataset_ref = DatasetRefService.create_dataset_ref(dataset)
+        document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, document_id)
+        document = DatasetRefService.get_document_by_ref(document_ref, session=session)
 
         if not document:
             raise NotFound("Document not found.")
 
-        if document.tenant_id != current_tenant_id:
-            raise Forbidden("No permission.")
-
         return document
 
     def get_batch_documents(
@@ -578,18 +580,33 @@ class DatasetDocumentListApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    @cloud_edition_billing_rate_limit_check("knowledge")
     @console_ns.response(204, "Documents deleted successfully")
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def delete(self, session: Session, dataset_id: UUID):
+    def delete(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+    ):
         dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
-        # check user's model setting
-        DatasetService.check_dataset_model_setting(dataset)
 
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
+
+        if not dify_config.RBAC_ENABLED:
+            try:
+                DatasetService.check_dataset_permission(dataset, current_user, session)
+            except services.errors.account.NoPermissionError as e:
+                raise Forbidden(str(e))
+
+        check_knowledge_rate_limit()
         try:
             document_ids = request.args.getlist("document_id")
             dataset_ref = DatasetRefService.create_dataset_ref(dataset)
@@ -1367,29 +1384,32 @@ class DocumentPauseApi(DocumentResource):
     @setup_required
     @login_required
     @account_initialization_required
-    @cloud_edition_billing_rate_limit_check("knowledge")
     @console_ns.response(204, "Document paused successfully")
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def patch(self, session: Session, dataset_id: UUID, document_id: UUID):
+    def patch(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+        document_id: UUID,
+    ):
         """pause document."""
         dataset_id_str = str(dataset_id)
         document_id_str = str(document_id)
 
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        if not dataset:
-            raise NotFound("Dataset not found.")
-
-        document = DocumentService.get_document(dataset.id, document_id_str, session=session)
-
-        # 404 if document not found
-        if document is None:
-            raise NotFound("Document Not Exists.")
+        document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
 
         # 403 if document is archived
         if DocumentService.check_archived(document):
             raise ArchivedDocumentImmutableError()
 
+        check_knowledge_rate_limit()
         try:
             # pause document
             DocumentService.pause_document(document, session)
@@ -1404,26 +1424,31 @@ class DocumentRecoverApi(DocumentResource):
     @setup_required
     @login_required
     @account_initialization_required
-    @cloud_edition_billing_rate_limit_check("knowledge")
     @console_ns.response(204, "Document resumed successfully")
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def patch(self, session: Session, dataset_id: UUID, document_id: UUID):
+    def patch(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+        document_id: UUID,
+    ):
         """recover document."""
         dataset_id_str = str(dataset_id)
         document_id_str = str(document_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        if not dataset:
-            raise NotFound("Dataset not found.")
-        document = DocumentService.get_document(dataset.id, document_id_str, session=session)
 
-        # 404 if document not found
-        if document is None:
-            raise NotFound("Document Not Exists.")
+        document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
 
         # 403 if document is archived
         if DocumentService.check_archived(document):
             raise ArchivedDocumentImmutableError()
+        check_knowledge_rate_limit()
         try:
             # pause document
             DocumentService.recover_document(document, session)
@@ -1438,21 +1463,40 @@ class DocumentRetryApi(DocumentResource):
     @setup_required
     @login_required
     @account_initialization_required
-    @cloud_edition_billing_rate_limit_check("knowledge")
     @console_ns.expect(console_ns.models[DocumentRetryPayload.__name__])
     @console_ns.response(204, "Documents retry started successfully")
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def post(self, session: Session, dataset_id: UUID):
+    def post(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+    ):
         """retry document."""
-        payload = DocumentRetryPayload.model_validate(console_ns.payload or {})
         dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        retry_documents = []
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if not dataset:
             raise NotFound("Dataset not found.")
-        documents = DocumentService.get_documents_by_ids(dataset.id, payload.document_ids, session)
+
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
+
+        if not dify_config.RBAC_ENABLED:
+            try:
+                DatasetService.check_dataset_permission(dataset, current_user, session)
+            except services.errors.account.NoPermissionError as e:
+                raise Forbidden(str(e))
+
+        payload = DocumentRetryPayload.model_validate(console_ns.payload or {})
+        documents = DocumentService.get_documents_by_ids(
+            DatasetRefService.create_dataset_ref(dataset), payload.document_ids, session
+        )
         documents_by_id = {document.id: document for document in documents}
+        retry_documents = []
         for document_id in payload.document_ids:
             try:
                 document = documents_by_id.get(document_id)
@@ -1473,6 +1517,7 @@ class DocumentRetryApi(DocumentResource):
                 logger.exception("Failed to retry document, document id: %s", document_id)
                 continue
         # retry document
+        check_knowledge_rate_limit()
         DocumentService.retry_document(dataset_id_str, retry_documents, session)
 
         return "", 204
@@ -1512,28 +1557,34 @@ class WebsiteDocumentSyncApi(DocumentResource):
     @login_required
     @account_initialization_required
     @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
+    @with_current_user
     @with_current_tenant_id
-    @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
+    @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def get(self, session: Session, current_tenant_id: str, dataset_id: UUID, document_id: UUID):
+    def get(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+        document_id: UUID,
+    ):
         """sync website document."""
         dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if not dataset:
             raise NotFound("Dataset not found.")
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
         document_id_str = str(document_id)
-        document = DocumentService.get_document(dataset.id, document_id_str, session=session)
-        if not document:
-            raise NotFound("Document not found.")
-        if document.tenant_id != current_tenant_id:
-            raise Forbidden("No permission.")
+        document = self.get_document(session, dataset.id, document_id_str, current_user, current_tenant_id)
         if document.data_source_type != "website_crawl":
             raise ValueError("Document is not a website document.")
         # 403 if document is archived
         if DocumentService.check_archived(document):
             raise ArchivedDocumentImmutableError()
         # sync document
-        DocumentService.sync_website_document(dataset_id_str, document, session)
+        DocumentService.sync_website_document(dataset, document, session)
 
         return SimpleResultResponse(result="success").model_dump(mode="json"), 200
 
@@ -1548,21 +1599,25 @@ class DocumentPipelineExecutionLogApi(DocumentResource):
     @setup_required
     @login_required
     @account_initialization_required
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
     @with_session(write=False)
-    def get(self, session: Session, dataset_id: UUID, document_id: UUID):
+    def get(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+        document_id: UUID,
+    ):
         dataset_id_str = str(dataset_id)
         document_id_str = str(document_id)
 
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
-        if not dataset:
-            raise NotFound("Dataset not found.")
-        document = DocumentService.get_document(dataset.id, document_id_str, session=session)
-        if not document:
-            raise NotFound("Document not found.")
+        document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
         log = session.scalar(
             select(DocumentPipelineExecutionLog)
-            .where(DocumentPipelineExecutionLog.document_id == document_id_str)
+            .where(DocumentPipelineExecutionLog.document_id == document.id)
             .order_by(DocumentPipelineExecutionLog.created_at.desc())
             .limit(1)
         )
@@ -1645,7 +1700,9 @@ class DocumentGenerateSummaryApi(Resource):
             raise ValueError("Summary index is not enabled for this dataset. Please enable it in the dataset settings.")
 
         # Verify all documents exist and belong to the dataset
-        documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, session)
+        documents = DocumentService.get_documents_by_ids(
+            DatasetRefService.create_dataset_ref(dataset), document_list, session
+        )
 
         if len(documents) != len(document_list):
             found_ids = {doc.id for doc in documents}

api/controllers/console/datasets/metadata.py

@@ -3,8 +3,10 @@ from uuid import UUID
 
 from flask_restx import Resource
 from sqlalchemy.orm import Session
-from werkzeug.exceptions import NotFound
+from werkzeug.exceptions import Forbidden, NotFound
 
+import services
+from configs import dify_config
 from controllers.common.controller_schemas import MetadataUpdatePayload
 from controllers.common.schema import register_response_schema_models, register_schema_models
 from controllers.common.session import with_session
@@ -35,6 +37,7 @@ from services.entities.knowledge_entities.knowledge_entities import (
     MetadataDetail,
     MetadataOperationData,
 )
+from services.errors.metadata import MetadataResourceNotFoundError
 from services.metadata_service import MetadataService
 
 register_schema_models(
@@ -87,13 +90,20 @@ class DatasetMetadataCreateApi(Resource):
     @console_ns.response(
         200, "Metadata retrieved successfully", console_ns.models[DatasetMetadataListResponse.__name__]
     )
+    @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
     @with_session(write=False)
-    def get(self, session: Session, dataset_id: UUID):
+    def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
         dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
+        if not dify_config.RBAC_ENABLED:
+            try:
+                DatasetService.check_dataset_permission(dataset, current_user, session)
+            except services.errors.account.NoPermissionError as e:
+                raise Forbidden(str(e))
         metadata = MetadataService.get_dataset_metadatas(dataset, session)
         return dump_response(DatasetMetadataListResponse, metadata), 200
 
@@ -124,14 +134,12 @@ class DatasetMetadataApi(Resource):
 
         dataset_id_str = str(dataset_id)
         metadata_id_str = str(metadata_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
-        metadata = MetadataService.update_metadata_name(
-            dataset_id_str, metadata_id_str, name, current_user, current_tenant_id, session=session
-        )
+        metadata = MetadataService.update_metadata_name(dataset, metadata_id_str, name, current_user, session=session)
         return dump_response(DatasetMetadataResponse, metadata), 200
 
     @setup_required
@@ -140,17 +148,25 @@ class DatasetMetadataApi(Resource):
     @enterprise_license_required
     @console_ns.response(204, "Metadata deleted successfully")
     @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
-    def delete(self, session: Session, current_user: Account, dataset_id: UUID, metadata_id: UUID):
+    def delete(
+        self,
+        session: Session,
+        current_tenant_id: str,
+        current_user: Account,
+        dataset_id: UUID,
+        metadata_id: UUID,
+    ):
         dataset_id_str = str(dataset_id)
         metadata_id_str = str(metadata_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
-        MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session)
+        MetadataService.delete_metadata(dataset, metadata_id_str, session)
         # Frontend callers only await success and invalidate metadata caches; no response body is consumed.
         return "", 204
 
@@ -208,18 +224,29 @@ class DocumentMetadataEditApi(Resource):
         204,
         "Documents metadata updated successfully",
     )
+    @console_ns.response(404, "Dataset, document, or metadata not found")
     @with_current_user
+    @with_current_tenant_id
     @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
     @with_session
     @model_validate(MetadataOperationData)
-    def post(self, req_data: MetadataOperationData, session: Session, current_user: Account, dataset_id: UUID):
-        dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+    def post(
+        self,
+        req_data: MetadataOperationData,
+        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.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
-        MetadataService.update_documents_metadata(dataset, req_data, current_user, session=session)
+        try:
+            MetadataService.update_documents_metadata(dataset, req_data, current_user, session=session)
+        except MetadataResourceNotFoundError as exc:
+            raise NotFound(str(exc)) from exc
 
         # Frontend callers only await success and invalidate caches; no response body is consumed.
         return "", 204

api/controllers/console/wraps.py

@@ -249,35 +249,35 @@ def cloud_edition_billing_knowledge_limit_check[**P, R](
     return interceptor
 
 
+def check_knowledge_rate_limit() -> None:
+    _, current_tenant_id = current_account_with_tenant()
+    knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(current_tenant_id)
+    if not knowledge_rate_limit.enabled:
+        return
+
+    current_time = int(time.time() * 1000)
+    key = f"rate_limit_{current_tenant_id}"
+    redis_client.zadd(key, {current_time: current_time})
+    redis_client.zremrangebyscore(key, 0, current_time - 60000)
+
+    if redis_client.zcard(key) > knowledge_rate_limit.limit:
+        db.session.add(  # guard-ignore: no-new-controller-sqlalchemy -- existing decorator audit write
+            RateLimitLog(
+                tenant_id=current_tenant_id,
+                subscription_plan=knowledge_rate_limit.subscription_plan,
+                operation="knowledge",
+            )
+        )
+        db.session.commit()
+        abort(403, "Sorry, you have reached the knowledge base request rate limit of your subscription.")
+
+
 def cloud_edition_billing_rate_limit_check[**P, R](resource: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
             if resource == "knowledge":
-                _, current_tenant_id = current_account_with_tenant()
-                knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(current_tenant_id)
-                if knowledge_rate_limit.enabled:
-                    current_time = int(time.time() * 1000)
-                    key = f"rate_limit_{current_tenant_id}"
-
-                    redis_client.zadd(key, {current_time: current_time})
-
-                    redis_client.zremrangebyscore(key, 0, current_time - 60000)
-
-                    request_count = redis_client.zcard(key)
-
-                    if request_count > knowledge_rate_limit.limit:
-                        # add ratelimit record
-                        rate_limit_log = RateLimitLog(
-                            tenant_id=current_tenant_id,
-                            subscription_plan=knowledge_rate_limit.subscription_plan,
-                            operation="knowledge",
-                        )
-                        db.session.add(rate_limit_log)
-                        db.session.commit()
-                        abort(
-                            403, "Sorry, you have reached the knowledge base request rate limit of your subscription."
-                        )
+                check_knowledge_rate_limit()
             return view(*args, **kwargs)
 
         return decorated

api/controllers/service_api/dataset/metadata.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, cast
 from uuid import UUID
 
 from flask_login import current_user
@@ -17,6 +17,7 @@ from fields.dataset_fields import (
     DatasetMetadataResponse,
 )
 from libs.helper import dump_response
+from models import Account
 from services.dataset_service import DatasetService
 from services.entities.knowledge_entities.knowledge_entities import (
     DocumentMetadataOperation,
@@ -24,6 +25,7 @@ from services.entities.knowledge_entities.knowledge_entities import (
     MetadataDetail,
     MetadataOperationData,
 )
+from services.errors.metadata import MetadataResourceNotFoundError
 from services.metadata_service import MetadataService
 
 BUILT_IN_METADATA_ACTION_PARAM = {
@@ -158,12 +160,14 @@ class DatasetMetadataServiceApi(DatasetApiResource):
 
         dataset_id_str = str(dataset_id)
         metadata_id_str = str(metadata_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
-        metadata = MetadataService.update_metadata_name(dataset_id_str, metadata_id_str, payload.name, session=session)
+        metadata = MetadataService.update_metadata_name(
+            dataset, metadata_id_str, payload.name, cast(Account, current_user), session=session
+        )
         return dump_response(DatasetMetadataResponse, metadata), 200
 
     @service_api_ns.doc(
@@ -194,12 +198,12 @@ class DatasetMetadataServiceApi(DatasetApiResource):
         """Delete metadata."""
         dataset_id_str = str(dataset_id)
         metadata_id_str = str(metadata_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, tenant_id, session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
-        MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session)
+        MetadataService.delete_metadata(dataset, metadata_id_str, session)
         return "", 204
 
 
@@ -297,7 +301,7 @@ class DocumentMetadataEditServiceApi(DatasetApiResource):
         responses={
             200: "Documents metadata updated successfully",
             401: "Unauthorized - invalid API token",
-            404: "Dataset not found",
+            404: "Dataset, document, or metadata not found",
         }
     )
     @service_api_ns.response(
@@ -309,14 +313,18 @@ class DocumentMetadataEditServiceApi(DatasetApiResource):
     @with_session
     def post(self, session: Session, tenant_id, dataset_id: UUID):
         """Update metadata for multiple documents."""
-        dataset_id_str = str(dataset_id)
-        dataset = DatasetService.get_dataset(dataset_id_str, session)
+        dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), str(tenant_id), session=session)
         if dataset is None:
             raise NotFound("Dataset not found.")
         DatasetService.check_dataset_permission(dataset, current_user, session)
 
         metadata_args = MetadataOperationData.model_validate(service_api_ns.payload or {})
 
-        MetadataService.update_documents_metadata(dataset, metadata_args, session=session)
+        try:
+            MetadataService.update_documents_metadata(
+                dataset, metadata_args, cast(Account, current_user), session=session
+            )
+        except MetadataResourceNotFoundError as exc:
+            raise NotFound(str(exc)) from exc
 
         return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200