Add throttling for Django REST Framework

We add two limits to the number of requests from authenticated users and non-authenticated users
to prevent the APIs from being overloaded.

Author: hoangpn

promgen/settings.py

@@ -199,6 +199,18 @@
     "DEFAULT_FILTER_BACKENDS": ("django_filters.rest_framework.DjangoFilterBackend",),
     "DEFAULT_SCHEMA_CLASS": "promgen.schemas.CustomSchema",
     "EXCEPTION_HANDLER": "promgen.middleware.custom_exception_handler",
+    "DEFAULT_THROTTLE_CLASSES": [
+        "rest_framework.throttling.AnonRateThrottle",
+        "rest_framework.throttling.UserRateThrottle",
+    ],
+    "DEFAULT_THROTTLE_RATES": {
+        # Limits the rate of API calls that may be made by a given user.
+        # The user id will be used as a unique cache key.
+        "user": "1000/day",
+        # Limits the rate of API calls that may be made by a anonymous users.
+        # The IP address of the request will be used as the unique cache key.
+        "anon": "100/day"
+    },
 }
 
 # If CELERY_BROKER_URL is set in our environment, then we configure celery as

promgen/tests/test_rest.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2018 LINE Corporation
 # These sources are released under the terms of the MIT license: see LICENSE
+import django.core.cache
 from django.contrib.auth.models import User, Permission
 from django.test import override_settings
 from django.urls import reverse
@@ -9,6 +10,11 @@
 
 
 class RestAPITest(tests.PromgenTest):
+    def setUp(self):
+        super().setUp()
+        # Clear the cache before each test to reset throttling
+        django.core.cache.cache.clear()
+
     @override_settings(PROMGEN=tests.SETTINGS)
     @override_settings(CELERY_TASK_ALWAYS_EAGER=True)
     def test_alert_blackhole(self):
@@ -1787,3 +1793,24 @@ def test_rest_shard(self):
         response = self.client.get(reverse("api-v2:shard-projects", args=[1]))
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.json(), expected)
+
+    @override_settings(PROMGEN=tests.SETTINGS)
+    def test_throttling(self):
+        # Check throttling for anonymous users
+        for _ in range(100):
+            response = self.client.get(reverse("api-v2:service-list"))
+            self.assertEqual(response.status_code, 200)
+        response = self.client.get(reverse("api-v2:service-list"))
+        self.assertEqual(response.status_code, 429)
+
+        # Check throttling for authenticated users
+        token = Token.objects.filter(user__username="demo").first().key
+        for _ in range(1000):
+            response = self.client.get(
+                reverse("api-v2:service-list"), HTTP_AUTHORIZATION=f"Token {token}"
+            )
+            self.assertEqual(response.status_code, 200)
+        response = self.client.get(
+            reverse("api-v2:service-list"), HTTP_AUTHORIZATION=f"Token {token}"
+        )
+        self.assertEqual(response.status_code, 429)