48 lines
No EOL
1.6 KiB
Python
48 lines
No EOL
1.6 KiB
Python
from rest_framework.throttling import UserRateThrottle
|
|
|
|
from apps.gooyal_oauth2.models import AccessToken
|
|
|
|
def get_application(request):
|
|
try:
|
|
application = request.auth.application
|
|
except:
|
|
application = None
|
|
return application
|
|
|
|
|
|
class TokenLimitThrottle:
|
|
def allow_request(self, request, view):
|
|
application = get_application(request)
|
|
|
|
if application:
|
|
if hasattr(view, "max_allowed_session"):
|
|
max_allowed_session = view.max_allowed_session
|
|
else:
|
|
max_allowed_session = application.max_allowed_session
|
|
|
|
if max_allowed_session:
|
|
session_count = AccessToken.objects.filter(application=application, user=request.user).count()
|
|
if max_allowed_session >= session_count:
|
|
return True
|
|
else:
|
|
active_tokens = AccessToken.objects.filter(
|
|
application=application, user=request.user
|
|
).order_by("created")[:max_allowed_session].values_list("token", flat=True)
|
|
if request.auth.token in active_tokens:
|
|
return True
|
|
else:
|
|
False
|
|
# raise ServiceUnavailable(code='max_allowed_session_reached')
|
|
|
|
else:
|
|
return True
|
|
|
|
def wait(self):
|
|
"""
|
|
Optionally, return a recommended number of seconds to wait before
|
|
the next request.
|
|
"""
|
|
return None
|
|
|
|
|
|
from rest_framework.exceptions import Throttled |