fix ignore
This commit is contained in:
+1
-1
@@ -36,7 +36,7 @@ pip-wheel-metadata/
|
||||
# Application runtime data
|
||||
app/temp_uploads/*
|
||||
storage_drive/
|
||||
storage/
|
||||
/storage/
|
||||
uploads/
|
||||
downloads/
|
||||
temp/
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
import logging
|
||||
import io
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
from app.core.settings import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# get_storage_client() was being called fresh on every single file operation
|
||||
# (every upload, download, stream, copy, ...), each paying for a DB round-trip
|
||||
# (TenantStorageConfig, sometimes SystemConfiguration too) plus a credential
|
||||
# decrypt and a new boto3 client construction — real, measured cost on a
|
||||
# remote DB, and a direct contributor to "opening any document takes a long
|
||||
# time." The (client, quarantine_bucket, clean_bucket) tuple is safe to reuse
|
||||
# across calls for the same tenant, so it's cached in-process with a short
|
||||
# TTL — long enough to eliminate the redundant work within a single document
|
||||
# open (which can call this several times), short enough that an admin
|
||||
# changing storage credentials in System Settings takes effect within a few
|
||||
# minutes without needing a server restart.
|
||||
_STORAGE_CLIENT_CACHE_TTL = 300 # seconds
|
||||
_storage_client_cache: dict = {}
|
||||
|
||||
|
||||
def invalidate_storage_client_cache(tenant_id=None) -> None:
|
||||
"""Drop cached storage client(s) so the next call rebuilds from the
|
||||
database — call this right after a tenant's storage config is
|
||||
created/updated/cleared, so admin credential changes take effect
|
||||
immediately instead of waiting out the TTL. Matches by string form
|
||||
since callers pass tenant_id as either a UUID or a str.
|
||||
"""
|
||||
if tenant_id is None:
|
||||
_storage_client_cache.clear()
|
||||
return
|
||||
for key in [k for k in _storage_client_cache if str(k[0]) == str(tenant_id)]:
|
||||
_storage_client_cache.pop(key, None)
|
||||
|
||||
class B2StorageClient:
|
||||
"""
|
||||
Backblaze B2 Storage Client using Boto3 (S3-compatible API).
|
||||
"""
|
||||
|
||||
def __init__(self, key_id: str, application_key: str, endpoint_url: str):
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
|
||||
region = "us-west-004"
|
||||
if "s3." in endpoint_url:
|
||||
parts = endpoint_url.split("s3.")
|
||||
if len(parts) > 1:
|
||||
region = parts[1].split(".")[0]
|
||||
|
||||
try:
|
||||
self.client = boto3.client(
|
||||
's3',
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=key_id,
|
||||
aws_secret_access_key=application_key,
|
||||
region_name=region,
|
||||
config=Config(signature_version='s3v4')
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize B2 client: {e}")
|
||||
raise
|
||||
|
||||
def put_object(self, Bucket: str, Key: str, Body: Any, ContentType: str = "application/octet-stream", Metadata: Optional[dict] = None, length: Optional[int] = None):
|
||||
try:
|
||||
if isinstance(Body, bytes):
|
||||
data = Body
|
||||
elif hasattr(Body, "read"):
|
||||
data = Body
|
||||
else:
|
||||
data = Body
|
||||
|
||||
if Metadata:
|
||||
Metadata = {str(k): str(v) for k, v in Metadata.items()}
|
||||
else:
|
||||
Metadata = {}
|
||||
|
||||
self.client.put_object(
|
||||
Bucket=Bucket,
|
||||
Key=Key,
|
||||
Body=data,
|
||||
ContentType=ContentType,
|
||||
Metadata=Metadata
|
||||
)
|
||||
return {"status": "success", "key": Key}
|
||||
except Exception as e:
|
||||
logger.error(f"B2 Put Error: {e}")
|
||||
raise
|
||||
|
||||
def get_object(self, Bucket: str, Key: str, Range: Optional[str] = None):
|
||||
try:
|
||||
kwargs = {"Bucket": Bucket, "Key": Key}
|
||||
if Range:
|
||||
kwargs["Range"] = Range
|
||||
response = self.client.get_object(**kwargs)
|
||||
return {
|
||||
"Body": response['Body'],
|
||||
"ContentType": response.get('ContentType'),
|
||||
"ContentLength": response.get('ContentLength'),
|
||||
"ContentRange": response.get('ContentRange'),
|
||||
"StatusCode": 206 if Range else 200,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"B2 Get Error: {e}")
|
||||
raise
|
||||
|
||||
def delete_object(self, Bucket: str, Key: str):
|
||||
try:
|
||||
self.client.delete_object(Bucket=Bucket, Key=Key)
|
||||
return {"status": "success"}
|
||||
except Exception as e:
|
||||
logger.error(f"B2 Delete Error: {e}")
|
||||
raise
|
||||
|
||||
def copy_object(self, Bucket: str, Key: str, CopySource: dict):
|
||||
try:
|
||||
self.client.copy_object(
|
||||
Bucket=Bucket,
|
||||
Key=Key,
|
||||
CopySource=CopySource
|
||||
)
|
||||
return {"status": "success"}
|
||||
except Exception as e:
|
||||
logger.error(f"B2 Copy Error: {e}")
|
||||
raise
|
||||
|
||||
def generate_presigned_url(self, ClientMethod: str, Params: Optional[dict] = None, ExpiresIn: int = 3600, user_id: Optional[int] = None, tenant_id: Optional[str] = None):
|
||||
try:
|
||||
params = Params or {}
|
||||
key = params.get("Key")
|
||||
|
||||
if user_id is not None and key:
|
||||
allowed_prefixes = [
|
||||
f"user_{user_id}/", f"users/{user_id}/",
|
||||
f"chatbot_documents/{user_id}/", f"doc_conversion/{user_id}/",
|
||||
]
|
||||
if tenant_id:
|
||||
allowed_prefixes.extend([
|
||||
f"tenants/{tenant_id}/user_{user_id}/", f"tenants/{tenant_id}/users/{user_id}/",
|
||||
f"tenants/{tenant_id}/chatbot_documents/{user_id}/", f"tenants/{tenant_id}/doc_conversion/{user_id}/",
|
||||
])
|
||||
if not any(key.startswith(p) for p in allowed_prefixes) and not any(f"/{p}" in f"/{key}" for p in allowed_prefixes):
|
||||
logger.warning(f"🚨 Security Alert: User {user_id} attempted to generate presigned URL for unauthorized key: {key}")
|
||||
raise PermissionError(f"Access denied: User {user_id} does not own resource '{key}'")
|
||||
|
||||
if ClientMethod in ["get_object", "put_object"]:
|
||||
return self.client.generate_presigned_url(
|
||||
ClientMethod=ClientMethod,
|
||||
Params=Params,
|
||||
ExpiresIn=ExpiresIn
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
if not isinstance(e, PermissionError):
|
||||
logger.error(f"B2 Presigned URL Error: {e}")
|
||||
raise
|
||||
|
||||
def list_objects(self, Bucket: str, Prefix: str = ""):
|
||||
try:
|
||||
response = self.client.list_objects_v2(Bucket=Bucket, Prefix=Prefix)
|
||||
objects = response.get('Contents', [])
|
||||
return [{"Key": obj['Key'], "Size": obj['Size']} for obj in objects]
|
||||
except Exception as e:
|
||||
logger.error(f"B2 List Error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def format_bytes(size: float) -> str:
|
||||
"""Format bytes to human readable string."""
|
||||
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
||||
if size < 1024.0:
|
||||
return f"{size:.2f} {unit}"
|
||||
size /= 1024.0
|
||||
return f"{size:.2f} PB"
|
||||
|
||||
|
||||
def generate_storage_key(tenant_id: str, user_id: int, filename: str) -> str:
|
||||
"""Generate a unique original storage key for an initial upload."""
|
||||
import uuid
|
||||
import os
|
||||
|
||||
tenant_part = f"tenants/{tenant_id}/" if tenant_id and tenant_id != 'None' else ""
|
||||
return f"{tenant_part}users/{user_id}/uploads/{uuid.uuid4()}_{filename}"
|
||||
|
||||
|
||||
def _derive_b2_endpoint(key_id: str) -> str:
|
||||
"""
|
||||
Derive the B2 S3 endpoint URL from the key ID.
|
||||
B2 key IDs start with a account cluster prefix.
|
||||
"""
|
||||
region_map = {
|
||||
"000": "us-west-000",
|
||||
"001": "us-west-001",
|
||||
"002": "us-west-002",
|
||||
"003": "us-west-004",
|
||||
"004": "us-west-004",
|
||||
"005": "us-east-005",
|
||||
}
|
||||
prefix = key_id[:3] if key_id and len(key_id) >= 3 else ""
|
||||
region = region_map.get(prefix, "us-west-004")
|
||||
|
||||
if key_id.startswith("003"):
|
||||
region = "us-west-004"
|
||||
elif key_id.startswith("005"):
|
||||
region = "us-east-005"
|
||||
|
||||
return f"https://s3.{region}.backblazeb2.com"
|
||||
|
||||
|
||||
def get_storage_client(db=None, tenant_id=None, provider: Optional[str] = None):
|
||||
"""
|
||||
Get storage client dynamically.
|
||||
1. Checks for tenant-specific configuration if tenant_id is provided.
|
||||
2. Falls back to global system configuration.
|
||||
Mandatory Backblaze B2.
|
||||
"""
|
||||
cache_key = (tenant_id, provider)
|
||||
cached = _storage_client_cache.get(cache_key)
|
||||
if cached is not None and time.time() - cached[3] < _STORAGE_CLIENT_CACHE_TTL:
|
||||
return cached[0], cached[1], cached[2]
|
||||
|
||||
if db is None:
|
||||
logger.error("get_storage_client called without DB session. Cannot fetch B2 config.")
|
||||
raise RuntimeError("Storage configuration requires a database session.")
|
||||
|
||||
from app.modules.configuration.models.system_configuration_model import SystemConfiguration
|
||||
from app.modules.tenant.models.tenant_storage_config_model import TenantStorageConfig
|
||||
from app.core.crypto import decrypt_data
|
||||
|
||||
if tenant_id:
|
||||
tenant_cfg = db.query(TenantStorageConfig).filter(
|
||||
TenantStorageConfig.tenant_id == tenant_id,
|
||||
TenantStorageConfig.is_active.is_(True)
|
||||
).first()
|
||||
|
||||
if tenant_cfg:
|
||||
try:
|
||||
key_id = (tenant_cfg.b2_key_id or "").strip()
|
||||
application_key = decrypt_data(tenant_cfg.encrypted_b2_application_key) if tenant_cfg.encrypted_b2_application_key else None
|
||||
|
||||
if not key_id or not application_key:
|
||||
logger.warning(f"[Storage] Tenant {tenant_id} configuration is incomplete (missing key or secret). Falling back to global.")
|
||||
else:
|
||||
endpoint = tenant_cfg.b2_endpoint
|
||||
if endpoint and endpoint.strip():
|
||||
endpoint = endpoint.strip()
|
||||
if not endpoint.startswith("http"):
|
||||
endpoint = f"https://{endpoint}"
|
||||
else:
|
||||
endpoint = _derive_b2_endpoint(key_id)
|
||||
|
||||
quarantine_bucket = tenant_cfg.b2_quarantine_bucket or "quarantine-bucket"
|
||||
clean_bucket = tenant_cfg.b2_clean_bucket or "clean-document-bucket"
|
||||
|
||||
logger.info(f"[Storage] Initializing B2 client for tenant {tenant_id} (Key ID: ...{key_id[-4:]})")
|
||||
|
||||
client = B2StorageClient(
|
||||
key_id=key_id,
|
||||
application_key=application_key,
|
||||
endpoint_url=endpoint
|
||||
)
|
||||
_storage_client_cache[cache_key] = (client, quarantine_bucket, clean_bucket, time.time())
|
||||
return client, quarantine_bucket, clean_bucket
|
||||
except Exception as e:
|
||||
logger.error(f"[Storage] Failed to initialize tenant storage for {tenant_id}: {e}. Falling back to global.")
|
||||
else:
|
||||
logger.debug(f"[Storage] No active tenant config for {tenant_id}, using global settings.")
|
||||
|
||||
config_items = db.query(SystemConfiguration).filter(SystemConfiguration.config_key.like("storage.%")).all()
|
||||
configs = {cfg.config_key: cfg for cfg in config_items}
|
||||
|
||||
def get_cfg_val(key, default=None):
|
||||
cfg = configs.get(key)
|
||||
return cfg.text_value if cfg and cfg.text_value is not None else default
|
||||
|
||||
key_id = get_cfg_val("storage.b2_key_id")
|
||||
enc_app_key = get_cfg_val("storage.b2_application_key")
|
||||
|
||||
if not key_id or not enc_app_key:
|
||||
logger.error("Backblaze B2 is not configured in system settings.")
|
||||
raise RuntimeError("Storage is not configured. Please set up Backblaze B2 in System Settings.")
|
||||
|
||||
try:
|
||||
key_id = key_id.strip()
|
||||
app_key = decrypt_data(enc_app_key).strip()
|
||||
|
||||
custom_endpoint = get_cfg_val("storage.b2_endpoint")
|
||||
if custom_endpoint and custom_endpoint.strip():
|
||||
endpoint = custom_endpoint.strip()
|
||||
if not endpoint.startswith("http"):
|
||||
endpoint = f"https://{endpoint}"
|
||||
else:
|
||||
endpoint = _derive_b2_endpoint(key_id)
|
||||
|
||||
quarantine = get_cfg_val("storage.b2_quarantine_bucket", "quarantine-bucket").strip()
|
||||
clean = get_cfg_val("storage.b2_clean_bucket", "clean-document-bucket").strip()
|
||||
|
||||
b2_client = B2StorageClient(key_id, app_key, endpoint_url=endpoint)
|
||||
logger.info(f"[Storage] Initializing global B2 client (Key: ...{key_id[-4:]})")
|
||||
_storage_client_cache[cache_key] = (b2_client, quarantine, clean, time.time())
|
||||
return b2_client, quarantine, clean
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize Global B2 client: {e}")
|
||||
raise RuntimeError(f"Failed to initialize storage client: {e}")
|
||||
@@ -0,0 +1,155 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.infrastructure.storage.local_storage_handler import get_storage_client
|
||||
from app.modules.drive.constants import PRESIGNED_URL_EXPIRY
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DriveStorageService:
|
||||
"""
|
||||
Thin wrapper around the Backblaze B2 client.
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session = None):
|
||||
self.db = db
|
||||
|
||||
def _get_client_and_bucket(self, tenant_id=None, provider=None):
|
||||
"""Returns (client, quarantine_bucket, clean_bucket)."""
|
||||
return get_storage_client(db=self.db, tenant_id=tenant_id, provider=provider)
|
||||
|
||||
async def upload(self, key: str, file_obj, content_type: str, length: Optional[int] = None, tenant_id=None, provider=None) -> str:
|
||||
"""Upload a file-like object to clean bucket under *key* (scan is done before upload)."""
|
||||
client, quarantine_bucket, clean_bucket = self._get_client_and_bucket(tenant_id, provider)
|
||||
|
||||
use_bucket = clean_bucket or quarantine_bucket
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: client.put_object(
|
||||
Bucket=use_bucket,
|
||||
Key=key,
|
||||
Body=file_obj,
|
||||
ContentType=content_type,
|
||||
length=length
|
||||
),
|
||||
)
|
||||
return "b2"
|
||||
|
||||
async def move_to_clean(self, key: str, quarantine_bucket: str, clean_bucket: str, tenant_id=None, provider=None) -> None:
|
||||
"""Copy from quarantine to clean bucket and delete the quarantine copy."""
|
||||
if quarantine_bucket == clean_bucket:
|
||||
return
|
||||
client, _, _ = self._get_client_and_bucket(tenant_id, provider)
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: client.copy_object(
|
||||
Bucket=clean_bucket,
|
||||
Key=key,
|
||||
CopySource={'Bucket': quarantine_bucket, 'Key': key}
|
||||
),
|
||||
)
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: client.delete_object(Bucket=quarantine_bucket, Key=key),
|
||||
)
|
||||
|
||||
async def copy(self, source_key: str, dest_key: str, tenant_id=None, provider=None, bucket: Optional[str] = None) -> None:
|
||||
"""Copy a single object within storage from *source_key* to *dest_key*."""
|
||||
client, quarantine_bucket, clean_bucket = self._get_client_and_bucket(tenant_id, provider)
|
||||
|
||||
use_bucket = bucket or clean_bucket or quarantine_bucket
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: client.copy_object(
|
||||
Bucket=use_bucket,
|
||||
Key=dest_key,
|
||||
CopySource={'Bucket': use_bucket, 'Key': source_key}
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
if not bucket and clean_bucket != quarantine_bucket:
|
||||
fallback_source = quarantine_bucket if use_bucket == clean_bucket else clean_bucket
|
||||
try:
|
||||
logger.info(f"[DriveStorage] Copy failed in {use_bucket}, trying fallback source {fallback_source}")
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: client.copy_object(
|
||||
Bucket=clean_bucket,
|
||||
Key=dest_key,
|
||||
CopySource={'Bucket': fallback_source, 'Key': source_key}
|
||||
),
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
async def delete(self, key: str, tenant_id=None, provider=None) -> None:
|
||||
"""Delete a single object from storage (tries clean then quarantine)."""
|
||||
if not key:
|
||||
return
|
||||
try:
|
||||
client, quarantine_bucket, clean_bucket = self._get_client_and_bucket(tenant_id, provider)
|
||||
if not client:
|
||||
return
|
||||
loop = asyncio.get_running_loop()
|
||||
for bucket in [clean_bucket, quarantine_bucket]:
|
||||
if not bucket:
|
||||
continue
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda b=bucket, k=key: client.delete_object(Bucket=b, Key=k),
|
||||
)
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.debug(f"[DriveStorage] Could not delete {key} from {bucket}: {exc}")
|
||||
logger.debug(f"[DriveStorage] Object {key} not found or already deleted from storage")
|
||||
except Exception as exc:
|
||||
logger.warning(f"[DriveStorage] Error during storage delete for {key}: {exc}")
|
||||
|
||||
def presigned_url(
|
||||
self,
|
||||
key: str,
|
||||
user_id: Optional[int] = None,
|
||||
expiry: int = PRESIGNED_URL_EXPIRY,
|
||||
method: str = "get_object",
|
||||
tenant_id=None,
|
||||
provider=None,
|
||||
bucket: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Return a presigned URL for the given object key."""
|
||||
client, quarantine_bucket, clean_bucket = self._get_client_and_bucket(tenant_id, provider)
|
||||
use_bucket = bucket or clean_bucket
|
||||
return client.generate_presigned_url(
|
||||
method,
|
||||
Params={"Bucket": use_bucket, "Key": key},
|
||||
ExpiresIn=expiry,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
def get_object(self, key: str, tenant_id=None, provider=None, bucket: Optional[str] = None, range_header: Optional[str] = None):
|
||||
"""Return the raw S3 response object (body, metadata). ``range_header``
|
||||
is the raw HTTP Range header value (e.g. "bytes=0-1023"), passed
|
||||
straight through to B2/S3 so callers can serve partial content."""
|
||||
client, quarantine_bucket, clean_bucket = self._get_client_and_bucket(tenant_id, provider)
|
||||
|
||||
if bucket:
|
||||
return client.get_object(Bucket=bucket, Key=key, Range=range_header)
|
||||
|
||||
try:
|
||||
return client.get_object(Bucket=clean_bucket, Key=key, Range=range_header)
|
||||
except Exception:
|
||||
if quarantine_bucket != clean_bucket:
|
||||
return client.get_object(Bucket=quarantine_bucket, Key=key, Range=range_header)
|
||||
raise
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import Any
|
||||
from app.modules.storage.services.storage_service import StorageService
|
||||
|
||||
|
||||
class StorageController:
|
||||
@staticmethod
|
||||
def get_storage_usage(user: Any, conn: Session):
|
||||
storage_service = StorageService(conn)
|
||||
return storage_service.get_storage_usage(user)
|
||||
|
||||
@staticmethod
|
||||
async def upload_file(user: Any, file, folder: str, conn: Session):
|
||||
storage_service = StorageService(conn)
|
||||
return await storage_service.upload_file(user, file)
|
||||
|
||||
@staticmethod
|
||||
def list_files(user: Any, limit: int, offset: int, conn: Session):
|
||||
storage_service = StorageService(conn)
|
||||
return storage_service.list_files(user, limit, offset)
|
||||
|
||||
@staticmethod
|
||||
def delete_file(user: Any, file_id: int, conn: Session):
|
||||
storage_service = StorageService(conn)
|
||||
return storage_service.delete_file(user, file_id)
|
||||
@@ -0,0 +1,49 @@
|
||||
import uuid
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
from sqlalchemy import (
|
||||
Column,
|
||||
Integer,
|
||||
String,
|
||||
BigInteger,
|
||||
TIMESTAMP,
|
||||
ForeignKey,
|
||||
Boolean,
|
||||
)
|
||||
from sqlalchemy.orm import relationship, Mapped, mapped_column
|
||||
from app.db.database import Base
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
class UserFile(Base):
|
||||
__tablename__ = "user_files"
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
tenant_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, index=True
|
||||
)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
||||
filename = Column(String, nullable=False)
|
||||
storage_provider = Column(String(50), nullable=False, default="b2")
|
||||
s3_key = Column(String, nullable=False)
|
||||
size = Column(BigInteger, default=0)
|
||||
content_type = Column(String, nullable=True)
|
||||
status = Column(String, default="uploaded")
|
||||
is_deleted = Column(Boolean, default=False)
|
||||
deleted_at = Column(TIMESTAMP, nullable=True)
|
||||
created_at = Column(TIMESTAMP, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
user = relationship("User", back_populates="user_files")
|
||||
|
||||
|
||||
class UserStorageUsage(Base):
|
||||
__tablename__ = "user_storage_usage"
|
||||
user_id = Column(Integer, ForeignKey("users.id"), primary_key=True, nullable=False)
|
||||
total_bytes_used = Column(BigInteger, default=0)
|
||||
total_files_count = Column(Integer, default=0)
|
||||
max_bytes_quota = Column(BigInteger, default=1073741824)
|
||||
last_updated = Column(
|
||||
TIMESTAMP,
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
onupdate=lambda: datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
user = relationship("User", back_populates="storage_usage")
|
||||
@@ -0,0 +1,287 @@
|
||||
import uuid
|
||||
from typing import List, Optional, Any
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import func, case, or_
|
||||
from app.modules.storage.models.storage_model import UserFile, UserStorageUsage
|
||||
from app.modules.drive.models.drive_model import DriveFile
|
||||
from app.modules.auth.models.user_model import User
|
||||
from app.modules.tenant.models.tenant_model import Tenant
|
||||
|
||||
|
||||
class StorageRepository:
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
|
||||
def get_user_file_usages(self, user_id: int) -> Any:
|
||||
return (
|
||||
self.db.query(
|
||||
func.coalesce(func.sum(UserFile.size), 0).label("total"),
|
||||
func.count(UserFile.id).label("count"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(UserFile.content_type == "application/pdf", UserFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("pdf_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(
|
||||
UserFile.content_type
|
||||
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
UserFile.size,
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("docx_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(UserFile.content_type == "text/html", UserFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("html_size"),
|
||||
)
|
||||
.filter(
|
||||
UserFile.user_id == user_id,
|
||||
or_(UserFile.is_deleted == False, UserFile.is_deleted == None),
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_drive_file_usages(self, user_id: int) -> Any:
|
||||
return (
|
||||
self.db.query(
|
||||
func.coalesce(func.sum(DriveFile.size), 0).label("total"),
|
||||
func.count(DriveFile.id).label("count"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(DriveFile.mime_type == "application/pdf", DriveFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("pdf_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(
|
||||
DriveFile.mime_type
|
||||
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
DriveFile.size,
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("docx_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(DriveFile.mime_type == "text/html", DriveFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("html_size"),
|
||||
)
|
||||
.filter(DriveFile.owner_id == user_id, DriveFile.is_trashed == False)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_tenant_user_file_usages(self, tenant_id: uuid.UUID) -> Any:
|
||||
return (
|
||||
self.db.query(
|
||||
func.coalesce(func.sum(UserFile.size), 0).label("total"),
|
||||
func.count(UserFile.id).label("count"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(UserFile.content_type == "application/pdf", UserFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("pdf_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(
|
||||
UserFile.content_type
|
||||
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
UserFile.size,
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("docx_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(UserFile.content_type == "text/html", UserFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("html_size"),
|
||||
)
|
||||
.filter(
|
||||
UserFile.tenant_id == tenant_id,
|
||||
or_(UserFile.is_deleted == False, UserFile.is_deleted == None),
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_tenant_drive_file_usages(self, tenant_id: uuid.UUID) -> Any:
|
||||
return (
|
||||
self.db.query(
|
||||
func.coalesce(func.sum(DriveFile.size), 0).label("total"),
|
||||
func.count(DriveFile.id).label("count"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(DriveFile.mime_type == "application/pdf", DriveFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("pdf_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(
|
||||
DriveFile.mime_type
|
||||
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
DriveFile.size,
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("docx_size"),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(DriveFile.mime_type == "text/html", DriveFile.size),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
).label("html_size"),
|
||||
)
|
||||
.filter(DriveFile.tenant_id == tenant_id, DriveFile.is_trashed == False)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_user(self, user_id: int) -> Optional[User]:
|
||||
return self.db.query(User).filter(User.id == user_id).first()
|
||||
|
||||
def get_tenant(self, tenant_id: uuid.UUID) -> Optional[Tenant]:
|
||||
return self.db.query(Tenant).filter(Tenant.id == tenant_id).first()
|
||||
|
||||
def create_user_file(self, user_file: UserFile) -> UserFile:
|
||||
self.db.add(user_file)
|
||||
self.db.flush()
|
||||
self.db.refresh(user_file)
|
||||
return user_file
|
||||
|
||||
def list_user_files(
|
||||
self, user_id: int, tenant_id: uuid.UUID, limit: int, offset: int
|
||||
) -> List[UserFile]:
|
||||
return (
|
||||
self.db.query(UserFile)
|
||||
.filter(
|
||||
UserFile.user_id == user_id,
|
||||
UserFile.tenant_id == tenant_id,
|
||||
or_(UserFile.is_deleted == False, UserFile.is_deleted == None),
|
||||
)
|
||||
.order_by(UserFile.created_at.desc())
|
||||
.limit(limit)
|
||||
.offset(offset)
|
||||
.all()
|
||||
)
|
||||
|
||||
def count_user_files(self, user_id: int, tenant_id: uuid.UUID) -> int:
|
||||
return (
|
||||
self.db.query(func.count(UserFile.id))
|
||||
.filter(
|
||||
UserFile.user_id == user_id,
|
||||
UserFile.tenant_id == tenant_id,
|
||||
or_(UserFile.is_deleted == False, UserFile.is_deleted == None),
|
||||
)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
def get_user_file_by_id(
|
||||
self, user_id: int, file_id: int, tenant_id: uuid.UUID
|
||||
) -> Optional[UserFile]:
|
||||
return (
|
||||
self.db.query(UserFile)
|
||||
.filter(
|
||||
UserFile.id == file_id,
|
||||
UserFile.user_id == user_id,
|
||||
UserFile.tenant_id == tenant_id,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_or_create_storage_usage(self, user_id: int) -> UserStorageUsage:
|
||||
"""
|
||||
Fetch the user's storage-usage row, creating it if it is missing.
|
||||
|
||||
The obvious version — SELECT, then INSERT if absent — is a race, and not
|
||||
a theoretical one: the B6.1 load harness hit it within seconds at ten
|
||||
concurrent users. Two requests both find no row, both insert, and the
|
||||
loser dies with `duplicate key value violates unique constraint
|
||||
"user_storage_usage_pkey"` — a 500 on a read-only endpoint, which is the
|
||||
kind of failure that only ever appears in production.
|
||||
|
||||
An `ON CONFLICT DO NOTHING` insert inside a savepoint lets the database
|
||||
arbitrate instead. The savepoint matters: without it a conflict would
|
||||
leave the surrounding transaction unusable even though nothing is wrong.
|
||||
"""
|
||||
usage = (
|
||||
self.db.query(UserStorageUsage)
|
||||
.filter(UserStorageUsage.user_id == user_id)
|
||||
.first()
|
||||
)
|
||||
if usage:
|
||||
return usage
|
||||
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
with self.db.begin_nested():
|
||||
self.db.execute(
|
||||
pg_insert(UserStorageUsage.__table__)
|
||||
.values(user_id=user_id)
|
||||
.on_conflict_do_nothing(index_elements=["user_id"])
|
||||
)
|
||||
|
||||
return (
|
||||
self.db.query(UserStorageUsage)
|
||||
.filter(UserStorageUsage.user_id == user_id)
|
||||
.one()
|
||||
)
|
||||
|
||||
def update_usage_stats(self, user_id: int, total_bytes: int, file_count: int):
|
||||
usage = self.get_or_create_storage_usage(user_id)
|
||||
usage.total_bytes_used = total_bytes
|
||||
usage.total_files_count = file_count
|
||||
usage.last_updated = func.now()
|
||||
self.db.flush()
|
||||
|
||||
def commit(self):
|
||||
self.db.flush()
|
||||
|
||||
def flush(self):
|
||||
self.db.flush()
|
||||
@@ -0,0 +1,45 @@
|
||||
from fastapi import APIRouter, Depends, UploadFile, File, Form, HTTPException
|
||||
from app.db.database import get_db
|
||||
from app.middleware.auth import get_current_user
|
||||
from app.modules.auth.models.user_model import User
|
||||
from app.modules.storage.controllers.storage_controller import StorageController
|
||||
from app.modules.storage.schemas.storage_schema import (
|
||||
StorageFilesResponse,
|
||||
StorageUsageResponse,
|
||||
)
|
||||
from app.core.schemas import SuccessOut
|
||||
|
||||
router = APIRouter(prefix="/storage", tags=["Storage"])
|
||||
|
||||
|
||||
@router.get("/usage", response_model=StorageUsageResponse)
|
||||
def get_usage(conn=Depends(get_db), user: User = Depends(get_current_user)):
|
||||
usage = StorageController.get_storage_usage(user, conn)
|
||||
return {"user_id": user.id, "tenant_id": user.tenant_id, "usage": usage}
|
||||
|
||||
|
||||
@router.post("/upload")
|
||||
async def upload_file(
|
||||
file: UploadFile = File(...),
|
||||
folder: str = Form(""),
|
||||
conn=Depends(get_db),
|
||||
user: User = Depends(get_current_user),
|
||||
):
|
||||
return await StorageController.upload_file(user, file, folder, conn)
|
||||
|
||||
|
||||
@router.get("/files", response_model=StorageFilesResponse)
|
||||
def list_files(
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
conn=Depends(get_db),
|
||||
user: User = Depends(get_current_user),
|
||||
):
|
||||
return StorageController.list_files(user, limit, offset, conn)
|
||||
|
||||
|
||||
@router.delete("/files/{file_id}", response_model=SuccessOut)
|
||||
def delete_file(
|
||||
file_id: int, conn=Depends(get_db), user: User = Depends(get_current_user)
|
||||
):
|
||||
return StorageController.delete_file(user, file_id, conn)
|
||||
@@ -0,0 +1,109 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
from uuid import UUID
|
||||
|
||||
|
||||
class UserFileOut(BaseModel):
|
||||
id: int
|
||||
user_id: int
|
||||
file_name: str
|
||||
s3_key: str
|
||||
file_size: int
|
||||
mime_type: Optional[str]
|
||||
is_deleted: bool
|
||||
deleted_at: Optional[datetime]
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class StorageUsageOut(BaseModel):
|
||||
total_bytes: int
|
||||
file_count: int
|
||||
last_updated: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class StorageStatsOut(BaseModel):
|
||||
user_id: int
|
||||
total_bytes: int
|
||||
file_count: int
|
||||
quota_bytes: int
|
||||
usage_percentage: float
|
||||
|
||||
|
||||
class StorageBreakdown(BaseModel):
|
||||
pdf: int = 0
|
||||
docx: int = 0
|
||||
html: int = 0
|
||||
other: int = 0
|
||||
|
||||
|
||||
class StorageAlerts(BaseModel):
|
||||
severity: Optional[str] = None
|
||||
ring_color: Optional[str] = None
|
||||
banner: Optional[str] = None
|
||||
center_message: Optional[str] = None
|
||||
current_threshold: Optional[float] = None
|
||||
|
||||
|
||||
class UserStorageStats(BaseModel):
|
||||
used_bytes: int = 0
|
||||
used_formatted: Optional[str] = None
|
||||
file_count: int = 0
|
||||
percent: float = 0.0
|
||||
breakdown: StorageBreakdown = StorageBreakdown()
|
||||
|
||||
|
||||
class StorageUsageDetail(BaseModel):
|
||||
"""The `usage` / `storage` object, shared by both endpoints."""
|
||||
|
||||
scope: Optional[str] = None
|
||||
tenant_id: Optional[str] = None
|
||||
used_bytes: int = 0
|
||||
used_formatted: Optional[str] = None
|
||||
total_bytes: int = 0
|
||||
total_formatted: Optional[str] = None
|
||||
quota_bytes: int = 0
|
||||
quota_formatted: Optional[str] = None
|
||||
remaining_bytes: int = 0
|
||||
remaining_formatted: Optional[str] = None
|
||||
file_count: int = 0
|
||||
percent: float = 0.0
|
||||
usage_percentage: float = 0.0
|
||||
breakdown: StorageBreakdown = StorageBreakdown()
|
||||
alerts: StorageAlerts = StorageAlerts()
|
||||
user_stats: Optional[UserStorageStats] = None
|
||||
|
||||
|
||||
class StorageUsageResponse(BaseModel):
|
||||
user_id: int
|
||||
tenant_id: Optional[UUID] = None
|
||||
usage: StorageUsageDetail
|
||||
|
||||
|
||||
class StorageFileSummary(BaseModel):
|
||||
"""One row of `GET /api/storage/files`, as built by StorageService."""
|
||||
|
||||
id: int
|
||||
name: Optional[str] = None
|
||||
size: int = 0
|
||||
size_formatted: Optional[str] = None
|
||||
mime_type: Optional[str] = None
|
||||
created_at: Optional[str] = None
|
||||
|
||||
|
||||
class Pagination(BaseModel):
|
||||
limit: int
|
||||
offset: int
|
||||
total_count: int
|
||||
|
||||
|
||||
class StorageFilesResponse(BaseModel):
|
||||
files: List[StorageFileSummary] = []
|
||||
pagination: Pagination
|
||||
storage: StorageUsageDetail
|
||||
@@ -0,0 +1,27 @@
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.orm import Session
|
||||
from app.db.database import get_db
|
||||
from app.modules.auth.models.user_model import User
|
||||
from app.modules.documents.models.document_model import Project
|
||||
from app.modules.drive.models.drive_model import DriveFile
|
||||
from app.middleware.auth import get_current_user
|
||||
|
||||
router = APIRouter(prefix="/stats", tags=["Stats"])
|
||||
|
||||
@router.get("/dashboard")
|
||||
async def get_dashboard_stats(
|
||||
current_user: User = Depends(get_current_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get dashboard statistics for the logged-in user.
|
||||
"""
|
||||
project_count = db.query(Project).filter(Project.user_id == current_user.id).count()
|
||||
file_count = db.query(DriveFile).filter(DriveFile.owner_id == current_user.id).count()
|
||||
|
||||
return {
|
||||
"projects_count": project_count,
|
||||
"files_count": file_count,
|
||||
"storage_used": 0,
|
||||
"subscription": current_user.subscription
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.mail import send_email
|
||||
from app.db.redis import redis_cache
|
||||
from app.infrastructure.storage.local_storage_handler import format_bytes
|
||||
from app.modules.auth.models.user_model import User
|
||||
from app.modules.notifications.services.notification_service import NotificationService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StorageAlertService:
|
||||
THRESHOLDS = (50, 75, 90, 100)
|
||||
REDIS_KEY_PREFIX = "storage_alert_thresholds"
|
||||
EMAIL_SUPPRESSION_HOURS = 48
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
self.notifications = NotificationService(db)
|
||||
|
||||
def process_usage(self, user: User, usage: Dict[str, Any]) -> None:
|
||||
if not user:
|
||||
return
|
||||
|
||||
if usage.get("quota_bytes") == -1:
|
||||
self._clear_threshold_state(user.id)
|
||||
return
|
||||
|
||||
percent = float(usage.get("usage_percentage") or 0)
|
||||
active_thresholds = self._get_active_thresholds(user.id)
|
||||
thresholds_to_activate: list[str] = []
|
||||
thresholds_to_clear: list[str] = []
|
||||
|
||||
for threshold in self.THRESHOLDS:
|
||||
key = str(threshold)
|
||||
if percent >= threshold:
|
||||
if key not in active_thresholds:
|
||||
if not self._dispatch_alert(user, usage, threshold):
|
||||
return
|
||||
thresholds_to_activate.append(key)
|
||||
elif key in active_thresholds:
|
||||
thresholds_to_clear.append(key)
|
||||
|
||||
for key in thresholds_to_activate:
|
||||
self._mark_threshold_active(user.id, key)
|
||||
|
||||
for key in thresholds_to_clear:
|
||||
self._clear_threshold(user.id, key)
|
||||
|
||||
def build_alert_context(self, user: User, usage: Dict[str, Any]) -> Dict[str, Any]:
|
||||
percent = float(usage.get("usage_percentage") or 0)
|
||||
current_threshold = self._current_threshold(percent)
|
||||
|
||||
return {
|
||||
"current_threshold": current_threshold,
|
||||
"severity": self._severity(percent),
|
||||
"ring_color": self._ring_color(percent),
|
||||
"center_message": self._center_message(percent),
|
||||
"banner": self._banner_payload(percent, current_threshold),
|
||||
}
|
||||
|
||||
def _dispatch_alert(self, user: User, usage: Dict[str, Any], threshold: int) -> bool:
|
||||
content = self._alert_content(user, usage, threshold)
|
||||
try:
|
||||
self.notifications.notify_storage_alert(
|
||||
user_id=user.id,
|
||||
title=content["in_app_title"],
|
||||
message=content["in_app_message"],
|
||||
tenant_id=user.tenant_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to create storage alert notification for user %s at %s%%: %s",
|
||||
user.id,
|
||||
threshold,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
if user.email and self._should_send_email(user):
|
||||
try:
|
||||
send_email(
|
||||
subject=content["email_subject"],
|
||||
recipient=user.email,
|
||||
template_name=content["template_name"],
|
||||
template_context=content["template_context"],
|
||||
db=self.db,
|
||||
tenant_id=user.tenant_id,
|
||||
user_id=user.id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to send storage alert email to %s for %s%% threshold: %s",
|
||||
user.email,
|
||||
threshold,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
def _alert_content(
|
||||
self, user: User, usage: Dict[str, Any], threshold: int
|
||||
) -> Dict[str, Any]:
|
||||
used_formatted = usage.get("used_formatted") or format_bytes(
|
||||
int(usage.get("used_bytes") or 0)
|
||||
)
|
||||
quota_formatted = usage.get("quota_formatted") or format_bytes(
|
||||
int(usage.get("quota_bytes") or 0)
|
||||
)
|
||||
config = {
|
||||
50: {
|
||||
"subject": "Storage usage is at 50%",
|
||||
"title": "Storage is halfway used",
|
||||
"message": "You have used 50% of your storage quota.",
|
||||
"template": "storage_alert_50.html",
|
||||
},
|
||||
75: {
|
||||
"subject": "Storage usage is at 75%",
|
||||
"title": "Storage usage is climbing",
|
||||
"message": "You have used 75% of your storage quota.",
|
||||
"template": "storage_alert_75.html",
|
||||
},
|
||||
90: {
|
||||
"subject": "Storage usage is at 90%",
|
||||
"title": "Storage is nearly full",
|
||||
"message": "You have used 90% of your storage quota.",
|
||||
"template": "storage_alert_90.html",
|
||||
},
|
||||
100: {
|
||||
"subject": "Storage usage has reached 100%",
|
||||
"title": "Storage has reached its quota",
|
||||
"message": "Your storage usage has reached 100% of the configured quota.",
|
||||
"template": "storage_alert_100.html",
|
||||
},
|
||||
}[threshold]
|
||||
|
||||
return {
|
||||
"email_subject": config["subject"],
|
||||
"in_app_title": config["title"],
|
||||
"in_app_message": config["message"],
|
||||
"template_name": config["template"],
|
||||
"template_context": {
|
||||
"user_name": user.name or user.email,
|
||||
"threshold": threshold,
|
||||
"used_formatted": used_formatted,
|
||||
"quota_formatted": quota_formatted,
|
||||
"percent_used": round(float(usage.get("usage_percentage") or 0), 1),
|
||||
"year": datetime.now(timezone.utc).year,
|
||||
},
|
||||
}
|
||||
|
||||
def _redis_key(self, user_id: int) -> str:
|
||||
return f"{self.REDIS_KEY_PREFIX}:{user_id}"
|
||||
|
||||
def _get_active_thresholds(self, user_id: int) -> set[str]:
|
||||
try:
|
||||
state = redis_cache.hgetall(self._redis_key(user_id))
|
||||
if not state:
|
||||
return set()
|
||||
active = set()
|
||||
for k, v in state.items():
|
||||
k_str = k.decode("utf-8") if isinstance(k, bytes) else str(k)
|
||||
v_str = v.decode("utf-8") if isinstance(v, bytes) else str(v)
|
||||
if v_str == "1":
|
||||
active.add(k_str)
|
||||
return active
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to load storage alert threshold state for user %s: %s",
|
||||
user_id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return set()
|
||||
|
||||
def _mark_threshold_active(self, user_id: int, threshold_key: str) -> None:
|
||||
try:
|
||||
redis_cache.hset(self._redis_key(user_id), threshold_key, "1")
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to persist storage alert threshold %s for user %s: %s",
|
||||
threshold_key,
|
||||
user_id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _clear_threshold(self, user_id: int, threshold_key: str) -> None:
|
||||
try:
|
||||
redis_cache.hdel(self._redis_key(user_id), threshold_key)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to clear storage alert threshold %s for user %s: %s",
|
||||
threshold_key,
|
||||
user_id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _clear_threshold_state(self, user_id: int) -> None:
|
||||
try:
|
||||
redis_cache.delete(self._redis_key(user_id))
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to clear storage alert state for user %s: %s",
|
||||
user_id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _should_send_email(self, user: User) -> bool:
|
||||
created_at = getattr(user, "created_at", None)
|
||||
if not created_at:
|
||||
return True
|
||||
if created_at.tzinfo is None:
|
||||
created_at = created_at.replace(tzinfo=timezone.utc)
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(hours=self.EMAIL_SUPPRESSION_HOURS)
|
||||
return created_at <= cutoff
|
||||
|
||||
def _current_threshold(self, percent: float) -> Optional[int]:
|
||||
crossed = [threshold for threshold in self.THRESHOLDS if percent >= threshold]
|
||||
return max(crossed) if crossed else None
|
||||
|
||||
def _severity(self, percent: float) -> str:
|
||||
if percent >= 90:
|
||||
return "critical"
|
||||
if percent >= 75:
|
||||
return "warning"
|
||||
return "normal"
|
||||
|
||||
def _ring_color(self, percent: float) -> str:
|
||||
if percent >= 90:
|
||||
return "red"
|
||||
if percent >= 75:
|
||||
return "amber"
|
||||
return "blue"
|
||||
|
||||
def _center_message(self, percent: float) -> Optional[str]:
|
||||
threshold = self._current_threshold(percent)
|
||||
messages = {
|
||||
75: "Storage needs attention",
|
||||
90: "Storage is nearly full",
|
||||
100: "Storage is at quota",
|
||||
}
|
||||
return messages.get(threshold)
|
||||
|
||||
def _banner_payload(
|
||||
self, percent: float, current_threshold: Optional[int]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if percent < 90 or not current_threshold:
|
||||
return None
|
||||
threshold = current_threshold
|
||||
banner_copy = {
|
||||
90: {
|
||||
"title": "Storage is nearly full",
|
||||
"message": "You have used 90% of your storage quota.",
|
||||
},
|
||||
100: {
|
||||
"title": "Storage has reached its quota",
|
||||
"message": "Your storage usage has reached 100% of the configured quota.",
|
||||
},
|
||||
}.get(threshold)
|
||||
if not banner_copy:
|
||||
return None
|
||||
|
||||
return {
|
||||
"id": threshold,
|
||||
"threshold": threshold,
|
||||
"title": banner_copy["title"],
|
||||
"message": banner_copy["message"],
|
||||
}
|
||||
@@ -0,0 +1,446 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import func
|
||||
|
||||
from app.modules.auth.models.user_model import User
|
||||
from app.modules.storage.models.storage_model import UserFile
|
||||
from app.modules.storage.repositories.storage_repository import StorageRepository
|
||||
from app.infrastructure.storage.local_storage_handler import (
|
||||
get_storage_client,
|
||||
format_bytes,
|
||||
generate_storage_key,
|
||||
)
|
||||
from app.modules.storage.services.storage_alert_service import StorageAlertService
|
||||
from app.core.clamav import (
|
||||
ClamAVError,
|
||||
MalwareDetectedError,
|
||||
scan_file_for_malware,
|
||||
)
|
||||
from app.core.transactions import compensate_on_rollback, defer_until_commit
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
DEFAULT_STORAGE_QUOTA_BYTES = 1024 * 1024 * 1024
|
||||
|
||||
|
||||
def _best_effort_delete(client, buckets, key: str) -> None:
|
||||
"""Remove `key` from the first bucket that has it. Never raises."""
|
||||
for bucket in buckets:
|
||||
try:
|
||||
client.delete_object(Bucket=bucket, Key=key)
|
||||
logger.info("Removed %s from [%s]", key, bucket)
|
||||
return
|
||||
except Exception: # noqa: BLE001 - try the next bucket
|
||||
continue
|
||||
logger.warning("Could not remove %s from any bucket", key)
|
||||
|
||||
|
||||
class StorageService:
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
self.storage_repo = StorageRepository(db)
|
||||
|
||||
@staticmethod
|
||||
def _build_usage_aggregates(res_total: Any, res_drive: Any) -> Dict[str, int]:
|
||||
total_bytes = int((res_total.total or 0) + (res_drive.total or 0))
|
||||
file_count = int((res_total.count or 0) + (res_drive.count or 0))
|
||||
|
||||
pdf_bytes = int((res_total.pdf_size or 0) + (res_drive.pdf_size or 0))
|
||||
docx_bytes = int((res_total.docx_size or 0) + (res_drive.docx_size or 0))
|
||||
html_bytes = int((res_total.html_size or 0) + (res_drive.html_size or 0))
|
||||
other_bytes = max(0, total_bytes - (pdf_bytes + docx_bytes + html_bytes))
|
||||
|
||||
return {
|
||||
"total_bytes": total_bytes,
|
||||
"file_count": file_count,
|
||||
"pdf_bytes": pdf_bytes,
|
||||
"docx_bytes": docx_bytes,
|
||||
"html_bytes": html_bytes,
|
||||
"other_bytes": other_bytes,
|
||||
}
|
||||
|
||||
def _resolve_quota(self, user: User) -> tuple[int, str]:
|
||||
"""
|
||||
Same number the drive path enforces, from the same place.
|
||||
|
||||
These two resolved the quota independently before, which is how they
|
||||
would eventually have disagreed — one reading the column, the other
|
||||
reading something else, with nothing to notice.
|
||||
"""
|
||||
if user.tenant_id:
|
||||
from app.modules.billing.services.entitlement_service import (
|
||||
EntitlementService,
|
||||
)
|
||||
|
||||
tenant = self.storage_repo.get_tenant(user.tenant_id)
|
||||
if tenant:
|
||||
return (
|
||||
int(EntitlementService(self.db).limit(tenant, "storage_bytes")),
|
||||
"tenant",
|
||||
)
|
||||
return DEFAULT_STORAGE_QUOTA_BYTES, "tenant"
|
||||
|
||||
usage_record = self.storage_repo.get_or_create_storage_usage(user.id)
|
||||
quota = usage_record.max_bytes_quota or DEFAULT_STORAGE_QUOTA_BYTES
|
||||
return int(quota), "user"
|
||||
|
||||
def _build_usage_response(
|
||||
self,
|
||||
user: User,
|
||||
own_aggregates: Dict[str, int],
|
||||
aggregates: Dict[str, int],
|
||||
limit: int,
|
||||
scope: str,
|
||||
) -> Dict[str, Any]:
|
||||
total_bytes = aggregates["total_bytes"]
|
||||
file_count = aggregates["file_count"]
|
||||
pdf_bytes = aggregates["pdf_bytes"]
|
||||
docx_bytes = aggregates["docx_bytes"]
|
||||
html_bytes = aggregates["html_bytes"]
|
||||
other_bytes = aggregates["other_bytes"]
|
||||
|
||||
if limit == -1:
|
||||
remaining_bytes = 0
|
||||
percent = 0
|
||||
quota_formatted = "Unlimited"
|
||||
remaining_formatted = "Unlimited"
|
||||
else:
|
||||
remaining_bytes = max(0, limit - total_bytes)
|
||||
percent = (total_bytes / limit) * 100 if limit > 0 else 0
|
||||
quota_formatted = format_bytes(limit)
|
||||
remaining_formatted = format_bytes(remaining_bytes)
|
||||
|
||||
user_used = own_aggregates["total_bytes"]
|
||||
user_percent = (user_used / limit * 100) if limit > 0 else 0
|
||||
|
||||
usage = {
|
||||
"total_bytes": total_bytes,
|
||||
"total_formatted": format_bytes(total_bytes),
|
||||
"used_bytes": total_bytes,
|
||||
"used_formatted": format_bytes(total_bytes),
|
||||
"file_count": file_count,
|
||||
"quota_bytes": limit,
|
||||
"quota_formatted": quota_formatted,
|
||||
"remaining_bytes": remaining_bytes,
|
||||
"remaining_formatted": remaining_formatted,
|
||||
"percent": percent,
|
||||
"usage_percentage": percent,
|
||||
"scope": scope,
|
||||
"tenant_id": str(user.tenant_id) if user.tenant_id else None,
|
||||
"breakdown": {
|
||||
"pdf": pdf_bytes,
|
||||
"docx": docx_bytes,
|
||||
"html": html_bytes,
|
||||
"other": other_bytes,
|
||||
},
|
||||
"user_stats": {
|
||||
"used_bytes": user_used,
|
||||
"used_formatted": format_bytes(user_used),
|
||||
"percent": user_percent,
|
||||
"file_count": own_aggregates["file_count"],
|
||||
"breakdown": {
|
||||
"pdf": own_aggregates["pdf_bytes"],
|
||||
"docx": own_aggregates["docx_bytes"],
|
||||
"html": own_aggregates["html_bytes"],
|
||||
"other": own_aggregates["other_bytes"],
|
||||
},
|
||||
}
|
||||
if user.tenant_id
|
||||
else None,
|
||||
}
|
||||
usage["alerts"] = self._build_alert_context(user, usage)
|
||||
return usage
|
||||
|
||||
def process_storage_alerts(self, user: User) -> Dict[str, Any]:
|
||||
usage = self.get_storage_usage(user)
|
||||
alert_service = StorageAlertService(self.db)
|
||||
|
||||
try:
|
||||
alert_service.process_usage(user, usage)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to process storage alerts for user %s: %s",
|
||||
user.id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
usage["alerts"] = self._build_alert_context(user, usage, alert_service)
|
||||
return usage
|
||||
|
||||
def _build_alert_context(
|
||||
self,
|
||||
user: User,
|
||||
usage: Dict[str, Any],
|
||||
alert_service: Optional[StorageAlertService] = None,
|
||||
) -> Dict[str, Any]:
|
||||
alert_service = alert_service or StorageAlertService(self.db)
|
||||
|
||||
try:
|
||||
return alert_service.build_alert_context(user, usage)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to build storage alert context for user %s: %s",
|
||||
user.id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return {
|
||||
"current_threshold": None,
|
||||
"severity": "normal",
|
||||
"ring_color": "blue",
|
||||
"center_message": None,
|
||||
"banner": None,
|
||||
"history": [],
|
||||
}
|
||||
|
||||
def _dispatch_storage_alerts(self, user: User) -> Dict[str, Any]:
|
||||
usage = self.get_storage_usage(user)
|
||||
|
||||
try:
|
||||
with self.db.begin_nested():
|
||||
usage = self.process_storage_alerts(user)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Storage alert dispatch failed for user %s: %s",
|
||||
user.id,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return usage
|
||||
|
||||
@staticmethod
|
||||
def invalidate_storage_cache(user_id: int):
|
||||
try:
|
||||
from app.db.redis import redis_cache
|
||||
from app.modules.drive.constants import CACHE_USER_STORAGE
|
||||
redis_cache.delete(CACHE_USER_STORAGE.format(user_id=user_id))
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to invalidate storage cache for user {user_id}: {e}")
|
||||
|
||||
def get_storage_usage(self, user: User) -> Dict[str, Any]:
|
||||
# 1. Try Redis cache lookup
|
||||
cache_key = None
|
||||
try:
|
||||
from app.db.redis import redis_cache
|
||||
from app.modules.drive.constants import CACHE_USER_STORAGE, CACHE_STORAGE_TTL_SECONDS
|
||||
cache_key = CACHE_USER_STORAGE.format(user_id=user.id)
|
||||
cached = redis_cache.get(cache_key)
|
||||
if cached:
|
||||
import json
|
||||
return json.loads(cached)
|
||||
except Exception as e:
|
||||
logger.debug(f"Redis cache lookup failed for storage usage: {e}")
|
||||
|
||||
# 2. Database calculation on miss
|
||||
own_storage_usage = self.storage_repo.get_user_file_usages(user.id)
|
||||
own_drive_usage = self.storage_repo.get_drive_file_usages(user.id)
|
||||
own_aggregates = self._build_usage_aggregates(own_storage_usage, own_drive_usage)
|
||||
self.storage_repo.update_usage_stats(
|
||||
user.id, own_aggregates["total_bytes"], own_aggregates["file_count"]
|
||||
)
|
||||
|
||||
if user.tenant_id:
|
||||
res_total = self.storage_repo.get_tenant_user_file_usages(user.tenant_id)
|
||||
res_drive = self.storage_repo.get_tenant_drive_file_usages(user.tenant_id)
|
||||
else:
|
||||
res_total = own_storage_usage
|
||||
res_drive = own_drive_usage
|
||||
|
||||
aggregates = self._build_usage_aggregates(res_total, res_drive)
|
||||
limit, scope = self._resolve_quota(user)
|
||||
result = self._build_usage_response(user, own_aggregates, aggregates, limit, scope)
|
||||
|
||||
# 3. Store in Redis
|
||||
if cache_key:
|
||||
try:
|
||||
import json
|
||||
redis_cache.set(cache_key, json.dumps(result), ex=CACHE_STORAGE_TTL_SECONDS)
|
||||
except Exception as e:
|
||||
logger.debug(f"Redis cache set failed for storage usage: {e}")
|
||||
|
||||
return result
|
||||
|
||||
async def upload_file(self, user: User, file: Any) -> Dict[str, Any]:
|
||||
"""
|
||||
Upload a file using the quarantine → scan → clean flow:
|
||||
1. Upload to quarantine bucket.
|
||||
2. Virus scan.
|
||||
3. On clean: copy to clean bucket, delete from quarantine.
|
||||
4. On malware: delete from quarantine, raise HTTP 400.
|
||||
"""
|
||||
try:
|
||||
file.file.seek(0, 2)
|
||||
size = file.file.tell()
|
||||
file.file.seek(0)
|
||||
except Exception:
|
||||
logger.warning("Could not determine size via seek for upload. Reading into memory as fallback.")
|
||||
content = await file.read()
|
||||
size = len(content)
|
||||
import io
|
||||
file_body = io.BytesIO(content)
|
||||
else:
|
||||
file_body = file.file
|
||||
|
||||
s3_key = generate_storage_key(
|
||||
str(user.tenant_id) if user.tenant_id else "", user.id, file.filename
|
||||
)
|
||||
mime = file.content_type or "application/octet-stream"
|
||||
|
||||
s3, quarantine_bucket, clean_bucket = get_storage_client(db=self.db, tenant_id=user.tenant_id)
|
||||
provider = "b2"
|
||||
|
||||
try:
|
||||
s3.put_object(
|
||||
Bucket=quarantine_bucket,
|
||||
Key=s3_key,
|
||||
Body=file_body,
|
||||
ContentType=mime,
|
||||
Metadata={"user_id": str(user.id)},
|
||||
length=size,
|
||||
)
|
||||
logger.info(f"Uploaded {s3_key} to quarantine bucket [{quarantine_bucket}]")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Storage Upload Error (quarantine): {e}")
|
||||
raise HTTPException(status_code=500, detail="Upload Failed")
|
||||
|
||||
try:
|
||||
if hasattr(file_body, 'seek'):
|
||||
file_body.seek(0)
|
||||
await asyncio.to_thread(scan_file_for_malware, file_body)
|
||||
except MalwareDetectedError as exc:
|
||||
try:
|
||||
s3.delete_object(Bucket=quarantine_bucket, Key=s3_key)
|
||||
except Exception:
|
||||
logger.error(f"Failed to clean up quarantine file {s3_key} after malware detection")
|
||||
raise HTTPException(status_code=400, detail=f"Virus detected: {exc.virus_name}")
|
||||
except ClamAVError as exc:
|
||||
logger.error(f"ClamAV scan failed for user {user.id}: {exc}")
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Malware scan service unavailable. Please try again.",
|
||||
)
|
||||
|
||||
if quarantine_bucket != clean_bucket:
|
||||
try:
|
||||
s3.copy_object(
|
||||
Bucket=clean_bucket,
|
||||
Key=s3_key,
|
||||
CopySource={"Bucket": quarantine_bucket, "Key": s3_key},
|
||||
)
|
||||
s3.delete_object(Bucket=quarantine_bucket, Key=s3_key)
|
||||
logger.info(f"Moved {s3_key} from quarantine [{quarantine_bucket}] to clean [{clean_bucket}]")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to move {s3_key} from quarantine to clean bucket: {e}. File stays in quarantine."
|
||||
)
|
||||
|
||||
compensate_on_rollback(
|
||||
self.db,
|
||||
lambda: _best_effort_delete(s3, [clean_bucket, quarantine_bucket], s3_key),
|
||||
)
|
||||
|
||||
user_file = UserFile(
|
||||
tenant_id=user.tenant_id,
|
||||
user_id=user.id,
|
||||
filename=file.filename,
|
||||
s3_key=s3_key,
|
||||
storage_provider=provider,
|
||||
size=size,
|
||||
content_type=mime,
|
||||
created_at=func.now(),
|
||||
)
|
||||
user_file = self.storage_repo.create_user_file(user_file)
|
||||
try:
|
||||
self.db.commit()
|
||||
self.invalidate_storage_cache(user.id)
|
||||
except Exception:
|
||||
self.db.rollback()
|
||||
raise
|
||||
|
||||
usage = self._dispatch_storage_alerts(user)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": "File uploaded successfully",
|
||||
"file": {
|
||||
"id": user_file.id,
|
||||
"name": user_file.filename,
|
||||
"size": user_file.size,
|
||||
"size_formatted": format_bytes(user_file.size),
|
||||
"mime_type": user_file.content_type,
|
||||
"created_at": (
|
||||
user_file.created_at.isoformat() if user_file.created_at else None
|
||||
),
|
||||
},
|
||||
"storage": usage,
|
||||
}
|
||||
|
||||
def list_files(self, user: User, limit: int, offset: int) -> Dict[str, Any]:
|
||||
rows = self.storage_repo.list_user_files(user.id, user.tenant_id, limit, offset)
|
||||
files = [
|
||||
{
|
||||
"id": r.id,
|
||||
"name": r.filename,
|
||||
"size": r.size,
|
||||
"size_formatted": format_bytes(r.size),
|
||||
"mime_type": r.content_type,
|
||||
"created_at": r.created_at.isoformat() if r.created_at else None,
|
||||
}
|
||||
for r in rows
|
||||
]
|
||||
|
||||
total_count = self.storage_repo.count_user_files(user.id, user.tenant_id)
|
||||
usage = self.get_storage_usage(user)
|
||||
|
||||
return {
|
||||
"files": files,
|
||||
"pagination": {
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"total_count": total_count,
|
||||
},
|
||||
"storage": usage,
|
||||
}
|
||||
|
||||
def delete_file(self, user: User, file_id: int) -> Dict[str, bool]:
|
||||
user_file = self.storage_repo.get_user_file_by_id(
|
||||
user.id, file_id, user.tenant_id
|
||||
)
|
||||
if user_file:
|
||||
s3_key = user_file.s3_key
|
||||
file_provider = getattr(user_file, "storage_provider", "b2")
|
||||
|
||||
user_file.is_deleted = True
|
||||
user_file.deleted_at = func.now()
|
||||
self.storage_repo.flush()
|
||||
|
||||
try:
|
||||
s3, quarantine_bucket, clean_bucket = get_storage_client(
|
||||
db=self.db,
|
||||
tenant_id=user.tenant_id,
|
||||
provider=file_provider,
|
||||
)
|
||||
defer_until_commit(
|
||||
self.db,
|
||||
lambda key=s3_key, client=s3, buckets=[clean_bucket, quarantine_bucket]:
|
||||
_best_effort_delete(client, buckets, key),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Could not resolve storage client for {s3_key}: {e}")
|
||||
|
||||
try:
|
||||
self.db.commit()
|
||||
self.invalidate_storage_cache(user.id)
|
||||
except Exception:
|
||||
self.db.rollback()
|
||||
raise
|
||||
|
||||
self._dispatch_storage_alerts(user)
|
||||
|
||||
return {"success": True}
|
||||
Reference in New Issue
Block a user