51 lines
1.4 KiB
Python
51 lines
1.4 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.
|
||
|
|
"""
|
||
|
|
from datetime import datetime
|
||
|
|
from zoneinfo import ZoneInfo
|
||
|
|
|
||
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
|
|
from sqlalchemy.orm import Session
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
from app.models import ApiQuotaUsage
|
||
|
|
|
||
|
|
_PACIFIC = ZoneInfo("America/Los_Angeles")
|
||
|
|
|
||
|
|
|
||
|
|
def pacific_today():
|
||
|
|
return datetime.now(_PACIFIC).date()
|
||
|
|
|
||
|
|
|
||
|
|
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, settings.quota_daily_budget - units_used_today(db))
|
||
|
|
|
||
|
|
|
||
|
|
def can_spend(db: Session, units: int) -> bool:
|
||
|
|
return remaining_today(db) >= units
|
||
|
|
|
||
|
|
|
||
|
|
def record_usage(db: Session, units: int) -> None:
|
||
|
|
"""Atomically add `units` to today's counter (upsert)."""
|
||
|
|
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)
|
||
|
|
db.commit()
|