peter
ad6ccbac05
Tokens, notifications, quota, debug sessions and the reset-link burn filtered their owner inline. Ten of the fifteen filters had no test that noticed one going missing; each now has a test where another account's rows must not count, show or change.
219 lines
8.7 KiB
Python
219 lines
8.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"
|
|
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 / log_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 log_action(user_id: int, action: str) -> None:
|
|
"""Record a zero-cost action event so per-user daily caps still count it.
|
|
|
|
`record_usage` only logs when units are actually charged; a scrape-based search spends no
|
|
quota yet must still be rate-limited per user, so it logs its event here directly.
|
|
|
|
Writes in its OWN session (like record_usage) — quota accounting is durable independent of the
|
|
caller's transaction and must never commit the caller's session (that would flush a job's
|
|
partial state mid-run, defeating its own rollback)."""
|
|
with SessionLocal() as qs:
|
|
qs.add(QuotaEvent(user_id=user_id, action=action, units=0))
|
|
qs.commit()
|
|
|
|
|
|
def actions_today(db: Session, user_id: int, action: str) -> int:
|
|
"""How many quota events of `action` this user has logged so far in the current Pacific
|
|
day — for per-user, per-action daily caps (e.g. the live-search limit). Counts events,
|
|
not units, so it only works for actions charged exactly once per user action."""
|
|
return (
|
|
db.scalar(
|
|
select(func.count())
|
|
.select_from(QuotaEvent)
|
|
.where(
|
|
owned_by(QuotaEvent.user_id, user_id),
|
|
QuotaEvent.action == action,
|
|
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()
|