43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
import newrelic.agent
|
|
import redis.exceptions
|
|
import werkzeug.exceptions
|
|
from limits.storage import RedisStorage
|
|
|
|
from app.log import LOG
|
|
|
|
lock_redis: Optional[RedisStorage] = None
|
|
|
|
|
|
def set_redis_concurrent_lock(redis: RedisStorage):
|
|
global lock_redis
|
|
lock_redis = redis
|
|
|
|
|
|
def check_bucket_limit(
|
|
lock_name: Optional[str] = None,
|
|
max_hits: int = 5,
|
|
bucket_seconds: int = 3600,
|
|
):
|
|
# Calculate current bucket time
|
|
int_time = int(datetime.utcnow().timestamp())
|
|
bucket_id = int_time - (int_time % bucket_seconds)
|
|
bucket_lock_name = f"bl:{lock_name}:{bucket_id}"
|
|
if not lock_redis:
|
|
return
|
|
try:
|
|
value = lock_redis.incr(bucket_lock_name, bucket_seconds)
|
|
if value > max_hits:
|
|
LOG.i(
|
|
f"Rate limit hit for {lock_name} (bucket id {bucket_id}) -> {value}/{max_hits}"
|
|
)
|
|
newrelic.agent.record_custom_event(
|
|
"BucketRateLimit",
|
|
{"lock_name": lock_name, "bucket_seconds": bucket_seconds},
|
|
)
|
|
raise werkzeug.exceptions.TooManyRequests()
|
|
except (redis.exceptions.RedisError, AttributeError):
|
|
LOG.e("Cannot connect to redis")
|