siftlode/backend/app/downloads/quota.py

"""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:
        # Floor max_concurrent at 1 (admin_set_quota already does on write; the sysconfig default's
        # min is 1). A stored 0 is nonsensical for a concurrency limit AND would diverge the two
        # "at limit" checks — the worker's claim-window exclusion vs at_concurrency_limit — since a
        # 0-limit user is "capped" with zero running jobs, which the SQL exclusion can't see. Keep
        # them consistent by never letting the effective limit drop below 1. (To BLOCK a user's
        # downloads, suspend them or set max_jobs=0, not max_concurrent=0.)
        return QuotaLimits(row.max_bytes, max(1, row.max_concurrent), row.max_jobs, row.unlimited)
    return QuotaLimits(
        max_bytes=sysconfig.get_int(db, "download_default_max_bytes"),
        max_concurrent=max(1, 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 check_upload(db: Session, user_id: int, add_bytes: int) -> None:
    """Like check_enqueue, but for an upload whose size is KNOWN at ingest: also reject when adding
    this file would push the footprint over max_bytes. A download can't pre-check its size (unknown
    until the worker finishes) so check_enqueue only compares the current footprint; an upload can,
    so it's charged against the cap before it lands. Raises QuotaExceeded(max_jobs|max_bytes)."""
    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)
    projected = footprint(db, user_id) + add_bytes
    if projected > lim.max_bytes:
        raise QuotaExceeded("max_bytes", lim.max_bytes, projected)


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,
    }