initial commit

.gitignore

@@ -0,0 +1 @@
+*.pyc

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"},
+]

pyproject.toml

@@ -0,0 +1,16 @@
+[tool.poetry]
+name = "ratelimit"
+version = "0.1.0"
+description = ""
+authors = ["James Turk <dev@jamesturk.net>"]
+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"

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:
+    """
+    <zone>:<key>:<hour><minute>         expires in 2 minutes
+    <zone>:<key>:<hour>                 expires in 2 hours
+    <zone>:<key>:<day>                  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")