diff --git a/core/pb_client.py b/core/pb_client.py index adf4b87..8cbfc45 100644 --- a/core/pb_client.py +++ b/core/pb_client.py @@ -75,6 +75,11 @@ async def _hdrs() -> dict[str, str]: return {"Authorization": await _ensure_auth()} +def _invalidate_token() -> None: + global _token_expiry + _token_expiry = 0.0 + + # --------------------------------------------------------------------------- # CRUD helpers # --------------------------------------------------------------------------- @@ -82,54 +87,70 @@ async def _hdrs() -> dict[str, str]: async def get_record(user_id: str) -> dict[str, Any] | None: """Fetch one economy record by Discord user_id. Returns None if not found.""" session = _get_session() - async with session.get( - f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", - params={"filter": f'user_id="{user_id}"', "perPage": 1}, - headers=await _hdrs(), - ) as resp: - resp.raise_for_status() - data = await resp.json() - items = data.get("items", []) - return items[0] if items else None + for attempt in range(2): + async with session.get( + f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", + params={"filter": f'user_id="{user_id}"', "perPage": 1}, + headers=await _hdrs(), + ) as resp: + if resp.status == 403 and attempt == 0: + _invalidate_token() + continue + resp.raise_for_status() + data = await resp.json() + items = data.get("items", []) + return items[0] if items else None async def create_record(record: dict[str, Any]) -> dict[str, Any]: """Create a new economy record. Returns the created record (includes PB id).""" session = _get_session() - async with session.post( - f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", - json=record, - headers=await _hdrs(), - ) as resp: - if resp.status not in (200, 201): - text = await resp.text() - raise RuntimeError(f"PocketBase create failed ({resp.status}): {text}") - return await resp.json() + for attempt in range(2): + async with session.post( + f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", + json=record, + headers=await _hdrs(), + ) as resp: + if resp.status == 403 and attempt == 0: + _invalidate_token() + continue + if resp.status not in (200, 201): + text = await resp.text() + raise RuntimeError(f"PocketBase create failed ({resp.status}): {text}") + return await resp.json() async def update_record(record_id: str, data: dict[str, Any]) -> dict[str, Any]: """PATCH an existing record by its PocketBase record id.""" session = _get_session() - async with session.patch( - f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records/{record_id}", - json=data, - headers=await _hdrs(), - ) as resp: - resp.raise_for_status() - return await resp.json() + for attempt in range(2): + async with session.patch( + f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records/{record_id}", + json=data, + headers=await _hdrs(), + ) as resp: + if resp.status == 403 and attempt == 0: + _invalidate_token() + continue + resp.raise_for_status() + return await resp.json() async def count_records() -> int: """Return the total number of records in the collection (single cheap request).""" session = _get_session() - async with session.get( - f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", - params={"perPage": 1, "page": 1}, - headers=await _hdrs(), - ) as resp: - resp.raise_for_status() - data = await resp.json() - return int(data.get("totalItems", 0)) + for attempt in range(2): + async with session.get( + f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", + params={"perPage": 1, "page": 1}, + headers=await _hdrs(), + ) as resp: + if resp.status == 403 and attempt == 0: + _invalidate_token() + continue + resp.raise_for_status() + data = await resp.json() + return int(data.get("totalItems", 0)) async def list_all_records(page_size: int = 500) -> list[dict[str, Any]]: @@ -137,18 +158,23 @@ async def list_all_records(page_size: int = 500) -> list[dict[str, Any]]: results: list[dict[str, Any]] = [] page = 1 session = _get_session() - hdrs = await _hdrs() while True: - async with session.get( - f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", - params={"perPage": page_size, "page": page}, - headers=hdrs, - ) as resp: - resp.raise_for_status() - data = await resp.json() - batch = data.get("items", []) - results.extend(batch) - if len(batch) < page_size: + hdrs = await _hdrs() + for attempt in range(2): + async with session.get( + f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", + params={"perPage": page_size, "page": page}, + headers=hdrs, + ) as resp: + if resp.status == 403 and attempt == 0: + _invalidate_token() + hdrs = await _hdrs() + continue + resp.raise_for_status() + data = await resp.json() + batch = data.get("items", []) + results.extend(batch) + if len(batch) < page_size: + return results + page += 1 break - page += 1 - return results