From 00485339149c9cd934584293893262da5c8fa977 Mon Sep 17 00:00:00 2001
From: Zack Fu Zi Xiang <speed2exe@live.com.sg>
Date: Thu, 14 Mar 2024 17:50:19 +0800
Subject: [PATCH] feat: total document usage

---
 admin_frontend/src/web_app.rs  |  2 --
 libs/database/src/workspace.rs | 22 ++++++++++++++++++++++
 src/api/workspace.rs           |  9 ++++++---
 src/biz/workspace/ops.rs       | 22 +++++++++++++++++++++-
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/admin_frontend/src/web_app.rs b/admin_frontend/src/web_app.rs
index 21b43fc3..67ddaade 100644
--- a/admin_frontend/src/web_app.rs
+++ b/admin_frontend/src/web_app.rs
@@ -2,14 +2,12 @@ use crate::error::WebAppError;
 use crate::ext::api::{
   get_user_workspace_count, get_user_workspace_limit, get_user_workspace_usages,
 };
-use crate::ext::entities::WorkspaceUsageLimits;
 use crate::session::UserSession;
 use askama::Template;
 use axum::extract::{Path, State};
 use axum::response::Result;
 use axum::{response::Html, routing::get, Router};
 use gotrue_entity::dto::User;
-use human_bytes::human_bytes;
 
 use crate::{templates, AppState};
 
diff --git a/libs/database/src/workspace.rs b/libs/database/src/workspace.rs
index 4c15d827..5f611c83 100644
--- a/libs/database/src/workspace.rs
+++ b/libs/database/src/workspace.rs
@@ -636,3 +636,25 @@ pub async fn select_permission_from_role_id(
   .await?;
   Ok(permission)
 }
+
+pub async fn select_workspace_total_collab_bytes(
+  pool: &PgPool,
+  workspace_id: &Uuid,
+) -> Result<i64, AppError> {
+  let sum = sqlx::query_scalar!(
+    r#"
+    SELECT SUM(len) FROM af_collab WHERE workspace_id = $1
+    "#,
+    workspace_id
+  )
+  .fetch_one(pool)
+  .await?;
+
+  match sum {
+    Some(s) => Ok(s),
+    None => Err(AppError::RecordNotFound(format!(
+      "Failed to get total collab bytes for workspace_id: {}",
+      workspace_id
+    ))),
+  }
+}
diff --git a/src/api/workspace.rs b/src/api/workspace.rs
index 33171b2c..b6d4761b 100644
--- a/src/api/workspace.rs
+++ b/src/api/workspace.rs
@@ -849,9 +849,12 @@ async fn get_workspace_usage_handler(
   workspace_id: web::Path<Uuid>,
   state: Data<AppState>,
 ) -> Result<Json<AppResponse<WorkspaceUsage>>> {
-  let res = WorkspaceUsage {
-    total_document_size: 19978, //  TODO
-  };
+  let res = biz::workspace::ops::get_workspace_document_total_bytes(
+    &state.pg_pool,
+    &user_uuid,
+    &workspace_id,
+  )
+  .await?;
   Ok(Json(AppResponse::Ok().with_data(res)))
 }
 
diff --git a/src/biz/workspace/ops.rs b/src/biz/workspace/ops.rs
index 7690d8f6..d210e71d 100644
--- a/src/biz/workspace/ops.rs
+++ b/src/biz/workspace/ops.rs
@@ -11,12 +11,14 @@ use database::user::select_uid_from_email;
 use database::workspace::{
   change_workspace_icon, delete_from_workspace, delete_workspace_members, get_invitation_by_id,
   insert_user_workspace, insert_workspace_invitation, rename_workspace, select_all_user_workspaces,
-  select_workspace, select_workspace_invitations_for_user, select_workspace_member_list,
+  select_user_is_workspace_owner, select_workspace, select_workspace_invitations_for_user,
+  select_workspace_member_list, select_workspace_total_collab_bytes,
   update_updated_at_of_workspace, update_workspace_invitation_set_status_accepted,
   upsert_workspace_member, upsert_workspace_member_with_txn,
 };
 use database_entity::dto::{
   AFAccessLevel, AFRole, AFWorkspace, AFWorkspaceInvitation, AFWorkspaceInvitationStatus,
+  WorkspaceUsage,
 };
 
 use gotrue::params::InviteUserParams;
@@ -366,3 +368,21 @@ pub async fn update_workspace_member(
 
   Ok(())
 }
+
+pub async fn get_workspace_document_total_bytes(
+  pg_pool: &PgPool,
+  user_uuid: &Uuid,
+  workspace_id: &Uuid,
+) -> Result<WorkspaceUsage, AppError> {
+  let is_owner = select_user_is_workspace_owner(pg_pool, user_uuid, workspace_id).await?;
+  if !is_owner {
+    return Err(AppError::UserUnAuthorized(
+      "User is not the owner of the workspace".to_string(),
+    ));
+  }
+
+  let byte_count = select_workspace_total_collab_bytes(pg_pool, workspace_id).await?;
+  Ok(WorkspaceUsage {
+    total_document_size: byte_count,
+  })
+}