hibiapi/utils/cache.py

import hashlib
from collections.abc import Awaitable
from datetime import timedelta
from functools import wraps
from typing import Any, Callable, Optional, TypeVar, cast

from cashews import Cache
from pydantic import BaseModel
from pydantic.decorator import ValidatedFunction

from .config import Config
from .log import logger

CACHE_CONFIG_KEY = "_cache_config"

AsyncFunc = Callable[..., Awaitable[Any]]
T_AsyncFunc = TypeVar("T_AsyncFunc", bound=AsyncFunc)


CACHE_ENABLED = Config["cache"]["enabled"].as_bool()
CACHE_DELTA = timedelta(seconds=Config["cache"]["ttl"].as_number())
CACHE_URI = Config["cache"]["uri"].as_str()
CACHE_CONTROLLABLE = Config["cache"]["controllable"].as_bool()

cache = Cache(name="hibiapi")
try:
    cache.setup(CACHE_URI)
except Exception as e:
    logger.warning(
        f"Cache URI <y>{CACHE_URI!r}</y> setup <r><b>failed</b></r>: "
        f"<r>{e!r}</r>, use memory backend instead."
    )


class CacheConfig(BaseModel):
    endpoint: AsyncFunc
    namespace: str
    enabled: bool = True
    ttl: timedelta = CACHE_DELTA

    @staticmethod
    def new(
        function: AsyncFunc,
        *,
        enabled: bool = True,
        ttl: timedelta = CACHE_DELTA,
        namespace: Optional[str] = None,
    ):
        return CacheConfig(
            endpoint=function,
            enabled=enabled,
            ttl=ttl,
            namespace=namespace or function.__qualname__,
        )


def cache_config(
    enabled: bool = True,
    ttl: timedelta = CACHE_DELTA,
    namespace: Optional[str] = None,
):
    def decorator(function: T_AsyncFunc) -> T_AsyncFunc:
        setattr(
            function,
            CACHE_CONFIG_KEY,
            CacheConfig.new(function, enabled=enabled, ttl=ttl, namespace=namespace),
        )
        return function

    return decorator


disable_cache = cache_config(enabled=False)


class CachedValidatedFunction(ValidatedFunction):
    def serialize(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> BaseModel:
        values = self.build_values(args=args, kwargs=kwargs)
        return self.model(**values)


def endpoint_cache(function: T_AsyncFunc) -> T_AsyncFunc:
    from .routing import request_headers, response_headers

    vf = CachedValidatedFunction(function, config={})
    config = cast(
        CacheConfig,
        getattr(function, CACHE_CONFIG_KEY, None) or CacheConfig.new(function),
    )

    config.enabled = CACHE_ENABLED and config.enabled

    @wraps(function)
    async def wrapper(*args, **kwargs):
        cache_policy = "public"

        if CACHE_CONTROLLABLE:
            cache_policy = request_headers.get().get("cache-control", cache_policy)

        if not config.enabled or cache_policy.casefold() == "no-store":
            return await vf.call(*args, **kwargs)

        key = (
            f"{config.namespace}:"
            + hashlib.md5(
                (model := vf.serialize(args=args, kwargs=kwargs))
                .json(exclude={"self"}, sort_keys=True, ensure_ascii=False)
                .encode()
            ).hexdigest()
        )

        response_header = response_headers.get()
        result: Optional[Any] = None

        if cache_policy.casefold() == "no-cache":
            await cache.delete(key)
        elif result := await cache.get(key):
            logger.debug(f"Request hit cache <b><e>{key}</e></b>")
            response_header.setdefault("X-Cache-Hit", key)

        if result is None:
            result = await vf.execute(model)
            await cache.set(key, result, expire=config.ttl)

        if (cache_remain := await cache.get_expire(key)) > 0:
            response_header.setdefault("Cache-Control", f"max-age={cache_remain}")

        return result

    return wrapper  # type:ignore