fastapi-best-architecture/backend/common/cache/decorator.py

import functools

from collections.abc import Callable, Sequence
from typing import Any, ParamSpec, TypeVar

from cachebox import make_hash_key
from msgspec import json

from backend.common.cache.local import local_cache_manager
from backend.common.cache.pubsub import cache_pubsub_manager
from backend.common.context import ctx
from backend.common.exception import errors
from backend.common.log import log
from backend.core.conf import settings
from backend.database.redis import redis_client
from backend.utils.serializers import select_columns_serialize, select_list_serialize

P = ParamSpec('P')
T = TypeVar('T')

# 哈希缓存键排除参数
_EXCLUDE_PARAMS = frozenset({'db', 'session', 'self', 'cls', 'request', 'response'})


def build_cache_key(
    name: str,
    key: str | None,
    key_builder: Callable[..., str] | None,
    *args: Any,
    **kwargs: Any,
) -> str:
    """构建缓存 Key"""
    if key_builder:
        return f'{name}:{key_builder(*args, **kwargs)}'

    if key:
        value = kwargs.get(key)
        if value is None:
            raise errors.ServerError(msg=f'缓存键构建失败，参数 "{key}" 不存在或值为空')
        return f'{name}:{value}'

    filtered = {k: v for k, v in kwargs.items() if k not in _EXCLUDE_PARAMS and v is not None}

    if filtered:
        hash_suffix = make_hash_key(*args, **kwargs)
        return f'{name}:{hash_suffix}'

    return name


def user_key_builder() -> str:
    """基于当前用户 ID 生成缓存 Key"""
    user_id = ctx.user_id
    if user_id is None:
        raise errors.ServerError(msg='用户缓存键构建失败')
    return str(user_id)


def _serialize_result(result: Any) -> bytes:
    """
    序列化缓存结果

    :param result: 需要进行序列化的结果
    :return:
    """
    # SQLAlchemy 查询表
    if hasattr(result, '__table__'):
        return json.encode(select_columns_serialize(result))

    # SQLAlchemy 查询列表
    if (
        isinstance(result, Sequence)
        and not isinstance(result, (str, bytes))
        and len(result) > 0
        and hasattr(result[0], '__table__')
    ):
        return json.encode(select_list_serialize(result))

    # 基本类型
    return json.encode(result)


def _deserialize_result(value: bytes) -> Any:
    """
    反序列化缓存结果

    :param value: 缓存结果
    :return:
    """
    try:
        return json.decode(value)
    except Exception:
        return value


def cached(  # noqa: C901
    name: str,
    *,
    key: str | None = None,
    key_builder: Callable[..., str] | None = None,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """
    缓存装饰器

    :param name: 缓存名称（通常为缓存 Key 前缀）
    :param key: 从方法参数中获取指定参数名的值作为缓存 Key，与 key_builder 互斥
    :param key_builder: 自定义 Key 生成函数，与 key 互斥
    :return:
    """
    if key is not None and key_builder is not None:
        raise errors.ServerError(msg='缓存 key 和 key_builder 不能同时使用')

    def decorator(func: Callable[P, T]) -> Callable[P, T]:  # noqa: C901
        @functools.wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            cache_key = build_cache_key(name, key, key_builder, *args, **kwargs)

            # L1: 本地缓存
            if settings.CACHE_LOCAL_ENABLED:
                local_value = local_cache_manager.get(cache_key)
                if local_value is not None:
                    return local_value

            # L2: Redis 缓存
            try:
                redis_value = await redis_client.get(cache_key)
                if redis_value is not None:
                    result = _deserialize_result(redis_value)
                    # 回填 L1
                    if settings.CACHE_LOCAL_ENABLED:
                        local_cache_manager.set(cache_key, result)
                    return result
            except Exception as e:
                log.warning(f'[Cache] GET error: {e}')

            # 缓存未命中
            result = await func(*args, **kwargs)

            if result is not None:
                try:
                    # 回填 L1
                    if settings.CACHE_LOCAL_ENABLED:
                        local_cache_manager.set(cache_key, result)

                    # 回填 L2
                    serialized_result = _serialize_result(result)
                    if settings.CACHE_REDIS_TTL:
                        await redis_client.setex(cache_key, settings.CACHE_REDIS_TTL, serialized_result)
                    else:
                        await redis_client.set(cache_key, serialized_result)
                except Exception as e:
                    log.warning(f'[Cache] SET error: {e}')

            return result

        return wrapper

    return decorator


def cache_invalidate(  # noqa: C901
    name: str,
    *,
    key: str | None = None,
    key_builder: Callable[..., str] | None = None,
    atomic: bool = True,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """
    缓存失效装饰器

    :param name: 缓存名称（通常为缓存 Key 前缀）
    :param key: 从方法参数中获取指定参数名的值作为缓存 Key，与 key_builder 互斥
    :param key_builder: 自定义 Key 生成函数，与 key 互斥
    :param atomic: 是否保证缓存原子性
    :return:
    """
    if key is not None and key_builder is not None:
        raise errors.ServerError(msg='缓存 key 和 key_builder 不能同时使用')

    def decorator(func: Callable[P, T]) -> Callable[P, T]:
        @functools.wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            result = await func(*args, **kwargs)

            # 尝试失效缓存
            invalidate_success = False
            invalidate_error = None

            try:
                invalidate_key = build_cache_key(name, key, key_builder, *args, **kwargs)

                # L2 缓存失效
                if invalidate_key == name:
                    await redis_client.delete(invalidate_key)
                else:
                    await redis_client.delete_prefix(invalidate_key)

                # L1 缓存失效
                if settings.CACHE_LOCAL_ENABLED:
                    if invalidate_key == name:
                        local_cache_manager.delete(invalidate_key)
                    else:
                        local_cache_manager.delete_prefix(invalidate_key)

                # 广播失效消息（通知其他节点清除本地缓存）
                if settings.CACHE_LOCAL_ENABLED:
                    if invalidate_key == name:
                        await cache_pubsub_manager.publish_invalidation(invalidate_key)
                    else:
                        await cache_pubsub_manager.publish_invalidation(invalidate_key, is_delete_prefix=True)

            except Exception as e:
                log.error(f'[Cache] INVALIDATE error: {e}')
                invalidate_error = e
            else:
                invalidate_success = True

            # 原子性检查
            if atomic and not invalidate_success:
                raise errors.ServerError(msg='缓存失效失败，数据可能不一致', data=invalidate_error)

            return result

        return wrapper

    return decorator