disable_rate_limit.py

"""Utilities for disabling Django Rest Framework rate limiting.

This is useful for performance tests in which we need to generate
a lot of traffic from a particular IP address.  By default,
Django Rest Framework uses the IP address to throttle traffic
for users who are not authenticated.

To disable rate limiting:

1) Decorate the Django Rest Framework APIView with `@can_disable_rate_limit`
2) In Django's admin interface, set `RateLimitConfiguration.enabled` to False.

Note: You should NEVER disable rate limiting in production.

"""
from functools import wraps
import logging
from rest_framework.views import APIView
from util.models import RateLimitConfiguration


LOGGER = logging.getLogger(__name__)


def _check_throttles_decorator(func):
    """Decorator for `APIView.check_throttles`.

    The decorated function will first check model-based config
    to see if rate limiting is disabled; if so, it skips
    the throttle check.  Otherwise, it calls the original
    function to enforce rate-limiting.

    Arguments:
        func (function): The function to decorate.

    Returns:
        The decorated function.

    """
    @wraps(func)
    def _decorated(*args, **kwargs):
        # Skip the throttle check entirely if we've disabled rate limiting.
        # Otherwise, perform the checks (as usual)
        if RateLimitConfiguration.current().enabled:
            return func(*args, **kwargs)
        else:
            msg = "Rate limiting is disabled because `RateLimitConfiguration` is not enabled."
            LOGGER.info(msg)
            return

    return _decorated


def can_disable_rate_limit(clz):
    """Class decorator that allows rate limiting to be disabled.

    Arguments:
        clz (class): The APIView subclass to decorate.

    Returns:
        class: the decorated class.

    Example Usage:
        >>> from rest_framework.views import APIView
        >>> @can_disable_rate_limit
        >>> class MyApiView(APIView):
        >>>     pass

    """
    # No-op if the class isn't a Django Rest Framework view.
    if not issubclass(clz, APIView):
        msg = (
            u"{clz} is not a Django Rest Framework APIView subclass."
        ).format(clz=clz)
        LOGGER.warning(msg)
        return clz

    # If we ARE explicitly disabling rate limiting,
    # modify the class to always allow requests.
    # Note that this overrides both rate limiting applied
    # for the particular view, as well as global rate limits
    # configured in Django settings.
    if hasattr(clz, 'check_throttles'):
        clz.check_throttles = _check_throttles_decorator(clz.check_throttles)

    return clz