四月二十一日

# -*- encoding:utf-8 -*-
# @time: 2023/4/21 15:41
# @author: ifeng
from django.core.cache import cache as default_cache
from rest_framework import exceptions
from rest_framework import status
from rest_framework.throttling import SimpleRateThrottle

class ThrottledException(exceptions.APIException):
status_code = status.HTTP_429_TOO_MANY_REQUESTS
default_code = 'throttled'

class MyRateThrottle(SimpleRateThrottle):
cache = default_cache # 访问记录存放在django的缓存中(需设置缓存)
scope = 'user' # 构造缓存中的key
cache_format = 'throttle_%(scope)s_%(ident)s'

# 设置其他访问评率, 例如: 一分钟允许访问10次
# 其他: 's': 'sec', 'm': 'min', 'h': 'hour', 'd': 'day'
THROTTLE_RATES = {'user': '10/m'}

def get_cache_key(self, request, view):
if request.user:
ident = request.user.id
else:
ident = self.get_ident(request) # 获取请求用户IP(request中找请求头)

# throttle_u # throttle_user_11.11.11.11ser_2

return self.cache_format % {'scope': self.scope, 'ident': ident}

def throttle_failure(self):
wait = self.wait()
detail = {
'code': 1005,
'data': '访问频率限制',
'detail': '需要等待 %s s才能访问' % (int(wait))
}
raise ThrottledException(detail)