siftlode/backend/app/quota.py

"""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: ``<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_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()