mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 02:28:30 +08:00
fix: fix miss create rbac binding (#39023)
This commit is contained in:
parent
62bdbc8628
commit
7dfd84472f
@ -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())
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user