django rest framework throttlers

throttling.py

import time

from rest_framework.throttling import SimpleRateThrottle

import redis


class ApiKeyRateThrottle(SimpleRateThrottle):
    scope = "api_key"

    def get_cache_key(self, request, view):
        if not request.user or not request.user.is_authenticated:
            return None
        return self.cache_format % {"scope": self.scope, "ident": request.user.pk}


class ApiKeyRedisSlidingWindowLogThrottle(ApiKeyRateThrottle):
    """
    Sliding window log with redis
    """

    def allow_request(self, request, view):
        """
        Implement the check to see if the request should be throttled.
        """
        if self.rate is None:
            return True

        self.key = self.get_cache_key(request, view)

        if self.key is None:
            return True

        redisclient: redis.Redis = self.cache._cache.get_client(self.key, write=True)

        now = self.timer()
        req_key = str(now)

        pipeline = redisclient.pipeline()

        pipeline.zremrangebyscore(self.key, 0, now - self.duration)
        pipeline.zadd(self.key, {req_key: now})
        pipeline.expire(self.key, self.duration)
        pipeline.zcard(self.key)

        _, _, _, num_requests_in_window = pipeline.execute()

        if num_requests_in_window > self.num_requests:
            # remove key, request was not processed and doesnt count
            redisclient.zrem(self.key, req_key)
            return False

        return True

    def wait(self):
        """
        Optionally, return a recommended number of seconds to wait before
        the next request.
        """
        return self.duration