fix: add throttling for influx query endpoints (#6453)

Author: matthewelwell

api/app/settings/common.py

@@ -320,6 +320,7 @@
         "mfa_code": "5/min",
         "invite": "10/min",
         "user": USER_THROTTLE_RATE,
+        "influx_query": "5/min",
     },
     "DEFAULT_FILTER_BACKENDS": ["django_filters.rest_framework.DjangoFilterBackend"],
     "DEFAULT_RENDERER_CLASSES": [

api/app/settings/test.py

@@ -11,6 +11,7 @@
     "invite": "10/min",
     "signup": "100/min",
     "user": "100000/day",
+    "influx_query": "50/min",
 }
 
 AWS_SSE_LOGS_BUCKET_NAME = "test_bucket"

api/app_analytics/throttles.py

@@ -0,0 +1,23 @@
+from django.views import View
+from rest_framework.request import Request
+from rest_framework.throttling import SimpleRateThrottle
+
+
+class InfluxQueryThrottle(SimpleRateThrottle):
+    scope = "influx_query"
+
+    def get_cache_key(self, request: Request, view: View) -> str:
+        """
+        Since we want to throttle requests across multiple views and viewsets that don't
+        necessarily have the option to define the scope themselves, we override the
+        get_cache_key method to ensure that the throttling logic behaves as expected.
+        """
+        if request.user and request.user.is_authenticated:
+            return self.cache_format % {
+                "scope": self.scope,
+                "ident": request.user.pk,
+            }
+        return self.cache_format % {  # pragma: no cover
+            "scope": self.scope,
+            "ident": self.get_ident(request),
+        }

api/app_analytics/views.py

@@ -4,7 +4,7 @@
 from common.core.utils import using_database_replica
 from drf_yasg.utils import swagger_auto_schema  # type: ignore[import-untyped]
 from rest_framework import status
-from rest_framework.decorators import api_view, permission_classes
+from rest_framework.decorators import api_view, permission_classes, throttle_classes
 from rest_framework.fields import IntegerField
 from rest_framework.generics import CreateAPIView
 from rest_framework.permissions import IsAuthenticated
@@ -20,6 +20,7 @@
 from app_analytics.mappers import (
     map_request_to_labels,
 )
+from app_analytics.throttles import InfluxQueryThrottle
 from environments.authentication import EnvironmentKeyAuthentication
 from environments.permissions.permissions import EnvironmentKeyPermissions
 from features.models import FeatureState
@@ -135,6 +136,7 @@ class SelfHostedTelemetryAPIView(CreateAPIView):  # type: ignore[type-arg]
 )
 @api_view(["GET"])
 @permission_classes([IsAuthenticated, UsageDataPermission])
+@throttle_classes([InfluxQueryThrottle])
 def get_usage_data_total_count_view(request: Request, organisation_pk: int) -> Response:
     organisation = using_database_replica(Organisation.objects).get(id=organisation_pk)
     count = get_total_events_count(organisation)
@@ -151,6 +153,7 @@ def get_usage_data_total_count_view(request: Request, organisation_pk: int) -> R
 )
 @api_view(["GET"])
 @permission_classes([IsAuthenticated, UsageDataPermission])
+@throttle_classes([InfluxQueryThrottle])
 def get_usage_data_view(request: Request, organisation_pk: int) -> Response:
     filters = UsageDataQuerySerializer(data=request.query_params)
     filters.is_valid(raise_exception=True)

api/organisations/views.py

@@ -21,6 +21,7 @@
     get_events_for_organisation,
     get_multiple_event_list_for_organisation,
 )
+from app_analytics.throttles import InfluxQueryThrottle
 from core.helpers import get_current_site_url
 from organisations.chargebee import webhook_event_types, webhook_handlers
 from organisations.exceptions import OrganisationHasNoPaidSubscription
@@ -164,6 +165,7 @@ def remove_users(self, request, pk):  # type: ignore[no-untyped-def]
     @action(
         detail=True,
         methods=["GET"],
+        throttle_classes=[InfluxQueryThrottle],
     )
     def usage(self, request, pk):  # type: ignore[no-untyped-def]
         organisation = self.get_object()
@@ -240,7 +242,12 @@ def get_hosted_page_url_for_subscription_upgrade(self, request, pk):  # type: ig
         operation_description="Please use /api/v1/organisations/{organisation_pk}/usage-data/",
         query_serializer=InfluxDataQuerySerializer(),
     )
-    @action(detail=True, methods=["GET"], url_path="influx-data")
+    @action(
+        detail=True,
+        methods=["GET"],
+        url_path="influx-data",
+        throttle_classes=[InfluxQueryThrottle],
+    )
     def get_influx_data(self, request, pk):  # type: ignore[no-untyped-def]
         filters = InfluxDataQuerySerializer(data=request.query_params)
         filters.is_valid(raise_exception=True)

api/tests/integration/app_analytics/test_influx_query_throttle.py

@@ -0,0 +1,35 @@
+from django.urls import reverse
+from pytest_mock import MockerFixture
+from rest_framework import status
+from rest_framework.test import APIClient
+
+
+def test_influx_data_endpoint_is_throttled(
+    admin_client: APIClient,
+    organisation: int,
+    mocker: MockerFixture,
+    reset_cache: None,
+) -> None:
+    # Given
+    mocker.patch(
+        "app_analytics.throttles.InfluxQueryThrottle.get_rate", return_value="1/minute"
+    )
+    mocker.patch(
+        "organisations.views.get_multiple_event_list_for_organisation", return_value=[]
+    )
+
+    url = reverse(
+        "api-v1:organisations:organisation-get-influx-data",
+        args=[organisation],
+    )
+
+    # When - first request should be successful
+    first_response = admin_client.get(url)
+    second_response = admin_client.get(url)
+
+    # Then
+    # The first response should have been successful
+    assert first_response.status_code == status.HTTP_200_OK
+
+    # But the second request should have been throttled
+    assert second_response.status_code == status.HTTP_429_TOO_MANY_REQUESTS