feat(webserver): add endpoint to calculate disk usage

ee/tabby-webserver/graphql/schema.graphql

@@ -1,3 +1,10 @@
+type StorageStats {
+  events: DirectoryStat!
+  indexedRepositories: DirectoryStat!
+  database: DirectoryStat!
+  models: DirectoryStat!
+}
+
 enum Language {
   RUST
   PYTHON
@@ -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!
+  storageStats: StorageStats!
 }
 
 input NetworkSettingInput {
@@ -280,6 +288,11 @@ type JobRunConnection {
   pageInfo: PageInfo!
 }
 
+type DirectoryStat {
+  filePaths: [String!]!
+  size: Float!
+}
+
 type RefreshTokenResponse {
   accessToken: String!
   refreshToken: String!

ee/tabby-webserver/src/path.rs

@@ -2,7 +2,7 @@
 
 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 StorageStats {
+    pub events: DirectoryStat,
+    pub indexed_repositories: DirectoryStat,
+    pub database: DirectoryStat,
+    pub models: DirectoryStat,
+}
+
+#[derive(GraphQLObject)]
+pub struct DirectoryStat {
+    pub file_paths: Vec<String>,
+    pub size: f64,
+}
+
+impl DirectoryStat {
+    pub fn combine(self, other: Self) -> Self {
+        DirectoryStat {
+            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 storage_stats(&self) -> Result<StorageStats>;
 }

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, StorageStats},
     auth::{
         JWTPayload, OAuthCredential, OAuthProvider, PasswordChangeInput, PasswordResetInput,
         RequestInvitationInput, RequestPasswordResetEmailInput, UpdateOAuthCredentialInput,
@@ -457,6 +457,12 @@ impl Query {
         )
         .await
     }
+
+    async fn storage_stats(ctx: &Context) -> Result<StorageStats> {
+        check_admin(ctx).await?;
+        let storage_stats = ctx.locator.analytic().storage_stats().await?;
+        Ok(storage_stats)
+    }
 }
 
 #[derive(GraphQLObject)]

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

@@ -1,14 +1,15 @@
-use std::sync::Arc;
+use std::{path::PathBuf, sync::Arc};
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use futures::Future;
 use juniper::ID;
 use tabby_db::DbConn;
 use tracing::warn;
 
 use super::AsRowid;
 use crate::schema::{
-    analytic::{AnalyticService, CompletionStats, Language},
+    analytic::{AnalyticService, CompletionStats, DirectoryStat, Language, StorageStats},
     Result,
 };
 
@@ -70,6 +71,41 @@
             .collect();
         Ok(stats)
     }
+
+    async fn storage_stats(&self) -> Result<StorageStats> {
+        Ok(StorageStats {
+            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
+fn recursive_dir_size(
+    path: PathBuf,
+) -> Box<dyn Future<Output = Result<DirectoryStat, anyhow::Error>> + Unpin + Send> {
+    Box::new(Box::pin(async move {
+        let mut size: f64 = 0.0;
+        if path.exists() {
+            let mut read_dir = tokio::fs::read_dir(path.clone()).await?;
+            while let Some(next) = read_dir.next_entry().await? {
+                let meta = next.metadata().await?;
+                if meta.is_dir() {
+                    size += recursive_dir_size(next.path()).await?.size;
+                } else {
+                    size += meta.len() as f64 / 1024.0;
+                }
+            }
+        }
+        Ok(DirectoryStat {
+            file_paths: vec![path.to_string_lossy().to_string()],
+            size,
+        })
+    }))
 }
 
 fn convert_ids(ids: Vec<ID>) -> Vec<i64> {