spb/ai-atlas

1"""Redis helpers: API cache (prefix `aia:api:<db namespace>:`), distributed locks (`aia:lock:`), heartbeats.23The API cache is namespaced by a hash of `DATABASE_URL`: two API processes pointing at different databases (dev on :8331, a production4copy on :8332) but sharing one Redis db must never serve each other's bodies or derived id sets (that is how `/frontier` once reported5frontier ids from another database and found no prices for them)."""6from __future__ import annotations78import asyncio9import hashlib10import logging11from collections.abc import AsyncIterator12from contextlib import asynccontextmanager13from typing import Any1415import orjson16from redis.asyncio import Redis1718from aiatlas.config import settings1920log = logging.getLogger(__name__)21_redis: Redis | None = None22_NS: str | None = None232425def redis() -> Redis:26    global _redis27    if _redis is None:28        _redis = Redis.from_url(settings.redis_url, decode_responses=False, socket_connect_timeout=3, socket_timeout=5)29    return _redis303132def namespace() -> str:33    """8-hex namespace derived from the database URL (host, port, database name) — same DB ⇒ same cache, different DB ⇒ disjoint keys."""34    global _NS35    if _NS is None:36        _NS = hashlib.sha1(settings.database_url.encode()).hexdigest()[:8]37    return _NS383940def api_key(key: str) -> str:41    return f"aia:api:{namespace()}:{key}"424344async def cache_get(key: str) -> Any | None:45    try:46        raw = await redis().get(api_key(key))47        return orjson.loads(raw) if raw else None48    except Exception:  # noqa: BLE00149        return None505152async def cache_set(key: str, value: Any, ttl_s: int = 300) -> None:53    try:54        await redis().set(api_key(key), orjson.dumps(value, default=str), ex=ttl_s)55    except Exception:  # noqa: BLE00156        pass575859async def cache_invalidate(prefix: str = "") -> int:60    """Delete `aia:api:<ns>:<prefix>*` — only this database's cache."""61    n = 062    try:63        async for key in redis().scan_iter(match=f"{api_key(prefix)}*", count=500):64            await redis().delete(key)65            n += 166    except Exception:  # noqa: BLE00167        pass68    return n697071@asynccontextmanager72async def lock(name: str, ttl_s: int = 3600) -> AsyncIterator[bool]:73    """Best-effort distributed lock. Yields False when someone else holds it (callers must skip)."""74    key = f"aia:lock:{name}"75    acquired = False76    try:77        acquired = bool(await redis().set(key, b"1", nx=True, ex=ttl_s))78    except Exception as exc:  # noqa: BLE00179        log.warning("redis unavailable, proceeding without lock", extra={"error": str(exc)})80        acquired = True81    try:82        yield acquired83    finally:84        if acquired:85            try:86                await redis().delete(key)87            except Exception:  # noqa: BLE00188                pass899091async def heartbeat(service: str, payload: dict[str, Any], ttl_s: int = 180) -> None:92    try:93        await redis().set(f"aia:heartbeat:{service}", orjson.dumps(payload, default=str), ex=ttl_s)94    except Exception:  # noqa: BLE00195        pass969798async def heartbeats() -> dict[str, Any]:99    out: dict[str, Any] = {}100    try:101        async for key in redis().scan_iter(match="aia:heartbeat:*"):102            raw = await redis().get(key)103            if raw:104                out[key.decode().rsplit(":", 1)[-1]] = orjson.loads(raw)105    except Exception:  # noqa: BLE001106        pass107    return out108109110async def close() -> None:111    global _redis112    if _redis is not None:113        await _redis.aclose()114        _redis = None115    await asyncio.sleep(0)116117118__all__ = ["api_key", "cache_get", "cache_invalidate", "cache_set", "close", "heartbeat", "heartbeats", "lock", "namespace", "redis"]119