136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
"""
|
|
Scoped Rate Throttling
|
|
Provides rate limiting with different scopes for different operation types
|
|
"""
|
|
from rest_framework.throttling import ScopedRateThrottle
|
|
from django.conf import settings
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DebugScopedRateThrottle(ScopedRateThrottle):
|
|
"""
|
|
Scoped rate throttle that can be bypassed in debug mode
|
|
|
|
Usage:
|
|
class MyViewSet(viewsets.ModelViewSet):
|
|
throttle_scope = 'planner'
|
|
throttle_classes = [DebugScopedRateThrottle]
|
|
"""
|
|
|
|
def allow_request(self, request, view):
|
|
"""
|
|
Check if request should be throttled
|
|
|
|
Bypasses throttling if:
|
|
- DEBUG mode is True
|
|
- IGNY8_DEBUG_THROTTLE environment variable is True
|
|
- User belongs to aws-admin or other system accounts
|
|
- User is admin/developer role
|
|
- Public blueprint list request with site filter (for Sites Renderer)
|
|
"""
|
|
# Check if throttling should be bypassed
|
|
debug_bypass = getattr(settings, 'DEBUG', False)
|
|
env_bypass = getattr(settings, 'IGNY8_DEBUG_THROTTLE', False)
|
|
|
|
# Bypass for public blueprint list requests (Sites Renderer fallback)
|
|
public_blueprint_bypass = False
|
|
if hasattr(view, 'action') and view.action == 'list':
|
|
if hasattr(request, 'query_params') and request.query_params.get('site'):
|
|
if not request.user or not hasattr(request.user, 'is_authenticated') or not request.user.is_authenticated:
|
|
public_blueprint_bypass = True
|
|
|
|
# Bypass for authenticated users (avoid user-facing 429s)
|
|
authenticated_bypass = False
|
|
if hasattr(request, 'user') and request.user and hasattr(request.user, 'is_authenticated') and request.user.is_authenticated:
|
|
authenticated_bypass = True # Do not throttle logged-in users
|
|
|
|
if debug_bypass or env_bypass or public_blueprint_bypass or authenticated_bypass:
|
|
# In debug mode or for system accounts, still set throttle headers but don't actually throttle
|
|
# This allows testing throttle headers without blocking requests
|
|
if hasattr(self, 'get_rate'):
|
|
# Set headers for debugging
|
|
self.scope = getattr(view, 'throttle_scope', None)
|
|
if self.scope:
|
|
# Get rate for this scope
|
|
rate = self.get_rate()
|
|
if rate:
|
|
# Parse rate (e.g., "10/min")
|
|
num_requests, duration = self.parse_rate(rate)
|
|
# Set headers
|
|
request._throttle_debug_info = {
|
|
'scope': self.scope,
|
|
'rate': rate,
|
|
'limit': num_requests,
|
|
'duration': duration
|
|
}
|
|
return True
|
|
|
|
# Normal throttling behavior
|
|
return super().allow_request(request, view)
|
|
|
|
def get_rate(self):
|
|
"""
|
|
Get rate for the current scope
|
|
"""
|
|
if not self.scope:
|
|
return None
|
|
|
|
# Get throttle rates from settings
|
|
throttle_rates = getattr(settings, 'REST_FRAMEWORK', {}).get('DEFAULT_THROTTLE_RATES', {})
|
|
|
|
# Get rate for this scope
|
|
rate = throttle_rates.get(self.scope)
|
|
|
|
# Fallback to default if scope not found
|
|
if not rate:
|
|
rate = throttle_rates.get('default', '100/min')
|
|
|
|
return rate
|
|
|
|
def parse_rate(self, rate):
|
|
"""
|
|
Parse rate string (e.g., "10/min") into (num_requests, duration)
|
|
|
|
Returns:
|
|
tuple: (num_requests, duration_in_seconds)
|
|
"""
|
|
if not rate:
|
|
return None, None
|
|
|
|
try:
|
|
num, period = rate.split('/')
|
|
num_requests = int(num)
|
|
|
|
# Parse duration
|
|
period = period.strip().lower()
|
|
if period == 'sec' or period == 's':
|
|
duration = 1
|
|
elif period == 'min' or period == 'm':
|
|
duration = 60
|
|
elif period == 'hour' or period == 'h':
|
|
duration = 3600
|
|
elif period == 'day' or period == 'd':
|
|
duration = 86400
|
|
else:
|
|
# Default to seconds
|
|
duration = 1
|
|
|
|
return num_requests, duration
|
|
except (ValueError, AttributeError):
|
|
# Invalid rate format, default to 100/min
|
|
logger.warning(f"Invalid rate format: {rate}, defaulting to 100/min")
|
|
return 100, 60
|
|
|
|
def throttle_success(self):
|
|
"""
|
|
Called when request is allowed
|
|
Sets throttle headers on response
|
|
"""
|
|
# This is called by DRF after allow_request returns True
|
|
# Headers are set automatically by ScopedRateThrottle
|
|
pass
|
|
|
|
|