├── .github └── workflows │ └── python-package.yml ├── .gitignore ├── .travis.yml ├── LICENSE ├── MANIFEST.in ├── Makefile ├── README.rst ├── redis_rate_limit └── __init__.py ├── requirements.txt ├── setup.cfg ├── setup.py └── tests ├── __init__.py └── rate_limit_test.py /.github/workflows/python-package.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: Run Tests 5 | 6 | on: 7 | push: 8 | branches: [ master ] 9 | pull_request: 10 | branches: [ master ] 11 | 12 | jobs: 13 | build: 14 | runs-on: ubuntu-20.04 15 | services: 16 | redis: 17 | image: redis 18 | ports: 19 | - 6379:6379 20 | strategy: 21 | fail-fast: false 22 | matrix: 23 | python-version: ["3.6", "3.7","3.8", "3.9", "3.10"] 24 | 25 | steps: 26 | - uses: actions/checkout@v2 27 | - name: Set up Python ${{ matrix.python-version }} 28 | uses: actions/setup-python@v2 29 | with: 30 | python-version: ${{ matrix.python-version }} 31 | - name: Install dependencies 32 | run: | 33 | python -m pip install --upgrade pip 34 | pip install -e ./ 35 | - name: Test 36 | run: | 37 | python tests/rate_limit_test.py 38 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyx 2 | *.pyc 3 | 4 | *.swp 5 | *.orig 6 | 7 | .idea/ 8 | 9 | # virtualenv 10 | venv/ 11 | env/ 12 | 13 | build/ 14 | dist/ 15 | 16 | *.egg** 17 | 18 | .tool-versions -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - "2.7" 4 | - "3.5" 5 | - "3.6" 6 | - "3.7" 7 | - "3.8" 8 | - "3.9" 9 | script: "make test" 10 | branches: 11 | only: 12 | - master 13 | services: 14 | - redis-server 15 | 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License 2 | 3 | Copyright (c) 2016, Evolux Sistemas Ltda. 4 | 5 | This project was originally created and maintained by Victor Torres. 6 | 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | 25 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst 2 | include requirements.txt 3 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | install: 2 | python setup.py install 3 | test: 4 | python tests/rate_limit_test.py 5 | 6 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | python-redis-rate-limit 2 | ======================= 3 | .. image:: https://github.com/EvoluxBR/python-redis-rate-limit/actions/workflows/python-package.yml/badge.svg 4 | :target: https://github.com/EvoluxBR/python-redis-rate-limit/actions/workflows/python-package.yml 5 | 6 | .. image:: https://img.shields.io/pypi/v/python-redis-rate-limit.svg 7 | :target: https://pypi.python.org/pypi/python-redis-rate-limit 8 | 9 | .. image:: https://img.shields.io/pypi/dm/python-redis-rate-limit.svg 10 | :target: https://pypi.python.org/pypi/python-redis-rate-limit 11 | 12 | 13 | This lib offers an abstraction of a Rate Limit algorithm implemented on top of 14 | Redis >= 2.6.0. 15 | 16 | Supported Python Versions: 2.7, 3.6+ 17 | 18 | Example: 10 requests per second 19 | 20 | .. code-block:: python 21 | 22 | from redis_rate_limit import RateLimit, TooManyRequests 23 | try: 24 | with RateLimit(resource='users_list', client='192.168.0.10', max_requests=10): 25 | return '200 OK' 26 | except TooManyRequests: 27 | return '429 Too Many Requests' 28 | 29 | Example: using as a decorator 30 | 31 | .. code-block:: python 32 | 33 | from redis_rate_limit import RateLimit, TooManyRequests 34 | 35 | @RateLimit(resource='users_list', client='192.168.0.10', max_requests=10) 36 | def list_users(): 37 | return '200 OK' 38 | 39 | try: 40 | return list_users() 41 | except TooManyRequests: 42 | return '429 Too Many Requests' 43 | 44 | Example: 600 requests per minute 45 | 46 | .. code-block:: python 47 | 48 | from redis_rate_limit import RateLimit, TooManyRequests 49 | try: 50 | with RateLimit(resource='users_list', client='192.168.0.10', max_requests=600, expire=60): 51 | return '200 OK' 52 | except TooManyRequests: 53 | return '429 Too Many Requests' 54 | 55 | Example: 100 requests per hour 56 | 57 | .. code-block:: python 58 | 59 | from redis_rate_limit import RateLimit, TooManyRequests 60 | try: 61 | with RateLimit(resource='users_list', client='192.168.0.10', max_requests=100, expire=3600): 62 | return '200 OK' 63 | except TooManyRequests: 64 | return '429 Too Many Requests' 65 | 66 | Example: you can also setup a factory to use it later 67 | 68 | .. code-block:: python 69 | 70 | from redis_rate_limit import RateLimiter, TooManyRequests 71 | limiter = RateLimiter(resource='users_list', max_requests=100, expire=3600) 72 | try: 73 | with limiter.limit(client='192.168.0.10'): 74 | return '200 OK' 75 | except TooManyRequests: 76 | return '429 Too Many Requests' 77 | 78 | Example: you can also pass an optional Redis Pool 79 | 80 | .. code-block:: python 81 | 82 | import redis 83 | from redis_rate_limit import RateLimit, TooManyRequests 84 | redis_pool = redis.ConnectionPool(host='127.0.0.1', port=6379, db=0) 85 | try: 86 | with RateLimit(resource='users_list', client='192.168.0.10', max_requests=10, redis_pool=redis_pool): 87 | return '200 OK' 88 | except TooManyRequests: 89 | return '429 Too Many Requests' 90 | -------------------------------------------------------------------------------- /redis_rate_limit/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import functools 5 | from hashlib import sha1 6 | from redis.exceptions import NoScriptError 7 | from redis import Redis, ConnectionPool 8 | 9 | 10 | # Adapted from http://redis.io/commands/incr#pattern-rate-limiter-2 11 | INCREMENT_SCRIPT = b""" 12 | local current 13 | current = tonumber(redis.call("incrby", KEYS[1], ARGV[2])) 14 | if current == tonumber(ARGV[2]) then 15 | redis.call("expire", KEYS[1], ARGV[1]) 16 | end 17 | return current 18 | """ 19 | INCREMENT_SCRIPT_HASH = sha1(INCREMENT_SCRIPT).hexdigest() 20 | 21 | REDIS_POOL = ConnectionPool(host='127.0.0.1', port=6379, db=0) 22 | 23 | class TooManyRequests(Exception): 24 | """ 25 | Occurs when the maximum number of requests is reached for a given resource 26 | of an specific user. 27 | """ 28 | pass 29 | 30 | 31 | class RateLimit(object): 32 | """ 33 | This class offers an abstraction of a Rate Limit algorithm implemented on 34 | top of Redis >= 2.6.0. 35 | """ 36 | def __init__(self, resource, client, max_requests, expire=None, redis_pool=REDIS_POOL): 37 | """ 38 | Class initialization method checks if the Rate Limit algorithm is 39 | actually supported by the installed Redis version and sets some 40 | useful properties. 41 | 42 | If Rate Limit is not supported, it raises an Exception. 43 | 44 | :param resource: resource identifier string (i.e. ‘user_pictures’) 45 | :param client: client identifier string (i.e. ‘192.168.0.10’) 46 | :param max_requests: integer (i.e. ‘10’) 47 | :param expire: seconds to wait before resetting counters (i.e. ‘60’) 48 | :param redis_pool: instance of redis.ConnectionPool. 49 | Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) 50 | """ 51 | self._redis = Redis(connection_pool=redis_pool) 52 | self._rate_limit_key = "rate_limit:{0}_{1}".format(resource, client) 53 | self._max_requests = max_requests 54 | self._expire = expire or 1 55 | 56 | def __call__(self, func): 57 | """ 58 | Returns a wrapped function that could be used as a decorator to 59 | rate limit resources. 60 | """ 61 | @functools.wraps(func) 62 | def wrapper(*args, **kwargs): 63 | with self: 64 | return func(*args, **kwargs) 65 | 66 | return wrapper 67 | 68 | def __enter__(self): 69 | return self.increment_usage() 70 | 71 | def __exit__(self, exc_type, exc_val, exc_tb): 72 | pass 73 | 74 | def get_usage(self): 75 | """ 76 | Returns actual resource usage by client. Note that it could be greater 77 | than the maximum number of requests set. 78 | 79 | :return: integer: current usage 80 | """ 81 | return int(self._redis.get(self._rate_limit_key) or 0) 82 | 83 | def get_wait_time(self): 84 | """ 85 | Returns estimated optimal wait time for subsequent requests. 86 | If limit has already been reached, return wait time until it gets reset. 87 | 88 | :return: float: wait time in seconds 89 | """ 90 | expire = self._redis.pttl(self._rate_limit_key) 91 | # Fallback if key has not yet been set or TTL can't be retrieved 92 | expire = expire / 1000.0 if expire > 0 else float(self._expire) 93 | if self.has_been_reached(): 94 | return expire 95 | else: 96 | return expire / (self._max_requests - self.get_usage()) 97 | 98 | def has_been_reached(self): 99 | """ 100 | Checks if Rate Limit has been reached. 101 | 102 | :return: bool: True if limit has been reached or False otherwise 103 | """ 104 | return self.get_usage() >= self._max_requests 105 | 106 | def increment_usage(self, increment_by=1): 107 | """ 108 | Calls a LUA script that should increment the resource usage by client. 109 | 110 | If the resource limit overflows the maximum number of requests, this 111 | method raises an Exception. 112 | 113 | :param increment_by: The count to increment the rate limiter by. 114 | This is by default 1, but higher values are provided for more flexible 115 | rate-limiting schemes. 116 | 117 | :return: integer: current usage 118 | """ 119 | if increment_by > self._max_requests: 120 | raise ValueError('increment_by {increment_by} overflows ' 121 | 'max_requests of {max_requests}' 122 | .format(increment_by=increment_by, 123 | max_requests=self._max_requests)) 124 | elif increment_by <= 0: 125 | raise ValueError('{increment_by} is not a valid increment, ' 126 | 'should be greater than or equal to zero.' 127 | .format(increment_by=increment_by)) 128 | 129 | try: 130 | current_usage = self._redis.evalsha( 131 | INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by) 132 | except NoScriptError: 133 | current_usage = self._redis.eval( 134 | INCREMENT_SCRIPT, 1, self._rate_limit_key, self._expire, increment_by) 135 | 136 | if int(current_usage) > self._max_requests: 137 | raise TooManyRequests() 138 | 139 | return current_usage 140 | 141 | def _reset(self): 142 | """ 143 | Deletes all keys that start with ‘rate_limit:’. 144 | """ 145 | matching_keys = self._redis.scan_iter(match='{0}*'.format('rate_limit:*')) 146 | for rate_limit_key in matching_keys: 147 | self._redis.delete(rate_limit_key) 148 | 149 | 150 | class RateLimiter(object): 151 | def __init__(self, resource, max_requests, expire=None, redis_pool=REDIS_POOL): 152 | """ 153 | Rate limit factory. Checks if RateLimit is supported when limit is called. 154 | :param resource: resource identifier string (i.e. ‘user_pictures’) 155 | :param max_requests: integer (i.e. ‘10’) 156 | :param expire: seconds to wait before resetting counters (i.e. ‘60’) 157 | :param redis_pool: instance of redis.ConnectionPool. 158 | Default: ConnectionPool(host='127.0.0.1', port=6379, db=0) 159 | """ 160 | self.resource = resource 161 | self.max_requests = max_requests 162 | self.expire = expire 163 | self.redis_pool = redis_pool 164 | 165 | def limit(self, client): 166 | """ 167 | :param client: client identifier string (i.e. ‘192.168.0.10’) 168 | """ 169 | return RateLimit( 170 | resource=self.resource, 171 | client=client, 172 | max_requests=self.max_requests, 173 | expire=self.expire, 174 | redis_pool=self.redis_pool, 175 | ) 176 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | redis 2 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal=1 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from setuptools import setup, find_packages 4 | 5 | 6 | with open('README.rst') as f: 7 | readme = f.read() 8 | 9 | with open('requirements.txt') as f: 10 | requires = f.readlines() 11 | 12 | setup( 13 | name='python-redis-rate-limit', 14 | version='0.0.10', 15 | description=u'Python Rate Limiter based on Redis.', 16 | long_description=readme, 17 | author=u'Victor Torres', 18 | author_email=u'vpaivatorres@gmail.com', 19 | url=u'https://github.com/evoluxbr/python-redis-rate-limit', 20 | license=u'MIT', 21 | packages=find_packages(exclude=('tests', 'docs')), 22 | classifiers=[ 23 | 'Development Status :: 4 - Beta', 24 | 'Intended Audience :: Developers', 25 | 'Programming Language :: Python', 26 | 'Programming Language :: Python :: 3.6', 27 | 'Programming Language :: Python :: 3.7', 28 | 'Programming Language :: Python :: 3.8', 29 | 'Programming Language :: Python :: 3.9', 30 | 'Programming Language :: Python :: 3.10', 31 | 'License :: OSI Approved :: MIT License' 32 | ], 33 | install_requires=requires 34 | ) 35 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/EvoluxBR/python-redis-rate-limit/97fc2ff5beaf2ee005740538b97eafb1597eae70/tests/__init__.py -------------------------------------------------------------------------------- /tests/rate_limit_test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import unittest 4 | import time 5 | from redis_rate_limit import RateLimit, RateLimiter, TooManyRequests 6 | 7 | 8 | class TestRedisRateLimit(unittest.TestCase): 9 | def setUp(self): 10 | """ 11 | Initialises Rate Limit class and delete all keys from Redis. 12 | """ 13 | self.rate_limit = RateLimit(resource='test', client='localhost', 14 | max_requests=10) 15 | self.rate_limit._reset() 16 | 17 | def _make_10_requests(self): 18 | """ 19 | Increments usage ten times. 20 | """ 21 | for x in range(0, 10): 22 | with self.rate_limit: 23 | pass 24 | 25 | def test_limit_10_max_request(self): 26 | """ 27 | Should raise TooManyRequests Exception when trying to increment for the 28 | eleventh time. 29 | """ 30 | self.assertEqual(self.rate_limit.get_usage(), 0) 31 | self.assertEqual(self.rate_limit.has_been_reached(), False) 32 | 33 | self._make_10_requests() 34 | self.assertEqual(self.rate_limit.get_usage(), 10) 35 | self.assertEqual(self.rate_limit.has_been_reached(), True) 36 | 37 | with self.assertRaises(TooManyRequests): 38 | with self.rate_limit: 39 | pass 40 | 41 | self.assertEqual(self.rate_limit.get_usage(), 11) 42 | self.assertEqual(self.rate_limit.has_been_reached(), True) 43 | 44 | def test_expire(self): 45 | """ 46 | Should not raise TooManyRequests Exception when trying to increment for 47 | the eleventh time after the expire time. 48 | """ 49 | self._make_10_requests() 50 | time.sleep(1) 51 | with self.rate_limit: 52 | pass 53 | 54 | def test_not_expired(self): 55 | """ 56 | Should raise TooManyRequests Exception when the expire time has not 57 | been reached yet. 58 | """ 59 | self.rate_limit = RateLimit(resource='test', client='localhost', 60 | max_requests=10, expire=2) 61 | self._make_10_requests() 62 | time.sleep(1) 63 | with self.assertRaises(TooManyRequests): 64 | with self.rate_limit: 65 | pass 66 | 67 | def test_limit_10_using_rate_limiter(self): 68 | """ 69 | Should raise TooManyRequests Exception when trying to increment for the 70 | eleventh time. 71 | """ 72 | self.rate_limit = RateLimiter(resource='test', max_requests=10, 73 | expire=2).limit(client='localhost') 74 | self.assertEqual(self.rate_limit.get_usage(), 0) 75 | self.assertEqual(self.rate_limit.has_been_reached(), False) 76 | 77 | self._make_10_requests() 78 | self.assertEqual(self.rate_limit.get_usage(), 10) 79 | self.assertEqual(self.rate_limit.has_been_reached(), True) 80 | 81 | with self.assertRaises(TooManyRequests): 82 | with self.rate_limit: 83 | pass 84 | 85 | self.assertEqual(self.rate_limit.get_usage(), 11) 86 | self.assertEqual(self.rate_limit.has_been_reached(), True) 87 | 88 | def test_wait_time_limit_reached(self): 89 | """ 90 | Should report wait time approximately equal to expire after reaching 91 | the limit without delay between requests. 92 | """ 93 | self.rate_limit = RateLimit(resource='test', client='localhost', 94 | max_requests=10, expire=1) 95 | self._make_10_requests() 96 | with self.assertRaises(TooManyRequests): 97 | with self.rate_limit: 98 | pass 99 | self.assertAlmostEqual(self.rate_limit.get_wait_time(), 1, places=2) 100 | 101 | def test_wait_time_limit_expired(self): 102 | """ 103 | Should report wait time equal to expire / max_requests before any 104 | requests were made and after the limit has expired. 105 | """ 106 | self.rate_limit = RateLimit(resource='test', client='localhost', 107 | max_requests=10, expire=1) 108 | self.assertEqual(self.rate_limit.get_wait_time(), 1./10) 109 | self._make_10_requests() 110 | time.sleep(1) 111 | self.assertEqual(self.rate_limit.get_wait_time(), 1./10) 112 | 113 | def test_context_manager_returns_usage(self): 114 | """ 115 | Should return the usage when used as a context manager. 116 | """ 117 | self.rate_limit = RateLimit(resource='test', client='localhost', 118 | max_requests=1, expire=1) 119 | with self.rate_limit as usage: 120 | self.assertEqual(usage, 1) 121 | 122 | def test_limit_10_using_as_decorator(self): 123 | """ 124 | Should raise TooManyRequests Exception when trying to increment for the 125 | eleventh time. 126 | """ 127 | self.assertEqual(self.rate_limit.get_usage(), 0) 128 | self.assertEqual(self.rate_limit.has_been_reached(), False) 129 | 130 | self._make_10_requests() 131 | self.assertEqual(self.rate_limit.get_usage(), 10) 132 | self.assertEqual(self.rate_limit.has_been_reached(), True) 133 | 134 | @self.rate_limit 135 | def limit_with_decorator(): 136 | pass 137 | 138 | with self.assertRaises(TooManyRequests): 139 | limit_with_decorator() 140 | 141 | self.assertEqual(self.rate_limit.get_usage(), 11) 142 | self.assertEqual(self.rate_limit.has_been_reached(), True) 143 | 144 | def test_increment_multiple(self): 145 | """ 146 | Test incrementing usage by a value > 1 147 | """ 148 | self.rate_limit.increment_usage(7) 149 | self.rate_limit.increment_usage(3) 150 | 151 | self.assertEqual(self.rate_limit.get_usage(), 10) 152 | self.assertEqual(self.rate_limit.has_been_reached(), True) 153 | 154 | with self.assertRaises(TooManyRequests): 155 | self.rate_limit.increment_usage(1) 156 | 157 | self.assertEqual(self.rate_limit.get_usage(), 11) 158 | self.assertEqual(self.rate_limit.has_been_reached(), True) 159 | 160 | with self.assertRaises(TooManyRequests): 161 | self.rate_limit.increment_usage(5) 162 | 163 | self.assertEqual(self.rate_limit.get_usage(), 16) 164 | self.assertEqual(self.rate_limit.has_been_reached(), True) 165 | 166 | def test_increment_multiple_too_much(self): 167 | """ 168 | Test that we cannot bulk-increment a value higher than 169 | the bucket limit. 170 | """ 171 | with self.assertRaises(ValueError): 172 | self.rate_limit.increment_usage(11) 173 | 174 | self.assertEqual(self.rate_limit.get_usage(), 0) 175 | self.assertEqual(self.rate_limit.has_been_reached(), False) 176 | 177 | def test_increment_by_zero(self): 178 | """ 179 | Should not allow increment by zero. 180 | """ 181 | self.assertEqual(self.rate_limit.get_usage(), 0) 182 | self.assertEqual(self.rate_limit.has_been_reached(), False) 183 | 184 | self.rate_limit.increment_usage(5) 185 | self.assertEqual(self.rate_limit.get_usage(), 5) 186 | self.assertEqual(self.rate_limit.has_been_reached(), False) 187 | 188 | with self.assertRaises(ValueError): 189 | self.rate_limit.increment_usage(0) 190 | 191 | self.assertEqual(self.rate_limit.get_usage(), 5) 192 | self.assertEqual(self.rate_limit.has_been_reached(), False) 193 | 194 | def test_increment_by_negative(self): 195 | """ 196 | Should not allow decrement the counter. 197 | """ 198 | self.assertEqual(self.rate_limit.get_usage(), 0) 199 | self.assertEqual(self.rate_limit.has_been_reached(), False) 200 | with self.assertRaises(ValueError): 201 | self.rate_limit.increment_usage(-5) 202 | 203 | self.assertEqual(self.rate_limit.get_usage(), 0) 204 | self.assertEqual(self.rate_limit.has_been_reached(), False) 205 | 206 | 207 | if __name__ == '__main__': 208 | unittest.main() 209 | --------------------------------------------------------------------------------