From 8cac1846bc3f0f13a15e3a6a0f2be3f69d99cee8 Mon Sep 17 00:00:00 2001 From: Christoph Buelter Date: Fri, 20 Dec 2024 13:26:01 +0100 Subject: [PATCH] Annotate request and exception with usage data This can e.g. be inspected by Django's exception handler to give an API user feedback about how long to wait until they may try again. --- django_ratelimit/decorators.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/django_ratelimit/decorators.py b/django_ratelimit/decorators.py index 40c9541..4d41214 100644 --- a/django_ratelimit/decorators.py +++ b/django_ratelimit/decorators.py @@ -5,7 +5,7 @@ from django_ratelimit import ALL, UNSAFE from django_ratelimit.exceptions import Ratelimited -from django_ratelimit.core import is_ratelimited +from django_ratelimit.core import get_usage __all__ = ['ratelimit'] @@ -16,14 +16,24 @@ def decorator(fn): @wraps(fn) def _wrapped(request, *args, **kw): old_limited = getattr(request, 'limited', False) - ratelimited = is_ratelimited(request=request, group=group, fn=fn, - key=key, rate=rate, method=method, - increment=True) + + usage = get_usage(request=request, group=group, fn=fn, + key=key, rate=rate, method=method, + increment=True) + if usage is None: + ratelimited = False + else: + ratelimited = usage['should_limit'] + request.limited = ratelimited or old_limited + request.usage = usage + if ratelimited and block: cls = getattr( settings, 'RATELIMIT_EXCEPTION_CLASS', Ratelimited) - raise (import_string(cls) if isinstance(cls, str) else cls)() + exc = (import_string(cls) if isinstance(cls, str) else cls)() + exc.usage = usage + raise exc return fn(request, *args, **kw) return _wrapped return decorator