peter
ed51f5be7e
Checked first and logged after search.list, requests sent together at the last free slot all passed, and each spent 100 units of the shared budget. Every live search now claims its slot (quota.claim_action) whatever the source, and gives it back when it showed nothing. The cap counts the zero-unit claims only, so search.list's own spend event is accounting, not a second slot. log_action, the unlocked writer, goes. Tests: the API source reads one page per request, is refused before the call when the budget cannot pay, counts one slot per search and refunds a failed one; the lock test now proves the count happens after the lock.
237 lines
9.7 KiB
Python
237 lines
9.7 KiB
Python
"""Central YouTube Data API quota guard.
|
|
|
|
The whole app shares one daily budget (the API resets at midnight US Pacific time).
|
|
Steady-state work (RSS detection is free; enrichment of new videos is cheap) should
|
|
always be allowed; expensive backfill must yield to the remaining budget.
|
|
"""
|
|
|
|
import contextlib
|
|
import contextvars
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any, cast
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from sqlalchemy import delete, func, select
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
from sqlalchemy.engine import CursorResult
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.userscope import owned_by
|
|
from app import logs, sysconfig
|
|
from app.db import SessionLocal
|
|
from app.models import ApiQuotaUsage, QuotaEvent
|
|
|
|
_PACIFIC = ZoneInfo("America/Los_Angeles")
|
|
|
|
|
|
class QuotaAction:
|
|
"""Canonical quota-attribution action keys (one source of truth).
|
|
|
|
Naming scheme: ``<entity>_<action>``, all snake_case, explicit pull/push direction.
|
|
The stored value is this machine key; the user-facing label is resolved from i18n
|
|
(frontend ``quotaActions.<key>``). Distinct from scheduler job ids and progress phase
|
|
labels — do not conflate.
|
|
"""
|
|
|
|
SUBSCRIPTIONS_PULL = "subscriptions_pull"
|
|
SUBSCRIPTIONS_RESYNC = "subscriptions_resync"
|
|
PLAYLISTS_PULL = "playlists_pull"
|
|
PLAYLISTS_PUSH = "playlists_push"
|
|
PLAYLISTS_DELETE = "playlists_delete"
|
|
VIDEOS_BACKFILL_RECENT = "videos_backfill_recent"
|
|
VIDEOS_BACKFILL_FULL = "videos_backfill_full"
|
|
VIDEOS_FEED_FALLBACK = "videos_feed_fallback"
|
|
VIDEOS_ENRICH = "videos_enrich"
|
|
VIDEOS_LOOKUP = "videos_lookup"
|
|
VIDEOS_SEARCH = "videos_search"
|
|
VIDEOS_RELATED = "videos_related"
|
|
CHANNELS_DISCOVER = "channels_discover"
|
|
CHANNELS_EXPLORE = "channels_explore"
|
|
CHANNELS_SUBSCRIBE = "channels_subscribe"
|
|
CHANNELS_UNSUBSCRIBE = "channels_unsubscribe"
|
|
MAINTENANCE_REVALIDATE = "maintenance_revalidate"
|
|
OTHER = "other"
|
|
|
|
|
|
# Request/job-scoped attribution for quota spend: who triggered it and what kind of work.
|
|
# Set at entry points (route handlers, scheduler jobs) via attribute(); read by
|
|
# record_usage. Default = background/system, generic action.
|
|
_actor_id: contextvars.ContextVar[int | None] = contextvars.ContextVar(
|
|
"quota_actor_id", default=None
|
|
)
|
|
_action: contextvars.ContextVar[str] = contextvars.ContextVar(
|
|
"quota_action", default=QuotaAction.OTHER
|
|
)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def attribute(actor_id: int | None, action: str):
|
|
"""Attribute any quota spent in this block to `actor_id` (None = system) under `action`."""
|
|
t1 = _actor_id.set(actor_id)
|
|
t2 = _action.set(action)
|
|
try:
|
|
yield
|
|
finally:
|
|
_actor_id.reset(t1)
|
|
_action.reset(t2)
|
|
|
|
|
|
class Measure:
|
|
"""Result handle from `measured()` — read .used / .remaining after the block exits."""
|
|
|
|
used = 0
|
|
remaining = 0
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def measured(db: Session, actor_id: int | None, action: str):
|
|
"""Like `attribute()`, but also reports how much quota the block spent. Yields a `Measure`
|
|
whose `.used` (units charged inside the block) and `.remaining` (today's remaining after)
|
|
are populated on exit — so manual sync routes don't each re-spell the before/after diff."""
|
|
m = Measure()
|
|
before = units_used_today(db)
|
|
with attribute(actor_id, action):
|
|
yield m
|
|
m.used = units_used_today(db) - before
|
|
m.remaining = remaining_today(db)
|
|
|
|
|
|
def pacific_today():
|
|
return datetime.now(_PACIFIC).date()
|
|
|
|
|
|
def gc(db: Session, retention_days: int) -> int:
|
|
"""Delete quota-history rows older than `retention_days` (<= 0 keeps them forever) and return how
|
|
many went. Flushes only; the cleanup job commits. Today's rows are never in reach: the per-user
|
|
search cap counts them, and the floor of two days outlasts any Pacific day (25 h at most)."""
|
|
if retention_days <= 0:
|
|
return 0
|
|
cutoff = datetime.now(timezone.utc) - timedelta(days=max(2, retention_days))
|
|
# A DELETE answers with a CursorResult; `execute` is typed as the general Result, without rowcount.
|
|
result = cast(
|
|
CursorResult[Any], db.execute(delete(QuotaEvent).where(QuotaEvent.created_at < cutoff))
|
|
)
|
|
return result.rowcount or 0
|
|
|
|
|
|
def pacific_day_start_utc() -> datetime:
|
|
"""UTC instant of the current Pacific day's midnight — aligns 'today' with the budget reset."""
|
|
start = datetime.now(_PACIFIC).replace(hour=0, minute=0, second=0, microsecond=0)
|
|
return start.astimezone(timezone.utc)
|
|
|
|
|
|
def units_used_today(db: Session) -> int:
|
|
# A scalar SELECT (not db.get) so it bypasses the caller session's identity map — quota writes
|
|
# commit in a SEPARATE session (see record_usage), and a cached ORM row would otherwise hide
|
|
# their effect from a caller that already loaded today's row (measured()'s before/after diff).
|
|
# Under READ COMMITTED a fresh SELECT sees the other session's committed writes.
|
|
return (
|
|
db.scalar(select(ApiQuotaUsage.units_used).where(ApiQuotaUsage.day == pacific_today())) or 0
|
|
)
|
|
|
|
|
|
def remaining_today(db: Session) -> int:
|
|
return max(0, sysconfig.get_int(db, "quota_daily_budget") - units_used_today(db))
|
|
|
|
|
|
def can_spend(db: Session, units: int) -> bool:
|
|
return remaining_today(db) >= units
|
|
|
|
|
|
def backfill_reserve_reached() -> bool:
|
|
"""True when today's remaining budget is at or below the backfill reserve — the shared yield
|
|
condition for every backfill/explore path (recent, deep, immediate full-history, explore, and
|
|
rolling revalidation).
|
|
|
|
Reads in its OWN short-lived session (like record_usage / claim_action), for two reasons: it
|
|
replaces the remaining/reserve comparison that was open-coded at every call site, and — critically
|
|
for the per-page guards in paging loops — a guard must not hold the caller's transaction open
|
|
across the network fetch that follows it. Quota is committed independently of any caller, so a
|
|
fresh session reads the same numbers."""
|
|
with SessionLocal() as qs:
|
|
return remaining_today(qs) <= sysconfig.get_int(qs, "backfill_quota_reserve")
|
|
|
|
|
|
def claim_action(user_id: int, action: str, limit: int) -> int | None:
|
|
"""Take one of this user's `limit` daily `action` slots, or None when they are all taken — the
|
|
check and the take in one transaction, under a lock per user and action. Checked first and
|
|
logged after the work, N requests sent together at the last free slot all passed the check
|
|
(the work in between is seconds of scraping). Returns the event's id, for `refund_action` when
|
|
the work then produced nothing.
|
|
|
|
The ONE way to count a capped action: a claim is a zero-unit event, and `actions_today` counts
|
|
those alone, so a call the action then pays for (search.list's own `record_usage` event, under
|
|
the same action) is accounting, never a second slot. Own session, committed at once — quota
|
|
accounting is durable whatever the caller's transaction does next, and must never commit it."""
|
|
with SessionLocal() as qs:
|
|
qs.execute(select(func.pg_advisory_xact_lock(func.hashtext(action), user_id)))
|
|
if actions_today(qs, user_id, action) >= limit:
|
|
return None
|
|
event = QuotaEvent(user_id=user_id, action=action, units=0)
|
|
qs.add(event)
|
|
qs.commit()
|
|
return event.id
|
|
|
|
|
|
def refund_action(event_id: int) -> None:
|
|
"""Give a slot `claim_action` took back: the work it was taken for failed outright, and a
|
|
request that showed the user nothing must not count against them."""
|
|
with SessionLocal() as qs:
|
|
qs.execute(delete(QuotaEvent).where(QuotaEvent.id == event_id))
|
|
qs.commit()
|
|
|
|
|
|
def actions_today(db: Session, user_id: int, action: str) -> int:
|
|
"""How many slots of `action` this user has claimed so far in the current Pacific day — for
|
|
per-user, per-action daily caps (the live search, the related lists). Counts the zero-unit
|
|
claim events (`claim_action`) only: `record_usage`'s events for units the action spent are not
|
|
slots (it logs only when units were charged)."""
|
|
return (
|
|
db.scalar(
|
|
select(func.count())
|
|
.select_from(QuotaEvent)
|
|
.where(
|
|
owned_by(QuotaEvent.user_id, user_id),
|
|
QuotaEvent.action == action,
|
|
QuotaEvent.units == 0,
|
|
QuotaEvent.created_at >= pacific_day_start_utc(),
|
|
)
|
|
)
|
|
or 0
|
|
)
|
|
|
|
|
|
def record_usage(units: int) -> None:
|
|
"""Atomically add `units` to today's counter (upsert) and log an attribution event.
|
|
|
|
Uses its OWN session, committed immediately, so the spend is DURABLE regardless of what the
|
|
calling job does next: the units were really consumed at YouTube, so a caller that later rolls
|
|
back (or is read-only and never commits) must NOT lose the record — under-counting the shared
|
|
daily budget risks overspending the real API quota. Equally, quota must never commit the
|
|
CALLER's session (the old behaviour) — that persisted a job's partial state mid-run behind its
|
|
own `except: db.rollback()`."""
|
|
if units <= 0:
|
|
return
|
|
day = pacific_today()
|
|
stmt = (
|
|
pg_insert(ApiQuotaUsage)
|
|
.values(day=day, units_used=units)
|
|
.on_conflict_do_update(
|
|
index_elements=[ApiQuotaUsage.day],
|
|
set_={"units_used": ApiQuotaUsage.units_used + units},
|
|
)
|
|
)
|
|
with SessionLocal() as qs:
|
|
qs.execute(stmt)
|
|
# Audit detail: who/what spent it (per-user attribution; NULL actor = system), and the
|
|
# background run it belongs to, so the run's row can say what it cost.
|
|
qs.add(
|
|
QuotaEvent(
|
|
user_id=_actor_id.get(),
|
|
action=_action.get(),
|
|
units=units,
|
|
job_run_id=logs.current_run_id(),
|
|
)
|
|
)
|
|
qs.commit()
|