diff --git a/api/commands/__init__.py b/api/commands/__init__.py index d62d0dbd7c..5483b24b20 100644 --- a/api/commands/__init__.py +++ b/api/commands/__init__.py @@ -14,6 +14,7 @@ from .plugin import ( setup_system_trigger_oauth_client, transform_datasource_credentials, ) +from .rbac import migrate_member_roles_to_rbac from .retention import ( archive_workflow_runs, clean_expired_messages, @@ -55,6 +56,7 @@ __all__ = [ "migrate_annotation_vector_database", "migrate_data_for_plugin", "migrate_knowledge_vector_database", + "migrate_member_roles_to_rbac", "migrate_oss", "old_metadata_migration", "remove_orphaned_files_on_storage", diff --git a/api/extensions/ext_commands.py b/api/extensions/ext_commands.py index fe95cc5816..70fb51038b 100644 --- a/api/extensions/ext_commands.py +++ b/api/extensions/ext_commands.py @@ -21,6 +21,7 @@ def init_app(app: DifyApp): install_plugins, install_rag_pipeline_plugins, migrate_data_for_plugin, + migrate_member_roles_to_rbac, migrate_oss, old_metadata_migration, remove_orphaned_files_on_storage, @@ -47,6 +48,7 @@ def init_app(app: DifyApp): upgrade_db, fix_app_site_missing, migrate_data_for_plugin, + migrate_member_roles_to_rbac, extract_plugins, extract_unique_plugins, install_plugins, diff --git a/api/services/account_service.py b/api/services/account_service.py index 6f7f6f7843..e9c08d704e 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -1147,6 +1147,18 @@ class TenantService: else: tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup) TenantService.create_tenant_member(tenant, account, role="owner") + if dify_config.RBAC_ENABLED: + resolved_role_id = AccountService.resolve_workspace_rbac_role_id( + tenant_id=str(tenant.id), + account_id=account.id, + role_identifier="所有者", + ) + RBACService.MemberRoles.replace( + tenant_id=str(tenant.id), + account_id=account.id, + member_account_id=account.id, + role_ids=[resolved_role_id], + ) account.current_tenant = tenant db.session.commit() tenant_was_created.send(tenant) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 25544fd11a..79cf4b6279 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -828,6 +828,40 @@ class TestTenantService: assert mock_target_join.role == "admin" self._assert_database_operations_called(mock_db) + def test_create_owner_tenant_if_not_exist_rbac_enabled_assigns_owner_role( + self, mock_db_dependencies, mock_external_service_dependencies + ): + mock_account = TestAccountAssociatedDataFactory.create_account_mock(account_id="user-rbac", name="RBAC User") + mock_external_service_dependencies[ + "feature_service" + ].get_system_features.return_value.is_allow_create_workspace = True + mock_external_service_dependencies[ + "feature_service" + ].get_system_features.return_value.license.workspaces.is_available.return_value = True + + mock_tenant = MagicMock() + mock_tenant.id = "tenant-rbac" + mock_tenant.name = "RBAC User's Workspace" + + with ( + patch("services.account_service.dify_config.RBAC_ENABLED", True), + patch("services.account_service.TenantService.create_tenant", return_value=mock_tenant), + patch("services.account_service.TenantService.create_tenant_member"), + patch("services.account_service.AccountService.resolve_workspace_rbac_role_id", return_value="rbac-owner-id"), + patch("services.account_service.RBACService") as mock_rbac_service, + patch("services.account_service.tenant_was_created.send"), + ): + mock_db_dependencies["db"].session.scalar.return_value = None + + TenantService.create_owner_tenant_if_not_exist(mock_account, is_setup=True) + + mock_rbac_service.MemberRoles.replace.assert_called_once_with( + tenant_id="tenant-rbac", + account_id="user-rbac", + member_account_id="user-rbac", + role_ids=["rbac-owner-id"], + ) + def test_admin_can_update_admin_member_role(self): """Test admin can update another non-owner member, including an admin.""" mock_tenant = MagicMock()