From e0ab3d45e592accc8739b6f0eca45dc286a64863 Mon Sep 17 00:00:00 2001 From: James Turk Date: Thu, 19 Nov 2020 15:29:56 -0500 Subject: [PATCH] remove zone concept, better handled as prefix --- README.md | 1 - rrl.py | 19 ++++++++----------- test_ratelimit.py | 29 +++++++++++++++-------------- 3 files changed, 23 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 6c719bf..b82e543 100644 --- a/README.md +++ b/README.md @@ -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.) diff --git a/rrl.py b/rrl.py index ae70243..fcc09d4 100644 --- a/rrl.py +++ b/rrl.py @@ -32,9 +32,9 @@ def _get_redis_connection() -> Redis: class RateLimiter: """ - :: expires in 2 minutes - :: expires in 2 hours - :: never expires + :: expires in 2 minutes + :: expires in 2 hours + :: 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)) diff --git a/test_ratelimit.py b/test_ratelimit.py index 84de977..3510f05 100644 --- a/test_ratelimit.py +++ b/test_ratelimit.py @@ -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))