"""Async PocketBase REST client for TipiLAN Bot. Handles admin authentication (auto-refreshed), and CRUD operations on the economy_users collection. Uses aiohttp, which discord.py already depends on. Environment variables (set in .env): PB_URL Base URL of PocketBase (default: http://127.0.0.1:8090) PB_ADMIN_EMAIL PocketBase admin e-mail PB_ADMIN_PASSWORD PocketBase admin password PB_ECONOMY_COLLECTION_DEV / PB_ECONOMY_COLLECTION_ECONOMY """ from __future__ import annotations import asyncio import logging import time from typing import Any import aiohttp import config class DatabaseError(Exception): """Raised when PocketBase is unreachable or returns an error.""" pass _log = logging.getLogger("tipiCOIN.pb") PB_URL = config.PB_URL PB_ADMIN_EMAIL = config.PB_ADMIN_EMAIL PB_ADMIN_PASSWORD = config.PB_ADMIN_PASSWORD ECONOMY_COLLECTION = config.PB_ECONOMY_COLLECTION _TIMEOUT = aiohttp.ClientTimeout(total=4) # --------------------------------------------------------------------------- # Persistent session (created once, reused for the lifetime of the process) # --------------------------------------------------------------------------- _session: aiohttp.ClientSession | None = None def _get_session() -> aiohttp.ClientSession: global _session if _session is None or _session.closed: _session = aiohttp.ClientSession(timeout=_TIMEOUT) return _session # --------------------------------------------------------------------------- # Auth token cache # --------------------------------------------------------------------------- _token: str = "" _token_expiry: float = 0.0 _auth_lock = asyncio.Lock() async def _ensure_auth() -> str: global _token, _token_expiry async with _auth_lock: if time.monotonic() < _token_expiry: return _token session = _get_session() try: async with session.post( f"{PB_URL}/api/collections/_superusers/auth-with-password", json={"identity": PB_ADMIN_EMAIL, "password": PB_ADMIN_PASSWORD}, ) as resp: if resp.status != 200: text = await resp.text() raise DatabaseError(f"PocketBase auth failed ({resp.status}): {text}") data = await resp.json() _token = data["token"] _token_expiry = time.monotonic() + 13 * 24 * 3600 # refresh well before expiry _log.debug("PocketBase admin token refreshed") except (aiohttp.ClientConnectorError, asyncio.TimeoutError) as e: raise DatabaseError(f"Database unavailable: {e}") from e return _token async def _hdrs() -> dict[str, str]: return {"Authorization": await _ensure_auth()} def _invalidate_token() -> None: global _token_expiry _token_expiry = 0.0 # --------------------------------------------------------------------------- # Request helper with auth-retry and error wrapping # --------------------------------------------------------------------------- async def _request(method: str, url: str, **kwargs: Any) -> Any: """Make an authenticated request, retrying once on 401/403 by re-authing. Returns the parsed JSON body. Raises DatabaseError on connection issues or non-2xx responses after retrying. """ session = _get_session() for attempt in range(2): kwargs["headers"] = await _hdrs() try: async with session.request(method, url, **kwargs) as resp: if resp.status in (401, 403) and attempt == 0: _invalidate_token() continue if not resp.ok: text = await resp.text() raise DatabaseError(f"Database unavailable: {resp.status}, {text}") return await resp.json() except (aiohttp.ClientConnectorError, asyncio.TimeoutError) as e: raise DatabaseError(f"Database unavailable: {e}") from e # --------------------------------------------------------------------------- # CRUD helpers # --------------------------------------------------------------------------- async def get_record(user_id: str) -> dict[str, Any] | None: """Fetch one economy record by Discord user_id. Returns None if not found.""" data = await _request( "GET", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", params={"filter": f'user_id="{user_id}"', "perPage": 1}, ) 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).""" return await _request( "POST", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", json=record, ) async def update_record(record_id: str, data: dict[str, Any]) -> dict[str, Any]: """PATCH an existing record by its PocketBase record id.""" return await _request( "PATCH", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records/{record_id}", json=data, ) async def get_collection_fields() -> set[str]: """Return the field names defined on the economy collection's schema.""" data = await _request("GET", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}") return {f["name"] for f in data.get("fields", [])} async def count_records() -> int: """Return the total number of records in the collection (single cheap request).""" data = await _request( "GET", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", params={"perPage": 1, "page": 1}, ) return int(data.get("totalItems", 0)) async def list_all_records(page_size: int = 500) -> list[dict[str, Any]]: """Fetch every record in the collection, handling PocketBase pagination.""" results: list[dict[str, Any]] = [] page = 1 while True: data = await _request( "GET", f"{PB_URL}/api/collections/{ECONOMY_COLLECTION}/records", params={"perPage": page_size, "page": page}, ) batch = data.get("items", []) results.extend(batch) if len(batch) < page_size: return results page += 1