diff --git a/poetry.lock b/poetry.lock index e3e07c9..a89a618 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,3 +1,126 @@ +[[package]] +category = "dev" +description = "Atomic file writes." +marker = "sys_platform == \"win32\"" +name = "atomicwrites" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +version = "1.4.0" + +[[package]] +category = "dev" +description = "Classes Without Boilerplate" +name = "attrs" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +version = "20.2.0" + +[package.extras] +dev = ["coverage (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "zope.interface", "sphinx", "sphinx-rtd-theme", "pre-commit"] +docs = ["sphinx", "sphinx-rtd-theme", "zope.interface"] +tests = ["coverage (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "zope.interface"] +tests_no_zope = ["coverage (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six"] + +[[package]] +category = "dev" +description = "Cross-platform colored terminal text." +marker = "sys_platform == \"win32\"" +name = "colorama" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" +version = "0.4.3" + +[[package]] +category = "dev" +description = "Let your Python tests travel through time" +name = "freezegun" +optional = false +python-versions = ">=3.5" +version = "1.0.0" + +[package.dependencies] +python-dateutil = ">=2.7" + +[[package]] +category = "dev" +description = "iniconfig: brain-dead simple config-ini parsing" +name = "iniconfig" +optional = false +python-versions = "*" +version = "1.0.1" + +[[package]] +category = "dev" +description = "Core utilities for Python packages" +name = "packaging" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +version = "20.4" + +[package.dependencies] +pyparsing = ">=2.0.2" +six = "*" + +[[package]] +category = "dev" +description = "plugin and hook calling mechanisms for python" +name = "pluggy" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +version = "0.13.1" + +[package.extras] +dev = ["pre-commit", "tox"] + +[[package]] +category = "dev" +description = "library with cross-python path, ini-parsing, io, code, log facilities" +name = "py" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +version = "1.9.0" + +[[package]] +category = "dev" +description = "Python parsing module" +name = "pyparsing" +optional = false +python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" +version = "2.4.7" + +[[package]] +category = "dev" +description = "pytest: simple powerful testing with Python" +name = "pytest" +optional = false +python-versions = ">=3.5" +version = "6.1.1" + +[package.dependencies] +atomicwrites = ">=1.0" +attrs = ">=17.4.0" +colorama = "*" +iniconfig = "*" +packaging = "*" +pluggy = ">=0.12,<1.0" +py = ">=1.8.2" +toml = "*" + +[package.extras] +checkqa_mypy = ["mypy (0.780)"] +testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "requests", "xmlschema"] + +[[package]] +category = "dev" +description = "Extensions to the standard Python datetime module" +name = "python-dateutil" +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" +version = "2.8.1" + +[package.dependencies] +six = ">=1.5" + [[package]] category = "main" description = "Python client for Redis key-value store" @@ -9,12 +132,80 @@ version = "3.5.3" [package.extras] hiredis = ["hiredis (>=0.1.3)"] +[[package]] +category = "dev" +description = "Python 2 and 3 compatibility utilities" +name = "six" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" +version = "1.15.0" + +[[package]] +category = "dev" +description = "Python Library for Tom's Obvious, Minimal Language" +name = "toml" +optional = false +python-versions = "*" +version = "0.10.1" + [metadata] -content-hash = "82ebb65843b5dc535b41ed7505e0166c19a5d34959a3aa3c76d31878cabca4a2" +content-hash = "694029b97dd478608ab5a39b9cf66955b18c7b0518f03d9a8cab63cf46612aaf" python-versions = "^3.8" [metadata.files] +atomicwrites = [ + {file = "atomicwrites-1.4.0-py2.py3-none-any.whl", hash = "sha256:6d1784dea7c0c8d4a5172b6c620f40b6e4cbfdf96d783691f2e1302a7b88e197"}, + {file = "atomicwrites-1.4.0.tar.gz", hash = "sha256:ae70396ad1a434f9c7046fd2dd196fc04b12f9e91ffb859164193be8b6168a7a"}, +] +attrs = [ + {file = "attrs-20.2.0-py2.py3-none-any.whl", hash = "sha256:fce7fc47dfc976152e82d53ff92fa0407700c21acd20886a13777a0d20e655dc"}, + {file = "attrs-20.2.0.tar.gz", hash = "sha256:26b54ddbbb9ee1d34d5d3668dd37d6cf74990ab23c828c2888dccdceee395594"}, +] +colorama = [ + {file = "colorama-0.4.3-py2.py3-none-any.whl", hash = "sha256:7d73d2a99753107a36ac6b455ee49046802e59d9d076ef8e47b61499fa29afff"}, + {file = "colorama-0.4.3.tar.gz", hash = "sha256:e96da0d330793e2cb9485e9ddfd918d456036c7149416295932478192f4436a1"}, +] +freezegun = [ + {file = "freezegun-1.0.0-py2.py3-none-any.whl", hash = "sha256:02b35de52f4699a78f6ac4518e4cd3390dddc43b0aeb978335a8f270a2d9668b"}, + {file = "freezegun-1.0.0.tar.gz", hash = "sha256:1cf08e441f913ff5e59b19cc065a8faa9dd1ddc442eaf0375294f344581a0643"}, +] +iniconfig = [ + {file = "iniconfig-1.0.1-py3-none-any.whl", hash = "sha256:80cf40c597eb564e86346103f609d74efce0f6b4d4f30ec8ce9e2c26411ba437"}, + {file = "iniconfig-1.0.1.tar.gz", hash = "sha256:e5f92f89355a67de0595932a6c6c02ab4afddc6fcdc0bfc5becd0d60884d3f69"}, +] +packaging = [ + {file = "packaging-20.4-py2.py3-none-any.whl", hash = "sha256:998416ba6962ae7fbd6596850b80e17859a5753ba17c32284f67bfff33784181"}, + {file = "packaging-20.4.tar.gz", hash = "sha256:4357f74f47b9c12db93624a82154e9b120fa8293699949152b22065d556079f8"}, +] +pluggy = [ + {file = "pluggy-0.13.1-py2.py3-none-any.whl", hash = "sha256:966c145cd83c96502c3c3868f50408687b38434af77734af1e9ca461a4081d2d"}, + {file = "pluggy-0.13.1.tar.gz", hash = "sha256:15b2acde666561e1298d71b523007ed7364de07029219b604cf808bfa1c765b0"}, +] +py = [ + {file = "py-1.9.0-py2.py3-none-any.whl", hash = "sha256:366389d1db726cd2fcfc79732e75410e5fe4d31db13692115529d34069a043c2"}, + {file = "py-1.9.0.tar.gz", hash = "sha256:9ca6883ce56b4e8da7e79ac18787889fa5206c79dcc67fb065376cd2fe03f342"}, +] +pyparsing = [ + {file = "pyparsing-2.4.7-py2.py3-none-any.whl", hash = "sha256:ef9d7589ef3c200abe66653d3f1ab1033c3c419ae9b9bdb1240a85b024efc88b"}, + {file = "pyparsing-2.4.7.tar.gz", hash = "sha256:c203ec8783bf771a155b207279b9bccb8dea02d8f0c9e5f8ead507bc3246ecc1"}, +] +pytest = [ + {file = "pytest-6.1.1-py3-none-any.whl", hash = "sha256:7a8190790c17d79a11f847fba0b004ee9a8122582ebff4729a082c109e81a4c9"}, + {file = "pytest-6.1.1.tar.gz", hash = "sha256:8f593023c1a0f916110285b6efd7f99db07d59546e3d8c36fc60e2ab05d3be92"}, +] +python-dateutil = [ + {file = "python-dateutil-2.8.1.tar.gz", hash = "sha256:73ebfe9dbf22e832286dafa60473e4cd239f8592f699aa5adaf10050e6e1823c"}, + {file = "python_dateutil-2.8.1-py2.py3-none-any.whl", hash = "sha256:75bb3f31ea686f1197762692a9ee6a7550b59fc6ca3a1f4b5d7e32fb98e2da2a"}, +] redis = [ {file = "redis-3.5.3-py2.py3-none-any.whl", hash = "sha256:432b788c4530cfe16d8d943a09d40ca6c16149727e4afe8c2c9d5580c59d9f24"}, {file = "redis-3.5.3.tar.gz", hash = "sha256:0e7e0cfca8660dea8b7d5cd8c4f6c5e29e11f31158c0b0ae91a397f00e5a05a2"}, ] +six = [ + {file = "six-1.15.0-py2.py3-none-any.whl", hash = "sha256:8b74bedcbbbaca38ff6d7491d76f2b06b3592611af620f8426e82dddb04a5ced"}, + {file = "six-1.15.0.tar.gz", hash = "sha256:30639c035cdb23534cd4aa2dd52c3bf48f06e5f4a941509c8bafd8ce11080259"}, +] +toml = [ + {file = "toml-0.10.1-py2.py3-none-any.whl", hash = "sha256:bda89d5935c2eac546d648028b9901107a595863cb36bae0c73ac804a9b4ce88"}, + {file = "toml-0.10.1.tar.gz", hash = "sha256:926b612be1e5ce0634a2ca03470f95169cf16f939018233a670519cb4ac58b0f"}, +] diff --git a/pyproject.toml b/pyproject.toml index 15825ea..c4d63ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,8 @@ python = "^3.8" redis = "^3.5.3" [tool.poetry.dev-dependencies] +pytest = "^6.1.1" +freezegun = "^1.0.0" [build-system] requires = ["poetry>=0.12"] diff --git a/ratelimit.py b/ratelimit.py index 8a3b37a..4bd3f75 100644 --- a/ratelimit.py +++ b/ratelimit.py @@ -12,7 +12,7 @@ class Tier: per_day: int -class RateLimitException(Exception): +class RateLimitExceeded(Exception): pass @@ -23,34 +23,56 @@ class RateLimiter: :: never expires """ - def __init__(self, prefix: str, tiers: typing.List[Tier]): + def __init__(self, tiers: typing.List[Tier], *, prefix="", use_redis_time=True): self.redis = Redis() self.tiers = {tier.name: tier for tier in tiers} self.prefix = prefix + self.use_redis_time = use_redis_time def check_limit(self, zone: str, key: str, tier_name: str): - timestamp = self.redis.time()[0] - now = datetime.datetime.fromtimestamp(timestamp) + if self.use_redis_time: + timestamp = self.redis.time()[0] + now = datetime.datetime.fromtimestamp(timestamp) + else: + now = datetime.datetime.utcnow() 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.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.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) + 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") + result = pipe.execute() + + # the result is pairs of results of incr and expire calls, so if all 3 limits are set + # it looks like [per_minute_calls, True, per_hour_calls, True, per_day_calls] + # we increment value_pos as we consume values so we know which location we're looking at + value_pos = 0 + if tier.per_minute: + if result[value_pos] > tier.per_minute: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_minute}/min: {result[value_pos]}" + ) + value_pos += 2 + if tier.per_hour: + if result[value_pos] > tier.per_hour: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_hour}/hour: {result[value_pos]}" + ) + value_pos += 2 + if tier.per_day: + if result[value_pos] > tier.per_day: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_day}/day: {result[value_pos]}" + ) + + return True diff --git a/test_ratelimit.py b/test_ratelimit.py new file mode 100644 index 0000000..bd15472 --- /dev/null +++ b/test_ratelimit.py @@ -0,0 +1,38 @@ +import pytest +from ratelimit import Tier, RateLimiter, RateLimitExceeded +from freezegun import freeze_time +from redis import Redis + +redis = Redis() +simple_minute_tier = Tier("minute", 10, 0, 0) +simple_hour_tier = Tier("hour", 0, 10, 0) +simple_daily_tier = Tier("day", 0, 0, 10) + + +@pytest.mark.parametrize( + "tier,reset_time", + [ + (simple_minute_tier, 60), + (simple_hour_tier, 3600), + (simple_daily_tier, 60 * 60 * 25), + ], +) +def test_check_limit_per_minute(tier, reset_time): + redis.flushall() + rl = RateLimiter(tiers=[tier], use_redis_time=False) + + count = 0 + with freeze_time() as frozen: + # don't loop infinitely if test is failing + while count < 20: + try: + rl.check_limit("test-zone", "test-key", tier.name) + count += 1 + except RateLimitExceeded as e: + print(e) + break + # assert that we broke after 10 + assert count == 10 + # resets after a given time + frozen.tick(reset_time) + assert rl.check_limit("test-zone", "test-key", tier.name)