fix: account has no owner workspace by member inviting (langgenius#2435)

api/controllers/console/auth/login.py

@@ -7,7 +7,7 @@
 from controllers.console.setup import setup_required
 from libs.helper import email
 from libs.password import valid_password
-from services.account_service import AccountService
+from services.account_service import AccountService, TenantService
 
 
 class LoginApi(Resource):
@@ -29,6 +29,8 @@ def post(self):
         except services.errors.account.AccountLoginError:
             return {'code': 'unauthorized', 'message': 'Invalid email or password'}, 401
 
+        TenantService.create_owner_tenant_if_not_exist(account)
+
         AccountService.update_last_login(account, request)
 
         # todo: return the user info

api/controllers/console/auth/oauth.py

@@ -10,7 +10,7 @@
 from extensions.ext_database import db
 from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo
 from models.account import Account, AccountStatus
-from services.account_service import AccountService, RegisterService
+from services.account_service import AccountService, RegisterService, TenantService
 
 from .. import api
 
@@ -76,6 +76,8 @@ def get(self, provider: str):
             account.initialized_at = datetime.utcnow()
             db.session.commit()
 
+        TenantService.create_owner_tenant_if_not_exist(account)
+
         AccountService.update_last_login(account, request)
 
         token = AccountService.get_account_jwt_token(account)

api/controllers/console/workspace/members.py

@@ -11,6 +11,7 @@
 from libs.login import login_required
 from models.account import Account
 from services.account_service import RegisterService, TenantService
+from services.errors.account import AccountAlreadyInTenantError
 
 account_fields = {
     'id': fields.String,
@@ -71,6 +72,13 @@ def post(self):
                     'email': invitee_email,
                     'url': f'{console_web_url}/activate?email={invitee_email}&token={token}'
                 })
+            except AccountAlreadyInTenantError:
+                invitation_results.append({
+                    'status': 'success',
+                    'email': invitee_email,
+                    'url': f'{console_web_url}/signin'
+                })
+                break
             except Exception as e:
                 invitation_results.append({
                     'status': 'failed',

api/services/account_service.py

@@ -1,5 +1,4 @@
 import base64
-import json
 import logging
 import secrets
 import uuid
@@ -36,15 +35,6 @@
 from tasks.mail_invite_member_task import send_invite_member_mail_task
 
 
-def _create_tenant_for_account(account) -> Tenant:
-    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
-
-    TenantService.create_tenant_member(tenant, account, role='owner')
-    account.current_tenant = tenant
-
-    return tenant
-
-
 class AccountService:
 
     @staticmethod
@@ -59,15 +49,14 @@ def load_user(user_id: str) -> Account:
         current_tenant = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
         if current_tenant:
             account.current_tenant_id = current_tenant.tenant_id
-            account.current_tenant_id = current_tenant.tenant_id
         else:
-            available_tenant = TenantAccountJoin.query.filter_by(account_id=account.id) \
+            available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
                 .order_by(TenantAccountJoin.id.asc()).first()
-            if not available_tenant:
-                raise Forbidden('No available tenant for the user.')
+            if not available_ta:
+                return None
 
-            account.current_tenant_id = available_tenant.tenant_id
-            available_tenant.current = True
+            account.current_tenant_id = available_ta.tenant_id
+            available_ta.current = True
             db.session.commit()
 
         if datetime.utcnow() - account.last_active_at > timedelta(minutes=10):
@@ -226,6 +215,21 @@ def create_tenant(name: str) -> Tenant:
         db.session.commit()
         return tenant
 
+    @staticmethod
+    def create_owner_tenant_if_not_exist(account: Account):
+        """Create owner tenant if not exist"""
+        available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
+            .order_by(TenantAccountJoin.id.asc()).first()
+
+        if available_ta:
+            return
+
+        tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+        TenantService.create_tenant_member(tenant, account, role='owner')
+        account.current_tenant = tenant
+        db.session.commit()
+        tenant_was_created.send(tenant)
+
     @staticmethod
     def create_tenant_member(tenant: Tenant, account: Account, role: str = 'normal') -> TenantAccountJoin:
         """Create tenant member"""
@@ -362,12 +366,6 @@ def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Accoun
             raise MemberNotInTenantError("Member not in tenant.")
 
         db.session.delete(ta)
-
-        account.initialized_at = None
-        account.status = AccountStatus.PENDING.value
-        account.password = None
-        account.password_salt = None
-
         db.session.commit()
 
     @staticmethod
@@ -418,12 +416,18 @@ def _get_invitation_token_key(cls, token: str) -> str:
         return f'member_invite:token:{token}'
 
     @classmethod
-    def register(cls, email, name, password: str = None, open_id: str = None, provider: str = None) -> Account:
+    def register(cls, email, name, password: str = None, open_id: str = None, provider: str = None,
+                 language: str = None, status: AccountStatus = None) -> Account:
         db.session.begin_nested()
         """Register account"""
         try:
-            account = AccountService.create_account(email, name, interface_language=languages[0], password=password)
-            account.status = AccountStatus.ACTIVE.value
+            account = AccountService.create_account(
+                email=email,
+                name=name,
+                interface_language=language if language else languages[0],
+                password=password
+            )
+            account.status = AccountStatus.ACTIVE.value if not status else status.value
             account.initialized_at = datetime.utcnow()
 
             if open_id is not None or provider is not None:
@@ -452,11 +456,12 @@ def invite_new_member(cls, tenant: Tenant, email: str, language: str, role: str
         if not account:
             TenantService.check_member_permission(tenant, inviter, None, 'add')
             name = email.split('@')[0]
-            account = AccountService.create_account(email, name, interface_language=language)
-            account.status = AccountStatus.PENDING.value
-            db.session.commit()
 
+            account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
+
+            # Create new tenant member for invited tenant
             TenantService.create_tenant_member(tenant, account, role)
+            TenantService.switch_tenant(account, tenant.id)
         else:
             TenantService.check_member_permission(tenant, inviter, account, 'add')
             ta = TenantAccountJoin.query.filter_by(