tipibot/core/pb_client.py

"""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