From cda20785eea345df9488021e3b8a4ed056bb3f7e Mon Sep 17 00:00:00 2001
From: GareArc <chen4851@purdue.edu>
Date: Sat, 21 Dec 2024 23:11:43 -0500
Subject: [PATCH] fix: refactor task

---
 api/tasks/delete_account_task.py | 114 ++++++++++++++++++-------------
 1 file changed, 68 insertions(+), 46 deletions(-)

diff --git a/api/tasks/delete_account_task.py b/api/tasks/delete_account_task.py
index bc13830f3a86ab..b6599b7395ddc9 100644
--- a/api/tasks/delete_account_task.py
+++ b/api/tasks/delete_account_task.py
@@ -17,53 +17,75 @@
 @shared_task(queue="dataset")
 def delete_account_task(account_id, reason: str):
     account = db.session.query(Account).filter(Account.id == account_id).first()
-    logger.info(click.style("Start delete account task.", fg="green"))
+    if not account:
+        logging.error(f"Account with ID {account_id} not found.")
+        return
+
+    logger.info(click.style(f"Start deletion task for account {account.email}.", fg="green"))
     start_at = time.perf_counter()
 
-    logger.info(f"Start deletion of account {account.email}.")
-    # find all tenants this account belongs to
-    tenant_account_joins = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.account_id == account.id).all()
-    for ta in tenant_account_joins:
-        try:
-            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 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)
-
-            # 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:
-            db.session.rollback()
-            logging.error(f"Failed to delete account {account.email}: {e}.")
-            send_deletion_fail_task.delay(account.interface_language, account.email)
-            return
-
-    send_deletion_success_task.delay(account.interface_language, account.email)
-    end_at = time.perf_counter()
-    logging.info(
-        click.style(
-            "Account deletion task completed: latency: {}".format(end_at - start_at), fg="green"
+    try:
+        _process_account_deletion(account, reason)
+        db.session.commit()
+        send_deletion_success_task.delay(account.interface_language, account.email)
+        logger.info(
+            click.style(
+                f"Account deletion task completed for {account.email}: latency: {time.perf_counter() - start_at}",
+                fg="green",
+            )
         )
+    except Exception as e:
+        db.session.rollback()
+        logging.error(f"Failed to delete account {account.email}: {e}.")
+        send_deletion_fail_task.delay(account.interface_language, account.email)
+        raise
+
+
+def _process_account_deletion(account, reason):
+    # Fetch all tenant-account associations
+    tenant_account_joins = db.session.query(TenantAccountJoin).filter(
+        TenantAccountJoin.account_id == account.id
+    ).all()
+
+    for ta in tenant_account_joins:
+        if ta.role == TenantAccountJoinRole.OWNER:
+            _handle_owner_tenant_deletion(ta)
+        else:
+            _remove_account_from_tenant(ta, account.email)
+
+    account_deletion_log = AccountDeletionLogService.create_account_deletion_log(
+        account.email, reason
     )
+    db.session.add(account_deletion_log)
+    db.session.delete(account)
+    logger.info(f"Account {account.email} successfully deleted.")
+
+
+def _handle_owner_tenant_deletion(ta):
+    """Handle deletion of a tenant where the account is an owner."""
+    tenant_id = ta.tenant_id
+
+    # Dismiss all tenant members
+    members_deleted = db.session.query(TenantAccountJoin).filter(
+        TenantAccountJoin.tenant_id == tenant_id
+    ).delete()
+    logger.info(f"Dismissed {members_deleted} members from tenant {tenant_id}.")
+
+    # Delete the tenant
+    db.session.query(Tenant).filter(Tenant.id == tenant_id).delete()
+    logger.info(f"Deleted tenant {tenant_id}.")
+
+    # Delete subscription
+    try:
+        BillingService.delete_tenant_customer(tenant_id)
+        logger.info(f"Subscription for tenant {tenant_id} deleted successfully.")
+    except Exception as e:
+        logger.error(f"Failed to delete subscription for tenant {tenant_id}: {e}.")
+        raise
+
+
+def _remove_account_from_tenant(ta, email):
+    """Remove the account from a tenant."""
+    tenant_id = ta.tenant_id
+    db.session.delete(ta)
+    logger.info(f"Removed account {email} from tenant {tenant_id}.")