From c8112601ac223083d08ad6a1a599f261fa9a4b0d Mon Sep 17 00:00:00 2001
From: GareArc <chen4851@purdue.edu>
Date: Sat, 21 Dec 2024 15:03:16 -0500
Subject: [PATCH] fix: db session error

---
 api/controllers/console/workspace/account.py |  5 +-
 api/services/account_service.py              |  1 -
 api/tasks/delete_account_task.py             | 50 ++++++++++----------
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py
index 15b4a5efe4ec99..412b8a8b212f54 100644
--- a/api/controllers/console/workspace/account.py
+++ b/api/controllers/console/workspace/account.py
@@ -22,7 +22,6 @@
 from services.account_service import AccountService
 from services.errors.account import \
     CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError
-from services.errors.account import RateLimitExceededError
 
 
 class AccountInitApi(Resource):
@@ -254,8 +253,8 @@ def get(self):
         try:
             token, code = AccountService.generate_account_deletion_verification_code(account)
             AccountService.send_account_delete_verification_email(account, code)
-        except RateLimitExceededError:
-            return {"result": "fail", "error": "Rate limit exceeded."}, 429
+        except Exception as e:
+            return {"result": "fail", "error": str(e)}, 429
 
         return {"result": "success", "data": token}
 
diff --git a/api/services/account_service.py b/api/services/account_service.py
index 69b88dc4eb5073..fe1ea9bc1b3de1 100644
--- a/api/services/account_service.py
+++ b/api/services/account_service.py
@@ -245,7 +245,6 @@ def generate_account_deletion_verification_code(account: Account) -> tuple[str,
         token = TokenManager.generate_token(
             account=account, token_type="account_deletion", additional_data={"code": code}
         )
-        logging.info(f"Account {account.id} generated account deletion verification code {code} with token {token}")
         return token, code
 
     @classmethod
diff --git a/api/tasks/delete_account_task.py b/api/tasks/delete_account_task.py
index ffea9d971e6491..bc13830f3a86ab 100644
--- a/api/tasks/delete_account_task.py
+++ b/api/tasks/delete_account_task.py
@@ -25,36 +25,38 @@ def delete_account_task(account_id, reason: str):
     tenant_account_joins = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.account_id == account.id).all()
     for ta in tenant_account_joins:
         try:
-            with db.session.begin():
-                if ta.role == TenantAccountJoinRole.OWNER:
-                    # dismiss all members of the tenant
-                    members = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.tenant_id == ta.tenant_id).delete()
-                    logging.info(f"Dismissed {members} members from tenant {ta.tenant_id}.")
+            if ta.role == TenantAccountJoinRole.OWNER:
+                # dismiss all members of the tenant
+                members = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.tenant_id == ta.tenant_id).delete()
+                logging.info(f"Dismissed {members} members from tenant {ta.tenant_id}.")
 
-                    # delete the tenant
-                    db.session.query(Tenant).filter(Tenant.id == ta.tenant_id).delete()
-                    logging.info(f"Deleted tenant {ta.tenant_id}.")
+                # delete the tenant
+                db.session.query(Tenant).filter(Tenant.id == ta.tenant_id).delete()
+                logging.info(f"Deleted tenant {ta.tenant_id}.")
 
-                    # delete subscription
-                    try:
-                        BillingService.delete_tenant_customer(ta.tenant_id)
-                    except Exception as e:
-                        logging.error(f"Failed to delete subscription for tenant {ta.tenant_id}: {e}.")
-                        raise
-                else:
-                    # remove the account from tenant
-                    db.session.delete(ta)
-                    logging.info(f"Removed account {account.email} from tenant {ta.tenant_id}.")
+                # delete subscription
+                try:
+                    BillingService.delete_tenant_customer(ta.tenant_id)
+                except Exception as e:
+                    logging.error(f"Failed to delete subscription for tenant {ta.tenant_id}: {e}.")
+                    raise
+            else:
+                # remove the account from tenant
+                db.session.delete(ta)
+                logging.info(f"Removed account {account.email} from tenant {ta.tenant_id}.")
 
-                # delete the account
-                db.session.delete(account)
+            # delete the account
+            db.session.delete(account)
 
-                # prepare account deletion log
-                account_deletion_log = AccountDeletionLogService.create_account_deletion_log(account.email, reason)
-                db.session.add(account_deletion_log)
+            # prepare account deletion log
+            account_deletion_log = AccountDeletionLogService.create_account_deletion_log(account.email, reason)
+            db.session.add(account_deletion_log)
+
+            db.session.commit()
 
         except Exception as e:
-            logging.error(f"Failed to delete account {account.email}.")
+            db.session.rollback()
+            logging.error(f"Failed to delete account {account.email}: {e}.")
             send_deletion_fail_task.delay(account.interface_language, account.email)
             return