fix: fix miss create rbac binding (#39023)

This commit is contained in:
wangxiaolei 2026-07-16 10:25:42 +08:00 committed by GitHub
parent 62bdbc8628
commit 7dfd84472f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 58 additions and 58 deletions

View File

@ -16,7 +16,6 @@ from controllers.console.auth.error import (
)
from controllers.console.error import AccountNotFound, EmailSendIpLimitError
from controllers.console.wraps import email_password_login_enabled, setup_required
from events.tenant_event import tenant_was_created
from extensions.ext_database import db
from libs.helper import EmailStr, extract_remote_ip
from libs.password import hash_password
@ -201,7 +200,4 @@ class ForgotPasswordResetApi(Resource):
not TenantService.get_join_tenants(account, session=db.session())
and FeatureService.get_system_features().is_allow_create_workspace
):
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(tenant, account, db.session(), role="owner")
account.set_current_tenant_with_session(tenant, session=db.session())
tenant_was_created.send(tenant)
TenantService.create_owner_tenant(account, session=db.session())

View File

@ -42,7 +42,6 @@ from controllers.console.wraps import (
setup_required,
with_current_user,
)
from events.tenant_event import tenant_was_created
from extensions.ext_database import db
from libs.helper import EmailStr, extract_remote_ip
from libs.helper import timezone as validate_timezone_string
@ -317,10 +316,7 @@ class EmailCodeLoginApi(Resource):
if not FeatureService.get_system_features().is_allow_create_workspace:
raise NotAllowedCreateWorkspace()
else:
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
account.set_current_tenant_with_session(new_tenant, session=db.session())
tenant_was_created.send(new_tenant)
TenantService.create_owner_tenant(account, session=db.session())
if account is None:
try:

View File

@ -11,7 +11,6 @@ from configs import dify_config
from constants.languages import languages
from controllers.common.fields import RedirectResponse
from controllers.common.schema import query_params_from_model, register_response_schema_model, register_schema_models
from events.tenant_event import tenant_was_created
from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
from libs.helper import extract_remote_ip
@ -282,10 +281,7 @@ def _generate_account(
if not FeatureService.get_system_features().is_allow_create_workspace:
raise WorkSpaceNotAllowedCreateError()
else:
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
account.set_current_tenant_with_session(new_tenant, session=db.session())
tenant_was_created.send(new_tenant)
TenantService.create_owner_tenant(account, session=db.session())
if not account:
normalized_email = user_info.email.lower()

View File

@ -47,10 +47,12 @@ class EnterpriseWorkspace(Resource):
if account is None:
return {"message": "owner account not found."}, 404
tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session())
TenantService.create_tenant_member(tenant, account, db.session(), role="owner")
tenant_was_created.send(tenant)
tenant = TenantService.create_owner_tenant(
account,
name=args.name,
is_from_dashboard=True,
session=db.session(),
)
resp = {
"id": tenant.id,

View File

@ -1297,7 +1297,7 @@ class TenantService:
def create_owner_tenant_if_not_exist(
account: Account, name: str | None = None, is_setup: bool | None = False, *, session: Session
):
"""Check if user have a workspace or not"""
"""Create an owner workspace only when the account has no membership."""
available_ta = session.scalar(
select(TenantAccountJoin)
.where(TenantAccountJoin.account_id == account.id)
@ -1308,8 +1308,28 @@ class TenantService:
if available_ta:
return
"""Create owner tenant if not exist"""
if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup:
TenantService.create_owner_tenant(account, name=name, is_setup=is_setup, session=session)
@staticmethod
def create_owner_tenant(
account: Account,
name: str | None = None,
is_setup: bool | None = False,
is_from_dashboard: bool | None = False,
*,
session: Session,
) -> Tenant:
"""Create an owner workspace and bind its owner RBAC role when enabled.
This is the single write path for a newly created workspace with an
owner. It persists the legacy membership before creating the matching
RBAC role binding, then makes the workspace current for the account.
"""
if (
not FeatureService.get_system_features().is_allow_create_workspace
and not is_setup
and not is_from_dashboard
):
raise WorkSpaceNotAllowedCreateError()
workspaces = FeatureService.get_system_features().license.workspaces
@ -1317,9 +1337,19 @@ class TenantService:
raise WorkspacesLimitExceededError()
if name:
tenant = TenantService.create_tenant(name=name, is_setup=is_setup, session=session)
tenant = TenantService.create_tenant(
name=name,
is_setup=is_setup,
is_from_dashboard=is_from_dashboard,
session=session,
)
else:
tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup, session=session)
tenant = TenantService.create_tenant(
name=f"{account.name}'s Workspace",
is_setup=is_setup,
is_from_dashboard=is_from_dashboard,
session=session,
)
TenantService.create_tenant_member(tenant, account, session, role="owner")
if dify_config.RBAC_ENABLED:
owner_role_id = AccountService._resolve_legacy_role_id(str(tenant.id), account.id, TenantAccountRole.OWNER)
@ -1333,6 +1363,7 @@ class TenantService:
account.set_current_tenant_with_session(tenant, session=session)
session.commit()
tenant_was_created.send(tenant)
return tenant
@staticmethod
def create_tenant_member(
@ -1984,10 +2015,7 @@ class RegisterService:
and FeatureService.get_system_features().license.workspaces.is_available()
):
try:
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=session)
TenantService.create_tenant_member(tenant, account, session, role="owner")
account.set_current_tenant_with_session(tenant, session=session)
tenant_was_created.send(tenant)
TenantService.create_owner_tenant(account, session=session)
except Exception:
_try_join_enterprise_default_workspace(str(account.id))
raise

View File

