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 import func, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.orm import Session

from app import 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_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 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 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). See [[quota-owns-its-session]]."""
    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(
                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()`. See [[quota-owns-its-session]]."""
    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).
        qs.add(QuotaEvent(user_id=_actor_id.get(), action=_action.get(), units=units))
        qs.commit()