From c54fcfb45dc957e089892ce60f9182aaf9945f88 Mon Sep 17 00:00:00 2001 From: Bowen Liang Date: Thu, 25 Apr 2024 18:20:08 +0800 Subject: [PATCH] extract enum type for tenant account role (#3788) --- api/controllers/console/workspace/members.py | 4 ++-- api/controllers/console/workspace/models.py | 5 +++-- api/models/account.py | 17 ++++++++++++++--- api/services/account_service.py | 6 +++--- api/services/billing_service.py | 4 ++-- api/tests/unit_tests/models/test_account.py | 12 ++++++++++++ 6 files changed, 36 insertions(+), 12 deletions(-) create mode 100644 api/tests/unit_tests/models/test_account.py diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index f40ccebf25..922b95cde5 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -9,7 +9,7 @@ from controllers.console.wraps import account_initialization_required, cloud_edi from extensions.ext_database import db from fields.member_fields import account_with_role_list_fields from libs.login import login_required -from models.account import Account +from models.account import Account, TenantAccountRole from services.account_service import RegisterService, TenantService from services.errors.account import AccountAlreadyInTenantError @@ -43,7 +43,7 @@ class MemberInviteEmailApi(Resource): invitee_emails = args['emails'] invitee_role = args['role'] interface_language = args['language'] - if invitee_role not in ['admin', 'normal']: + if invitee_role not in [TenantAccountRole.ADMIN, TenantAccountRole.NORMAL]: return {'code': 'invalid-role', 'message': 'Invalid role'}, 400 inviter = current_user diff --git a/api/controllers/console/workspace/models.py b/api/controllers/console/workspace/models.py index 5745c0d408..23239b1902 100644 --- a/api/controllers/console/workspace/models.py +++ b/api/controllers/console/workspace/models.py @@ -11,6 +11,7 @@ from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.errors.validate import CredentialsValidateFailedError from core.model_runtime.utils.encoders import jsonable_encoder from libs.login import login_required +from models.account import TenantAccountRole from services.model_provider_service import ModelProviderService @@ -94,7 +95,7 @@ class ModelProviderModelApi(Resource): @login_required @account_initialization_required def post(self, provider: str): - if current_user.current_tenant.current_role not in ['admin', 'owner']: + if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role): raise Forbidden() tenant_id = current_user.current_tenant_id @@ -125,7 +126,7 @@ class ModelProviderModelApi(Resource): @login_required @account_initialization_required def delete(self, provider: str): - if current_user.current_tenant.current_role not in ['admin', 'owner']: + if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role): raise Forbidden() tenant_id = current_user.current_tenant_id diff --git a/api/models/account.py b/api/models/account.py index 7854e3f63e..d8e587c90c 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -100,10 +100,11 @@ class Account(UserMixin, db.Model): return db.session.query(ai).filter( ai.account_id == self.id ).all() + # check current_user.current_tenant.current_role in ['admin', 'owner'] @property def is_admin_or_owner(self): - return self._current_tenant.current_role in ['admin', 'owner'] + return TenantAccountRole.is_privileged_role(self._current_tenant.current_role) class TenantStatus(str, enum.Enum): @@ -111,6 +112,16 @@ class TenantStatus(str, enum.Enum): ARCHIVE = 'archive' +class TenantAccountRole(str, enum.Enum): + OWNER = 'owner' + ADMIN = 'admin' + NORMAL = 'normal' + + @staticmethod + def is_privileged_role(role: str) -> bool: + return role and role in {TenantAccountRole.ADMIN, TenantAccountRole.OWNER} + + class Tenant(db.Model): __tablename__ = 'tenants' __table_args__ = ( @@ -132,11 +143,11 @@ class Tenant(db.Model): Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id ).all() - + @property def custom_config_dict(self) -> dict: return json.loads(self.custom_config) if self.custom_config else {} - + @custom_config_dict.setter def custom_config_dict(self, value: dict): self.custom_config = json.dumps(value) diff --git a/api/services/account_service.py b/api/services/account_service.py index 64fe3a4f0f..7551c9cb4b 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -344,9 +344,9 @@ class TenantService: def check_member_permission(tenant: Tenant, operator: Account, member: Account, action: str) -> None: """Check member permission""" perms = { - 'add': ['owner', 'admin'], - 'remove': ['owner'], - 'update': ['owner'] + 'add': [TenantAccountRole.OWNER, TenantAccountRole.ADMIN], + 'remove': [TenantAccountRole.OWNER], + 'update': [TenantAccountRole.OWNER] } if action not in ['add', 'remove', 'update']: raise InvalidActionError("Invalid action.") diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 37d362d083..4e791d075b 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -3,7 +3,7 @@ import os import requests from extensions.ext_database import db -from models.account import TenantAccountJoin +from models.account import TenantAccountJoin, TenantAccountRole class BillingService: @@ -74,5 +74,5 @@ class BillingService: TenantAccountJoin.account_id == current_user.id ).first() - if join.role not in ['owner', 'admin']: + if TenantAccountRole.is_privileged_role(join.role): raise ValueError('Only team owner or team admin can perform this action') diff --git a/api/tests/unit_tests/models/test_account.py b/api/tests/unit_tests/models/test_account.py new file mode 100644 index 0000000000..ddb4e8cb75 --- /dev/null +++ b/api/tests/unit_tests/models/test_account.py @@ -0,0 +1,12 @@ +from models.account import TenantAccountRole + + +def test_account_is_privileged_role() -> None: + assert TenantAccountRole.ADMIN == 'admin' + assert TenantAccountRole.OWNER == 'owner' + assert TenantAccountRole.NORMAL == 'normal' + + assert TenantAccountRole.is_privileged_role(TenantAccountRole.ADMIN) + assert TenantAccountRole.is_privileged_role(TenantAccountRole.OWNER) + assert not TenantAccountRole.is_privileged_role(TenantAccountRole.NORMAL) + assert not TenantAccountRole.is_privileged_role('')