"""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 / PB_ECONOMY_COLLECTION_LAN """ from __future__ import annotations import asyncio import logging import time from typing import Any import aiohttp import config _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=10) # --------------------------------------------------------------------------- # 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() 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 RuntimeError(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") return _token async def _hdrs() -> dict[str, str]: return {"Authorization": await _ensure_auth()} def _escape_filter_value(value: str) -> str: return value.replace("\\", "\\\\").replace('"', '\\"') # --------------------------------------------------------------------------- # 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.""" return await get_first_record( ECONOMY_COLLECTION, f'user_id="{_escape_filter_value(user_id)}"', ) async def get_first_record(collection: str, filter_expr: str) -> dict[str, Any] | None: """Fetch one record from any collection by a PocketBase filter expression.""" session = _get_session() async with session.get( f"{PB_URL}/api/collections/{collection}/records", params={"filter": filter_expr, "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 async def get_record_by_field(collection: str, field: str, value: str) -> dict[str, Any] | None: """Fetch one record where `field` exactly equals `value`.""" escaped = _escape_filter_value(value) return await get_first_record(collection, f'{field}="{escaped}"') 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 create_record_in(ECONOMY_COLLECTION, record) async def create_record_in(collection: str, record: dict[str, Any]) -> dict[str, Any]: """Create a new record in any collection. Returns the created record.""" session = _get_session() async with session.post( f"{PB_URL}/api/collections/{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() 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 update_record_in(ECONOMY_COLLECTION, record_id, data) async def update_record_in(collection: str, record_id: str, data: dict[str, Any]) -> dict[str, Any]: """PATCH an existing record in any collection by its PocketBase record id.""" session = _get_session() async with session.patch( f"{PB_URL}/api/collections/{collection}/records/{record_id}", json=data, headers=await _hdrs(), ) as resp: 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).""" return await count_records_in(ECONOMY_COLLECTION) async def count_records_in(collection: str) -> int: """Return the total number of records in any collection.""" session = _get_session() async with session.get( f"{PB_URL}/api/collections/{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)) async def list_all_records(page_size: int = 500) -> list[dict[str, Any]]: """Fetch every record in the collection, handling PocketBase pagination.""" return await list_all_records_in(ECONOMY_COLLECTION, page_size=page_size) async def list_all_records_in(collection: str, page_size: int = 500) -> list[dict[str, Any]]: """Fetch every record in any collection, handling PocketBase pagination.""" results: list[dict[str, Any]] = [] page = 1 session = _get_session() hdrs = await _hdrs() while True: async with session.get( f"{PB_URL}/api/collections/{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: break page += 1 return results async def upsert_record_by_field( collection: str, field: str, value: str, data: dict[str, Any], ) -> tuple[dict[str, Any], bool]: """Create or update a record. Returns (record, created).""" existing = await get_record_by_field(collection, field, value) if existing: return await update_record_in(collection, existing["id"], data), False return await create_record_in(collection, data), True async def get_collection(collection: str) -> dict[str, Any] | None: """Fetch collection metadata, returning None if it doesn't exist.""" session = _get_session() async with session.get( f"{PB_URL}/api/collections/{collection}", headers=await _hdrs(), ) as resp: if resp.status == 404: return None resp.raise_for_status() return await resp.json() async def create_collection(payload: dict[str, Any]) -> dict[str, Any]: """Create a PocketBase collection from a full collection payload.""" session = _get_session() async with session.post( f"{PB_URL}/api/collections", json=payload, headers=await _hdrs(), ) as resp: if resp.status not in (200, 201): text = await resp.text() raise RuntimeError(f"PocketBase collection create failed ({resp.status}): {text}") return await resp.json() async def ensure_collection(collection: str, payload: dict[str, Any]) -> bool: """Create `collection` when missing. Returns True if created.""" if await get_collection(collection): return False await create_collection(payload) return True