From 6d0a7cc95b19779efbb8c29961bef34a4c4c2630 Mon Sep 17 00:00:00 2001
From: Zack Fu Zi Xiang <speed2exe@live.com.sg>
Date: Wed, 4 Sep 2024 12:00:34 +0800
Subject: [PATCH] fix: delete user should only remove owned workspaces

---
 ...4cce9a0dabccfff03bea125dfb729a9de2033.json | 22 ++++++++++++++++++
 libs/database/src/workspace.rs                | 19 +++++++++++++++
 src/biz/user/user_delete.rs                   | 10 ++++----
 tests/user/delete.rs                          | 23 +++++++++++++++++++
 4 files changed, 69 insertions(+), 5 deletions(-)
 create mode 100644 .sqlx/query-57ef3a6399630ec23dd238a97eb4cce9a0dabccfff03bea125dfb729a9de2033.json

diff --git a/.sqlx/query-57ef3a6399630ec23dd238a97eb4cce9a0dabccfff03bea125dfb729a9de2033.json b/.sqlx/query-57ef3a6399630ec23dd238a97eb4cce9a0dabccfff03bea125dfb729a9de2033.json
new file mode 100644
index 00000000..887c0a79
--- /dev/null
+++ b/.sqlx/query-57ef3a6399630ec23dd238a97eb4cce9a0dabccfff03bea125dfb729a9de2033.json
@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n      SELECT workspace_id\n      FROM af_workspace\n      where OWNER_UID = (SELECT uid FROM public.af_user WHERE uuid = $1)\n    ",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "workspace_id",
+        "type_info": "Uuid"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid"
+      ]
+    },
+    "nullable": [
+      false
+    ]
+  },
+  "hash": "57ef3a6399630ec23dd238a97eb4cce9a0dabccfff03bea125dfb729a9de2033"
+}
diff --git a/libs/database/src/workspace.rs b/libs/database/src/workspace.rs
index 2b638f4f..870a5932 100644
--- a/libs/database/src/workspace.rs
+++ b/libs/database/src/workspace.rs
@@ -691,6 +691,25 @@ pub async fn select_all_user_workspaces<'a, E: Executor<'a, Database = Postgres>
   Ok(workspaces)
 }
 
+/// Returns a list of workspace ids that the user is owner of.
+#[inline]
+pub async fn select_user_owned_workspaces_id<'a, E: Executor<'a, Database = Postgres>>(
+  executor: E,
+  user_uuid: &Uuid,
+) -> Result<Vec<Uuid>, AppError> {
+  let workspace_ids = sqlx::query_scalar!(
+    r#"
+      SELECT workspace_id
+      FROM af_workspace
+      where OWNER_UID = (SELECT uid FROM public.af_user WHERE uuid = $1)
+    "#,
+    user_uuid
+  )
+  .fetch_all(executor)
+  .await?;
+  Ok(workspace_ids)
+}
+
 pub async fn select_member_count_for_workspaces<'a, E: Executor<'a, Database = Postgres>>(
   executor: E,
   workspace_ids: &[Uuid],
diff --git a/src/biz/user/user_delete.rs b/src/biz/user/user_delete.rs
index 01bd6ec8..84fb905e 100644
--- a/src/biz/user/user_delete.rs
+++ b/src/biz/user/user_delete.rs
@@ -1,11 +1,11 @@
 use std::sync::Arc;
 
-use crate::biz::workspace::ops::get_all_user_workspaces;
 use crate::state::GoTrueAdmin;
 use crate::{biz::workspace::ops::delete_workspace_for_user, config::config::AppleOAuthSetting};
 use app_error::ErrorCode;
 use authentication::jwt::Authorization;
 use database::file::s3_client_impl::S3BucketStorage;
+use database::workspace::select_user_owned_workspaces_id;
 use gotrue::params::AdminDeleteUserParams;
 use secrecy::{ExposeSecret, Secret};
 use shared_entity::response::AppResponseError;
@@ -48,14 +48,14 @@ pub async fn delete_user(
     .await
     .map_err(AppResponseError::from)?;
 
-  // spawn tasks to delete all user workspace and object storage
-  let user_workspaces = get_all_user_workspaces(pg_pool, &user_uuid, false).await?;
+  // spawn tasks to delete all workspaces owned by the user
+  let workspace_ids = select_user_owned_workspaces_id(pg_pool, &user_uuid).await?;
   let mut tasks = vec![];
-  for workspace in user_workspaces {
+  for workspace_id in workspace_ids {
     let cloned_pg_pool = pg_pool.clone();
     tasks.push(tokio::spawn(delete_workspace_for_user(
       cloned_pg_pool,
-      workspace.workspace_id,
+      workspace_id,
       bucket_storage.clone(),
     )));
   }
diff --git a/tests/user/delete.rs b/tests/user/delete.rs
index fe451a52..550dcf8d 100644
--- a/tests/user/delete.rs
+++ b/tests/user/delete.rs
@@ -1,3 +1,4 @@
+use client_api::entity::AFRole;
 use client_api_test::*;
 use gotrue::params::{AdminDeleteUserParams, AdminUserParams};
 
@@ -30,6 +31,28 @@ async fn user_delete_self() {
   }
 }
 
+/// Scenario:
+/// - User1 owns WorkspaceA
+/// - User1 invites User2 to WorkspaceA
+/// - User2 deletes itself
+/// - WorkspaceA should still exist
+#[tokio::test]
+async fn user_delete_self_shared_workspace() {
+  let user_1 = TestClient::new_user_without_ws_conn().await;
+  let workspace_a = user_1.workspace_id().await;
+  let user_2 = TestClient::new_user_without_ws_conn().await;
+  user_1
+    .invite_and_accepted_workspace_member(&workspace_a, &user_2, AFRole::Member)
+    .await
+    .unwrap();
+  user_2.api_client.delete_user().await.unwrap();
+  let user_1_workspaces = user_1.api_client.get_workspaces().await.unwrap();
+  let _workspace_a = user_1_workspaces
+    .into_iter()
+    .find(|w| w.workspace_id.to_string() == workspace_a)
+    .unwrap();
+}
+
 #[tokio::test]
 async fn admin_delete_create_same_user_hard() {
   let (client, user) = generate_unique_registered_user_client().await;