diff --git a/appbase/decorators.py b/appbase/decorators.py index 133d3c1..b548432 100644 --- a/appbase/decorators.py +++ b/appbase/decorators.py @@ -1,22 +1,23 @@ import datetime +from flask import request try: import _pickle as pickle except: import pickle import redis -import sys import traceback from functools import wraps from appbase.helpers import notify_dev, make_key_from_params +from appbase.errors import TooManyRequestsError from settings import REDIS_HOST, REDIS_PORT, REDIS_DB #TODO: Move below settings to settings/converge -NOTIFICATION_GAP = datetime.timedelta(0, (10*60)) +NOTIFICATION_GAP = datetime.timedelta(0, (10 * 60)) CACHE_TTL = 3 * 7 * 24 * 60 * 60 -# Not using decode_responses=True, +# Not using decode_responses=True, # because pickle generates binary serialization format which doesn't have decode method rconn = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB) @@ -25,7 +26,6 @@ def failsafe(f): """This decorator can be used for making any function fail safe. Basically, at every call it would save the output against a key, uniquely created by function name and args, kw passed to the function. - And in case of error it will return the previously saved output and will notify dev about error. """ f.last_notified_at = datetime.datetime.now() - NOTIFICATION_GAP @@ -46,3 +46,26 @@ def wrapper(*args, **kw): result = pickle.loads(dump) if dump else None return result return wrapper + + +def api_rate_limiter(max_calls, duration): + def decorator(f): + @wraps(f) + def wrapper(*args, **kw): + ip = request.remote_addr + key = 'arl:{}:{}'.format( + ip, + make_key_from_params(f.__name__, args, kw, strict=True) + ) + cnt = rconn.incr(key) + if cnt > max_calls: + raise TooManyRequestsError( + data={ + 'request': f.__name__, + 'duration': duration, + 'ip': ip}) + elif cnt == 1: + rconn.expire(key, duration) + return f(*args, **kw) + return wrapper + return decorator \ No newline at end of file diff --git a/appbase/errors.py b/appbase/errors.py index d94baa6..c3a06f4 100644 --- a/appbase/errors.py +++ b/appbase/errors.py @@ -32,3 +32,8 @@ class InvalidSessionError(BaseError): class ConflictError(BaseError): code = 409 msg = 'Duplicate resource' + + +class TooManyRequestsError(BaseError): + code = 429 + msg = 'Too many requests'