DRF-Throttle元件原始碼分析及改編原始碼

harry6發表於2024-10-27
元件原始碼

1. 限流元件原始碼分析

注意：以下原始碼為方便理解已進行簡化，只保留了與限流元件相關的程式碼。註釋前的數字代表執行順序。

image
image

SimpleRateThrottle類的原始碼解析：
image

2. 實踐：編寫一個限流類

from rest_framework.throttling import SimpleRateThrottle
from django.core.cache import cache as default_cache

class IpThrottle(SimpleRateThrottle):
    scope = "ip"  # 節流類的名稱
    # THROTTLE_RATES = {"x1": "5/m"}   類中未定義THROTTLE_RATES則去settings中獲取
    cache = default_cache  # 節流類指定寫入的快取

    def get_cache_key(self, request, view):
        ident = self.get_ident(request)  # 獲取請求使用者IP（去request中找請求頭）
        return self.cache_format % {'scope': self.scope, 'ident': ident}

class UserThrottle(SimpleRateThrottle):
    scope = "user"  # 節流類的名稱
    # THROTTLE_RATES = {"x1": "5/m"}   類中未定義THROTTLE_RATES則去settings中獲取
    cache = default_cache  # 節流類指定寫入的快取

    def get_cache_key(self, request, view):
        ident = request.user.pk  # 使用者ID
        return self.cache_format % {'scope': self.scope, 'ident': ident}

3. 原始碼改編

  • 重寫限流的時長規則：原本只支援 1分鐘x次，重寫後支援 x分鐘x次：
def parse_rate(self, rate):
	'''重寫獲取時長的parse_rate方法，原本只能寫1/m或2/m，現在支援：1/5m, 3/10m的格式'''
	if rate is None:
		return (None, None)
	num, period = rate.split('/')  # "1/5m"
	num_requests = int(num)
	duration = {'s': 1, 'm': 60, 'h': 3600, 'd': 86400}[period[-1]]
	count = int(period[0:-1])
	return (num_requests, duration * count)



  • 請求操作失敗時，不計入限流次數
# 思路：在檢視的操作成功後再進行計入操作（throttle_success）
# 限流類中： 重寫throttle_success操作，並把計次的操作提取出來，自定義一個函式，在檢視操作成功後才呼叫該函式進行計次
class MyThrottle(SimpleRateThrottle):
	
	def throttle_success(self):
		# self.history.insert(0, self.now)
		# self.cache.set(self.key, self.history, self.duration)
		return True

	def done(self):
		""" 檢視函式執行成功後，再呼叫throttle_success裡的方法"""
		self.history.insert(0, self.now)
		self.cache.set(self.key, self.history, self.duration)


# 檢視類中：
class MyView(APIView):
	def post(request):
	"""中間程式碼省略"""
	# 成功操作後才記錄一次限流   
	for throttle in self.get_throttles():
		throttle.done()



  • 自定義限流類的報錯資訊
# 在限流類中：
# 自定義丟擲報錯的異常類
class ThrottledException(exceptions.APIException):
	status_code = status.HTTP_429_TOO_MANY_REQUESTS
	default_code = 'throttled'


#重寫throttle_failure方法，自定義報錯資訊並直接丟擲異常；
def throttle_failure(self):
	wait = self.wait()
	detail = {
		"code": return_code.TOO_MANY_REQUESTS,
		"data": "訪問頻率限制",
		'detail': "需等待{}秒後才能訪問".format(int(wait))
	}
	raise ThrottledException(detail)

相關文章