From a706c312b3ef58f4508a58b29e230097e28ae2e4 Mon Sep 17 00:00:00 2001
From: Bartosz Sypytkowski <b.sypytkowski@gmail.com>
Date: Mon, 24 Jun 2024 06:22:13 +0200
Subject: [PATCH] chore: making embeddings code compile

---
 services/appflowy-collaborate/src/collab/queue.rs | 5 ++++-
 src/api/workspace.rs                              | 4 ++--
 src/biz/user/user_init.rs                         | 2 ++
 tests/collab/collab_curd_test.rs                  | 3 +++
 tests/collab/storage_test.rs                      | 2 ++
 tests/sql_test/workspace_test.rs                  | 1 +
 6 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/services/appflowy-collaborate/src/collab/queue.rs b/services/appflowy-collaborate/src/collab/queue.rs
index 8dfa675c..c3ccba8f 100644
--- a/services/appflowy-collaborate/src/collab/queue.rs
+++ b/services/appflowy-collaborate/src/collab/queue.rs
@@ -122,6 +122,7 @@ impl StorageQueue {
       workspace_id: workspace_id.to_string(),
       object_id: params.object_id.clone(),
       collab_type: params.collab_type.clone(),
+      embeddings: params.embeddings.clone(),
     };
 
     self
@@ -419,7 +420,7 @@ async fn write_pending_to_disk(
           object_id: meta.object_id.clone(),
           collab_type: meta.collab_type.clone(),
           encode_collab_v1,
-          embeddings,
+          embeddings: meta.embeddings.clone(),
         })
       } else {
         None
@@ -526,6 +527,8 @@ pub struct PendingWriteMeta {
   pub workspace_id: String,
   pub object_id: String,
   pub collab_type: CollabType,
+  #[serde(default)]
+  pub embeddings: Option<AFCollabEmbeddings>,
 }
 
 impl From<&PendingWriteMeta> for QueryCollab {
diff --git a/src/api/workspace.rs b/src/api/workspace.rs
index 8b659c2f..cbf6a852 100644
--- a/src/api/workspace.rs
+++ b/src/api/workspace.rs
@@ -609,7 +609,7 @@ async fn batch_create_collab_handler(
     {
       state
         .collab_access_control_storage
-        .insert_new_collab(&workspace_id, &uid, params, None)
+        .insert_new_collab(&workspace_id, &uid, params)
         .await?;
 
       state
@@ -854,7 +854,7 @@ async fn update_collab_handler(
   let (params, workspace_id) = create_params.split();
   state
     .collab_access_control_storage
-    .insert_or_update_collab(&workspace_id, &uid, params, None, false)
+    .insert_or_update_collab(&workspace_id, &uid, params, false)
     .await?;
   Ok(AppResponse::Ok().into())
 }
diff --git a/src/biz/user/user_init.rs b/src/biz/user/user_init.rs
index e42081c3..e7c11bfc 100644
--- a/src/biz/user/user_init.rs
+++ b/src/biz/user/user_init.rs
@@ -66,6 +66,7 @@ where
           object_id,
           encoded_collab_v1,
           collab_type: template.object_type,
+          embeddings: None,
         },
         txn,
       )
@@ -104,6 +105,7 @@ async fn create_workspace_database_collab(
         object_id: object_id.to_string(),
         encoded_collab_v1,
         collab_type,
+        embeddings: None,
       },
       txn,
     )
diff --git a/tests/collab/collab_curd_test.rs b/tests/collab/collab_curd_test.rs
index f784902c..502dff19 100644
--- a/tests/collab/collab_curd_test.rs
+++ b/tests/collab/collab_curd_test.rs
@@ -87,6 +87,7 @@ async fn batch_insert_collab_success_test() {
       object_id: Uuid::new_v4().to_string(),
       encoded_collab_v1: mock_encoded_collab_v1[i].encode_to_bytes().unwrap(),
       collab_type: CollabType::Unknown,
+      embeddings: None,
     })
     .collect::<Vec<_>>();
 
@@ -182,6 +183,7 @@ async fn create_collab_compatibility_with_json_params_test() {
       object_id: object_id.clone(),
       encoded_collab_v1: encoded_collab.encode_to_bytes().unwrap(),
       collab_type: CollabType::Unknown,
+      embeddings: None,
     },
     workspace_id: workspace_id.clone(),
   };
@@ -238,6 +240,7 @@ async fn batch_create_collab_compatibility_with_uncompress_params_test() {
       object_id: object_id.clone(),
       encoded_collab_v1: encoded_collab.encode_to_bytes().unwrap(),
       collab_type: CollabType::Unknown,
+      embeddings: None,
     }],
   }
   .to_bytes()
diff --git a/tests/collab/storage_test.rs b/tests/collab/storage_test.rs
index 79d8df7a..c68eb051 100644
--- a/tests/collab/storage_test.rs
+++ b/tests/collab/storage_test.rs
@@ -468,6 +468,7 @@ async fn simulate_small_data_set_write(pool: PgPool) {
       object_id: format!("object_id_{}", i),
       collab_type: CollabType::Unknown,
       encoded_collab_v1: encode_collab.encode_to_bytes().unwrap(),
+      embeddings: None,
     };
     cloned_storage_queue
       .push(
@@ -543,6 +544,7 @@ async fn simulate_large_data_set_write(pool: PgPool) {
     object_id: uuid::Uuid::new_v4().to_string(),
     collab_type: CollabType::Unknown,
     encoded_collab_v1: origin_encode_collab.encode_to_bytes().unwrap(),
+    embeddings: None,
   };
   storage_queue
     .push(&user.workspace_id, &user.uid, &params, WritePriority::Low)
diff --git a/tests/sql_test/workspace_test.rs b/tests/sql_test/workspace_test.rs
index 41971193..544e74ca 100644
--- a/tests/sql_test/workspace_test.rs
+++ b/tests/sql_test/workspace_test.rs
@@ -29,6 +29,7 @@ async fn insert_collab_sql_test(pool: PgPool) {
       object_id,
       collab_type: CollabType::Unknown,
       encoded_collab_v1,
+      embeddings: None,
     };
     insert_into_af_collab(&mut txn, &user.uid, &user.workspace_id, &params)
       .await