From 7dfd84472f3de34764e345779b3b4a81b3986c6e Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Thu, 16 Jul 2026 10:25:42 +0800 Subject: [PATCH] fix: fix miss create rbac binding (#39023) --- .../console/auth/forgot_password.py | 6 +-- api/controllers/console/auth/login.py | 6 +-- api/controllers/console/auth/oauth.py | 6 +-- .../inner_api/workspace/workspace.py | 10 ++-- api/services/account_service.py | 46 +++++++++++++++---- .../controllers/console/auth/test_oauth.py | 11 +---- .../inner_api/workspace/test_workspace.py | 14 +++--- .../services/test_account_service.py | 17 ++----- 8 files changed, 58 insertions(+), 58 deletions(-) diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index 8ea15e1ee5a..2c04dd397ef 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -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()) diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 0497cfd03cf..e022c9dd28b 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -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: diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 46d4eff01ec..2160f3e38ec 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -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() diff --git a/api/controllers/inner_api/workspace/workspace.py b/api/controllers/inner_api/workspace/workspace.py index b3a571112f6..7767653202a 100644 --- a/api/controllers/inner_api/workspace/workspace.py +++ b/api/controllers/inner_api/workspace/workspace.py @@ -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, diff --git a/api/services/account_service.py b/api/services/account_service.py index 83716915a98..85816d36106 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -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 diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py index ced586a815c..d681bcfdce0 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py @@ -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) diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py index bda25bb2fa8..35b18ae7d20 100644 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py @@ -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): diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 763f8c0bd23..1239e1e03cb 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -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