test: switch account

api/controllers/console/auth/login.py

@@ -30,10 +30,10 @@ class LoginApi(Resource):
         except services.errors.account.AccountLoginError:
             return {'code': 'unauthorized', 'message': 'Invalid email or password'}, 401
 
-        try:
-            TenantService.switch_tenant(account)
-        except Exception:
-            pass
+        # try:
+        #     TenantService.switch_tenant(account)
+        # except Exception:
+        #     pass
 
         AccountService.update_last_login(account, request)
 

api/migrations/versions/16830a790f0f_.py

+"""empty message
+
+Revision ID: 16830a790f0f
+Revises: 380c6aa5a70d
+Create Date: 2024-02-01 08:21:31.111119
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '16830a790f0f'
+down_revision = '380c6aa5a70d'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tenant_account_joins', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('current', sa.Boolean(), server_default=sa.text('false'), nullable=False))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tenant_account_joins', schema=None) as batch_op:
+        batch_op.drop_column('current')
+
+    # ### end Alembic commands ###

api/models/account.py

@@ -155,6 +155,7 @@ class TenantAccountJoin(db.Model):
     id = db.Column(UUID, server_default=db.text('uuid_generate_v4()'))
     tenant_id = db.Column(UUID, nullable=False)
     account_id = db.Column(UUID, nullable=False)
+    current = db.Column(db.Boolean, nullable=False, server_default=db.text('false'))
     role = db.Column(db.String(16), nullable=False, server_default='normal')
     invited_by = db.Column(UUID, nullable=True)
     created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))

api/services/account_service.py

@@ -45,11 +45,20 @@ class AccountService:
 
         if account.status in [AccountStatus.BANNED.value, AccountStatus.CLOSED.value]:
             raise Forbidden('Account is banned or closed.')
-
-        tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id).first()
-        if not tenant_account_join:
+        
+        # init owner's tenant
+        tenant_owner = TenantAccountJoin.query.filter_by(account_id=account.id, role='owner').first()
+        if not tenant_owner:
             _create_tenant_for_account(account)
-        # Update last_active_at if more than 10 minutes have passed
+        
+        current_tenant = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
+        if current_tenant:
+            account.current_tenant_id = current_tenant.tenant_id
+        else:
+            account.current_tenant_id = tenant_owner.tenant_id
+            tenant_owner.current = True
+            db.session.commit()
+       
         if datetime.utcnow() - account.last_active_at > timedelta(minutes=10):
             account.last_active_at = datetime.utcnow()
             db.session.commit()
@@ -247,14 +256,16 @@ class TenantService:
     @staticmethod
     def switch_tenant(account: Account, tenant_id: int = None) -> None:
         """Switch the current workspace for the account"""
-        if not tenant_id:
-            tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id).first()
-        else:
-            tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id, tenant_id=tenant_id).first()
+
+        TenantAccountJoin.query.filter_by(account_id=account.id).update({'current': False})
+        tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id, tenant_id=tenant_id).first()
 
         # Check if the tenant exists and the account is a member of the tenant
         if not tenant_account_join:
             raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.")
+        else:
+            tenant_account_join.current = True
+            db.session.commit()
 
         # Set the current tenant for the account
         account.current_tenant_id = tenant_account_join.tenant_id