django-rest-framework中如何自定义一个限流类

 1. 自定义一个限流类

class VisitThrottle(BaseThrottle):
    """自定义限流控制类, 60秒内只能访问3次"""
    def __init__(self):
        self.history = None

    def allow_request(self, request, view):
        # 获取用户的Ip地址, 此处的get_ident方法是父类中已经帮我们写好的获取远程ip地址
        remote_addr = self.get_ident(request)
        ctime = time.time()

        # 从缓存中获取字典数据
        VISIT_RECORD = cache.get("VISIT_RECORD", {})

        if remote_addr not in VISIT_RECORD:
            VISIT_RECORD[remote_addr] = [ctime]
            print(VISIT_RECORD)
            # 更新缓存中的数据
            cache.set("VISIT_RECORD", VISIT_RECORD)
            return True

        history = VISIT_RECORD.get(remote_addr)
        self.history = history

        while history and history[0] < ctime - 60:
            history.pop(0)

        if len(history) < 3:
            # 千万要注意这里：如果历史记录小于3条，一定要将当前访问时间追加进去，然后在返回True
            history.append(ctime)
            # 更新缓存中的数据
            cache.set("VISIT_RECORD", VISIT_RECORD)
            print(VISIT_RECORD)
            return True

        return False

    def wait(self):
        """还需要等多少秒才能访问"""
        ctime = time.time()
        wait_time = 60 - (ctime - self.history[0])
        return wait_time

2. 给登录认证接口增加限流配置

class AuthView(APIView):
    """用于用户登录认证，认证成功创建Token"""
    # 取消认证
    authentication_classes = []
    permission_classes = []
    throttle_classes = [VisitThrottle]

    def post(self, request):
        username = request._request.POST.get("username", "")
        password = request._request.POST.get("password", "")
        try:
            user = UserInfo.objects.get(username=username, password=password)
        except UserInfo.DoesNotExist:
            res = {"status": 1000, "msg": "用户名或密码错误"}
        else:
            # 创建唯一的用户token
            token = md5(user)
            # 没有就创建，有了就更新
            UserToken.objects.update_or_create(user=user, defaults={"token": token})
            res = {"status": 1001, "msg": "认证成功", "token": token}
        return JsonResponse(res)

3. 访问频率限制类可以简写成下面这种方式

from rest_framework.throttling import SimpleRateThrottle


class VisitThrottle(SimpleRateThrottle):
    # 匿名用户访问频率限制
    scope = "mayanan"

    def get_cache_key(self, request, view):
        return self.get_ident(request)


class UserThrottle(SimpleRateThrottle):
    # 登录用户访问频率限制
    scope = "mayananUser"

    def get_cache_key(self, request, view):
        return request.user.username

然后在全局配置文件中这样写：

# rest-framework配置
REST_FRAMEWORK = {
    "DEFAULT_THROTTLE_CLASSES": [
        "app01.my_throttle.UserThrottle",
    ],
    "DEFAULT_THROTTLE_RATES": {
        "mayanan": "3/m",
        "mayananUser": "5/m",
    },

因为我们全局配置的限流使用的是登录用户访问频率限制，所以如果想对单个类视图使用匿名用户访问频率限制，需要这样写

class AuthView(APIView):
    """用于用户登录认证，认证成功创建Token"""
    # 取消认证
    authentication_classes = []
    # 取消权限
    permission_classes = []
    # 设置匿名用户访问频率限制
    throttle_classes = [VisitThrottle]