read redis config from environment
This commit is contained in:
parent
4d5b4b8346
commit
5844268ebe
10
ratelimit.py
10
ratelimit.py
@ -1,3 +1,4 @@
|
|||||||
|
import os
|
||||||
from redis import Redis
|
from redis import Redis
|
||||||
import datetime
|
import datetime
|
||||||
import typing
|
import typing
|
||||||
@ -16,6 +17,13 @@ class RateLimitExceeded(Exception):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _get_redis_connection():
|
||||||
|
host = os.environ.get("RRL_REDIS_HOST", "localhost")
|
||||||
|
port = os.environ.get("RRL_REDIS_PORT", 6379)
|
||||||
|
db = os.environ.get("RRL_REDIS_DB", 0)
|
||||||
|
return Redis(host=host, port=port, db=db)
|
||||||
|
|
||||||
|
|
||||||
class RateLimiter:
|
class RateLimiter:
|
||||||
"""
|
"""
|
||||||
<zone>:<key>:<hour><minute> expires in 2 minutes
|
<zone>:<key>:<hour><minute> expires in 2 minutes
|
||||||
@ -24,7 +32,7 @@ class RateLimiter:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, tiers: typing.List[Tier], *, prefix="", use_redis_time=True):
|
def __init__(self, tiers: typing.List[Tier], *, prefix="", use_redis_time=True):
|
||||||
self.redis = Redis()
|
self.redis = _get_redis_connection()
|
||||||
self.tiers = {tier.name: tier for tier in tiers}
|
self.tiers = {tier.name: tier for tier in tiers}
|
||||||
self.prefix = prefix
|
self.prefix = prefix
|
||||||
self.use_redis_time = use_redis_time
|
self.use_redis_time = use_redis_time
|
||||||
|
@ -1,9 +1,8 @@
|
|||||||
import pytest
|
import pytest
|
||||||
from ratelimit import Tier, RateLimiter, RateLimitExceeded
|
from ratelimit import Tier, RateLimiter, RateLimitExceeded, _get_redis_connection
|
||||||
from freezegun import freeze_time
|
from freezegun import freeze_time
|
||||||
from redis import Redis
|
|
||||||
|
|
||||||
redis = Redis()
|
redis = _get_redis_connection()
|
||||||
simple_minute_tier = Tier("10/minute", 10, 0, 0)
|
simple_minute_tier = Tier("10/minute", 10, 0, 0)
|
||||||
simple_hour_tier = Tier("10/hour", 0, 10, 0)
|
simple_hour_tier = Tier("10/hour", 0, 10, 0)
|
||||||
simple_daily_tier = Tier("10/day", 0, 0, 10)
|
simple_daily_tier = Tier("10/day", 0, 0, 10)
|
||||||
|
Loading…
Reference in New Issue
Block a user