mirror of
https://github.com/langgenius/dify.git
synced 2026-09-19 10:11:30 +08:00
fix: fix create dataset with document miss rbac operation (#40599)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
autofix-ci[bot]
parent
6899c7720f
commit
eeaaf11a51
@@ -55,14 +55,17 @@ from libs.helper import dump_response, to_timestamp
|
||||
from libs.login import login_required
|
||||
from libs.pagination import paginate_query
|
||||
from models import Account, Document, DocumentSegment, UploadFile
|
||||
from models.dataset import DocumentPipelineExecutionLog
|
||||
from models.dataset import DatasetPermissionEnum, DocumentPipelineExecutionLog
|
||||
from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus
|
||||
from services.dataset_ref_service import DatasetRefService
|
||||
from services.dataset_service import DatasetService, DocumentService
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings
|
||||
from services.entities.knowledge_entities.knowledge_entities import KnowledgeConfig, ProcessRule, RetrievalModel
|
||||
from services.file_service import FileService
|
||||
from services.vector_space_admission_service import get_vector_space_admission_error_fields
|
||||
from tasks.generate_summary_index_task import generate_summary_index_task
|
||||
from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task
|
||||
|
||||
from ..app.error import (
|
||||
ProviderModelCurrentlyNotSupportError,
|
||||
@@ -679,6 +682,21 @@ class DatasetInitApi(Resource):
|
||||
except ModelCurrentlyNotSupportError:
|
||||
raise ProviderModelCurrentlyNotSupportError()
|
||||
|
||||
if dify_config.RBAC_ENABLED:
|
||||
dataset.permission = DatasetPermissionEnum.ALL_TEAM
|
||||
else:
|
||||
dataset.permission = DatasetPermissionEnum.ONLY_ME
|
||||
session.flush()
|
||||
|
||||
if dify_config.RBAC_ENABLED:
|
||||
enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist(
|
||||
current_tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
|
||||
return dump_response(
|
||||
DatasetAndDocumentResponse,
|
||||
{"dataset": dataset, "documents": document_responses(documents, session=session), "batch": batch},
|
||||
|
||||
@@ -44,10 +44,11 @@ from controllers.console.datasets.error import (
|
||||
)
|
||||
from core.entities.knowledge_entities import IndexingEstimate
|
||||
from core.rag.index_processor.constant.index_type import IndexStructureType
|
||||
from models.dataset import Dataset
|
||||
from models.dataset import Dataset, DatasetPermissionEnum
|
||||
from models.dataset import Document as DatasetDocument
|
||||
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
|
||||
from services.dataset_ref_service import DatasetRef, DocumentRef
|
||||
from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings
|
||||
from services.vector_space_admission_service import (
|
||||
VECTOR_SPACE_ADMISSION_ERROR_CODE,
|
||||
format_vector_space_admission_error,
|
||||
@@ -501,6 +502,7 @@ class TestDatasetInitApi:
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch("controllers.console.datasets.datasets_document.dify_config.RBAC_ENABLED", True),
|
||||
patch(
|
||||
"controllers.console.datasets.datasets_document.DocumentService.document_create_args_validate",
|
||||
return_value=None,
|
||||
@@ -509,6 +511,12 @@ class TestDatasetInitApi:
|
||||
"controllers.console.datasets.datasets_document.DocumentService.save_document_without_dataset_id",
|
||||
return_value=(created_dataset, [created_document], "batch-init"),
|
||||
),
|
||||
patch(
|
||||
"controllers.console.datasets.datasets_document.enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist"
|
||||
) as replace_whitelist,
|
||||
patch(
|
||||
"controllers.console.datasets.datasets_document.initialize_created_app_rbac_access_task"
|
||||
) as initialize_rbac_task,
|
||||
):
|
||||
response = method(api, session, tenant_id, user)
|
||||
assert response["dataset"]["id"] == "ds-1"
|
||||
@@ -516,6 +524,14 @@ class TestDatasetInitApi:
|
||||
assert response["documents"][0]["data_source_info"] == {}
|
||||
assert response["documents"][0]["doc_metadata"] == []
|
||||
assert response["batch"] == "batch-init"
|
||||
assert created_dataset.permission == DatasetPermissionEnum.ALL_TEAM
|
||||
replace_whitelist.assert_called_once_with(
|
||||
tenant_id,
|
||||
user.id,
|
||||
created_dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_rbac_task.delay.assert_called_once_with(tenant_id, user.id, dataset_id=created_dataset.id)
|
||||
|
||||
|
||||
class TestDocumentResource:
|
||||
|
||||
Reference in New Issue
Block a user