extract enum type for tenant account role (#3788)

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

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

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__ = (

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.")

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')

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('')