@ -656,10 +656,8 @@ class TestAccountGeneration:
@patch("controllers.console.auth.oauth.TenantService")
@patch("controllers.console.auth.oauth.FeatureService")
@patch("controllers.console.auth.oauth.AccountService")
@patch("controllers.console.auth.oauth.tenant_was_created")
def test_should_create_workspace_for_account_without_tenant(
self,
mock_event: MagicMock,
mock_account_service: MagicMock,
mock_feature_service: MagicMock,
mock_tenant_service: MagicMock,
@ -672,16 +670,9 @@ class TestAccountGeneration:
mock_tenant_service.get_join_tenants.return_value = []
mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
mock_new_tenant = MagicMock()
mock_tenant_service.create_tenant.return_value = mock_new_tenant
with app.test_request_context(headers={"Accept-Language": "en-US,en;q=0.9"}):
result, oauth_new_user = _generate_account("github", user_info)
assert result == mock_account
assert oauth_new_user is False
mock_tenant_service.create_tenant.assert_called_once_with("Test User's Workspace", session=ANY)
mock_tenant_service.create_tenant_member.assert_called_once_with(
mock_new_tenant, mock_account, ANY, role="owner"
)
mock_event.send.assert_called_once_with(mock_new_tenant)
mock_tenant_service.create_owner_tenant.assert_called_once_with(mock_account, session=ANY)

View File

@ -84,10 +84,9 @@ class TestEnterpriseWorkspace:
assert hasattr(api_instance, "post")
assert callable(api_instance.post)
@patch("controllers.inner_api.workspace.workspace.tenant_was_created")
@patch("controllers.inner_api.workspace.workspace.TenantService")
@patch("controllers.inner_api.workspace.workspace.db")
def test_post_creates_workspace_with_owner(self, mock_db, mock_tenant_svc, mock_event, api_instance, app: Flask):
def test_post_creates_workspace_with_owner(self, mock_db, mock_tenant_svc, api_instance, app: Flask):
"""Test that post() creates a workspace and assigns the owner account"""
# Arrange
mock_account = MagicMock()
@ -102,7 +101,7 @@ class TestEnterpriseWorkspace:
mock_tenant.status = TenantStatus.NORMAL
mock_tenant.created_at = now
mock_tenant.updated_at = now
mock_tenant_svc.create_tenant.return_value = mock_tenant
mock_tenant_svc.create_owner_tenant.return_value = mock_tenant
# Act — unwrap to bypass auth/setup decorators (tested in test_auth_wraps.py)
unwrapped_post = inspect.unwrap(api_instance.post)
@ -115,11 +114,12 @@ class TestEnterpriseWorkspace:
assert result["message"] == "enterprise workspace created."
assert result["tenant"]["id"] == "tenant-id"
assert result["tenant"]["name"] == "My Workspace"
mock_tenant_svc.create_tenant.assert_called_once_with("My Workspace", is_from_dashboard=True, session=ANY)
mock_tenant_svc.create_tenant_member.assert_called_once_with(
mock_tenant, mock_account, mock_db.session(), role="owner"
mock_tenant_svc.create_owner_tenant.assert_called_once_with(
mock_account,
name="My Workspace",
is_from_dashboard=True,
session=ANY,
)
mock_event.send.assert_called_once_with(mock_tenant)
@patch("controllers.inner_api.workspace.workspace.db")
def test_post_returns_404_when_owner_not_found(self, mock_db, api_instance, app: Flask):

View File

@ -1008,7 +1008,7 @@ class TestTenantService:
assert target_join.role == TenantAccountRole.ADMIN
@pytest.mark.parametrize("sqlite_session", [(TenantAccountJoin,)], indirect=True)
def test_create_owner_tenant_if_not_exist_rbac_enabled_assigns_owner_role(
def test_create_owner_tenant_rbac_enabled_assigns_owner_role(
self, sqlite_session: Session, mock_external_service_dependencies
):
mock_account = TestAccountAssociatedDataFactory.create_account_mock(account_id="user-rbac", name="RBAC User")
@ -1034,7 +1034,7 @@ class TestTenantService:
patch("services.account_service.RBACService") as mock_rbac_service,
patch("services.account_service.tenant_was_created.send"),
):
TenantService.create_owner_tenant_if_not_exist(mock_account, is_setup=True, session=sqlite_session)
TenantService.create_owner_tenant(mock_account, is_setup=True, session=sqlite_session)
mock_rbac_service.MemberRoles.replace.assert_called_once_with(
tenant_id="tenant-rbac",
@ -1476,16 +1476,9 @@ class TestRegisterService:
with patch("services.account_service.AccountService.create_account") as mock_create_account:
mock_create_account.return_value = mock_account
# Mock TenantService.create_tenant and create_tenant_member
with (
patch("services.account_service.TenantService.create_tenant") as mock_create_tenant,
patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
patch("services.account_service.tenant_was_created") as mock_event,
patch("services.account_service.TenantService.create_owner_tenant") as mock_create_owner_tenant,
):
mock_tenant = MagicMock()
mock_tenant.id = "tenant-456"
mock_create_tenant.return_value = mock_tenant
# Execute test
result = RegisterService.register(
email="test@example.com",
@ -1508,9 +1501,7 @@ class TestRegisterService:
timezone=None,
session=sqlite_session,
)
mock_create_tenant.assert_called_once_with("Test User's Workspace", session=sqlite_session)
mock_create_member.assert_called_once_with(mock_tenant, mock_account, sqlite_session, role="owner")
mock_event.send.assert_called_once_with(mock_tenant)
mock_create_owner_tenant.assert_called_once_with(mock_account, session=sqlite_session)
def test_register_calls_default_workspace_join_when_enterprise_enabled(
self, sqlite_session: Session, mock_external_service_dependencies, monkeypatch: pytest.MonkeyPatch