OpenConceptLab/ocl_issues#2121 | API Throttling framework

Author: snyaggarwal

core/client_configs/views.py

@@ -7,6 +7,7 @@
 from core.client_configs.serializers import ClientConfigSerializer, ClientConfigTemplateSerializer
 from core.common.views import BaseAPIView
 from .models import ClientConfig
+from ..common.throttling import ThrottleUtil
 
 
 class ClientConfigBaseView(generics.GenericAPIView):
@@ -15,6 +16,9 @@ class ClientConfigBaseView(generics.GenericAPIView):
     queryset = ClientConfig.objects.filter(is_active=True)
     serializer_class = ClientConfigSerializer
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
 
 class ClientConfigView(ClientConfigBaseView, RetrieveAPIView, UpdateAPIView, DestroyAPIView):
     def perform_destroy(self, instance: ClientConfig):
@@ -50,6 +54,9 @@ def post(self, request, *args, **kwargs):  # pylint: disable=unused-argument
 class ResourceTemplatesView(ListAPIView, CreateAPIView):
     serializer_class = ClientConfigTemplateSerializer
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get_queryset(self):
         user = self.request.user
         return ClientConfig.objects.filter(

core/collections/views.py

@@ -55,6 +55,7 @@
     canonical_url_param
 from core.common.tasks import add_references, export_collection, delete_collection, index_expansion_concepts, \
     index_expansion_mappings
+from core.common.throttling import ThrottleUtil
 from core.common.utils import compact_dict_by_values, parse_boolean_query_param
 from core.common.views import BaseAPIView, BaseLogoView, ConceptContainerExtraRetrieveUpdateDestroyView
 from core.concepts.documents import ConceptDocument
@@ -1103,6 +1104,9 @@ class CollectionClientConfigsView(CollectionBaseView, ResourceClientConfigsView)
 class ReferenceExpressionResolveView(APIView):
     serializer_class = ReferenceExpressionResolveSerializer
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get_results(self):
         data = self.request.data
         if not isinstance(data, list):

core/common/signals.py

@@ -1,9 +1,10 @@
 from django.db.models.signals import pre_save, post_save
 from django.dispatch import receiver
+from pydash import get
 
 from core.common.models import BaseModel
 from core.orgs.models import Organization
-from core.users.models import UserProfile
+from core.users.models import UserProfile, UserRateLimit
 
 
 @receiver(pre_save)
@@ -15,6 +16,8 @@ def stamp_uri(sender, instance, **kwargs):  # pylint: disable=unused-argument
 @receiver(post_save, sender=Organization)
 @receiver(post_save, sender=UserProfile)
 def propagate_owner_status(sender, instance=None, created=False, **kwargs):  # pylint: disable=unused-argument
+    if created and instance.__class__ == UserProfile and not get(instance, 'api_rate_limit'):
+        UserRateLimit(user=instance).save()
     if created and instance.__class__ == Organization and instance.id != 1:
         instance.record_create_event()
     if not created and instance:

core/common/throttling.py

@@ -0,0 +1,55 @@
+from django.conf import settings
+from pydash import get
+from rest_framework.throttling import UserRateThrottle
+
+
+class GuestMinuteThrottle(UserRateThrottle):
+    scope = 'guest_minute'
+
+
+class GuestDayThrottle(UserRateThrottle):
+    scope = 'guest_day'
+
+
+class LiteMinuteThrottle(UserRateThrottle):
+    scope = 'lite_minute'
+
+
+class LiteDayThrottle(UserRateThrottle):
+    scope = 'lite_day'
+
+
+class PremiumDayThrottle(UserRateThrottle):
+    scope = 'premium_day'
+
+
+class PremiumMinuteThrottle(UserRateThrottle):
+    scope = 'premium_minute'
+
+
+class ThrottleUtil:
+    @staticmethod
+    def get_limit_remaining(throttle, request, view):
+        key = throttle.get_cache_key(request, view)
+        if key is not None:
+            history = throttle.cache.get(key, None)
+            if history is None:
+                return None
+            while history and history[-1] <= throttle.timer() - throttle.duration:
+                history.pop()
+            remaining = throttle.num_requests - len(history)
+        else:
+            remaining = 'unlimited'
+
+        return remaining
+
+    @staticmethod
+    def get_throttles_by_user_plan(user):
+        if not settings.ENABLE_THROTTLING:
+            return []
+        # order is important, first one has to be minute throttle
+        if get(user, 'api_rate_limit.is_premium'):
+            return [PremiumMinuteThrottle(), PremiumDayThrottle()]
+        if get(user, 'api_rate_limit.is_guest') or not get(user, 'is_authenticated'):
+            return [GuestMinuteThrottle(), GuestDayThrottle()]
+        return [LiteMinuteThrottle(), LiteDayThrottle()]

core/common/views.py

@@ -30,6 +30,7 @@
 from core.common.search import CustomESSearch
 from core.common.serializers import RootSerializer
 from core.common.swagger_parameters import all_resource_query_param
+from core.common.throttling import ThrottleUtil
 from core.common.utils import compact_dict_by_values, to_snake_case, parse_updated_since_param, \
     to_int, get_falsy_values, get_truthy_values, format_url_for_search
 from core.concepts.permissions import CanViewParentDictionary, CanEditParentDictionary
@@ -59,6 +60,9 @@ class BaseAPIView(generics.GenericAPIView, PathWalkerMixin):
     facet_class = None
     total_count = 0
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def has_no_kwargs(self):
         return len(self.kwargs.values()) == 0
 
@@ -913,6 +917,9 @@ def __set_params(self):
 class SourceChildExtrasBaseView:
     default_qs_sort_attr = '-created_at'
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get_object(self):
         queryset = self.get_queryset()
 
@@ -980,6 +987,9 @@ class APIVersionView(APIView):  # pragma: no cover
     permission_classes = (AllowAny,)
     swagger_schema = None
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @staticmethod
     def get(_):
         return Response(__version__)
@@ -989,6 +999,9 @@ class ChangeLogView(APIView):  # pragma: no cover
     permission_classes = (AllowAny, )
     swagger_schema = None
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @staticmethod
     def get(_):
         resp = requests.get('https://raw.githubusercontent.com/OpenConceptLab/oclapi2/master/changelog.md')
@@ -1043,6 +1056,9 @@ def post(self, request, *args, **kwargs):  # pylint: disable=unused-argument
 class FeedbackView(APIView):  # pragma: no cover
     permission_classes = (AllowAny, )
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @staticmethod
     @swagger_auto_schema(request_body=openapi.Schema(
         type=openapi.TYPE_OBJECT,
@@ -1141,6 +1157,9 @@ class AbstractChecksumView(APIView):
     permission_classes = (IsAuthenticated,)
     smart = False
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(
         manual_parameters=[all_resource_query_param],
         request_body=openapi.Schema(

core/concepts/views.py

@@ -27,6 +27,7 @@
     cascade_levels_param, cascade_direction_param, cascade_view_hierarchy, return_map_types_param,
     omit_if_exists_in_param, equivalency_map_types_param, search_from_latest_repo_header)
 from core.common.tasks import delete_concept, make_hierarchy
+from core.common.throttling import ThrottleUtil
 from core.common.utils import to_parent_uri_from_kwargs, generate_temp_version, get_truthy_values, to_int
 from core.common.views import SourceChildCommonBaseView, SourceChildExtrasView, \
     SourceChildExtraRetrieveUpdateDestroyView, BaseAPIView
@@ -752,6 +753,9 @@ class ConceptsHierarchyAmendAdminView(APIView):  # pragma: no cover
     swagger_schema = None
     permission_classes = (IsAdminUser, )
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @staticmethod
     def post(request):
         concept_map = request.data

core/fhir/views.py

@@ -5,13 +5,17 @@
 from rest_framework.views import APIView
 
 from core import settings
+from core.common.throttling import ThrottleUtil
 
 logger = logging.getLogger('oclapi')
 
 
 class CapabilityStatementView(APIView):
     permission_classes = (AllowAny,)
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get(self, request):
         mode = request.query_params.get('mode')
         fhir_base_url = f"{settings.API_BASE_URL}/fhir"

core/importers/views.py

@@ -20,6 +20,7 @@
 from rest_framework.views import APIView
 
 from core.common.constants import DEPRECATED_API_HEADER
+from core.common.throttling import ThrottleUtil
 from core.common.views import BaseAPIView
 from core.common.tasks import bulk_import_new
 from core.common.swagger_parameters import update_if_exists_param, task_param, result_param, username_param, \
@@ -71,6 +72,9 @@ def import_response(request, import_queue, data, threads=None, inline=False, dep
 
 
 class ImportRetrieveDestroyMixin(BaseAPIView):
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get_serializer_class(self):
         if self.request.GET.get('task'):
             return TaskDetailSerializer
@@ -141,6 +145,9 @@ class BulkImportParallelInlineView(APIView):
     permission_classes = (IsAuthenticated, )
     deprecated = True
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     def get_parsers(self):
         if 'application/json' in [self.request.META.get('CONTENT_TYPE')]:
             return [JSONParser()]
@@ -225,6 +232,9 @@ class BulkImportFileUploadView(APIView):  # pragma: no cover
     parser_classes = (MultiPartParser, )
     deprecated = True
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(
         manual_parameters=[update_if_exists_param, file_upload_param],
         deprecated=True
@@ -255,6 +265,9 @@ class BulkImportFileURLView(APIView):  # pragma: no cover
     parser_classes = (MultiPartParser, )
     deprecated = True
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(
         manual_parameters=[update_if_exists_param, file_url_param],
         deprecated=True
@@ -336,6 +349,9 @@ class BulkImportInlineView(APIView):  # pragma: no cover
     parser_classes = (MultiPartParser, FormParser)
     deprecated = True
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(
         manual_parameters=[update_if_exists_param, file_url_param, file_upload_param],
         deprecated=True

core/indexes/views.py

@@ -9,6 +9,7 @@
 
 from core.common.swagger_parameters import apps_param, ids_param, resources_body_param, uri_param, filter_param
 from core.common.tasks import rebuild_indexes, populate_indexes, batch_index_resources
+from core.common.throttling import ThrottleUtil
 from core.common.utils import get_resource_class_from_resource_name
 from core.tasks.models import Task
 
@@ -18,6 +19,9 @@ class BaseESIndexView(APIView):  # pragma: no cover
     parser_classes = (MultiPartParser,)
     task = None
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(manual_parameters=[apps_param])
     def post(self, request):
         apps = request.data.get('apps', None)
@@ -49,6 +53,9 @@ class ResourceIndexView(APIView):
     permission_classes = (IsAdminUser,)
     parser_classes = (MultiPartParser,)
 
+    def get_throttles(self):
+        return ThrottleUtil.get_throttles_by_user_plan(self.request.user)
+
     @swagger_auto_schema(manual_parameters=[ids_param, uri_param, filter_param, resources_body_param])
     def post(self, _, resource):
         model = get_resource_class_from_resource_name(resource)

core/integration_tests/tests_users.py

@@ -785,6 +785,76 @@ def test_put_204(self):
         self.assertTrue(inactive_user.is_active)
 
 
+class UserRateLimitViewTest(OCLAPITestCase):
+    def setUp(self):
+        super().setUp()
+        self.superuser = UserProfile.objects.get(username='ocladmin')
+
+    def test_put_bad_request(self):
+        user = UserProfileFactory()
+        self.assertEqual(user.api_rate_limit.rate_plan, 'lite')
+
+        response = self.client.put(
+            f'/users/{user.username}/rate-limit/',
+            {},
+            HTTP_AUTHORIZATION='Token ' + self.superuser.get_token(),
+            format='json'
+        )
+        self.assertEqual(response.status_code, 400)
+        self.assertEqual(
+            response.data,
+            {'detail': ErrorDetail(string='"rate_plan" needs to be one of "guest", "lite" or "premium"', code='bad_request')}
+        )
+
+        user.refresh_from_db()
+        self.assertEqual(user.api_rate_limit.rate_plan, 'lite')
+
+        response = self.client.put(
+            f'/users/{user.username}/rate-limit/',
+            {'rate_plan': 'blah'},
+            HTTP_AUTHORIZATION='Token ' + self.superuser.get_token(),
+            format='json'
+        )
+        self.assertEqual(response.status_code, 400)
+        self.assertEqual(response.data, {'rate_plan': ["Value 'blah' is not a valid choice."]})
+
+        self.assertEqual(user.api_rate_limit.rate_plan, 'lite')
+
+        response = self.client.put(
+            f'/users/{user.username}/rate-limit/',
+            {'rate_plan': 'blah'},
+            HTTP_AUTHORIZATION='Token ' + user.get_token(),
+            format='json'
+        )
+        self.assertEqual(response.status_code, 403)
+
+    def test_put_204(self):
+        user = UserProfileFactory()
+        self.assertEqual(user.api_rate_limit.rate_plan, 'lite')
+
+        response = self.client.put(
+            f'/users/{user.username}/rate-limit/',
+            {'rate_plan': 'guest'},
+            HTTP_AUTHORIZATION='Token ' + self.superuser.get_token(),
+            format='json'
+        )
+        self.assertEqual(response.status_code, 204)
+
+        user.refresh_from_db()
+        self.assertEqual(user.api_rate_limit.rate_plan, 'guest')
+
+        response = self.client.put(
+            f'/users/{user.username}/rate-limit/',
+            {'rate_plan': 'premium'},
+            HTTP_AUTHORIZATION='Token ' + self.superuser.get_token(),
+            format='json'
+        )
+        self.assertEqual(response.status_code, 204)
+
+        user.refresh_from_db()
+        self.assertEqual(user.api_rate_limit.rate_plan, 'premium')
+
+
 class UserStaffToggleViewTest(OCLAPITestCase):
     def setUp(self):
         super().setUp()