在drf原生的流量控制中,只能支持1/m,1/s,1/h,1/d 等控制规则,不能实现1/5m,1/10s 等更加精细的控制。要是5分钟访问一次或多次,只能是通过12/h 这种来控制,不太友好。
重写SimpleRateThrottle 类的parse_rate方法,来满足开发的实际需求:
from rest_framework.throttling import SimpleRateThrottle
#重写了限制频率的方法,原始只能支持1/s,m,h,d , 不能支持5分钟一次,10秒中一次
class NewSimpleRateThrottle(SimpleRateThrottle):
def parse_rate(self, rate):
"""
这个方法,就是将2/10m 提取出2 和 10m
计算出用户需要等待多少秒,才能再次访问的
"""
if rate is None:
return (None, None)
num, period = rate.split('/') # 2/10m 10分钟两次请求
dic = {'s': 1, 'm': 60, 'h': 3600, 'd': 86400}
#1、计算需要等待的总秒数
if len(period)>1:
#这里,支持了1/5m,1/10s 等频率规则
key = period[-1] #必须是s,m,h,d
if key not in ('s','m','h','d'):
raise Exception(f'流量控制设置的rate有问题,不能是{rate},格式:1/10m')
try:
period_num = int(period.split(key)[0])
except Exception:
raise Exception(f'流量控制设置的rate有问题,不能是{rate},格式: 2/10m')
#求出需要的等待的总时长,展示到前端时,还需要除以int(num)
duration = dic[key]*period_num
elif len(period) == 1:
if period not in ('s','m','h','d'):
raise Exception(f'流量控制设置的rate有问题,不能是{rate},格式:3/10m')
#这里是对1/s,2/m,等的支持
duration = dic[period]
else:
#排除掉 1/ 2/ 等错误的频率规则
raise Exception(f'流量控制设置的rate有问题,不能是{rate},格式:4/m')
#2、除以整个数值,就是真实等待秒数
try:
num_requests = int(num)
except Exception:
raise Exception(f'流量控制设置的rate有问题,不能是{rate},格式: 5/m')
return (num_requests, duration)
#对请求源的IP进行限制
class IPThrottle(NewSimpleRateThrottle):
# 在settings中配置频率时使用到的关键字,有scope来指定
scope = 'IP'
def get_cache_key(self, request, view):
# 这里return什么,就以什么作为限制,上面的scope本来就是用于这里去配置文件中获取rate的
return request.META.get('REMOTE_ADDR')+self.scope
# 可以在这里设置访问频率
def get_rate(self):
return '10/m'
#对请求的用户进行限制
class UserThrottle(NewSimpleRateThrottle):
scope = 'UserID'
def get_cache_key(self, request, view):
# 这里return什么,就以什么作为限制,上面的scope本来就是用于这里去配置文件中获取rate的
user_pk = request.user.pk
return user_pk+self.scope
# 可以在这里设置访问频率,就无需到settings.py中设置了
def get_rate(self):
return '30/m'
#用户请求发送验证码的接口进行限制
class PhoneNumberThrottle(NewSimpleRateThrottle):
#唯一标识的一部分
scope = 'phoneNumber'
#以什么作为唯一标识
def get_cache_key(self, request, view):
# 这里return什么,就以什么作为限制,上面的scope本来是用于这里去配置文件中获取rate的
phone = request.data.get('phone')
return phone+self.scope
def get_rate(self):
#设置1分钟一次
return '1/5m'
REST_FRAMEWORK = {
#匿名用户:设置纯净版时,要设置好匿名用户
'UNAUTHENTICATED_USER':None,
#4、全局的使用的频率限制
'DEFAULT_THROTTLE_CLASSES': (
'base.throttle.IPThrottle', #使用的频率了
),
#4.1、对于scope=IP的限制
'DEFAULT_THROTTLE_RATES': {
'IP': '60/m'
# 一个ip一分钟只能访问60次 ,原始的get_rate 就是通过scope在这里获取到60/m,在设置时
},
}
在业务场景中,如短信服务,一般都是3分钟或5分钟或10分钟只能访问一次,可以有效的防止短信可使用条数被消耗掉。所以重写parse_rate 方法是很有必要,可以根据业务场景随意的设置控制的规则。