241 lines
8.3 KiB
Python
241 lines
8.3 KiB
Python
"""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
|