peter
e9c60f72cd
Capture's routes (downloads, files, public watch links) and the token list no longer build SQL, call the session or set a model's fields: reads are downloads/queries, writes the unit's service (service, folders, links, quota), and each write flushes for the handler to commit. Whose a row is goes through userscope (owned_by, owned_row), which owned_or_404 now uses too. The fence holds Capture below its routes as well: it reads only Catalog's and Plex's tables, writes no other unit's, filters owners through userscope and commits only in its jobs and the upload ingest. A test per service use case, among them that the tray lists only the caller's files.
248 lines
10 KiB
Python
248 lines
10 KiB
Python
"""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"). An
|
|
upload counts from the moment it lands, ready or not: its original is on disk while its conversion
|
|
waits, and stays there when the conversion fails (for Retry) — bytes a ready-only sum never saw.
|
|
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).
|
|
|
|
DOWNLOAD jobs additionally consult `remaining_bytes` in the worker, because enqueue can only compare
|
|
the footprint as it stood when the job was created and earlier jobs may have filled the cap since.
|
|
That is a per-job ceiling, not an aggregate ledger: concurrent jobs of one user each read the same
|
|
remainder (in-flight bytes are not `ready`, so `footprint` cannot see them), so a wave of
|
|
max_concurrent downloads can still land over the cap — the deliberate trade, since a download's size
|
|
is unknown until it finishes. CLIP jobs consult it before their fetch or encode and again with
|
|
their output's real size before it is placed. Edit jobs do not consult it at all (tracked as SEC-19).
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from sqlalchemy import and_, func, or_, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app import sysconfig
|
|
from app.downloads.status import JobKind, JobStatus
|
|
from app.downloads.models import DownloadJob, DownloadQuota, MediaAsset
|
|
from app.userscope import owned_by
|
|
|
|
# Jobs a user is considered to be "holding" (occupy a job slot / footprint). Terminal failures
|
|
# (error, canceled) don't count.
|
|
_HOLDING = (JobStatus.QUEUED, JobStatus.RUNNING, JobStatus.PAUSED, JobStatus.DONE)
|
|
_RUNNING = (JobStatus.QUEUED, JobStatus.RUNNING)
|
|
# An asset whose file is on disk: a ready one, or an upload in any state (its original lands at
|
|
# ingest, see the module docstring). Every byte sum — a user's footprint, the GC's total-cap pass,
|
|
# the admin's storage view — reads it, so they all count the same disk.
|
|
ON_DISK = or_(
|
|
MediaAsset.status == "ready",
|
|
and_(MediaAsset.source_kind == "upload", MediaAsset.rel_path.is_not(None)),
|
|
)
|
|
|
|
|
|
@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 — the worker leaves a suspended user's queued jobs unclaimed — 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 assets this user's jobs hold on disk: the ready ones their active
|
|
jobs reference, and every upload of theirs — including one whose conversion failed, which keeps
|
|
its original (see the module docstring)."""
|
|
held_assets = (
|
|
select(DownloadJob.asset_id)
|
|
.where(
|
|
owned_by(DownloadJob.user_id, user_id),
|
|
or_(
|
|
DownloadJob.status.in_(_HOLDING),
|
|
and_(DownloadJob.job_kind == JobKind.UPLOAD, DownloadJob.status == JobStatus.ERROR),
|
|
),
|
|
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), ON_DISK
|
|
)
|
|
).scalar()
|
|
or 0
|
|
)
|
|
|
|
|
|
def _count(db: Session, user_id: int, states) -> int:
|
|
return (
|
|
db.execute(
|
|
select(func.count())
|
|
.select_from(DownloadJob)
|
|
.where(owned_by(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, (JobStatus.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 remaining_bytes(db: Session, user_id: int) -> int | None:
|
|
"""Storage this user may still spend on ONE download: ``None`` for an unlimited user, else the
|
|
cap minus the current footprint, floored at 0.
|
|
|
|
Zero is a real answer, not "no limit" — it means the cap is already met and the caller must
|
|
REFUSE rather than proceed unbounded. `check_enqueue` cannot make that refusal on the caller's
|
|
behalf: it ran when the job was queued, while the user was still under the cap, and the jobs
|
|
that filled it have landed since. Keeping the arithmetic here (rather than inline in the worker)
|
|
is what keeps every max_bytes decision in this module."""
|
|
lim = resolve(db, user_id)
|
|
if lim.unlimited:
|
|
return None
|
|
return max(0, lim.max_bytes - footprint(db, user_id))
|
|
|
|
|
|
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 fits_after_conversion(db: Session, user_id: int, old_bytes: int, new_bytes: int) -> bool:
|
|
"""Whether an upload converted from `old_bytes` (what it is charged at now) to `new_bytes` keeps
|
|
its owner within the cap. A re-encode can come out larger than the original, and nothing else
|
|
re-checks it once the upload was let in."""
|
|
if new_bytes <= old_bytes: # a smaller file never breaks a cap the larger one was let in under
|
|
return True
|
|
lim = resolve(db, user_id)
|
|
if lim.unlimited:
|
|
return True
|
|
return footprint(db, user_id) - old_bytes + new_bytes <= lim.max_bytes
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
|
|
def set_custom(
|
|
db: Session,
|
|
user_id: int,
|
|
*,
|
|
max_bytes: int | None = None,
|
|
max_concurrent: int | None = None,
|
|
max_jobs: int | None = None,
|
|
unlimited: bool | None = None,
|
|
) -> None:
|
|
"""The admin's own limits for a user: starts from what applies now, then each value given (a
|
|
size floored at 0, a count at 1). The handler commits."""
|
|
row = db.get(DownloadQuota, user_id)
|
|
if row is None:
|
|
cur = resolve(db, user_id)
|
|
row = DownloadQuota(
|
|
user_id=user_id,
|
|
max_bytes=cur.max_bytes,
|
|
max_concurrent=cur.max_concurrent,
|
|
max_jobs=cur.max_jobs,
|
|
unlimited=cur.unlimited,
|
|
)
|
|
db.add(row)
|
|
if max_bytes is not None:
|
|
row.max_bytes = max(0, max_bytes)
|
|
if max_concurrent is not None:
|
|
row.max_concurrent = max(1, max_concurrent)
|
|
if max_jobs is not None:
|
|
row.max_jobs = max(1, max_jobs)
|
|
if unlimited is not None:
|
|
row.unlimited = unlimited
|
|
db.flush()
|
|
|
|
|
|
def reset_custom(db: Session, user_id: int) -> bool:
|
|
"""Back to the defaults; False when the user had no limits of their own."""
|
|
row = db.get(DownloadQuota, user_id)
|
|
if row is None:
|
|
return False
|
|
db.delete(row)
|
|
db.flush()
|
|
return True
|