diff --git a/backend/app/quota.py b/backend/app/quota.py index 66d4fa4d..13ec2e66 100644 --- a/backend/app/quota.py +++ b/backend/app/quota.py @@ -143,7 +143,7 @@ def backfill_reserve_reached() -> bool: 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 + Reads in its OWN short-lived session (like record_usage / claim_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 @@ -152,26 +152,17 @@ def backfill_reserve_reached() -> bool: 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 claim_action(user_id: int, action: str, limit: int) -> int | None: """Take one of this user's `limit` daily `action` slots, or None when they are all taken — the check and the take in one transaction, under a lock per user and action. Checked first and logged after the work, N requests sent together at the last free slot all passed the check (the work in between is seconds of scraping). Returns the event's id, for `refund_action` when - the work then produced nothing. Own session, as `log_action`.""" + the work then produced nothing. + + The ONE way to count a capped action: a claim is a zero-unit event, and `actions_today` counts + those alone, so a call the action then pays for (search.list's own `record_usage` event, under + the same action) is accounting, never a second slot. Own session, committed at once — quota + accounting is durable whatever the caller's transaction does next, and must never commit it.""" with SessionLocal() as qs: qs.execute(select(func.pg_advisory_xact_lock(func.hashtext(action), user_id))) if actions_today(qs, user_id, action) >= limit: @@ -191,9 +182,10 @@ def refund_action(event_id: int) -> None: 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.""" + """How many slots of `action` this user has claimed so far in the current Pacific day — for + per-user, per-action daily caps (the live search, the related lists). Counts the zero-unit + claim events (`claim_action`) only: `record_usage`'s events for units the action spent are not + slots (it logs only when units were charged).""" return ( db.scalar( select(func.count()) @@ -201,6 +193,7 @@ def actions_today(db: Session, user_id: int, action: str) -> int: .where( owned_by(QuotaEvent.user_id, user_id), QuotaEvent.action == action, + QuotaEvent.units == 0, QuotaEvent.created_at >= pacific_day_start_utc(), ) ) diff --git a/backend/app/routes/search.py b/backend/app/routes/search.py index 06741528..91a25442 100644 --- a/backend/app/routes/search.py +++ b/backend/app/routes/search.py @@ -142,27 +142,20 @@ def search_youtube( daily_limit = sysconfig.get_int(db, "search_daily_limit_per_user") if daily_limit <= 0: raise refused(409, "search_disabled") - claimed: int | None = None - if source == "api": - # The one search.list call logs this action's event itself (record_usage, attributed - # below), so here the cap is only checked: a claimed slot would count it twice. - if quota.actions_today(db, user.id, quota.QuotaAction.VIDEOS_SEARCH) >= daily_limit: - _cap_spent("search", daily_limit) - raise refused(429, "search_daily_limit", limit=daily_limit) - # The 100-unit budget pre-check only applies to the API source. - if not quota.can_spend(db, SEARCH_COST): - log.info( - "a live search was refused: the shared YouTube budget is spent", - extra=logs.event("search.refused", reason="shared_quota", cost=SEARCH_COST), - ) - raise refused(429, "quota_spent") - else: - # The check and the take are one step (`quota.claim_action`); a search that fails - # outright gives the slot back below. - claimed = quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, daily_limit) - if claimed is None: - _cap_spent("search", daily_limit) - raise refused(429, "search_daily_limit", limit=daily_limit) + # The 100-unit budget pre-check only applies to the API source. + if source == "api" and not quota.can_spend(db, SEARCH_COST): + log.info( + "a live search was refused: the shared YouTube budget is spent", + extra=logs.event("search.refused", reason="shared_quota", cost=SEARCH_COST), + ) + raise refused(429, "quota_spent") + # The check and the take are one step (`quota.claim_action`), whichever the source: checked + # first and logged after, requests sent together at the last free slot all passed — and with + # the API source each of them spent 100 units. A search that fails outright gives it back below. + claimed = quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, daily_limit) + if claimed is None: + _cap_spent("search", daily_limit) + raise refused(429, "search_daily_limit", limit=daily_limit) # Channels this user has blocked are dropped before ingestion — never shown, never stored. blocked = queries.blocked_ids(db, user) @@ -230,8 +223,7 @@ def search_youtube( ) except Exception: # A search that showed the user nothing does not count against them. - if claimed is not None: - quota.refund_action(claimed) + quota.refund_action(claimed) raise ordered = collected[:limit] diff --git a/backend/tests/test_owner_scoping.py b/backend/tests/test_owner_scoping.py index 8ab99547..c85a69de 100644 --- a/backend/tests/test_owner_scoping.py +++ b/backend/tests/test_owner_scoping.py @@ -75,9 +75,9 @@ def test_trimming_read_notifications_trims_only_mine(db, make_user): def test_todays_actions_count_only_mine(db, make_user): me, other = make_user(), make_user() - quota.log_action(me.id, quota.QuotaAction.VIDEOS_SEARCH) - quota.log_action(other.id, quota.QuotaAction.VIDEOS_SEARCH) - quota.log_action(other.id, quota.QuotaAction.VIDEOS_SEARCH) + quota.claim_action(me.id, quota.QuotaAction.VIDEOS_SEARCH, 5) + quota.claim_action(other.id, quota.QuotaAction.VIDEOS_SEARCH, 5) + quota.claim_action(other.id, quota.QuotaAction.VIDEOS_SEARCH, 5) assert quota.actions_today(db, me.id, quota.QuotaAction.VIDEOS_SEARCH) == 1 diff --git a/backend/tests/test_quota_ownership.py b/backend/tests/test_quota_ownership.py index a951ef98..36a269ee 100644 --- a/backend/tests/test_quota_ownership.py +++ b/backend/tests/test_quota_ownership.py @@ -44,15 +44,15 @@ def test_reads_see_separate_session_spend_mid_transaction(db): assert after - before == 7 -def test_log_action_survives_caller_rollback(db): +def test_a_claimed_slot_survives_caller_rollback(db): user = User(email="searcher@example.com") db.add(user) db.commit() # a committed user to attribute the (free) action to db.add(User(email="pending@example.com")) # uncommitted caller work - quota.log_action(user.id, quota.QuotaAction.VIDEOS_SEARCH) + assert quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, 5) is not None db.rollback() - # The zero-cost action event survived (per-user daily caps must still count it) ... + # The zero-cost claim event survived (per-user daily caps must still count it) ... events = db.execute(select(QuotaEvent)).scalars().all() assert len(events) == 1 assert events[0].user_id == user.id diff --git a/backend/tests/test_related_route.py b/backend/tests/test_related_route.py index a0f8aa29..c1665b78 100644 --- a/backend/tests/test_related_route.py +++ b/backend/tests/test_related_route.py @@ -336,10 +336,11 @@ def test_the_last_slot_is_taken_under_a_lock(db, make_user): t.start() t.join(timeout=1) assert t.is_alive() # it waits for the one holding the lock + # ...and counts only once it has the lock: the one holding it takes the last slot meanwhile. + db.add(QuotaEvent(user_id=user.id, action=action, units=0)) db.commit() t.join(timeout=10) - assert got and got[0] is not None - assert quota.claim_action(user.id, action, 1) is None # the one slot is taken + assert got == [None] def _videos_list_fails(monkeypatch, *, on_call: int = 1, reason: str = "quotaExceeded") -> None: diff --git a/backend/tests/test_search_cap.py b/backend/tests/test_search_cap.py index 447db142..f2cca60a 100644 --- a/backend/tests/test_search_cap.py +++ b/backend/tests/test_search_cap.py @@ -1,9 +1,11 @@ -"""`GET /api/search/youtube` — the per-user daily cap of the live search (scrape source). +"""`GET /api/search/youtube` — the per-user daily cap of the live search, both sources. -The scrape is replaced by pages the test hands out, and YouTube's videos.list by an offline enrich. -What is held: every request takes a slot, a cursor's too; a search that failed outright gives its -slot back; a zero limit turns live search off, which says so instead of "try again tomorrow"; and a -lookup that fails says why instead of answering "no results".""" +The scrape and search.list are replaced by pages the test hands out, and YouTube's videos.list by an +offline enrich. What is held: every request takes a slot, a cursor's too; a search that failed +outright gives its slot back; a zero limit turns live search off, which says so instead of "try +again tomorrow"; a lookup that fails says why instead of answering "no results"; and the API source +— 100 units a page — reads one page per request, is refused before the call when the shared budget +cannot pay for it, and counts one slot per search though search.list logs its own event.""" import pytest @@ -42,6 +44,19 @@ class _FakeYT: def get_channels(self, ids): return [] + # search.list: the pages the test hands out, each logged as the real client does (100 units, + # under the action the route attributes it to). + api_pages: list = [] + api_calls = 0 + + def search_videos(self, q, page_token=None, order="relevance"): + _FakeYT.api_calls += 1 + quota.record_usage(100) + page = _FakeYT.api_pages.pop(0) + if isinstance(page, Exception): + raise page + return page + def get_videos(self, ids): if _FakeYT.fail_videos: raise YouTubeError("videos -> 403", reason="quotaExceeded", path="videos") @@ -74,6 +89,8 @@ def pages(monkeypatch, db): monkeypatch.setattr(search_routes.search_scrape, "search", search) monkeypatch.setattr(search_routes, "YouTubeClient", _FakeYT) monkeypatch.setattr(_FakeYT, "fail_videos", False) + monkeypatch.setattr(_FakeYT, "api_pages", []) + monkeypatch.setattr(_FakeYT, "api_calls", 0) return state @@ -130,3 +147,78 @@ def test_a_lookup_that_fails_says_why_not_no_results(db, make_user, as_user, pag assert resp.status_code == 422 assert resp.json()["detail"] == {"code": "youtube", "reason": "quota"} assert quota.actions_today(db, user.id, ACTION) == 0 + + +def _api(db) -> None: + sysconfig.set_value(db, "search_source", "api") + sysconfig.set_value(db, "quota_daily_budget", 10000) + + +def test_the_api_source_reads_one_page_per_request(db, make_user, as_user, pages): + # Each page is 100 units: a short page is answered as it is, with its cursor, never paged on. + _api(db) + _FakeYT.api_pages = [{"items": [_item(1)], "next_page_token": "nextpage"}, _page(2)] + resp = as_user(make_user()).get("/api/search/youtube?q=cats&limit=20") + assert resp.status_code == 200 + assert _FakeYT.api_calls == 1 + assert resp.json()["next_cursor"] == "nextpage" + + +def test_the_api_source_counts_one_slot_per_search(db, make_user, as_user, pages): + # search.list logs its own event under the action (record_usage): accounting, not a second slot. + user = make_user() + _api(db) + sysconfig.set_value(db, "search_daily_limit_per_user", 2) + _FakeYT.api_pages = [_page(1), _page(2), _page(3)] + client = as_user(user) + assert client.get("/api/search/youtube?q=cats").status_code == 200 + assert quota.actions_today(db, user.id, ACTION) == 1 + assert client.get("/api/search/youtube?q=dogs").status_code == 200 + refused = client.get("/api/search/youtube?q=birds") + assert refused.status_code == 429 + assert refused.json()["detail"]["reason"] == "search_daily_limit" + assert _FakeYT.api_calls == 2 + + +def test_the_api_source_takes_its_slot_under_the_lock(db, make_user, as_user, pages, monkeypatch): + user = make_user() + _api(db) + claimed: list = [] + real = quota.claim_action + + def claim(uid, action, limit): + claimed.append(action) + return real(uid, action, limit) + + monkeypatch.setattr(search_routes.quota, "claim_action", claim) + _FakeYT.api_pages = [_page(1)] + assert as_user(user).get("/api/search/youtube?q=cats").status_code == 200 + assert claimed == [ACTION] + + +def test_the_api_source_is_refused_before_the_call_when_the_budget_cannot_pay( + db, make_user, as_user, pages +): + user = make_user() + _api(db) + sysconfig.set_value(db, "quota_daily_budget", 50) + _FakeYT.api_pages = [_page(1)] + resp = as_user(user).get("/api/search/youtube?q=cats") + assert resp.status_code == 429 + assert resp.json()["detail"]["reason"] == "quota_spent" + assert _FakeYT.api_calls == 0 + assert quota.actions_today(db, user.id, ACTION) == 0 + + +def test_an_api_search_that_failed_gives_its_slot_back(db, make_user, as_user, pages): + user = make_user() + _api(db) + sysconfig.set_value(db, "search_daily_limit_per_user", 1) + _FakeYT.api_pages = [ + YouTubeError("search -> 403", reason="quotaExceeded", path="search"), + _page(1), + ] + client = as_user(user) + assert client.get("/api/search/youtube?q=cats").status_code == 422 + assert quota.actions_today(db, user.id, ACTION) == 0 + assert client.get("/api/search/youtube?q=cats").status_code == 200