A.Mohammadi
A.Mohammadi

Reputation: 364

How to apply throttling just to create action in DRF viewsets?

I have a viewset and I want to apply throttling to only create action of that viewset and I don't want it to be applied to update, destroy, retrieve and etc...

class UserViewSet(viewsets.ModelViewSet):
    # should only be applied to the create action
    throttle_classes = [SomeThrottle] 
    ...

Upvotes: 8

Views: 1631

Answers (3)

Pablo Abdelhay
Pablo Abdelhay

Reputation: 1008

I built a Mixin for convenience. You can use it on your ViewSet and set the throttle_action_classes attribute as a dict of throttle_classes, where the key is the action and the value is the list of throttle_classes for that action.

class ThrottlePerActionMixin(object):
    """
    This mixing let you define throttles_classes per actions by passing a dict to throttle_action_classes attribute.
    The throttle_classes on throttle_action_classes will be added to the default's ViewSet throttle_classes.
    eg.:
    class MyViewSet(ModelViewSet):
        throttle_action_classes = {
            'create': [CustomThrottleClass],
            'retrieve': [AnonThrottleClass]
        }
    """

    def __init__(self, *args, **kwargs):
        if not getattr(self, 'throttle_action_classes', None) or \
            isinstance(self.throttle_action_classes, (list, tuple)):
            raise Exception("Must set throttle_action_classes as a dict. "
                            "eg.: throttle_action_classes = {'create': [CustomThrottle]}")
        super().__init__(*args, **kwargs)

    def get_throttles(self):
        throttles = super().get_throttles()
        for action, throttle_classes_for_action in self.throttle_action_classes.items():
            if self.action == action:
                if not isinstance(throttle_classes_for_action, (list, tuple)):
                    throttle_classes_for_action = [throttle_classes_for_action]
                for throttle_class in throttle_classes_for_action:
                    throttles.append(throttle_class())  # The throttles_class needs to be instantiated.
        return throttles

Upvotes: 1

C.K.
C.K.

Reputation: 5498

if use ScopedRateThrottle instead of Custom throttles

# settings

    'DEFAULT_THROTTLE_CLASSES': [
        'rest_framework.throttling.ScopedRateThrottle',
    ],
    'DEFAULT_THROTTLE_RATES': {
        'submit': '1/minute',
    }

class SomeViewSet(mixins.CreateModelMixin, GenericViewSet)
    ....

    def get_throttles(self):
        if self.action == 'create':
            self.throttle_scope = 'submit'
        return super().get_throttles()

Upvotes: 5

Abdul Aziz Barkat
Abdul Aziz Barkat

Reputation: 21787

As described in Introspecting ViewSet actions [DRF docs] you can inspect the action attribute to set the throttle_classes based on the current action by overriding the get_throttles method:

class UserViewSet(viewsets.ModelViewSet):
    def get_throttles(self):
        if self.action == 'create':
            throttle_classes = [SomeThrottle]
        else:
            throttle_classes = []  # No throttle for other actions
        return [throttle() for throttle in throttle_classes]

Upvotes: 10

Related Questions