commit 4fc2a898c0f9f494078bb2e279c05ebd9890d11e Author: James Turk Date: Mon Oct 12 13:42:58 2020 -0400 initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0d20b64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +*.pyc diff --git a/poetry.lock b/poetry.lock new file mode 100644 index 0000000..e3e07c9 --- /dev/null +++ b/poetry.lock @@ -0,0 +1,20 @@ +[[package]] +category = "main" +description = "Python client for Redis key-value store" +name = "redis" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" +version = "3.5.3" + +[package.extras] +hiredis = ["hiredis (>=0.1.3)"] + +[metadata] +content-hash = "82ebb65843b5dc535b41ed7505e0166c19a5d34959a3aa3c76d31878cabca4a2" +python-versions = "^3.8" + +[metadata.files] +redis = [ + {file = "redis-3.5.3-py2.py3-none-any.whl", hash = "sha256:432b788c4530cfe16d8d943a09d40ca6c16149727e4afe8c2c9d5580c59d9f24"}, + {file = "redis-3.5.3.tar.gz", hash = "sha256:0e7e0cfca8660dea8b7d5cd8c4f6c5e29e11f31158c0b0ae91a397f00e5a05a2"}, +] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..15825ea --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,16 @@ +[tool.poetry] +name = "ratelimit" +version = "0.1.0" +description = "" +authors = ["James Turk "] +license = "MIT" + +[tool.poetry.dependencies] +python = "^3.8" +redis = "^3.5.3" + +[tool.poetry.dev-dependencies] + +[build-system] +requires = ["poetry>=0.12"] +build-backend = "poetry.masonry.api" diff --git a/ratelimit.py b/ratelimit.py new file mode 100644 index 0000000..8a3b37a --- /dev/null +++ b/ratelimit.py @@ -0,0 +1,56 @@ +from redis import Redis +import datetime +import typing +from dataclasses import dataclass + + +@dataclass +class Tier: + name: str + per_minute: int + per_hour: int + per_day: int + + +class RateLimitException(Exception): + pass + + +class RateLimiter: + """ + :: expires in 2 minutes + :: expires in 2 hours + :: never expires + """ + + def __init__(self, prefix: str, tiers: typing.List[Tier]): + self.redis = Redis() + self.tiers = {tier.name: tier for tier in tiers} + self.prefix = prefix + + def check_limit(self, zone: str, key: str, tier_name: str): + timestamp = self.redis.time()[0] + now = datetime.datetime.fromtimestamp(timestamp) + tier = self.tiers[tier_name] + + pipe = self.redis.pipeline() + + if tier.per_minute: + minute_key = f"{self.prefix}:{zone}:{key}:m{now.minute}" + calls = pipe.incr(minute_key) + pipe.expire(minute_key, 60) + if calls > tier.per_minute: + raise RateLimitException(f"exceeded limit of {tier.per_minute}/min") + if tier.per_hour: + hour_key = f"{self.prefix}:{zone}:{key}:h{now.hour}" + calls = pipe.incr(hour_key) + pipe.expire(hour_key, 3600) + if calls > tier.per_hour: + raise RateLimitException(f"exceeded limit of {tier.per_hour}/hour") + if tier.per_day: + day = now.strftime("%Y%m%d") + day_key = f"{self.prefix}:{zone}:{key}:d{day}" + calls = pipe.incr(day_key) + # do not expire day keys for now, useful for metrics + if calls > tier.per_day: + raise RateLimitException(f"exceeded limit of {tier.per_day}/day")