"""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, timezone from zoneinfo import ZoneInfo from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session from app import sysconfig from app.models import ApiQuotaUsage, QuotaEvent _PACIFIC = ZoneInfo("America/Los_Angeles") class QuotaAction: """Canonical quota-attribution action keys (one source of truth). Naming scheme: ``_``, all snake_case, explicit pull/push direction. The stored value is this machine key; the user-facing label is resolved from i18n (frontend ``quotaActions.``). 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_ENRICH = "videos_enrich" VIDEOS_LOOKUP = "videos_lookup" VIDEOS_SEARCH = "videos_search" CHANNELS_DISCOVER = "channels_discover" 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 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: row = db.get(ApiQuotaUsage, pacific_today()) return row.units_used if row else 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 log_action(db: Session, 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.""" db.add(QuotaEvent(user_id=user_id, action=action, units=0)) db.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.""" from sqlalchemy import func, select # local import: keep the module's import head lean return ( db.scalar( select(func.count()) .select_from(QuotaEvent) .where( QuotaEvent.user_id == user_id, QuotaEvent.action == action, QuotaEvent.created_at >= pacific_day_start_utc(), ) ) or 0 ) def record_usage(db: Session, units: int) -> None: """Atomically add `units` to today's counter (upsert) and log an attribution event.""" 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}, ) ) db.execute(stmt) # Audit detail: who/what spent it (per-user attribution; NULL actor = system). db.add(QuotaEvent(user_id=_actor_id.get(), action=_action.get(), units=units)) db.commit()