python-memcache responds to incr on non-existent keys raising a ValueError. Cmemcache...

[plcapi.git] / aspects / ratelimitaspects.py

diff --git a/aspects/ratelimitaspects.py b/aspects/ratelimitaspects.py
index 78548c2..8c28dd5 100644 (file)
--- a/aspects/ratelimitaspects.py
+++ b/aspects/ratelimitaspects.py
@@ -23,6 +23,8 @@ class BaseRateLimit(object):
         self.requests = 50 # Number of allowed requests in that time period
         self.expire_after = (self.minutes + 1) * 60
 
+        self.whitelist = []
+
     def before(self, wobj, data, *args, **kwargs):
         # ratelimit_128.112.139.115_201011091532 = 1
         # ratelimit_128.112.139.115_201011091533 = 14
@@ -33,7 +35,7 @@ class BaseRateLimit(object):
         api_method_name = wobj.name
         api_method_source = wobj.source
 
-        if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP:
+        if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP or api_method_source[0] in self.whitelist:
             return
 
         mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
@@ -43,8 +45,11 @@ class BaseRateLimit(object):
         keys_to_check = ["%s_%s_%s" % (self.prefix, api_method_source[0], (now - timedelta(minutes = minute)).strftime("%Y%m%d%H%M")) for minute in range(self.minutes + 1)]
 
         try:
-            mc.incr(current_key)
+            value = mc.incr(current_key)
         except ValueError:
+            value = None
+
+        if value == None:
             mc.set(current_key, 1, time=self.expire_after)
 
         result = mc.get_multi(keys_to_check)