FastAPI 接入 FastAPI-Limiter 以及使用 Redis 进行限流指南

fastapi-limiter 更新至 v0.2.0 后与之前不太不一样，这里以接入 Redis 进行示例说明

安装：

pip install redis[hiredis]==7.4.0
pip install fastapi-limiter

创建 Redis 管理器

def _get_redis_pool(use_cache_db: bool = True):
    uri = settings.REDIS_CACHE_URI if use_cache_db else settings.REDIS_URI

    return redis.ConnectionPool.from_url(uri.encoded_string(), decode_responses=True)


redis_client = redis.Redis(connection_pool=_get_redis_pool(False))
redis_cache_client = redis.Redis(connection_pool=_get_redis_pool())

注意这里我使用了一个参数来区别 Redis 的 DB，实际上这里的 URI 来自 pydantic，示例如下：

@computed_field
@property
def REDIS_URI(self) -> RedisDsn:
    return RedisDsn.build(
        scheme="redis",
        host=self.REDIS_HOST,
        port=self.REDIS_PORT,
        password=self.REDIS_PASSWORD or None,
        path=str(self.REDIS_DB)
    )

@computed_field
@property
def REDIS_CACHE_URI(self) -> RedisDsn:
    return RedisDsn.build(
        scheme="redis",
        host=self.REDIS_HOST,
        port=self.REDIS_PORT,
        password=self.REDIS_PASSWORD or None,
        path=str(self.REDIS_CACHE_DB)
    )

创建 pyrate_limiter 限制器并使用 Redis

因为 fastapi-limiter 基于 pyrate-limiter 实现，这里我们需要先创建限制器，然后提供给 fastapi_limiter

from pyrate_limiter import Rate, Duration, Limiter, RedisBucket

# 用于区分全局和业务的不同限制
_GLOBAL_BUCKET_KEY = "global-rate-limit"
_BIZ_BUCKET_KEY = "biz-rate-limit"

# 自定义默认限流
_DEFAULT_RATES = [
    Rate(60, Duration.MINUTE),          # 60 次每分钟
    Rate(60 * 5, Duration.MINUTE * 10), # 300 次每十分钟
    Rate(1000, Duration.HOUR)           # 1000 次每小时
]


def get_redis_limiter(rates: list[Rate] | None = None, is_global: bool = False) -> Limiter:
    if is_global:
        bucket_key = _GLOBAL_BUCKET_KEY

        if not rates:
            rates = _DEFAULT_RATES
    else:
        if not rates:
            raise ValueError("Biz limit rate rules cannot be None")

        bucket_key = _BIZ_BUCKET_KEY
    
    # 使用我们已有 redis 管理器中的客户端
    bucket = RedisBucket.init(rates, redis_client, bucket_key)

    return Limiter(bucket)

依赖注入方法

这里可在前置路由依赖中先校验 jwt 后将用户编号或者 ID 放到 request.state 中，随后依赖该 get_rate_limiter 方法

须实现用户唯一标识才可用于判断是否登录并用来实现精确流控

from fastapi import Request
from fastapi_limiter.depends import RateLimiter
from fastapi_limiter.identifier import default_identifier
from pyrate_limiter import Rate


async def get_identifier(request: Request) -> str:
    # 自定义实现逻辑
    user_no = getattr(request.state, "user_no", None)

    if user_no:
        return f"user_{user_no}:{request.scope["path"]}"

    return await default_identifier(request)


def get_rate_limiter(rules: list[Rate] | None = None, is_global: bool = False) -> RateLimiter:
    return RateLimiter(get_redis_limiter(rules, is_global), get_identifier)

路由全局限流

此处示例，假设 get_current_user 方法中实现了对 request.state 的用户唯一身份设置，就能轻松区分是否已登录用户并分别限流

get_current_user 中应当实现对 jwt 的解码，拿到用户唯一编号后进行数据库查询，确认无误后设置到 request.state（参考逻辑）

from fastapi import APIRouter, Depends
from app.api.deps import get_current_user
from app.infrastructure.limiter.api_rate_limiter import get_rate_limiter

api = APIRouter(dependencies=[Depends(get_current_user), Depends(get_rate_limiter(is_global=True))])

路由单独限流

由于 FastAPI 的 dependencies 是叠加的，所以上方 RedisBucket 的 key 需要区分

否则在例如发送短信等接口上会因为全局限流已经写入了一条 Redis 记录，导致设置 1 分钟 1 次 的限流不可用，会直接触发 429 请求频繁

这里例如我们有一个发送短信接口：

_SEND_CODE_RATES = [
    Rate(1, Duration.MINUTE), # 1 次每分钟
    Rate(5, Duration.HOUR),   # 5 次每小时
    Rate(10, Duration.DAY)    # 10 次每天
]

@router.post("/send-code", dependencies=[Depends(get_rate_limiter(_SEND_CODE_RATES))])
async def send_code():
    pass

在触发一次短信发送后，会往 Redis 写入两条记录如下：

ZSET Key: global-rate-limit
Value: user_123456:/api/send-code:8:1:adee7ca88b:1

ZSET Key: biz-rate-limit
Value: user_123456:/api/send-code:8:2:7d2ecd023e:1