youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-15 09:27:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support dataset permission migrate to rbac (#38166)

Browse Source
This commit is contained in:
wangxiaolei
2026-06-30 14:43:46 +08:00
committed by GareArc
parent c49b4e20bf
commit d2e1f08026
9 changed files with 420 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
api/commands/__init__.py
View File

@ -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",

2
api/commands/data_migrate.py
View File

@ -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)

228
api/commands/rbac.py
View File

@ -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",
)
)

10
api/controllers/console/workspace/rbac.py
View File

@ -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):

4
api/core/rbac/__init__.py
View File

@ -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"]

8
api/core/rbac/entities.py
View File

@ -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``.

16
api/services/enterprise/rbac_service.py
View File

@ -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):

168
api/tests/unit_tests/commands/test_legacy_model_type_migration.py
View File

@ -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,

2
api/tests/unit_tests/controllers/console/workspace/test_rbac.py
View File

@ -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"})