Merge pull request #140 from gbw/fix/storage-usage-calculation
fix: correct storage usage calculation by using direct SQL query
This commit is contained in:
Dionisio Pozo
@@ -1,7 +1,8 @@
|
||||
use crate::application::ports::auth_ports::UserStoragePort;
|
||||
use crate::application::ports::storage_ports::{FileReadPort, StorageUsagePort};
|
||||
use crate::application::ports::storage_ports::StorageUsagePort;
|
||||
use crate::common::errors::DomainError;
|
||||
use async_trait::async_trait;
|
||||
use sqlx::PgPool;
|
||||
use std::sync::Arc;
|
||||
use tokio::task;
|
||||
use tracing::{debug, error, info};
|
||||
@@ -11,20 +12,20 @@ use tracing::{debug, error, info};
|
||||
*
|
||||
* This service is responsible for calculating how much storage each user
|
||||
* is using and updating this information in the user records.
|
||||
*
|
||||
* Storage usage is calculated directly from the `storage.files` table
|
||||
* by summing file sizes for each user (using the `user_id` column).
|
||||
*/
|
||||
pub struct StorageUsageService {
|
||||
file_repository: Arc<dyn FileReadPort>,
|
||||
pool: Arc<PgPool>,
|
||||
user_repository: Arc<dyn UserStoragePort>,
|
||||
}
|
||||
|
||||
impl StorageUsageService {
|
||||
/// Creates a new storage usage service
|
||||
pub fn new(
|
||||
file_repository: Arc<dyn FileReadPort>,
|
||||
user_repository: Arc<dyn UserStoragePort>,
|
||||
) -> Self {
|
||||
pub fn new(pool: Arc<PgPool>, user_repository: Arc<dyn UserStoragePort>) -> Self {
|
||||
Self {
|
||||
file_repository,
|
||||
pool,
|
||||
user_repository,
|
||||
}
|
||||
}
|
||||
@@ -33,12 +34,8 @@ impl StorageUsageService {
|
||||
pub async fn update_user_storage_usage(&self, user_id: &str) -> Result<i64, DomainError> {
|
||||
info!("Updating storage usage for user: {}", user_id);
|
||||
|
||||
// Get user's home folder pattern
|
||||
let user = self.user_repository.get_user_by_id(user_id).await?;
|
||||
let username = user.username();
|
||||
|
||||
// Calculate storage usage for this user
|
||||
let total_usage = self.calculate_user_storage_usage(username).await?;
|
||||
// Calculate storage usage directly from database
|
||||
let total_usage = self.calculate_user_storage_usage(user_id).await?;
|
||||
|
||||
// Update the user's storage usage in the database
|
||||
self.user_repository
|
||||
@@ -53,101 +50,34 @@ impl StorageUsageService {
|
||||
Ok(total_usage)
|
||||
}
|
||||
|
||||
/// Calculates a user's storage usage based on their home folder
|
||||
async fn calculate_user_storage_usage(&self, username: &str) -> Result<i64, DomainError> {
|
||||
debug!("Calculating storage for user: {}", username);
|
||||
/// Calculates a user's storage usage by summing all their file sizes.
|
||||
/// Uses a direct SQL query for O(1) performance.
|
||||
async fn calculate_user_storage_usage(&self, user_id: &str) -> Result<i64, DomainError> {
|
||||
debug!("Calculating storage for user: {}", user_id);
|
||||
|
||||
// First, try to find the user's home folder
|
||||
// List all folders to locate the user's folder
|
||||
let all_folders = self
|
||||
.file_repository
|
||||
.list_files(None)
|
||||
.await
|
||||
.map_err(|e| DomainError::internal_error("File repository", e.to_string()))?;
|
||||
// Direct SQL query to sum all file sizes for this user
|
||||
// This is much more efficient than recursively walking folders
|
||||
// Note: user_id is stored as varchar, not uuid, so we bind it directly as text
|
||||
let total_size: i64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COALESCE(SUM(size), 0)::bigint
|
||||
FROM storage.files
|
||||
WHERE user_id = $1 AND NOT is_trashed
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(self.pool.as_ref())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
DomainError::internal_error("StorageUsage", format!("Failed to calculate usage: {e}"))
|
||||
})?;
|
||||
|
||||
// Find the user's home folder (named "My Folder - {username}")
|
||||
let home_folder_name = format!("My Folder - {}", username);
|
||||
debug!("Looking for home folder: {}", home_folder_name);
|
||||
debug!(
|
||||
"Calculated storage for user {}: {} bytes",
|
||||
user_id, total_size
|
||||
);
|
||||
|
||||
let mut total_usage: i64 = 0;
|
||||
let mut home_folder_id = None;
|
||||
|
||||
// Find the home folder ID
|
||||
for folder in &all_folders {
|
||||
if folder.name() == home_folder_name {
|
||||
home_folder_id = Some(folder.id().to_string());
|
||||
debug!(
|
||||
"Found home folder for user {}: ID={}",
|
||||
username,
|
||||
folder.id()
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// If we found the home folder, calculate total size
|
||||
if let Some(folder_id) = home_folder_id {
|
||||
// Calculate recursively
|
||||
total_usage = self.calculate_folder_size(&folder_id).await?;
|
||||
} else {
|
||||
// If no home folder found, just return 0
|
||||
debug!("No home folder found for user: {}", username);
|
||||
}
|
||||
|
||||
Ok(total_usage)
|
||||
}
|
||||
|
||||
/// Recursively calculates the size of a folder and all its contents
|
||||
async fn calculate_folder_size(&self, folder_id: &str) -> Result<i64, DomainError> {
|
||||
// Implementation with explicit boxing to handle recursion in async functions
|
||||
async fn inner_calculate_size(
|
||||
repo: Arc<dyn FileReadPort>,
|
||||
folder_id: &str,
|
||||
) -> Result<i64, DomainError> {
|
||||
let mut total_size: i64 = 0;
|
||||
|
||||
// Get files directly in this folder
|
||||
let files = repo
|
||||
.list_files(Some(folder_id))
|
||||
.await
|
||||
.map_err(|e| DomainError::internal_error("File repository", e.to_string()))?;
|
||||
|
||||
// Sum the size of all files
|
||||
for file in &files {
|
||||
// Skip subdirectories at this level - we'll process them separately
|
||||
if file.mime_type() == "directory" || file.mime_type() == "application/directory" {
|
||||
// Recursively calculate subfolder size with explicit boxing
|
||||
let subfolder_id = file.id().to_string(); // Create owned copy
|
||||
let repo_clone = repo.clone(); // Clone the repository
|
||||
|
||||
// Use Box::pin to handle recursive async call
|
||||
let subfolder_size_future =
|
||||
Box::pin(inner_calculate_size(repo_clone, &subfolder_id));
|
||||
|
||||
match subfolder_size_future.await {
|
||||
Ok(size) => {
|
||||
total_size += size;
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Error calculating size for subfolder {}: {}",
|
||||
subfolder_id, e
|
||||
);
|
||||
// Continue with other folders even if one fails
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Add file size to total
|
||||
total_size += file.size() as i64;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(total_size)
|
||||
}
|
||||
|
||||
// Start the calculation with a clone of our repository reference
|
||||
let repo_clone = Arc::clone(&self.file_repository);
|
||||
inner_calculate_size(repo_clone, folder_id).await
|
||||
Ok(total_size)
|
||||
}
|
||||
|
||||
/// Calculates and updates storage usage for a user identified by username.
|
||||
@@ -161,7 +91,7 @@ impl StorageUsageService {
|
||||
let user_id = user.id().to_string();
|
||||
|
||||
// Reuse the existing calculation logic
|
||||
let total_usage = self.calculate_user_storage_usage(username).await?;
|
||||
let total_usage = self.calculate_user_storage_usage(&user_id).await?;
|
||||
|
||||
// Update the user's storage usage in the database
|
||||
self.user_repository
|
||||
@@ -287,7 +217,7 @@ impl StorageUsagePort for StorageUsageService {
|
||||
impl Clone for StorageUsageService {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
file_repository: Arc::clone(&self.file_repository),
|
||||
pool: Arc::clone(&self.pool),
|
||||
user_repository: Arc::clone(&self.user_repository),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user