feat(webserver): add endpoint to calculate disk usage

ee/tabby-webserver/Cargo.toml

@@ -51,10 +51,12 @@ validator = { version = "0.16.1", features = ["derive"] }
 regex.workspace = true
 tabby-search = { path = "../tabby-search" }
 octocrab = "0.38.0"
+fs_extra = "1.3.0"
 
 [dev-dependencies]
 assert_matches = "1.5.0"
 tokio = { workspace = true, features = ["macros"] }
 tabby-db = { path = "../../ee/tabby-db", features = ["testutils"] }
+tabby-common = { path = "../../crates/tabby-common", features = [ "testutils" ] }
 serial_test = { workspace = true }
 temp_testdir = { workspace = true }

ee/tabby-webserver/graphql/schema.graphql

@@ -91,6 +91,13 @@ type UserConnection {
   pageInfo: PageInfo!
 }
 
+type DiskUsageStats {
+  events: DiskUsage!
+  indexedRepositories: DiskUsage!
+  database: DiskUsage!
+  models: DiskUsage!
+}
+
 input UpdateOAuthCredentialInput {
   provider: OAuthProvider!
   clientId: String!
@@ -258,6 +265,7 @@ type Query {
   dailyStatsInPastYear(users: [ID!]): [CompletionStats!]!
   dailyStats(start: DateTimeUtc!, end: DateTimeUtc!, users: [ID!], languages: [Language!]): [CompletionStats!]!
   userEvents(after: String, before: String, first: Int, last: Int, users: [ID!], start: DateTimeUtc!, end: DateTimeUtc!): UserEventConnection!
+  diskUsageStats: DiskUsageStats!
 }
 
 input NetworkSettingInput {
@@ -275,6 +283,11 @@ type UserEdge {
   cursor: String!
 }
 
+type DiskUsage {
+  filePaths: [String!]!
+  size: Float!
+}
+
 type JobRunConnection {
   edges: [JobRunEdge!]!
   pageInfo: PageInfo!

ee/tabby-webserver/src/path.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use tabby_common::path::tabby_root;
 
-fn tabby_ee_root() -> PathBuf {
+pub fn tabby_ee_root() -> PathBuf {
     tabby_root().join("ee")
 }
 

ee/tabby-webserver/src/schema/analytic.rs

@@ -8,6 +8,33 @@
 
 use crate::schema::Result;
 
+#[derive(GraphQLObject)]
+pub struct DiskUsageStats {
+    pub events: DiskUsage,
+    pub indexed_repositories: DiskUsage,
+    pub database: DiskUsage,
+    pub models: DiskUsage,
+}
+
+#[derive(GraphQLObject)]
+pub struct DiskUsage {
+    pub file_paths: Vec<String>,
+    pub size: f64,
+}
+
+impl DiskUsage {
+    pub fn combine(self, other: Self) -> Self {
+        DiskUsage {
+            size: self.size + other.size,
+            file_paths: self
+                .file_paths
+                .into_iter()
+                .chain(other.file_paths)
+                .collect(),
+        }
+    }
+}
+
 #[derive(GraphQLObject, Debug, Clone)]
 pub struct CompletionStats {
     pub start: DateTime<Utc>,
@@ -101,4 +128,6 @@
         users: Vec<ID>,
         languages: Vec<Language>,
     ) -> Result<Vec<CompletionStats>>;
+
+    async fn disk_usage_stats(&self) -> Result<DiskUsageStats>;
 }

ee/tabby-webserver/src/schema/mod.rs

@@ -29,7 +29,7 @@ use validator::{Validate, ValidationErrors};
 use worker::{Worker, WorkerService};
 
 use self::{
-    analytic::{AnalyticService, CompletionStats},
+    analytic::{AnalyticService, CompletionStats, DiskUsageStats},
     auth::{
         JWTPayload, OAuthCredential, OAuthProvider, PasswordChangeInput, PasswordResetInput,
         RequestInvitationInput, RequestPasswordResetEmailInput, UpdateOAuthCredentialInput,
@@ -457,6 +457,12 @@ impl Query {
         )
         .await
     }
+
+    async fn disk_usage_stats(ctx: &Context) -> Result<DiskUsageStats> {
+        check_admin(ctx).await?;
+        let storage_stats = ctx.locator.analytic().disk_usage_stats().await?;
+        Ok(storage_stats)
+    }
 }
 
 #[derive(GraphQLObject)]

ee/tabby-webserver/src/service/analytic.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{path::PathBuf, sync::Arc};
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
@@ -8,7 +8,7 @@
 
 use super::AsRowid;
 use crate::schema::{
-    analytic::{AnalyticService, CompletionStats, Language},
+    analytic::{AnalyticService, CompletionStats, DiskUsage, DiskUsageStats, Language},
     Result,
 };
 
@@ -70,6 +70,35 @@
             .collect();
         Ok(stats)
     }
+
+    async fn disk_usage_stats(&self) -> Result<DiskUsageStats> {
+        Ok(DiskUsageStats {
+            events: recursive_dir_size(tabby_common::path::events_dir()).await?,
+            indexed_repositories: recursive_dir_size(tabby_common::path::dataset_dir())
+                .await?
+                .combine(recursive_dir_size(tabby_common::path::index_dir()).await?),
+            database: recursive_dir_size(crate::path::tabby_ee_root()).await?,
+            models: recursive_dir_size(tabby_common::path::models_dir()).await?,
+        })
+    }
+}
+
+/// Calculate the size of a directory in kilobytes recursively
+async fn recursive_dir_size(path: PathBuf) -> Result<DiskUsage, anyhow::Error> {
+    let path_str = path.to_string_lossy().to_string();
+
+    let size = if path.exists() {
+        tokio::task::spawn_blocking(|| async { fs_extra::dir::get_size(path) })
+            .await?
+            .await?
+    } else {
+        0
+    };
+
+    Ok(DiskUsage {
+        file_paths: vec![path_str],
+        size: size as f64 / 1024.0,
+    })
 }
 
 fn convert_ids(ids: Vec<ID>) -> Vec<i64> {
@@ -91,6 +120,8 @@
 #[cfg(test)]
 mod tests {
     use chrono::{Days, Duration};
+    use tabby_common::path::set_tabby_root;
+    use temp_testdir::TempDir;
 
     use super::*;
     use crate::service::AsID;
@@ -297,4 +328,31 @@
             .unwrap()
             .is_empty());
     }
+
+    #[tokio::test]
+    async fn test_disk_usage() {
+        let tmp_dir = TempDir::default();
+        set_tabby_root(tmp_dir.to_path_buf());
+
+        tokio::fs::create_dir_all(tabby_common::path::models_dir())
+            .await
+            .unwrap();
+
+        tokio::fs::write(
+            tabby_common::path::models_dir().join("testfile"),
+            "0".repeat(1024).as_bytes(),
+        )
+        .await
+        .unwrap();
+
+        let db = DbConn::new_in_memory().await.unwrap();
+        let service = new_analytic_service(db);
+
+        let disk_usage = service.disk_usage_stats().await.unwrap();
+
+        assert_eq!(disk_usage.events.size, 0.0);
+        assert_eq!(disk_usage.indexed_repositories.size, 0.0);
+        assert_eq!(disk_usage.database.size, 0.0);
+        assert_eq!(disk_usage.models.size, 1.0);
+    }
 }