diff --git a/api/auth.py b/api/auth.py index 7642e09..8722d3d 100644 --- a/api/auth.py +++ b/api/auth.py @@ -1,10 +1,10 @@ from typing import Optional from fastapi import Header, HTTPException, Depends from sqlalchemy.orm.exc import NoResultFound -from rrl import RateLimiter, Tier, RateLimitExceeded from .db import SessionLocal, get_db, models +from .rate_limiter import V3RateLimiter, Tier, RateLimitExceeded -limiter = RateLimiter( +limiter = V3RateLimiter( prefix="v3", tiers=[ Tier("default", 10, 0, 250), @@ -37,7 +37,7 @@ def apikey_auth( .one() ) try: - limiter.check_limit(provided_apikey, key.api_tier) + limiter.check_limit_and_increment_counters(provided_apikey, key.api_tier) except RateLimitExceeded as e: raise HTTPException(429, detail=str(e)) except ValueError: diff --git a/api/rate_limiter.py b/api/rate_limiter.py new file mode 100644 index 0000000..3cd4a0e --- /dev/null +++ b/api/rate_limiter.py @@ -0,0 +1,153 @@ +import os +import datetime +import typing +from dataclasses import dataclass +from redis import Redis + + +@dataclass +class Tier: + name: str + per_minute: int + per_hour: int + per_day: int + + +@dataclass +class DailyUsage: + date: datetime.date + calls: int + + +class RateLimitExceeded(Exception): + pass + + +def _get_redis_connection() -> Redis: + host = os.environ.get("RRL_REDIS_HOST", "localhost") + port = int(os.environ.get("RRL_REDIS_PORT", 6379)) + db = int(os.environ.get("RRL_REDIS_DB", 0)) + return Redis(host=host, port=port, db=db) + + +class V3RateLimiter: + """ + :: expires in 2 minutes + :: expires in 2 hours + :: never expires + """ + + def __init__( + self, + tiers: typing.List[Tier], + *, + prefix: str = "", + use_redis_time: bool = True, + track_daily_usage: bool = True, + ): + self.redis = _get_redis_connection() + self.tiers = {tier.name: tier for tier in tiers} + self.prefix = prefix + self.use_redis_time = use_redis_time + self.track_daily_usage = track_daily_usage + + def check_limit_and_increment_counters(self, key: str, tier_name: str) -> bool: + try: + tier = self.tiers[tier_name] + except KeyError: + raise ValueError(f"unknown tier: {tier_name}") + if self.use_redis_time: + timestamp = self.redis.time()[0] + now = datetime.datetime.fromtimestamp(timestamp) + else: + now = datetime.datetime.utcnow() + + # check AND increment usage counters + pipe = self.redis.pipeline() + day = now.strftime("%Y%m%d") + day_key = f"{self.prefix}:{key}:d{day}" + day_requests_key = f"{self.prefix}:{key}:dr{day}" + if tier.per_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}:{key}:h{now.hour}" + pipe.incr(hour_key) + pipe.expire(hour_key, 3600) + if tier.per_day or self.track_daily_usage: + # Keep separate day and day-requests keys + # day key: used for aggregate usage tracking, so we want to limit this to + # track allowed requests the user has used + # day-requests key: tracking how many TOTAL (incl blocked) requests made + pipe.incr(day_key) + pipe.incr(day_requests_key) + # keep data around for usage tracking + if not self.track_daily_usage: + pipe.expire(day_key, 86400) + pipe.expire(day_requests_key, 86400) + result = pipe.execute() + + # parse redis pipeline results + # 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_allowed_calls, per_day_raw_calls] + # we increment value_pos as we consume values so we know which location we're looking at + value_pos = 0 + minute_calls = hour_calls = day_calls = 0 + minute_exceeded = hour_exceeded = day_exceeded = False + if tier.per_minute: + minute_calls = result[value_pos] + if result[value_pos] > tier.per_minute: + minute_exceeded = True + value_pos += 2 + if tier.per_hour: + hour_calls = result[value_pos] + if result[value_pos] > tier.per_hour: + hour_exceeded = True + value_pos += 2 + if tier.per_day: + # report back the # of raw requests, not just allowed requests + day_calls = result[value_pos + 1] + if result[value_pos] > tier.per_day: + day_exceeded = True + # daily usage numbers are used to report overall usage + # so actually want to decrement back to the prior value + # otherwise the usage count for the day will include all *blocked* requests + self.redis.decr(day_key) + + # Raise appropriate exception if limit exceeded + if minute_exceeded: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_minute}/min: {minute_calls}" + ) + if hour_exceeded: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_hour}/hour: {hour_calls}" + ) + if day_exceeded: + raise RateLimitExceeded( + f"exceeded limit of {tier.per_day}/day: {day_calls}" + ) + + return True + + def get_usage_since( + self, + key: str, + start: datetime.date, + end: typing.Optional[datetime.date] = None, + ) -> typing.List[DailyUsage]: + if not self.track_daily_usage: + raise RuntimeError("track_daily_usage is not enabled") + if not end: + end = datetime.date.today() + days = [] + day = start + while day <= end: + days.append(day) + day += datetime.timedelta(days=1) + 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/poetry.lock b/poetry.lock index 818a051..1ca542e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -59,6 +59,17 @@ files = [ [package.extras] tests = ["mypy (>=0.800)", "pytest", "pytest-asyncio"] +[[package]] +name = "async-timeout" +version = "5.0.1" +description = "Timeout context manager for asyncio programs" +optional = false +python-versions = ">=3.8" +files = [ + {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"}, + {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"}, +] + [[package]] name = "atomicwrites" version = "1.4.1" @@ -1817,6 +1828,7 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1824,8 +1836,16 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1842,6 +1862,7 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1849,6 +1870,7 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -1856,17 +1878,22 @@ files = [ [[package]] name = "redis" -version = "3.5.3" -description = "Python client for Redis key-value store" +version = "6.4.0" +description = "Python client for Redis database and key-value store" optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" +python-versions = ">=3.9" files = [ - {file = "redis-3.5.3-py2.py3-none-any.whl", hash = "sha256:432b788c4530cfe16d8d943a09d40ca6c16149727e4afe8c2c9d5580c59d9f24"}, - {file = "redis-3.5.3.tar.gz", hash = "sha256:0e7e0cfca8660dea8b7d5cd8c4f6c5e29e11f31158c0b0ae91a397f00e5a05a2"}, + {file = "redis-6.4.0-py3-none-any.whl", hash = "sha256:f0544fa9604264e9464cdf4814e7d4830f74b165d52f2a330a760a88dd248b7f"}, + {file = "redis-6.4.0.tar.gz", hash = "sha256:b01bc7282b8444e28ec36b261df5375183bb47a07eb9c603f284e89cbc5ef010"}, ] +[package.dependencies] +async-timeout = {version = ">=4.0.3", markers = "python_full_version < \"3.11.3\""} + [package.extras] -hiredis = ["hiredis (>=0.1.3)"] +hiredis = ["hiredis (>=3.2.0)"] +jwt = ["pyjwt (>=2.9.0)"] +ocsp = ["cryptography (>=36.0.1)", "pyopenssl (>=20.0.1)", "requests (>=2.31.0)"] [[package]] name = "requests" @@ -1906,20 +1933,6 @@ idna = {version = "*", optional = true, markers = "extra == \"idna2008\""} [package.extras] idna2008 = ["idna"] -[[package]] -name = "rrl" -version = "0.3.1" -description = "simple redis rate limiting" -optional = false -python-versions = ">=3.7,<4.0" -files = [ - {file = "rrl-0.3.1-py3-none-any.whl", hash = "sha256:4ac0f1373600ba2d1c4bce05938cd97122e21d8dad81a72ed2108f8f52abc97d"}, - {file = "rrl-0.3.1.tar.gz", hash = "sha256:cfbeb818198bb53c5a3c1cb8df44e2be6049d8484a589d848f678613bc0ecaa3"}, -] - -[package.dependencies] -redis = ">=3.5.3,<4.0.0" - [[package]] name = "s3transfer" version = "0.6.0" @@ -2640,4 +2653,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "cba3a6f334f071ecee3b17a9828a8e5013e35a184d694a22f8a87e9be3b17c9f" +content-hash = "0c087ff50f1c89aebdc922f204a0d9fa8f7240e95fa6424d587a8758393731db" diff --git a/pyproject.toml b/pyproject.toml index 5ed14aa..7d2d9ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "openstates-api" -version = "3.0.0" +version = "3.0.1 " description = "Open States API v3" authors = ["James Turk "] license = "MIT" @@ -15,7 +15,7 @@ gunicorn = "^20.0.4" sentry-sdk = "^1.0.0" pybase62 = "^0.4.3" python-slugify = "^4.0.1" -rrl = "^0.3.1" +redis = "^6.4.0" prometheus-fastapi-instrumentator = "^5.8.2" fastapi = {extras = ["all"], version = "^0.87.0"}