From c4017370bf52f95e18343611b44e9ab01af4a480 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Tue, 2 Jan 2024 23:09:46 -0500
Subject: [PATCH] return bookmark count after adding or removing a bookmark

---
 .../get_total_bookmark_count_for_entity.rs    | 61 +++++++++++++++++++
 .../src/queries/user_bookmarks/mod.rs         |  9 +--
 .../create_user_bookmark_handler.rs           | 12 ++++
 3 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs

diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs b/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs
new file mode 100644
index 0000000000..214686c8c4
--- /dev/null
+++ b/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs
@@ -0,0 +1,61 @@
+use anyhow::anyhow;
+use log::warn;
+use sqlx::{Executor, MySql};
+use sqlx::pool::PoolConnection;
+
+use errors::AnyhowResult;
+
+use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken;
+
+pub struct BookmarkCount {
+  pub total_count: usize,
+}
+
+pub async fn get_total_bookmark_count_for_entity(
+  user_bookmark_entity_token: &UserBookmarkEntityToken,
+  mysql_connection: &mut PoolConnection<MySql>,
+) -> AnyhowResult<BookmarkCount> {
+
+  let (entity_type, entity_token) = user_bookmark_entity_token.get_composite_keys();
+
+  let maybe_result = sqlx::query_as!(
+      RawRecord,
+        r#"
+SELECT
+    COUNT(*) as total_count
+FROM
+    user_bookmarks
+WHERE
+    entity_type = ?
+    AND entity_token = ?
+    AND deleted_at IS NULL
+        "#,
+      entity_type,
+      entity_token
+    )
+      .fetch_one(&mut **mysql_connection)
+      .await;
+
+  match maybe_result {
+    Err(err) => match err {
+      sqlx::Error::RowNotFound => Ok(BookmarkCount { total_count: 0 }),
+      _ => {
+        warn!("get_total_bookmark_count db error: {:?}", err);
+        Err(anyhow!("error with query: {:?}", err))
+      }
+    },
+    Ok(result) => Ok(result.into_public_type()),
+  }
+}
+
+pub struct RawRecord {
+  total_count: i64,
+}
+
+impl RawRecord {
+  pub fn into_public_type(self) -> BookmarkCount {
+    BookmarkCount {
+      total_count: self.total_count as usize,
+    }
+  }
+}
diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs b/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs
index df6d7f415c..da3ed97d84 100644
--- a/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs
+++ b/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs
@@ -1,7 +1,8 @@
-pub mod delete_user_bookmark;
 pub mod create_user_bookmark;
-pub mod user_bookmark_entity_token;
+pub mod delete_user_bookmark;
+pub mod get_total_bookmark_count_for_entity;
 pub mod get_user_bookmark;
-pub mod list_user_bookmarks_for_entity;
 pub mod list_user_bookmarks;
-pub mod list_user_bookmarks_result;
\ No newline at end of file
+pub mod list_user_bookmarks_for_entity;
+pub mod list_user_bookmarks_result;
+pub mod user_bookmark_entity_token;
\ No newline at end of file
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs
index 62709c8fae..f08114fde1 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs
@@ -14,6 +14,7 @@ use utoipa::ToSchema;
 
 use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType;
 use mysql_queries::queries::user_bookmarks::create_user_bookmark::{create_user_bookmark, CreateUserBookmarkArgs};
+use mysql_queries::queries::user_bookmarks::get_total_bookmark_count_for_entity::get_total_bookmark_count_for_entity;
 use mysql_queries::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
@@ -39,6 +40,9 @@ pub struct CreateUserBookmarkRequest {
 pub struct CreateUserBookmarkSuccessResponse {
   pub success: bool,
   pub user_bookmark_token: UserBookmarkToken,
+
+  /// This is the new bookmark count (across all users) for the entity in question.
+  pub new_bookmark_count_for_entity: usize,
 }
 
 #[derive(Debug, ToSchema)]
@@ -155,9 +159,17 @@ pub async fn create_user_bookmark_handler(
     }
   };
 
+  let count = get_total_bookmark_count_for_entity(&entity_token, &mut mysql_connection)
+      .await
+      .map_err(|err| {
+        warn!("error getting updated bookmark count: {:?}", err);
+        CreateUserBookmarkError::ServerError
+      })?;
+
   let response = CreateUserBookmarkSuccessResponse {
     success: true,
     user_bookmark_token,
+    new_bookmark_count_for_entity: count.total_count,
   };
 
   let body = serde_json::to_string(&response)