jpt-archive/rrl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

remove zone concept, better handled as prefix

Browse Source
This commit is contained in:
James Turk 2020-11-19 15:29:56 -05:00
parent 79f799bd7b
commit e0ab3d45e5
3 changed files with 23 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
README.md
View File

@ -36,7 +36,6 @@ limiter = RateLimiter(tiers=[bronze, silver])
Then to apply limiting, you'll call the `check_limit` function, which takes three parameters:
* `zone` - Limits are considered unique per zone. So if you want calls against your geocoding API to count against a different limit than your user API you could pass those as unique zones.
* `key` - A unique-per user key, often the user's API key or username. (Note: `rrl` does not know if a key is valid or not, that validation should be in your application and usually occur before the call to `check_limit`.)
* `tier_name` - The name of one of the tiers as specified when instantiating the `RateLimiter` class. (Note: `rrl` does not have a concept of which users are in which tier, that logic should be handled by your key validation code.)

19
rrl.py
View File

@ -32,9 +32,9 @@ def _get_redis_connection() -> Redis:
class RateLimiter:
"""
<zone>:<key>:<hour><minute> expires in 2 minutes
<zone>:<key>:<hour> expires in 2 hours
<zone>:<key>:<day> never expires
<prefix>:<key>:<hour><minute> expires in 2 minutes
<prefix>:<key>:<hour> expires in 2 hours
<prefix>:<key>:<day> never expires
"""
def __init__(
@ -51,7 +51,7 @@ class RateLimiter:
self.use_redis_time = use_redis_time
self.track_daily_usage = track_daily_usage
def check_limit(self, zone: str, key: str, tier_name: str) -> bool:
def check_limit(self, key: str, tier_name: str) -> bool:
try:
tier = self.tiers[tier_name]
except KeyError:
@ -64,16 +64,16 @@ class RateLimiter:
pipe = self.redis.pipeline()
if tier.per_minute:
minute_key = f"{self.prefix}:{zone}:{key}:m{now.minute}"
minute_key = f"{self.prefix}:{key}:m{now.minute}"
pipe.incr(minute_key)
pipe.expire(minute_key, 60)
if tier.per_hour:
hour_key = f"{self.prefix}:{zone}:{key}:h{now.hour}"
hour_key = f"{self.prefix}:{key}:h{now.hour}"
pipe.incr(hour_key)
pipe.expire(hour_key, 3600)
if tier.per_day or self.track_daily_usage:
day = now.strftime("%Y%m%d")
day_key = f"{self.prefix}:{zone}:{key}:d{day}"
day_key = f"{self.prefix}:{key}:d{day}"
pipe.incr(day_key)
# keep data around for usage tracking
if not self.track_daily_usage:
@ -106,7 +106,6 @@ class RateLimiter:
def get_usage_since(
self,
zone: str,
key: str,
start: datetime.date,
end: typing.Optional[datetime.date] = None,
@ -120,9 +119,7 @@ class RateLimiter:
while day <= end:
days.append(day)
day += datetime.timedelta(days=1)
day_keys = [
f"{self.prefix}:{zone}:{key}:d{day.strftime('%Y%m%d')}" for day in days
]
day_keys = [f"{self.prefix}:{key}:d{day.strftime('%Y%m%d')}" for day in days]
return [
DailyUsage(d, int(calls.decode()) if calls else 0)
for d, calls in zip(days, self.redis.mget(day_keys))

29
test_ratelimit.py
View File

@ -31,7 +31,7 @@ def test_check_limit_per_minute(tier, reset_time):
# don't loop infinitely if test is failing
while count < 20:
try:
rl.check_limit("test-zone", "test-key", tier.name)
rl.check_limit("test-key", tier.name)
count += 1
except RateLimitExceeded as e:
print(e)
@ -40,7 +40,7 @@ def test_check_limit_per_minute(tier, reset_time):
assert count == 10
# resets after a given time
frozen.tick(reset_time)
assert rl.check_limit("test-zone", "test-key", tier.name)
assert rl.check_limit("test-key", tier.name)
def test_using_redis_time():
@ -51,7 +51,7 @@ def test_using_redis_time():
count = 0
while count < 20:
try:
rl.check_limit("test-zone", "test-key", simple_daily_tier.name)
rl.check_limit("test-key", simple_daily_tier.name)
count += 1
except RateLimitExceeded:
break
@ -63,19 +63,20 @@ def test_invalid_tier():
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
with pytest.raises(ValueError):
rl.check_limit("test-zone", "test-key", "non-existent-tier")
rl.check_limit("test-key", "non-existent-tier")
def test_multiple_zones():
def test_multiple_prefix():
redis.flushall()
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
rl1 = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True, prefix="zone1")
rl2 = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True, prefix="zone2")
# don't loop infinitely if test is failing
count = 0
while count < 20:
try:
rl.check_limit("zone1", "test-key", simple_daily_tier.name)
rl.check_limit("zone2", "test-key", simple_daily_tier.name)
rl1.check_limit("test-key", simple_daily_tier.name)
rl2.check_limit("test-key", simple_daily_tier.name)
count += 1
except RateLimitExceeded:
break
@ -90,8 +91,8 @@ def test_multiple_keys():
count = 0
while count < 20:
try:
rl.check_limit("zone", "test-key1", simple_daily_tier.name)
rl.check_limit("zone", "test-key2", simple_daily_tier.name)
rl.check_limit("test-key1", simple_daily_tier.name)
rl.check_limit("test-key2", simple_daily_tier.name)
count += 1
except RateLimitExceeded:
break
@ -108,10 +109,10 @@ def test_get_daily_usage():
for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"):
for _ in range(n):
rl.check_limit("zone", "test-key", unlimited_tier.name)
rl.check_limit("test-key", unlimited_tier.name)
with freeze_time("2020-01-15"):
usage = rl.get_usage_since("zone", "test-key", datetime.date(2020, 1, 1))
usage = rl.get_usage_since("test-key", datetime.date(2020, 1, 1))
assert usage[0] == DailyUsage(datetime.date(2020, 1, 1), 1)
assert usage[3] == DailyUsage(datetime.date(2020, 1, 4), 4)
assert usage[8] == DailyUsage(datetime.date(2020, 1, 9), 9)
@ -130,8 +131,8 @@ def test_get_daily_usage_untracked():
for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"):
for _ in range(n):
rl.check_limit("zone", "test-key", unlimited_tier.name)
rl.check_limit("test-key", unlimited_tier.name)
# values would be incorrect (likely zero), warn the caller
with pytest.raises(RuntimeError):
rl.get_usage_since("zone", "test-key", datetime.date(2020, 1, 1))
rl.get_usage_since("test-key", datetime.date(2020, 1, 1))