From 2d1f0982e0242724d65de0b6c8d45e27130d618e Mon Sep 17 00:00:00 2001
From: Branson Stephens <branson.stephens@ligo.org>
Date: Thu, 9 Jul 2015 14:47:35 -0500
Subject: [PATCH] Added throttling classes and some basic rates.

---
 gracedb/api.py       |  7 ++++++
 gracedb/throttles.py | 53 ++++++++++++++++++++++++++++++++++++++++++++
 settings/default.py  | 13 ++++++++++-
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 gracedb/throttles.py

diff --git a/gracedb/api.py b/gracedb/api.py
index 76d1577a1..198fdb398 100644
--- a/gracedb/api.py
+++ b/gracedb/api.py
@@ -31,6 +31,8 @@ from forms import CreateEventForm
 from permission_utils import user_has_perm, filter_events_for_user
 from guardian.models import GroupObjectPermission
 
+from throttles import EventCreationThrottle, AnnotationThrottle
+
 from alert import issueAlertForUpdate
 from buildVOEvent import buildVOEvent, VOEventBuilderException
 
@@ -384,6 +386,7 @@ class EventList(APIView):
     permission_classes = (IsAuthenticated,IsAuthorizedForPipeline)
     parser_classes = (parsers.MultiPartParser,)
     renderer_classes = (JSONRenderer, BrowsableAPIRenderer, LigoLwRenderer, TSVRenderer,)
+    throttle_classes = (EventCreationThrottle,)
 
     def get(self, request, *args, **kwargs):
 
@@ -729,6 +732,7 @@ class EventLogList(APIView):
     """
     authentication_classes = (LigoAuthentication,)
     permission_classes = (IsAuthenticated,IsAuthorizedForEvent,)
+    throttle_classes = (AnnotationThrottle,)
 
     @event_and_auth_required
     def get(self, request, event):
@@ -847,6 +851,7 @@ class EMBBEventLogList(APIView):
     """
     authentication_classes = (LigoAuthentication,)
     permission_classes = (IsAuthenticated,IsAuthorizedForEvent,)
+    throttle_classes = (AnnotationThrottle,)
 
     @event_and_auth_required
     def get(self, request, event):
@@ -917,6 +922,7 @@ class EMObservationList(APIView):
     """
     authentication_classes = (LigoAuthentication,)
     permission_classes = (IsAuthenticated,IsAuthorizedForEvent,)
+    throttle_classes = (AnnotationThrottle,)
 
     @event_and_auth_required
     def get(self, request, event):
@@ -1667,6 +1673,7 @@ class VOEventList(APIView):
     """
     authentication_classes = (LigoAuthentication,)
     permission_classes = (IsAuthenticated,IsAuthorizedForEvent,)
+    throttle_classes = (AnnotationThrottle,)
 
     @event_and_auth_required
     def get(self, request, event):
diff --git a/gracedb/throttles.py b/gracedb/throttles.py
new file mode 100644
index 000000000..597db4b73
--- /dev/null
+++ b/gracedb/throttles.py
@@ -0,0 +1,53 @@
+from rest_framework.throttling import UserRateThrottle
+
+class PostOrPutUserRateThrottle(UserRateThrottle):
+
+    def allow_request(self, request, view):
+        """
+        This is mostly copied from the Rest Framework's SimpleRateThrottle
+        except we now pass the request to throttle_success
+        """
+        if self.rate is None:
+            return True
+
+        self.key = self.get_cache_key(request, view)
+        if self.key is None:
+            return True
+
+        self.history = self.cache.get(self.key, [])
+        self.now = self.timer()
+
+        # Drop any requests from the history which have now passed the
+        # throttle duration
+        while self.history and self.history[-1] <= self.now - self.duration:
+            self.history.pop()
+        if len(self.history) >= self.num_requests:
+            return self.throttle_failure()
+        return self.throttle_success(request)
+
+    def throttle_success(self, request):
+        """
+        Inserts the current request's timestamp along with the key
+        into the cache. Except we only do this if the request is a
+        writing method (POST or PUT). That's why we needed the request.
+        """
+        if request.method in ['POST', 'PUT']:
+            self.history.insert(0, self.now)
+        self.cache.set(self.key, self.history, self.duration)
+        return True
+
+    def wait(self):
+        """
+        The HTTPError exception includes a little message with the recommended
+        wait time. However, this doesn't seem to work very well with fractional
+        seconds. Returning 'None' will prevent it from trying to recommend a 
+        wait time.
+        """
+        return None
+
+class EventCreationThrottle(PostOrPutUserRateThrottle):
+    scope = 'event_creation'
+
+class AnnotationThrottle(PostOrPutUserRateThrottle):
+    scope = 'annotation'
+
diff --git a/settings/default.py b/settings/default.py
index b1a11945b..46d4d5d24 100644
--- a/settings/default.py
+++ b/settings/default.py
@@ -105,6 +105,13 @@ DATABASES = {
     }
 }
 
+CACHES = {
+    'default': {
+        'BACKEND': 'django.core.cache.backends.memcached.MemcachedCache',
+        'LOCATION': '127.0.0.1:11211',
+    }
+}
+
 # SkyAlert
 
 SKYALERT_IVORN_PATTERN = "ivo://gwnet/gcn_sender#%s"
@@ -305,7 +312,11 @@ INSTALLED_APPS = (
 )
 
 REST_FRAMEWORK = {
-    'PAGINATE_BY': 10
+    'PAGINATE_BY': 10,
+    'DEFAULT_THROTTLE_RATES': {
+        'event_creation': '5/second',
+        'annotation'    : '10/second',
+    },
 }
 
 
-- 
GitLab