From 2028f4bffd1c3c06b2968fb583d5e99496354557 Mon Sep 17 00:00:00 2001
From: Zack <speed2exe@live.com.sg>
Date: Mon, 11 Nov 2024 16:02:03 +0800
Subject: [PATCH] feat: use batch insert collab for workspace creation (#977)

* feat: use batch insert collab for workspace creation

* fix: separate user creation from initial workspace creation

* feat: separate workspace creation and collab initialization
---
 libs/database/src/workspace.rs |  4 ++--
 src/biz/user/user_init.rs      | 26 +++++++++++---------------
 src/biz/user/user_verify.rs    | 19 ++++++++++++++-----
 src/biz/workspace/ops.rs       | 10 +++++-----
 4 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/libs/database/src/workspace.rs b/libs/database/src/workspace.rs
index ca50c218..2f033755 100644
--- a/libs/database/src/workspace.rs
+++ b/libs/database/src/workspace.rs
@@ -42,7 +42,7 @@ pub async fn delete_from_workspace(pg_pool: &PgPool, workspace_id: &Uuid) -> Res
 
 #[inline]
 pub async fn insert_user_workspace(
-  tx: &mut Transaction<'_, sqlx::Postgres>,
+  pg_pool: &PgPool,
   user_uuid: &Uuid,
   workspace_name: &str,
   is_initialized: bool,
@@ -73,7 +73,7 @@ pub async fn insert_user_workspace(
     workspace_name,
     is_initialized,
   )
-  .fetch_one(tx.deref_mut())
+  .fetch_one(pg_pool)
   .await?;
 
   Ok(workspace)
diff --git a/src/biz/user/user_init.rs b/src/biz/user/user_init.rs
index d16ee380..2b18ebc7 100644
--- a/src/biz/user/user_init.rs
+++ b/src/biz/user/user_init.rs
@@ -37,6 +37,7 @@ where
     .await?;
 
   let mut database_records = vec![];
+  let mut collab_params = Vec::with_capacity(templates.len());
   for template in templates {
     let template_id = template.template_id;
     let (view_id, object_id) = match &template_id {
@@ -53,21 +54,12 @@ where
       .encoded_collab
       .encode_to_bytes()
       .map_err(|err| AppError::Internal(anyhow::Error::from(err)))?;
-
-    collab_storage
-      .insert_new_collab_with_transaction(
-        &workspace_id,
-        &uid,
-        CollabParams {
-          object_id: object_id.clone(),
-          encoded_collab_v1: encoded_collab_v1.into(),
-          collab_type: object_type.clone(),
-          embeddings: None,
-        },
-        txn,
-        "initialize workspace for user",
-      )
-      .await?;
+    collab_params.push(CollabParams {
+      object_id: object_id.clone(),
+      encoded_collab_v1: encoded_collab_v1.into(),
+      collab_type: object_type.clone(),
+      embeddings: None,
+    });
 
     // push the database record
     if object_type == CollabType::Database {
@@ -81,6 +73,10 @@ where
     }
   }
 
+  collab_storage
+    .batch_insert_new_collab(&workspace_id, &uid, collab_params)
+    .await?;
+
   // Create a workspace database object for given user
   // The database_storage_id is auto-generated when the workspace is created. So, it should be available
   if let Some(database_storage_id) = row.database_storage_id.as_ref() {
diff --git a/src/biz/user/user_verify.rs b/src/biz/user/user_verify.rs
index 602e83d1..4896db3a 100644
--- a/src/biz/user/user_verify.rs
+++ b/src/biz/user/user_verify.rs
@@ -22,6 +22,7 @@ pub async fn verify_token(access_token: &str, state: &AppState) -> Result<bool,
   let user_uuid = uuid::Uuid::parse_str(&user.id)?;
   let name = name_from_user_metadata(&user.user_metadata);
 
+  // Create new user if it doesn't exist
   let mut txn = state
     .pg_pool
     .begin()
@@ -41,24 +42,32 @@ pub async fn verify_token(access_token: &str, state: &AppState) -> Result<bool,
       .workspace_access_control
       .insert_role(&new_uid, &workspace_id, AFRole::Owner)
       .await?;
+    // Need to commit the transaction for the record in `af_user` to be inserted
+    // so that `initialize_workspace_for_user` will be able to find the user
+    txn
+      .commit()
+      .await
+      .context("fail to commit transaction to verify token")?;
 
     // Create a workspace with the GetStarted template
+    let mut txn2 = state.pg_pool.begin().await?;
     initialize_workspace_for_user(
       new_uid,
       &user_uuid,
       &workspace_row,
-      &mut txn,
+      &mut txn2,
       vec![GettingStartedTemplate],
       &state.collab_access_control_storage,
     )
     .await?;
+    txn2
+      .commit()
+      .await
+      .context("fail to commit transaction to initialize workspace")?;
   } else {
     trace!("user already exists:{},{}", user.id, user.email);
   }
-  txn
-    .commit()
-    .await
-    .context("fail to commit transaction to verify token")?;
+
   Ok(is_new)
 }
 
diff --git a/src/biz/workspace/ops.rs b/src/biz/workspace/ops.rs
index 249add08..67419b23 100644
--- a/src/biz/workspace/ops.rs
+++ b/src/biz/workspace/ops.rs
@@ -77,14 +77,14 @@ pub async fn create_empty_workspace(
   user_uid: i64,
   workspace_name: &str,
 ) -> Result<AFWorkspace, AppResponseError> {
-  let mut txn = pg_pool.begin().await?;
-  let new_workspace_row = insert_user_workspace(&mut txn, user_uuid, workspace_name, false).await?;
+  let new_workspace_row = insert_user_workspace(pg_pool, user_uuid, workspace_name, false).await?;
   workspace_access_control
     .insert_role(&user_uid, &new_workspace_row.workspace_id, AFRole::Owner)
     .await?;
   let workspace_id = new_workspace_row.workspace_id.to_string();
 
   // create CollabType::Folder
+  let mut txn = pg_pool.begin().await?;
   create_workspace_collab(
     user_uid,
     &workspace_id,
@@ -130,14 +130,14 @@ pub async fn create_workspace_for_user(
   user_uid: i64,
   workspace_name: &str,
 ) -> Result<AFWorkspace, AppResponseError> {
-  let mut txn = pg_pool.begin().await?;
-  let new_workspace_row = insert_user_workspace(&mut txn, user_uuid, workspace_name, true).await?;
+  let new_workspace_row = insert_user_workspace(pg_pool, user_uuid, workspace_name, true).await?;
 
   workspace_access_control
     .insert_role(&user_uid, &new_workspace_row.workspace_id, AFRole::Owner)
     .await?;
 
   // add create initial collab for user
+  let mut txn = pg_pool.begin().await?;
   initialize_workspace_for_user(
     user_uid,
     user_uuid,
@@ -147,9 +147,9 @@ pub async fn create_workspace_for_user(
     collab_storage,
   )
   .await?;
+  txn.commit().await?;
 
   let new_workspace = AFWorkspace::try_from(new_workspace_row)?;
-  txn.commit().await?;
   Ok(new_workspace)
 }