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: 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`.) * `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.) * `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: class RateLimiter:
""" """
<zone>:<key>:<hour><minute> expires in 2 minutes <prefix>:<key>:<hour><minute> expires in 2 minutes
<zone>:<key>:<hour> expires in 2 hours <prefix>:<key>:<hour> expires in 2 hours
<zone>:<key>:<day> never expires <prefix>:<key>:<day> never expires
""" """
def __init__( def __init__(
@ -51,7 +51,7 @@ class RateLimiter:
self.use_redis_time = use_redis_time self.use_redis_time = use_redis_time
self.track_daily_usage = track_daily_usage 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: try:
tier = self.tiers[tier_name] tier = self.tiers[tier_name]
except KeyError: except KeyError:
@ -64,16 +64,16 @@ class RateLimiter:
pipe = self.redis.pipeline() pipe = self.redis.pipeline()
if tier.per_minute: 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.incr(minute_key)
pipe.expire(minute_key, 60) pipe.expire(minute_key, 60)
if tier.per_hour: 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.incr(hour_key)
pipe.expire(hour_key, 3600) pipe.expire(hour_key, 3600)
if tier.per_day or self.track_daily_usage: if tier.per_day or self.track_daily_usage:
day = now.strftime("%Y%m%d") 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) pipe.incr(day_key)
# keep data around for usage tracking # keep data around for usage tracking
if not self.track_daily_usage: if not self.track_daily_usage:
@ -106,7 +106,6 @@ class RateLimiter:
def get_usage_since( def get_usage_since(
self, self,
zone: str,
key: str, key: str,
start: datetime.date, start: datetime.date,
end: typing.Optional[datetime.date] = None, end: typing.Optional[datetime.date] = None,
@ -120,9 +119,7 @@ class RateLimiter:
while day <= end: while day <= end:
days.append(day) days.append(day)
day += datetime.timedelta(days=1) day += datetime.timedelta(days=1)
day_keys = [ day_keys = [f"{self.prefix}:{key}:d{day.strftime('%Y%m%d')}" for day in days]
f"{self.prefix}:{zone}:{key}:d{day.strftime('%Y%m%d')}" for day in days
]
return [ return [
DailyUsage(d, int(calls.decode()) if calls else 0) DailyUsage(d, int(calls.decode()) if calls else 0)
for d, calls in zip(days, self.redis.mget(day_keys)) 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 # don't loop infinitely if test is failing
while count < 20: while count < 20:
try: try:
rl.check_limit("test-zone", "test-key", tier.name) rl.check_limit("test-key", tier.name)
count += 1 count += 1
except RateLimitExceeded as e: except RateLimitExceeded as e:
print(e) print(e)
@ -40,7 +40,7 @@ def test_check_limit_per_minute(tier, reset_time):
assert count == 10 assert count == 10
# resets after a given time # resets after a given time
frozen.tick(reset_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(): def test_using_redis_time():
@ -51,7 +51,7 @@ def test_using_redis_time():
count = 0 count = 0
while count < 20: while count < 20:
try: try:
rl.check_limit("test-zone", "test-key", simple_daily_tier.name) rl.check_limit("test-key", simple_daily_tier.name)
count += 1 count += 1
except RateLimitExceeded: except RateLimitExceeded:
break break
@ -63,19 +63,20 @@ def test_invalid_tier():
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True) rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
with pytest.raises(ValueError): 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() 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 # don't loop infinitely if test is failing
count = 0 count = 0
while count < 20: while count < 20:
try: try:
rl.check_limit("zone1", "test-key", simple_daily_tier.name) rl1.check_limit("test-key", simple_daily_tier.name)
rl.check_limit("zone2", "test-key", simple_daily_tier.name) rl2.check_limit("test-key", simple_daily_tier.name)
count += 1 count += 1
except RateLimitExceeded: except RateLimitExceeded:
break break
@ -90,8 +91,8 @@ def test_multiple_keys():
count = 0 count = 0
while count < 20: while count < 20:
try: try:
rl.check_limit("zone", "test-key1", simple_daily_tier.name) rl.check_limit("test-key1", simple_daily_tier.name)
rl.check_limit("zone", "test-key2", simple_daily_tier.name) rl.check_limit("test-key2", simple_daily_tier.name)
count += 1 count += 1
except RateLimitExceeded: except RateLimitExceeded:
break break
@ -108,10 +109,10 @@ def test_get_daily_usage():
for n in range(1, 10): for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"): with freeze_time(f"2020-01-0{n}"):
for _ in range(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"): 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[0] == DailyUsage(datetime.date(2020, 1, 1), 1)
assert usage[3] == DailyUsage(datetime.date(2020, 1, 4), 4) assert usage[3] == DailyUsage(datetime.date(2020, 1, 4), 4)
assert usage[8] == DailyUsage(datetime.date(2020, 1, 9), 9) 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): for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"): with freeze_time(f"2020-01-0{n}"):
for _ in range(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 # values would be incorrect (likely zero), warn the caller
with pytest.raises(RuntimeError): 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))