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/admin.py

@@ -172,3 +172,8 @@ def has_add_permission(self, request, obj=None):
 
     def has_change_permission(self, request, obj=None):
         return False
+
+
+@admin.register(models.SiteConfiguration)
+class SiteConfigurationAdmin(admin.ModelAdmin):
+    list_display = ("key", "value")

promgen/migrations/0025_siteconfiguration.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.11 on 2025-04-03 10:19
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("promgen", "0024_alter_farm_name_alter_project_name_alter_rule_name_and_more"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="SiteConfiguration",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                ("key", models.CharField(max_length=128, unique=True)),
+                ("value", models.JSONField(default=dict)),
+            ],
+        ),
+    ]

promgen/models.py

@@ -617,3 +617,11 @@ class Meta:
         ordering = ["shard", "host"]
         unique_together = (("host", "port"),)
         verbose_name_plural = "prometheis"
+
+
+class SiteConfiguration(models.Model):
+    key = models.CharField(max_length=128, unique=True)
+    value = models.JSONField(default=dict)
+
+    def __str__(self):
+        return f"{self.key}={self.value}"

promgen/settings.py

@@ -197,6 +197,14 @@
     "DEFAULT_FILTER_BACKENDS": ("django_filters.rest_framework.DjangoFilterBackend",),
     "DEFAULT_SCHEMA_CLASS": "promgen.schemas.CustomSchema",
     "EXCEPTION_HANDLER": "promgen.middleware.custom_exception_handler",
+    "DEFAULT_THROTTLE_CLASSES": [
+        "promgen.util.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",
+    },
 }
 
 # If CELERY_BROKER_URL is set in our environment, then we configure celery as

promgen/signals.py

@@ -333,3 +333,12 @@ def add_default_project_subscription(instance, created, **kwargs):
             value=instance.owner.username,
             defaults={"owner": instance.owner},
         )
+
+
+@receiver(post_save, sender=models.SiteConfiguration)
+@skip_raw
+def clear_cache(*, sender, instance, **kwargs):
+    # We need to clear our cache when we change our configuration
+    # so that we can pick up the new settings
+    if instance.key == "THROTTLE_RATES":
+        cache.clear()

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):
@@ -2284,3 +2290,31 @@ def test_rest_shard(self):
         )
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.json(), expected)
+
+    @override_settings(PROMGEN=tests.SETTINGS)
+    def test_throttling(self):
+        # 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)
+
+        # Check changing rate
+        models.SiteConfiguration.objects.get_or_create(
+            key="THROTTLE_RATES", value={"user": "3/day"}
+        )
+        for _ in range(3):
+            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)

promgen/util.py

@@ -8,6 +8,9 @@
 from django.conf import settings
 from django.db.models import F
 from django.http import HttpResponse
+from rest_framework import throttling
+
+from promgen import models
 
 # Wrappers around request api to ensure we always attach our user agent
 # https://github.com/requests/requests/blob/master/requests/api.py
@@ -143,3 +146,11 @@ def proxy_error(response: requests.Response) -> HttpResponse:
 get.__doc__ = requests.get.__doc__
 post.__doc__ = requests.post.__doc__
 delete.__doc__ = requests.delete.__doc__
+
+
+class UserRateThrottle(throttling.UserRateThrottle):
+    def get_rate(self):
+        rate = models.SiteConfiguration.objects.filter(key="THROTTLE_RATES").first()
+        if rate and rate.value["user"]:
+            return rate.value["user"]
+        return super().get_rate()