"""Per-user download limits + footprint accounting. A user's *footprint* is the total size of the ready assets their library references (cache hits count — it's "how much disk you're responsible for", not "how much you personally downloaded"). Limits come from a per-user DownloadQuota row if the admin set one, else the sysconfig defaults. `unlimited` (e.g. the admin's own account) bypasses the byte cap. Enforcement split: max_jobs + max_bytes are checked at enqueue (they bound holdings); per-user max_concurrent is enforced by the worker at claim time (it bounds simultaneous downloads). """ from dataclasses import dataclass from sqlalchemy import func, select from sqlalchemy.orm import Session from app import sysconfig from app.models import DownloadJob, DownloadQuota, MediaAsset # Jobs a user is considered to be "holding" (occupy a job slot / footprint). Terminal failures # (error, canceled) don't count. _HOLDING = ("queued", "running", "paused", "done") _RUNNING = ("queued", "running") @dataclass class QuotaLimits: max_bytes: int max_concurrent: int max_jobs: int unlimited: bool class QuotaExceeded(Exception): """Raised by check_enqueue; `reason` is one of max_jobs|max_bytes (i18n key on the client).""" def __init__(self, reason: str, limit: int, current: int): self.reason = reason self.limit = limit self.current = current super().__init__(f"download quota exceeded: {reason} ({current}/{limit})") def resolve(db: Session, user_id: int) -> QuotaLimits: row = db.get(DownloadQuota, user_id) if row is not None: return QuotaLimits(row.max_bytes, row.max_concurrent, row.max_jobs, row.unlimited) return QuotaLimits( max_bytes=sysconfig.get_int(db, "download_default_max_bytes"), max_concurrent=sysconfig.get_int(db, "download_default_max_concurrent"), max_jobs=sysconfig.get_int(db, "download_default_max_jobs"), unlimited=False, ) def footprint(db: Session, user_id: int) -> int: """Total bytes of the distinct ready assets this user's active jobs reference.""" held_assets = ( select(DownloadJob.asset_id) .where( DownloadJob.user_id == user_id, DownloadJob.status.in_(_HOLDING), DownloadJob.asset_id.is_not(None), ) .distinct() ) return int( db.execute( select(func.coalesce(func.sum(MediaAsset.size_bytes), 0)).where( MediaAsset.id.in_(held_assets), MediaAsset.status == "ready" ) ).scalar() or 0 ) def _count(db: Session, user_id: int, states) -> int: return ( db.execute( select(func.count()) .select_from(DownloadJob) .where(DownloadJob.user_id == user_id, DownloadJob.status.in_(states)) ).scalar() or 0 ) def active_jobs(db: Session, user_id: int) -> int: return _count(db, user_id, _HOLDING) def running_jobs(db: Session, user_id: int) -> int: return _count(db, user_id, _RUNNING) def at_concurrency_limit(db: Session, user_id: int) -> bool: """Worker-side gate: is this user already downloading their max? Unlimited users never are.""" lim = resolve(db, user_id) if lim.unlimited: return False running = _count(db, user_id, ("running",)) return running >= lim.max_concurrent def check_enqueue(db: Session, user_id: int) -> None: """Raise QuotaExceeded if the user can't take another job (holdings caps). Concurrency is enforced later by the worker, so a big queue is allowed; it just drains max_concurrent-wide.""" lim = resolve(db, user_id) if lim.unlimited: return n = active_jobs(db, user_id) if n >= lim.max_jobs: raise QuotaExceeded("max_jobs", lim.max_jobs, n) used = footprint(db, user_id) if used >= lim.max_bytes: raise QuotaExceeded("max_bytes", lim.max_bytes, used) def usage(db: Session, user_id: int) -> dict: """Compact usage snapshot for the UI (footprint + counts vs limits).""" lim = resolve(db, user_id) return { "footprint_bytes": footprint(db, user_id), "active_jobs": active_jobs(db, user_id), "running_jobs": running_jobs(db, user_id), "max_bytes": lim.max_bytes, "max_concurrent": lim.max_concurrent, "max_jobs": lim.max_jobs, "unlimited": lim.unlimited, }