diff --git a/.gitignore b/.gitignore index 67179d7..1232e91 100644 --- a/.gitignore +++ b/.gitignore @@ -36,7 +36,7 @@ pip-wheel-metadata/ # Application runtime data app/temp_uploads/* storage_drive/ -storage/ +/storage/ uploads/ downloads/ temp/ diff --git a/app/infrastructure/storage/local_storage_handler.py b/app/infrastructure/storage/local_storage_handler.py new file mode 100644 index 0000000..41d4b5a --- /dev/null +++ b/app/infrastructure/storage/local_storage_handler.py @@ -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}") \ No newline at end of file diff --git a/app/modules/drive/storage/drive_storage_service.py b/app/modules/drive/storage/drive_storage_service.py new file mode 100644 index 0000000..efbc9e5 --- /dev/null +++ b/app/modules/drive/storage/drive_storage_service.py @@ -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 + diff --git a/app/modules/storage/controllers/storage_controller.py b/app/modules/storage/controllers/storage_controller.py new file mode 100644 index 0000000..d83a0af --- /dev/null +++ b/app/modules/storage/controllers/storage_controller.py @@ -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) diff --git a/app/modules/storage/models/storage_model.py b/app/modules/storage/models/storage_model.py new file mode 100644 index 0000000..9bb15bf --- /dev/null +++ b/app/modules/storage/models/storage_model.py @@ -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") diff --git a/app/modules/storage/repositories/storage_repository.py b/app/modules/storage/repositories/storage_repository.py new file mode 100644 index 0000000..42b13d2 --- /dev/null +++ b/app/modules/storage/repositories/storage_repository.py @@ -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() diff --git a/app/modules/storage/routes/storage_routes.py b/app/modules/storage/routes/storage_routes.py new file mode 100644 index 0000000..6d958db --- /dev/null +++ b/app/modules/storage/routes/storage_routes.py @@ -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) diff --git a/app/modules/storage/schemas/storage_schema.py b/app/modules/storage/schemas/storage_schema.py new file mode 100644 index 0000000..67b88b4 --- /dev/null +++ b/app/modules/storage/schemas/storage_schema.py @@ -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 diff --git a/app/modules/storage/services/stats.py b/app/modules/storage/services/stats.py new file mode 100644 index 0000000..660df03 --- /dev/null +++ b/app/modules/storage/services/stats.py @@ -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 + } diff --git a/app/modules/storage/services/storage_alert_service.py b/app/modules/storage/services/storage_alert_service.py new file mode 100644 index 0000000..1bb978c --- /dev/null +++ b/app/modules/storage/services/storage_alert_service.py @@ -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"], + } diff --git a/app/modules/storage/services/storage_service.py b/app/modules/storage/services/storage_service.py new file mode 100644 index 0000000..56bdaea --- /dev/null +++ b/app/modules/storage/services/storage_service.py @@ -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}