DRF-Throttle组件源码分析及改编源码

生活在水面的鱼

1. 限流组件源码分析

注意：以下源码为方便理解已进行简化，只保留了与限流组件相关的代码。注释前的数字代表执行顺序。

SimpleRateThrottle类的源码解析：

2. 实践：编写一个限流类

1. from rest_framework.throttling import SimpleRateThrottle
2. from django.core.cache import cache as default_cache
4. class IpThrottle(SimpleRateThrottle):
5.     scope = "ip"  # 节流类的名称
6.     # THROTTLE_RATES = {"x1": "5/m"}   类中未定义THROTTLE_RATES则去settings中获取
7.     cache = default_cache  # 节流类指定写入的缓存
9.     def get_cache_key(self, request, view):
10.         ident = self.get_ident(request)  # 获取请求用户IP（去request中找请求头）
11.         return self.cache_format % {'scope': self.scope, 'ident': ident}
13. class UserThrottle(SimpleRateThrottle):
14.     scope = "user"  # 节流类的名称
15.     # THROTTLE_RATES = {"x1": "5/m"}   类中未定义THROTTLE_RATES则去settings中获取
16.     cache = default_cache  # 节流类指定写入的缓存
18.     def get_cache_key(self, request, view):
19.         ident = request.user.pk  # 用户ID
20.         return self.cache_format % {'scope': self.scope, 'ident': ident}

复制代码

3. 源码改编

• 重写限流的时长规则：原本只支持 1分钟x次，重写后支持 x分钟x次：
  

1. def parse_rate(self, rate):
2.         '''重写获取时长的parse_rate方法，原本只能写1/m或2/m，现在支持：1/5m, 3/10m的格式'''
3.         if rate is None:
4.                 return (None, None)
5.         num, period = rate.split('/')  # "1/5m"
6.         num_requests = int(num)
7.         duration = {'s': 1, 'm': 60, 'h': 3600, 'd': 86400}[period[-1]]
8.         count = int(period[0:-1])
9.         return (num_requests, duration * count)

复制代码

• 请求操作失败时，不计入限流次数
  

1. # 思路：在视图的操作成功后再进行计入操作（throttle_success）
2. # 限流类中： 重写throttle_success操作，并把计次的操作提取出来，自定义一个函数，在视图操作成功后才调用该函数进行计次
3. class MyThrottle(SimpleRateThrottle):
4.        
5.         def throttle_success(self):
6.                 # self.history.insert(0, self.now)
7.                 # self.cache.set(self.key, self.history, self.duration)
8.                 return True
10.         def done(self):
11.                 """ 视图函数执行成功后，再调用throttle_success里的方法"""
12.                 self.history.insert(0, self.now)
13.                 self.cache.set(self.key, self.history, self.duration)
16. # 视图类中：
17. class MyView(APIView):
18.         def post(request):
19.         """中间代码省略"""
20.         # 成功操作后才记录一次限流   
21.         for throttle in self.get_throttles():
22.                 throttle.done()

复制代码

• 自定义限流类的报错信息
  

1. # 在限流类中：
2. # 自定义抛出报错的异常类
3. class ThrottledException(exceptions.APIException):
4.         status_code = status.HTTP_429_TOO_MANY_REQUESTS
5.         default_code = 'throttled'
8. #重写throttle_failure方法，自定义报错信息并直接抛出异常；
9. def throttle_failure(self):
10.         wait = self.wait()
11.         detail = {
12.                 "code": return_code.TOO_MANY_REQUESTS,
13.                 "data": "访问频率限制",
14.                 'detail': "需等待{}秒后才能访问".format(int(wait))
15.         }
16.         raise ThrottledException(detail)

复制代码

来源:https://www.cnblogs.com/harry6/p/18508725
免责声明：由于采集信息均来自互联网，如果侵犯了您的权益，请联系我们【E-Mail:cb@itdo.tech】 我们会及时删除侵权内容，谢谢合作！