From d2e1f080269e4cccaf482737dcfecf4f0e67a72a Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Tue, 30 Jun 2026 14:43:46 +0800 Subject: [PATCH] feat: support dataset permission migrate to rbac (#38166) --- api/commands/__init__.py | 3 +- api/commands/data_migrate.py | 2 + api/commands/rbac.py | 228 +++++++++++++++++- api/controllers/console/workspace/rbac.py | 10 +- api/core/rbac/__init__.py | 4 +- api/core/rbac/entities.py | 8 + api/services/enterprise/rbac_service.py | 16 +- .../test_legacy_model_type_migration.py | 168 +++++++++++++ .../console/workspace/test_rbac.py | 2 +- 9 files changed, 420 insertions(+), 21 deletions(-) diff --git a/api/commands/__init__.py b/api/commands/__init__.py index e4207bea74e..1e40d3f7928 100644 --- a/api/commands/__init__.py +++ b/api/commands/__init__.py @@ -22,7 +22,7 @@ from .plugin import ( setup_system_trigger_oauth_client, transform_datasource_credentials, ) -from .rbac import migrate_member_roles_to_rbac +from .rbac import migrate_dataset_permissions_to_rbac, migrate_member_roles_to_rbac from .retention import ( archive_workflow_runs, archive_workflow_runs_plan, @@ -76,6 +76,7 @@ __all__ = [ "legacy_model_types", "migrate_annotation_vector_database", "migrate_data_for_plugin", + "migrate_dataset_permissions_to_rbac", "migrate_knowledge_vector_database", "migrate_member_roles_to_rbac", "migrate_oss", diff --git a/api/commands/data_migrate.py b/api/commands/data_migrate.py index 4a71d864b26..35c2be999b7 100644 --- a/api/commands/data_migrate.py +++ b/api/commands/data_migrate.py @@ -7,6 +7,7 @@ from typing import cast import click +from commands.rbac import migrate_dataset_permissions_to_rbac from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType from services.legacy_model_type_migration import ( @@ -177,3 +178,4 @@ def legacy_model_types( data_migrate.add_command(legacy_model_types) +data_migrate.add_command(migrate_dataset_permissions_to_rbac) diff --git a/api/commands/rbac.py b/api/commands/rbac.py index 630304b67e3..0793d11cbb2 100644 --- a/api/commands/rbac.py +++ b/api/commands/rbac.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor, as_completed @@ -8,8 +9,11 @@ from sqlalchemy import select from configs import dify_config from core.db.session_factory import session_factory -from models import TenantAccountJoin, TenantAccountRole -from services.enterprise.rbac_service import ListOption, RBACService +from core.rbac import RBACResourceWhitelistScope +from models import Dataset, DatasetPermission, DatasetPermissionEnum, TenantAccountJoin, TenantAccountRole +from services.enterprise.rbac_service import ListOption, RBACService, ReplaceMemberBindings, ReplaceUserAccessPolicies + +_RBAC_DEFAULT_ACCESS_POLICY_ID = "default" _LEGACY_ROLE_TO_BUILTIN_TAG = { TenantAccountRole.OWNER.value: "owner", @@ -258,3 +262,223 @@ def migrate_member_roles_to_rbac( fg="green", ) ) + + +def _dataset_permission_enum(permission: DatasetPermissionEnum | str | None) -> DatasetPermissionEnum: + if permission is None: + return DatasetPermissionEnum.ONLY_ME + try: + return DatasetPermissionEnum(permission) + except ValueError as exc: + raise ValueError(f"Unsupported legacy dataset permission: {permission}") from exc + + +def _rbac_dataset_scope_for_legacy_permission(permission: DatasetPermissionEnum) -> RBACResourceWhitelistScope: + if permission is DatasetPermissionEnum.ALL_TEAM: + return RBACResourceWhitelistScope.ALL + if permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.PARTIAL_TEAM}: + return RBACResourceWhitelistScope.SPECIFIC + raise ValueError(f"Unsupported legacy dataset permission: {permission}") + + +def _emit_dataset_permission_migration_event(payload: dict[str, object]) -> None: + click.echo(json.dumps(payload, sort_keys=True)) + + +@click.command( + "rbac-migrate-dataset-permissions", + help=( + "Migrate legacy dataset permission scopes and partial members into RBAC dataset access bindings. " + "Side effect: replacing each dataset whitelist clears existing per-user policy bindings; " + "the command then recreates legacy partial-member default bindings." + ), +) +@click.option("--tenant-id", help="Only migrate datasets in a single workspace.") +@click.option("--dataset-id", help="Only migrate a single dataset.") +@click.option("--batch-size", default=500, show_default=True, type=click.IntRange(min=1)) +@click.option( + "--dry-run/--apply", + default=True, + show_default=True, + help="Preview the migration without writing RBAC bindings. Use --apply to write changes.", +) +def migrate_dataset_permissions_to_rbac( + tenant_id: str | None, + dataset_id: str | None, + batch_size: int, + dry_run: bool, +) -> None: + """Backfill RBAC dataset access config from legacy `Dataset.permission`. + + Legacy mapping: + - all_team_members -> RBAC dataset whitelist scope "all" + - partial_members -> RBAC dataset whitelist scope "specific" plus each partial member gets the + virtual default policy + - only_me -> RBAC dataset whitelist scope "specific" with no member policy bindings + + The command replaces each dataset's RBAC whitelist scope first. RBAC clears + existing per-user policy bindings during that replace, then this command + recreates the legacy partial-member default bindings. Re-running it is + therefore idempotent for a dataset's current legacy configuration. + """ + click.echo(click.style("Starting RBAC dataset permission migration.", fg="green")) + + scanned_count = 0 + scope_migrated_count = 0 + user_policy_migrated_count = 0 + partial_dataset_count = 0 + + last_dataset_id: str | None = None + while True: + with session_factory.create_session() as session: + stmt = ( + select(Dataset.id, Dataset.tenant_id, Dataset.permission, Dataset.created_by) + .order_by(Dataset.id.asc()) + .limit(batch_size) + ) + if tenant_id: + stmt = stmt.where(Dataset.tenant_id == tenant_id) + if dataset_id: + stmt = stmt.where(Dataset.id == dataset_id) + if last_dataset_id: + stmt = stmt.where(Dataset.id > last_dataset_id) + + dataset_rows = list(session.execute(stmt).all()) + if not dataset_rows: + break + + dataset_ids = [str(row.id) for row in dataset_rows] + partial_members_by_dataset_id: dict[str, list[str]] = {item: [] for item in dataset_ids} + permission_rows = session.execute( + select(DatasetPermission.dataset_id, DatasetPermission.account_id).where( + DatasetPermission.dataset_id.in_(dataset_ids) + ) + ).all() + for row in permission_rows: + partial_members_by_dataset_id[str(row.dataset_id)].append(str(row.account_id)) + + for dataset in dataset_rows: + workspace_id = str(dataset.tenant_id) + current_dataset_id = str(dataset.id) + operator_account_id = str(dataset.created_by) + permission_value = _dataset_permission_enum(dataset.permission) + scope = _rbac_dataset_scope_for_legacy_permission(permission_value) + partial_member_ids = sorted(set(partial_members_by_dataset_id[current_dataset_id])) + should_bind_partial_members = permission_value is DatasetPermissionEnum.PARTIAL_TEAM + + click.echo( + f"tenant={workspace_id} dataset={current_dataset_id} " + f"operator={operator_account_id} " + f"legacy_permission={permission_value} -> rbac_scope={scope} " + f"partial_members={len(partial_member_ids) if should_bind_partial_members else 0}" + ) + + scanned_count += 1 + replace_whitelist_payload = ReplaceMemberBindings(scope=scope) + if dry_run: + _emit_dataset_permission_migration_event( + { + "event": "dataset_permission_migration_proposed_change", + "action": "replace_whitelist", + "dry_run": True, + "tenant_id": workspace_id, + "dataset_id": current_dataset_id, + "operator_account_id": operator_account_id, + "before": { + "legacy_dataset_permission": permission_value.value, + "legacy_partial_member_ids": partial_member_ids if should_bind_partial_members else [], + }, + "after": { + "rbac_whitelist_scope": scope.value, + }, + "call": { + "method": "RBACService.DatasetAccess.replace_whitelist", + "kwargs": { + "tenant_id": workspace_id, + "account_id": operator_account_id, + "dataset_id": current_dataset_id, + "payload": replace_whitelist_payload.model_dump(mode="json"), + }, + }, + } + ) + if not dry_run: + RBACService.DatasetAccess.replace_whitelist( + tenant_id=workspace_id, + account_id=operator_account_id, + dataset_id=current_dataset_id, + payload=replace_whitelist_payload, + ) + scope_migrated_count += 1 + + if should_bind_partial_members: + partial_dataset_count += 1 + for member_account_id in partial_member_ids: + replace_user_access_policies_payload = ReplaceUserAccessPolicies( + access_policy_ids=[_RBAC_DEFAULT_ACCESS_POLICY_ID], + ) + if dry_run: + _emit_dataset_permission_migration_event( + { + "event": "dataset_permission_migration_proposed_change", + "action": "replace_user_access_policies", + "dry_run": True, + "tenant_id": workspace_id, + "dataset_id": current_dataset_id, + "operator_account_id": operator_account_id, + "target_account_id": member_account_id, + "before": { + "legacy_dataset_permission": permission_value.value, + "legacy_partial_member_id": member_account_id, + }, + "after": { + "rbac_user_access_policy_ids": [_RBAC_DEFAULT_ACCESS_POLICY_ID], + }, + "call": { + "method": "RBACService.DatasetAccess.replace_user_access_policies", + "kwargs": { + "tenant_id": workspace_id, + "account_id": operator_account_id, + "dataset_id": current_dataset_id, + "target_account_id": member_account_id, + "payload": replace_user_access_policies_payload.model_dump(mode="json"), + }, + }, + } + ) + continue + RBACService.DatasetAccess.replace_user_access_policies( + tenant_id=workspace_id, + account_id=operator_account_id, + dataset_id=current_dataset_id, + target_account_id=member_account_id, + payload=replace_user_access_policies_payload, + ) + user_policy_migrated_count += 1 + + last_dataset_id = dataset_ids[-1] + + if dataset_id: + break + + if scanned_count == 0: + click.echo(click.style("No datasets found for migration.", fg="yellow")) + return + + if dry_run: + click.echo( + click.style( + f"Dry run completed. Scanned {scanned_count} datasets; " + f"{partial_dataset_count} partial-member datasets would be migrated.", + fg="yellow", + ) + ) + else: + click.echo( + click.style( + "RBAC dataset permission migration completed. " + f"Scanned {scanned_count} datasets, migrated {scope_migrated_count} scopes, " + f"wrote {user_policy_migrated_count} user default-policy bindings.", + fg="green", + ) + ) diff --git a/api/controllers/console/workspace/rbac.py b/api/controllers/console/workspace/rbac.py index c3a3420b908..be1783dddd0 100644 --- a/api/controllers/console/workspace/rbac.py +++ b/api/controllers/console/workspace/rbac.py @@ -1,6 +1,5 @@ from __future__ import annotations -from enum import StrEnum from typing import Any from flask import request @@ -14,6 +13,7 @@ from controllers.common.schema import register_response_schema_models from controllers.console import console_ns from controllers.console.wraps import RBACPermission, RBACResourceScope, rbac_permission_required from core.db.session_factory import session_factory +from core.rbac import RBACResourceWhitelistScope from libs.login import current_account_with_tenant, login_required from models import Account from services.enterprise import rbac_service as svc @@ -511,14 +511,8 @@ class RBACAccessPolicyBindingUnlockApi(Resource): # --------------------------------------------------------------------------- -class _AccessScope(StrEnum): - ALL = "all" - SPECIFIC = "specific" - ONLY_ME = "only_me" - - class _ResourceAccessScopeRequest(BaseModel): - scope: _AccessScope + scope: RBACResourceWhitelistScope class _ReplaceBindingsRequest(BaseModel): diff --git a/api/core/rbac/__init__.py b/api/core/rbac/__init__.py index 495ac90f492..e8cf30e393f 100644 --- a/api/core/rbac/__init__.py +++ b/api/core/rbac/__init__.py @@ -1,3 +1,3 @@ -from core.rbac.entities import RBACPermission, RBACResourceScope +from core.rbac.entities import RBACPermission, RBACResourceScope, RBACResourceWhitelistScope -__all__ = ["RBACPermission", "RBACResourceScope"] +__all__ = ["RBACPermission", "RBACResourceScope", "RBACResourceWhitelistScope"] diff --git a/api/core/rbac/entities.py b/api/core/rbac/entities.py index 16a05f13111..200c33e9424 100644 --- a/api/core/rbac/entities.py +++ b/api/core/rbac/entities.py @@ -13,6 +13,14 @@ class RBACResourceScope(StrEnum): WORKSPACE = "workspace" +class RBACResourceWhitelistScope(StrEnum): + """Whitelist scopes accepted by RBAC app and dataset access config APIs.""" + + ALL = "all" + SPECIFIC = "specific" + ONLY_ME = "only_me" + + class RBACPermission(StrEnum): """Permission points (RBAC scenes) checked by ``rbac_permission_required``. diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index ff790b1690c..93eb1c1ff98 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -12,6 +12,7 @@ from sqlalchemy.exc import SQLAlchemyError from configs import dify_config from core.db.session_factory import session_factory +from core.rbac import RBACResourceWhitelistScope from models import TenantAccountJoin, TenantAccountRole from services.enterprise.base import EnterpriseRequest @@ -649,17 +650,18 @@ class ReplaceRoleBindings(_RBACModel): class ReplaceMemberBindings(_RBACModel): - scope: str = "specific" + scope: RBACResourceWhitelistScope = RBACResourceWhitelistScope.SPECIFIC @field_validator("scope") @classmethod - def _normalize_scope(cls, value: Any) -> str: + def _normalize_scope(cls, value: Any) -> RBACResourceWhitelistScope: scope = str(value or "").strip().lower() - if scope in {"", "specific"}: - return "specific" - if scope in {"all", "only_me"}: - return scope - raise ValueError(f"invalid scope: {value}") + if scope == "": + return RBACResourceWhitelistScope.SPECIFIC + try: + return RBACResourceWhitelistScope(scope) + except ValueError as exc: + raise ValueError(f"invalid scope: {value}") from exc class DeleteMemberBindings(_RBACModel): diff --git a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py index 7eead948c1c..e14c8ed3243 100644 --- a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py +++ b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py @@ -336,6 +336,174 @@ def _insert_load_balancing_model_config( ) +def test_data_migrate_group_registers_dataset_permission_rbac_migration(command_module) -> None: + command = command_module.data_migrate.commands["rbac-migrate-dataset-permissions"] + + assert command is command_module.migrate_dataset_permissions_to_rbac + assert "operator_account_id" not in {param.name for param in command.params} + + +def test_dataset_permission_rbac_migration_help_mentions_binding_clear_side_effect(command_module) -> None: + result = CliRunner().invoke( + command_module.data_migrate, + ["rbac-migrate-dataset-permissions", "--help"], + ) + + assert result.exit_code == 0 + normalized_output = " ".join(result.output.split()) + assert "clears existing per-user policy bindings" in normalized_output + assert "recreates legacy partial-member default bindings" in normalized_output + + +def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scopes() -> None: + rbac_module = importlib.import_module("commands.rbac") + + assert ( + rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.ALL_TEAM) + is rbac_module.RBACResourceWhitelistScope.ALL + ) + assert ( + rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.PARTIAL_TEAM) + is rbac_module.RBACResourceWhitelistScope.SPECIFIC + ) + assert rbac_module._dataset_permission_enum("partial_members") is rbac_module.DatasetPermissionEnum.PARTIAL_TEAM + assert rbac_module._dataset_permission_enum(None) is rbac_module.DatasetPermissionEnum.ONLY_ME + + +def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator( + command_module, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rbac_module = importlib.import_module("commands.rbac") + dataset_row = SimpleNamespace( + id="dataset-1", + tenant_id="tenant-1", + permission="only_me", + created_by="creator-account-1", + ) + execute_results = [[dataset_row], [], []] + calls: list[dict[str, object]] = [] + session_closed = False + + class FakeExecuteResult: + def __init__(self, rows: list[object]) -> None: + self._rows = rows + + def all(self) -> list[object]: + return self._rows + + class FakeSession: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback) -> None: + nonlocal session_closed + session_closed = True + pass + + def execute(self, stmt): + return FakeExecuteResult(execute_results.pop(0)) + + class FakeSessionFactory: + @staticmethod + def create_session() -> FakeSession: + return FakeSession() + + def fake_replace_whitelist(**kwargs): + assert session_closed is True + calls.append(kwargs) + + monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory) + monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist) + + command_module.migrate_dataset_permissions_to_rbac.callback( + tenant_id=None, + dataset_id=None, + batch_size=500, + dry_run=False, + ) + + assert calls[0]["tenant_id"] == "tenant-1" + assert calls[0]["account_id"] == "creator-account-1" + assert calls[0]["dataset_id"] == "dataset-1" + assert calls[0]["payload"].scope is rbac_module.RBACResourceWhitelistScope.SPECIFIC + + +def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes( + command_module, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rbac_module = importlib.import_module("commands.rbac") + dataset_row = SimpleNamespace( + id="dataset-1", + tenant_id="tenant-1", + permission="partial_members", + created_by="creator-account-1", + ) + permission_row = SimpleNamespace(dataset_id="dataset-1", account_id="member-account-1") + execute_results = [[dataset_row], [permission_row], []] + + class FakeExecuteResult: + def __init__(self, rows: list[object]) -> None: + self._rows = rows + + def all(self) -> list[object]: + return self._rows + + class FakeSession: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback) -> None: + pass + + def execute(self, stmt): + return FakeExecuteResult(execute_results.pop(0)) + + class FakeSessionFactory: + @staticmethod + def create_session() -> FakeSession: + return FakeSession() + + monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: pytest.fail("dry-run must not replace whitelist"), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: pytest.fail("dry-run must not replace user access policies"), + ) + + result = CliRunner().invoke( + command_module.data_migrate, + ["rbac-migrate-dataset-permissions", "--dry-run"], + ) + + assert result.exit_code == 0 + events = [json.loads(line) for line in result.output.splitlines() if line.startswith("{")] + assert [event["action"] for event in events] == ["replace_whitelist", "replace_user_access_policies"] + assert events[0]["before"] == { + "legacy_dataset_permission": "partial_members", + "legacy_partial_member_ids": ["member-account-1"], + } + assert events[0]["after"] == {"rbac_whitelist_scope": "specific"} + assert events[0]["call"] == { + "method": "RBACService.DatasetAccess.replace_whitelist", + "kwargs": { + "tenant_id": "tenant-1", + "account_id": "creator-account-1", + "dataset_id": "dataset-1", + "payload": {"scope": "specific"}, + }, + } + assert events[1]["target_account_id"] == "member-account-1" + assert events[1]["after"] == {"rbac_user_access_policy_ids": ["default"]} + assert events[1]["call"]["kwargs"]["payload"] == {"access_policy_ids": ["default"]} + + def test_data_migrate_command_defaults_output_to_stdout_stream( command_module, monkeypatch: pytest.MonkeyPatch, diff --git a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py b/api/tests/unit_tests/controllers/console/workspace/test_rbac.py index 2960bfef324..92e819f73bf 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_rbac.py @@ -136,7 +136,7 @@ class TestPydanticModels: def test_resource_access_scope_defaults_empty_account_ids(self): parsed = rbac_mod._ResourceAccessScopeRequest.model_validate({"scope": "specific"}) - assert parsed.scope is rbac_mod._AccessScope.SPECIFIC + assert parsed.scope is rbac_mod.RBACResourceWhitelistScope.SPECIFIC def test_resource_access_scope_coerce_null_account_ids(self): rbac_mod._ResourceAccessScopeRequest.model_validate({"scope": "all"})