ratelimit-io 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Bogdan
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: ratelimit-io
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Flexible bidirectional rate limiting library with redis backend
|
|
5
|
+
Home-page: https://github.com/bagowix/ratelimit-io
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: asyncio,rate-limiting,I/O,redis,aioredis
|
|
8
|
+
Author: Galushko Bogdan
|
|
9
|
+
Author-email: galushko355@gmail.com
|
|
10
|
+
Requires-Python: >=3.8,<4.0
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Requires-Dist: redis (>=5.2.1,<6.0.0)
|
|
20
|
+
Project-URL: Repository, https://github.com/bagowix/ratelimit-io
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# RatelimitIO
|
|
24
|
+
|
|
25
|
+
A Python library for rate limiting, supporting both synchronous and asynchronous paradigms. Built on Redis, this library provides decorators, context managers, and easy integration with APIs to manage incoming and outgoing request limits effectively.
|
|
26
|
+
|
|
27
|
+
#### Project Information
|
|
28
|
+
[](https://github.com/bagowix/ratelimit-io/actions/workflows/actions.yml)
|
|
29
|
+
[](https://github.com/bagowix/ratelimit-io/blob/main/coverage.json)
|
|
30
|
+
[](https://pypi.org/project/ratelimit-io/)
|
|
31
|
+
[](LICENSE)
|
|
32
|
+
|
|
33
|
+
---
|
|
34
|
+
|
|
35
|
+
## Features
|
|
36
|
+
|
|
37
|
+
- **Synchronous and Asynchronous Support**: Works seamlessly in both paradigms.
|
|
38
|
+
- **Redis Backend**: Leverages Redis for fast and scalable rate limiting.
|
|
39
|
+
- **Flexible API**:
|
|
40
|
+
- Use as a **decorator** for methods or functions.
|
|
41
|
+
- Use as a **context manager**.
|
|
42
|
+
- Integrate directly into API clients or middlewares.
|
|
43
|
+
- **Customizable Rate Limits**: Specify limits per key, time period, and requests.
|
|
44
|
+
- **Robust Lua Script**: Efficient and atomic rate limiting logic.
|
|
45
|
+
|
|
46
|
+
---
|
|
47
|
+
|
|
48
|
+
## Installation
|
|
49
|
+
|
|
50
|
+
Install via pip:
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
pip install ratelimit-io
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## Quick Start
|
|
57
|
+
|
|
58
|
+
### Using as a Synchronous Decorator
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
from ratelimit_io import RatelimitIO, LimitSpec
|
|
62
|
+
from redis import Redis
|
|
63
|
+
|
|
64
|
+
redis_client = Redis(host="localhost", port=6379)
|
|
65
|
+
limiter = RatelimitIO(backend=redis_client)
|
|
66
|
+
|
|
67
|
+
@limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123")
|
|
68
|
+
def fetch_data():
|
|
69
|
+
return "Request succeeded!"
|
|
70
|
+
|
|
71
|
+
# Use the decorated function
|
|
72
|
+
fetch_data()
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
### Using as a Asynchronous Decorator
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
from ratelimit_io import RatelimitIO, LimitSpec
|
|
79
|
+
from redis.asyncio import Redis as AsyncRedis
|
|
80
|
+
|
|
81
|
+
async_redis_client = AsyncRedis(host="localhost", port=6379)
|
|
82
|
+
async_limiter = RatelimitIO(backend=async_redis_client)
|
|
83
|
+
|
|
84
|
+
@async_limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123")
|
|
85
|
+
async def fetch_data():
|
|
86
|
+
return "Request succeeded!"
|
|
87
|
+
|
|
88
|
+
# Use the decorated function
|
|
89
|
+
await fetch_data()
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
### Using as a Synchronous Context Manager
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
with limiter:
|
|
96
|
+
limiter.wait("user:456", LimitSpec(requests=5, seconds=10))
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
### Using as a Asynchronous Context Manager
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
async with limiter:
|
|
103
|
+
await limiter.a_wait("user:456", LimitSpec(requests=5, seconds=10))
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
## License
|
|
107
|
+
|
|
108
|
+
[MIT](https://github.com/bagowix/ratelimit-io/blob/main/LICENSE)
|
|
109
|
+
|
|
110
|
+
## Contribution
|
|
111
|
+
|
|
112
|
+
Contributions are welcome! Follow the Contribution Guide for details.
|
|
113
|
+
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
# RatelimitIO
|
|
2
|
+
|
|
3
|
+
A Python library for rate limiting, supporting both synchronous and asynchronous paradigms. Built on Redis, this library provides decorators, context managers, and easy integration with APIs to manage incoming and outgoing request limits effectively.
|
|
4
|
+
|
|
5
|
+
#### Project Information
|
|
6
|
+
[](https://github.com/bagowix/ratelimit-io/actions/workflows/actions.yml)
|
|
7
|
+
[](https://github.com/bagowix/ratelimit-io/blob/main/coverage.json)
|
|
8
|
+
[](https://pypi.org/project/ratelimit-io/)
|
|
9
|
+
[](LICENSE)
|
|
10
|
+
|
|
11
|
+
---
|
|
12
|
+
|
|
13
|
+
## Features
|
|
14
|
+
|
|
15
|
+
- **Synchronous and Asynchronous Support**: Works seamlessly in both paradigms.
|
|
16
|
+
- **Redis Backend**: Leverages Redis for fast and scalable rate limiting.
|
|
17
|
+
- **Flexible API**:
|
|
18
|
+
- Use as a **decorator** for methods or functions.
|
|
19
|
+
- Use as a **context manager**.
|
|
20
|
+
- Integrate directly into API clients or middlewares.
|
|
21
|
+
- **Customizable Rate Limits**: Specify limits per key, time period, and requests.
|
|
22
|
+
- **Robust Lua Script**: Efficient and atomic rate limiting logic.
|
|
23
|
+
|
|
24
|
+
---
|
|
25
|
+
|
|
26
|
+
## Installation
|
|
27
|
+
|
|
28
|
+
Install via pip:
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
pip install ratelimit-io
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
## Quick Start
|
|
35
|
+
|
|
36
|
+
### Using as a Synchronous Decorator
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
from ratelimit_io import RatelimitIO, LimitSpec
|
|
40
|
+
from redis import Redis
|
|
41
|
+
|
|
42
|
+
redis_client = Redis(host="localhost", port=6379)
|
|
43
|
+
limiter = RatelimitIO(backend=redis_client)
|
|
44
|
+
|
|
45
|
+
@limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123")
|
|
46
|
+
def fetch_data():
|
|
47
|
+
return "Request succeeded!"
|
|
48
|
+
|
|
49
|
+
# Use the decorated function
|
|
50
|
+
fetch_data()
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
### Using as a Asynchronous Decorator
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
from ratelimit_io import RatelimitIO, LimitSpec
|
|
57
|
+
from redis.asyncio import Redis as AsyncRedis
|
|
58
|
+
|
|
59
|
+
async_redis_client = AsyncRedis(host="localhost", port=6379)
|
|
60
|
+
async_limiter = RatelimitIO(backend=async_redis_client)
|
|
61
|
+
|
|
62
|
+
@async_limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123")
|
|
63
|
+
async def fetch_data():
|
|
64
|
+
return "Request succeeded!"
|
|
65
|
+
|
|
66
|
+
# Use the decorated function
|
|
67
|
+
await fetch_data()
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
### Using as a Synchronous Context Manager
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
with limiter:
|
|
74
|
+
limiter.wait("user:456", LimitSpec(requests=5, seconds=10))
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### Using as a Asynchronous Context Manager
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
async with limiter:
|
|
81
|
+
await limiter.a_wait("user:456", LimitSpec(requests=5, seconds=10))
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
## License
|
|
85
|
+
|
|
86
|
+
[MIT](https://github.com/bagowix/ratelimit-io/blob/main/LICENSE)
|
|
87
|
+
|
|
88
|
+
## Contribution
|
|
89
|
+
|
|
90
|
+
Contributions are welcome! Follow the Contribution Guide for details.
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
[tool.poetry]
|
|
2
|
+
name = "ratelimit-io"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Flexible bidirectional rate limiting library with redis backend"
|
|
5
|
+
authors = ["Galushko Bogdan <galushko355@gmail.com>"]
|
|
6
|
+
license = "MIT"
|
|
7
|
+
readme = "README.md"
|
|
8
|
+
packages = [{ include = "ratelimit_io"}]
|
|
9
|
+
|
|
10
|
+
package-mode = true
|
|
11
|
+
|
|
12
|
+
repository = "https://github.com/bagowix/ratelimit-io"
|
|
13
|
+
|
|
14
|
+
keywords = [
|
|
15
|
+
"asyncio",
|
|
16
|
+
"rate-limiting",
|
|
17
|
+
"I/O",
|
|
18
|
+
"redis",
|
|
19
|
+
"aioredis",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[tool.poetry.dependencies]
|
|
23
|
+
python = "^3.8"
|
|
24
|
+
redis = "^5.2.1"
|
|
25
|
+
|
|
26
|
+
[tool.poetry.group.dev.dependencies]
|
|
27
|
+
mypy = "^1.14.1"
|
|
28
|
+
ruff = "^0.9.1"
|
|
29
|
+
pytest = "^8.3.4"
|
|
30
|
+
pytest-cov = "^5.0.0"
|
|
31
|
+
pytest-asyncio = "^0.18.3"
|
|
32
|
+
pre-commit = "^3.5"
|
|
33
|
+
fakeredis = "^2.26.2"
|
|
34
|
+
|
|
35
|
+
[tool.ruff]
|
|
36
|
+
line-length = 79
|
|
37
|
+
target-version = "py38"
|
|
38
|
+
|
|
39
|
+
[tool.ruff.lint]
|
|
40
|
+
ignore = [
|
|
41
|
+
"PT001",
|
|
42
|
+
"PT023",
|
|
43
|
+
"SIM103",
|
|
44
|
+
]
|
|
45
|
+
select = [
|
|
46
|
+
"B", # flake8-bugbear,
|
|
47
|
+
"C4", # flake8-comprehensions
|
|
48
|
+
"E", # pycodestyle error
|
|
49
|
+
"F", # pyflakes
|
|
50
|
+
"SIM", # flake8-simplify
|
|
51
|
+
"T20", # flake8-print
|
|
52
|
+
"W", # pycodestyle warning
|
|
53
|
+
"UP", # pyupgrade
|
|
54
|
+
"I", # isort
|
|
55
|
+
]
|
|
56
|
+
|
|
57
|
+
[tool.ruff.lint.mccabe]
|
|
58
|
+
max-complexity = 12
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
[tool.ruff.lint.per-file-ignores]
|
|
62
|
+
"__init__.py" = [
|
|
63
|
+
"F401",
|
|
64
|
+
"F403",
|
|
65
|
+
]
|
|
66
|
+
"test_*.py" = [
|
|
67
|
+
"S101",
|
|
68
|
+
"E501",
|
|
69
|
+
"ANN201",
|
|
70
|
+
]
|
|
71
|
+
|
|
72
|
+
[tool.ruff.lint.isort]
|
|
73
|
+
known-first-party = [
|
|
74
|
+
"ratelimit_io",
|
|
75
|
+
]
|
|
76
|
+
known-third-party = [
|
|
77
|
+
"asyncio",
|
|
78
|
+
"hashlib",
|
|
79
|
+
"redis",
|
|
80
|
+
"pytest",
|
|
81
|
+
"pytest-asyncio",
|
|
82
|
+
]
|
|
83
|
+
section-order = [
|
|
84
|
+
"future",
|
|
85
|
+
"standard-library",
|
|
86
|
+
"third-party",
|
|
87
|
+
"first-party",
|
|
88
|
+
"local-folder",
|
|
89
|
+
]
|
|
90
|
+
force-single-line = true
|
|
91
|
+
|
|
92
|
+
[tool.ruff.format]
|
|
93
|
+
quote-style = "double"
|
|
94
|
+
|
|
95
|
+
[tool.pytest.ini_options]
|
|
96
|
+
asyncio_mode = "auto"
|
|
97
|
+
|
|
98
|
+
[build-system]
|
|
99
|
+
requires = ["poetry-core<=1.8.5"]
|
|
100
|
+
build-backend = "poetry.core.masonry.api"
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
from typing import Optional
|
|
2
|
+
from typing import Union
|
|
3
|
+
|
|
4
|
+
from redis import Redis
|
|
5
|
+
from redis.asyncio import Redis as AsyncRedis
|
|
6
|
+
from typing_extensions import TypeAlias
|
|
7
|
+
|
|
8
|
+
from ratelimit_io.rate_limit import RatelimitIO # noqa
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class RatelimitIOError(Exception):
|
|
12
|
+
"""Raised when the rate limit is exceeded."""
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class LimitSpec:
|
|
16
|
+
"""
|
|
17
|
+
Specifies the number of requests allowed in a time frame.
|
|
18
|
+
|
|
19
|
+
Attributes:
|
|
20
|
+
requests (int): Maximum number of requests.
|
|
21
|
+
seconds (Optional[int]): Time frame in seconds.
|
|
22
|
+
minutes (Optional[int]): Time frame in minutes.
|
|
23
|
+
hours (Optional[int]): Time frame in hours.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
requests: int,
|
|
29
|
+
seconds: Optional[int] = None,
|
|
30
|
+
minutes: Optional[int] = None,
|
|
31
|
+
hours: Optional[int] = None,
|
|
32
|
+
):
|
|
33
|
+
if requests <= 0:
|
|
34
|
+
raise ValueError("Requests must be greater than 0.")
|
|
35
|
+
|
|
36
|
+
self.requests = requests
|
|
37
|
+
self.seconds = seconds
|
|
38
|
+
self.minutes = minutes
|
|
39
|
+
self.hours = hours
|
|
40
|
+
|
|
41
|
+
if self.total_seconds() == 0:
|
|
42
|
+
raise ValueError(
|
|
43
|
+
"At least one time frame "
|
|
44
|
+
"(seconds, minutes, or hours) must be provided."
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
def total_seconds(self) -> int:
|
|
48
|
+
"""
|
|
49
|
+
Calculates the total time frame in seconds.
|
|
50
|
+
|
|
51
|
+
Returns:
|
|
52
|
+
int: Total time in seconds.
|
|
53
|
+
"""
|
|
54
|
+
total = 0
|
|
55
|
+
if self.seconds:
|
|
56
|
+
total += self.seconds
|
|
57
|
+
if self.minutes:
|
|
58
|
+
total += self.minutes * 60
|
|
59
|
+
if self.hours:
|
|
60
|
+
total += self.hours * 3600
|
|
61
|
+
return total
|
|
62
|
+
|
|
63
|
+
def __str__(self) -> str:
|
|
64
|
+
return f"{self.requests}/{self.total_seconds()}s"
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
RedisBackend: TypeAlias = Union[Redis, AsyncRedis]
|
|
@@ -0,0 +1,396 @@
|
|
|
1
|
+
import time
|
|
2
|
+
from typing import Callable
|
|
3
|
+
from typing import Optional
|
|
4
|
+
from typing import Tuple
|
|
5
|
+
from typing import Union
|
|
6
|
+
|
|
7
|
+
import asyncio
|
|
8
|
+
import hashlib
|
|
9
|
+
from redis import Redis
|
|
10
|
+
from redis.asyncio import Redis as AsyncRedis
|
|
11
|
+
from redis.exceptions import NoScriptError
|
|
12
|
+
from typing_extensions import TypeAlias
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class RatelimitIOError(Exception):
|
|
16
|
+
"""Raised when the rate limit is exceeded."""
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class LimitSpec:
|
|
20
|
+
"""
|
|
21
|
+
Specifies the number of requests allowed in a time frame.
|
|
22
|
+
|
|
23
|
+
Attributes:
|
|
24
|
+
requests (int): Maximum number of requests.
|
|
25
|
+
seconds (Optional[int]): Time frame in seconds.
|
|
26
|
+
minutes (Optional[int]): Time frame in minutes.
|
|
27
|
+
hours (Optional[int]): Time frame in hours.
|
|
28
|
+
"""
|
|
29
|
+
|
|
30
|
+
def __init__(
|
|
31
|
+
self,
|
|
32
|
+
requests: int,
|
|
33
|
+
seconds: Optional[int] = None,
|
|
34
|
+
minutes: Optional[int] = None,
|
|
35
|
+
hours: Optional[int] = None,
|
|
36
|
+
):
|
|
37
|
+
if requests <= 0:
|
|
38
|
+
raise ValueError("Requests must be greater than 0.")
|
|
39
|
+
|
|
40
|
+
self.requests = requests
|
|
41
|
+
self.seconds = seconds
|
|
42
|
+
self.minutes = minutes
|
|
43
|
+
self.hours = hours
|
|
44
|
+
|
|
45
|
+
if self.total_seconds() == 0:
|
|
46
|
+
raise ValueError(
|
|
47
|
+
"At least one time frame "
|
|
48
|
+
"(seconds, minutes, or hours) must be provided."
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
def total_seconds(self) -> int:
|
|
52
|
+
"""
|
|
53
|
+
Calculates the total time frame in seconds.
|
|
54
|
+
|
|
55
|
+
Returns:
|
|
56
|
+
int: Total time in seconds.
|
|
57
|
+
"""
|
|
58
|
+
total = 0
|
|
59
|
+
if self.seconds:
|
|
60
|
+
total += self.seconds
|
|
61
|
+
if self.minutes:
|
|
62
|
+
total += self.minutes * 60
|
|
63
|
+
if self.hours:
|
|
64
|
+
total += self.hours * 3600
|
|
65
|
+
return total
|
|
66
|
+
|
|
67
|
+
def __str__(self) -> str:
|
|
68
|
+
return f"{self.requests}/{self.total_seconds()}s"
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
RedisBackend: TypeAlias = Union[Redis, AsyncRedis]
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
class RatelimitIO:
|
|
75
|
+
"""Rate limiter for managing incoming and outgoing request limits."""
|
|
76
|
+
|
|
77
|
+
def __init__(
|
|
78
|
+
self,
|
|
79
|
+
backend: RedisBackend,
|
|
80
|
+
base_url: Optional[str] = None,
|
|
81
|
+
base_limit: Optional[LimitSpec] = None,
|
|
82
|
+
):
|
|
83
|
+
"""
|
|
84
|
+
Initializes the RatelimitIO instance.
|
|
85
|
+
|
|
86
|
+
Args:
|
|
87
|
+
backend (Redis | AsyncRedis): Redis backend instance.
|
|
88
|
+
base_url (Optional[str]): Base URL for outgoing request limits.
|
|
89
|
+
base_limit (Optional[LimitSpec]): Default rate
|
|
90
|
+
limit for the base URL.
|
|
91
|
+
"""
|
|
92
|
+
if not isinstance(backend, (Redis, AsyncRedis)):
|
|
93
|
+
raise RuntimeError("Unsupported Redis backend.")
|
|
94
|
+
|
|
95
|
+
self.backend = backend
|
|
96
|
+
self.base_url = base_url
|
|
97
|
+
self.base_limit = base_limit
|
|
98
|
+
|
|
99
|
+
self._lua_script = b"""
|
|
100
|
+
local current = redis.call("GET", KEYS[1])
|
|
101
|
+
local limit = tonumber(ARGV[1])
|
|
102
|
+
local ttl = tonumber(ARGV[2])
|
|
103
|
+
|
|
104
|
+
if current and tonumber(current) >= limit then
|
|
105
|
+
return 0
|
|
106
|
+
else
|
|
107
|
+
local new_count = redis.call("INCR", KEYS[1])
|
|
108
|
+
if new_count == 1 then
|
|
109
|
+
redis.call("EXPIRE", KEYS[1], ttl)
|
|
110
|
+
end
|
|
111
|
+
return 1
|
|
112
|
+
end
|
|
113
|
+
"""
|
|
114
|
+
self._lua_script_hash = hashlib.sha1(self._lua_script).hexdigest()
|
|
115
|
+
self._script_loaded = False
|
|
116
|
+
|
|
117
|
+
def __call__(
|
|
118
|
+
self, limit_spec: LimitSpec, unique_key: Optional[str] = None
|
|
119
|
+
) -> Callable:
|
|
120
|
+
"""
|
|
121
|
+
Decorator for applying rate limits to functions.
|
|
122
|
+
|
|
123
|
+
Args:
|
|
124
|
+
limit_spec (LimitSpec): Rate limit specification.
|
|
125
|
+
unique_key (Optional[str]): Optional unique key for rate limiting.
|
|
126
|
+
|
|
127
|
+
Returns:
|
|
128
|
+
Callable: Decorated function.
|
|
129
|
+
"""
|
|
130
|
+
|
|
131
|
+
def decorator(func: Callable) -> Callable:
|
|
132
|
+
async def async_wrapper(*args, **kwargs):
|
|
133
|
+
key = unique_key or kwargs.get("ip", "unknown_ip")
|
|
134
|
+
try:
|
|
135
|
+
await self.a_wait(f"incoming:{key}", limit_spec)
|
|
136
|
+
except RatelimitIOError as exc:
|
|
137
|
+
raise RuntimeError(
|
|
138
|
+
f"Rate limit exceeded in {func.__name__}: {str(exc)}"
|
|
139
|
+
) from exc
|
|
140
|
+
return await func(*args, **kwargs)
|
|
141
|
+
|
|
142
|
+
def sync_wrapper(*args, **kwargs):
|
|
143
|
+
key = unique_key or kwargs.get("ip", "unknown_ip")
|
|
144
|
+
try:
|
|
145
|
+
self.wait(f"incoming:{key}", limit_spec)
|
|
146
|
+
except RatelimitIOError as exc:
|
|
147
|
+
raise RuntimeError(
|
|
148
|
+
f"Rate limit exceeded in {func.__name__}: {str(exc)}"
|
|
149
|
+
) from exc
|
|
150
|
+
return func(*args, **kwargs)
|
|
151
|
+
|
|
152
|
+
return (
|
|
153
|
+
async_wrapper
|
|
154
|
+
if asyncio.iscoroutinefunction(func)
|
|
155
|
+
else sync_wrapper
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
return decorator
|
|
159
|
+
|
|
160
|
+
def __enter__(self) -> "RatelimitIO":
|
|
161
|
+
"""Synchronous context manager."""
|
|
162
|
+
self._ensure_script_loaded_sync()
|
|
163
|
+
return self
|
|
164
|
+
|
|
165
|
+
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
166
|
+
"""No cleanup is necessary for this context manager."""
|
|
167
|
+
pass
|
|
168
|
+
|
|
169
|
+
async def __aenter__(self) -> "RatelimitIO":
|
|
170
|
+
"""Context manager for async operations."""
|
|
171
|
+
await self._ensure_script_loaded_async()
|
|
172
|
+
return self
|
|
173
|
+
|
|
174
|
+
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
|
175
|
+
"""No cleanup is necessary for this context manager."""
|
|
176
|
+
pass
|
|
177
|
+
|
|
178
|
+
def wait(
|
|
179
|
+
self,
|
|
180
|
+
key: Optional[str] = None,
|
|
181
|
+
limit_spec: Optional[LimitSpec] = None,
|
|
182
|
+
max_wait_time: float = 10.0,
|
|
183
|
+
backoff_start: float = 0.01,
|
|
184
|
+
backoff_max: float = 0.1,
|
|
185
|
+
) -> None:
|
|
186
|
+
"""
|
|
187
|
+
Synchronous rate limiting. Waits if the limit is exceeded.
|
|
188
|
+
|
|
189
|
+
Args:
|
|
190
|
+
key (Optional[str]): Unique identifier for the rate limit.
|
|
191
|
+
Defaults to `base_url`.
|
|
192
|
+
limit_spec (Optional[LimitSpec]): Rate specification.
|
|
193
|
+
Defaults to `base_limit`.
|
|
194
|
+
max_wait_time (float): Maximum wait time in seconds
|
|
195
|
+
before raising an error.
|
|
196
|
+
backoff_start (float): Initial backoff delay in seconds.
|
|
197
|
+
backoff_max (float): Maximum backoff delay in seconds.
|
|
198
|
+
|
|
199
|
+
Raises:
|
|
200
|
+
RatelimitIOError: If the rate limit is exceeded and
|
|
201
|
+
max wait time is reached.
|
|
202
|
+
"""
|
|
203
|
+
self._ensure_script_loaded_sync()
|
|
204
|
+
|
|
205
|
+
key, limit_spec = self._prepare_key_and_limit(key, limit_spec)
|
|
206
|
+
start_time = time.time()
|
|
207
|
+
backoff = backoff_start
|
|
208
|
+
|
|
209
|
+
while not self._enforce_limit_sync(key, limit_spec):
|
|
210
|
+
if time.time() - start_time > max_wait_time:
|
|
211
|
+
raise RatelimitIOError(
|
|
212
|
+
f"Rate limit exceeded for {key}, wait time exceeded."
|
|
213
|
+
)
|
|
214
|
+
time.sleep(backoff)
|
|
215
|
+
backoff = min(backoff * 2, backoff_max)
|
|
216
|
+
|
|
217
|
+
async def a_wait(
|
|
218
|
+
self,
|
|
219
|
+
key: Optional[str] = None,
|
|
220
|
+
limit_spec: Optional[LimitSpec] = None,
|
|
221
|
+
max_wait_time: float = 10.0,
|
|
222
|
+
backoff_start: float = 0.01,
|
|
223
|
+
backoff_max: float = 0.1,
|
|
224
|
+
) -> None:
|
|
225
|
+
"""
|
|
226
|
+
Asynchronous rate limiting. Waits if the limit is exceeded.
|
|
227
|
+
|
|
228
|
+
Args:
|
|
229
|
+
key (Optional[str]): Unique identifier for the rate limit.
|
|
230
|
+
Defaults to `base_url`.
|
|
231
|
+
limit_spec (Optional[LimitSpec]): Rate specification.
|
|
232
|
+
Defaults to `base_limit`.
|
|
233
|
+
max_wait_time (float): Maximum wait time in seconds
|
|
234
|
+
before raising an error.
|
|
235
|
+
backoff_start (float): Initial backoff delay in seconds.
|
|
236
|
+
backoff_max (float): Maximum backoff delay in seconds.
|
|
237
|
+
|
|
238
|
+
Raises:
|
|
239
|
+
RatelimitIOError: If the rate limit is exceeded and
|
|
240
|
+
max wait time is reached.
|
|
241
|
+
"""
|
|
242
|
+
await self._ensure_script_loaded_async()
|
|
243
|
+
|
|
244
|
+
key, limit_spec = self._prepare_key_and_limit(key, limit_spec)
|
|
245
|
+
start_time = time.time()
|
|
246
|
+
backoff = backoff_start
|
|
247
|
+
|
|
248
|
+
while not await self._enforce_limit_async(key, limit_spec):
|
|
249
|
+
if time.time() - start_time > max_wait_time:
|
|
250
|
+
raise RatelimitIOError(
|
|
251
|
+
f"Rate limit exceeded for {key}, wait time exceeded."
|
|
252
|
+
)
|
|
253
|
+
await asyncio.sleep(backoff)
|
|
254
|
+
backoff = min(backoff * 2, backoff_max)
|
|
255
|
+
|
|
256
|
+
def _prepare_key_and_limit(
|
|
257
|
+
self, key: Optional[str] = None, limit_spec: Optional[LimitSpec] = None
|
|
258
|
+
) -> Tuple[str, LimitSpec]:
|
|
259
|
+
"""
|
|
260
|
+
Prepares the key and limit specification, falling back to
|
|
261
|
+
defaults if needed.
|
|
262
|
+
|
|
263
|
+
Args:
|
|
264
|
+
key (Optional[str]): Unique identifier for the rate limit.
|
|
265
|
+
limit_spec (Optional[LimitSpec]): Limit specification.
|
|
266
|
+
|
|
267
|
+
Returns:
|
|
268
|
+
tuple[str, LimitSpec]: Prepared key and rate specification.
|
|
269
|
+
|
|
270
|
+
Raises:
|
|
271
|
+
ValueError: If neither key nor base settings are provided.
|
|
272
|
+
"""
|
|
273
|
+
if not (limit_spec or self.base_limit):
|
|
274
|
+
raise ValueError("limit_spec or self.base_limit must be provided.")
|
|
275
|
+
|
|
276
|
+
requests_for_key: str = (
|
|
277
|
+
str(limit_spec.requests)
|
|
278
|
+
if limit_spec
|
|
279
|
+
else str(self.base_limit.requests) # type: ignore
|
|
280
|
+
)
|
|
281
|
+
time_for_key: str = (
|
|
282
|
+
str(limit_spec.total_seconds())
|
|
283
|
+
if limit_spec
|
|
284
|
+
else str(self.base_limit.total_seconds()) # type: ignore
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
key = key or (
|
|
288
|
+
f"outgoing:ratelimit-io:{self.base_url}:"
|
|
289
|
+
f"requests:{requests_for_key}:"
|
|
290
|
+
f"time:{time_for_key}"
|
|
291
|
+
)
|
|
292
|
+
limit_spec = limit_spec or self.base_limit
|
|
293
|
+
|
|
294
|
+
if not key or not limit_spec:
|
|
295
|
+
raise ValueError("Key and limit_spec must be provided.")
|
|
296
|
+
return key, limit_spec
|
|
297
|
+
|
|
298
|
+
def _enforce_limit_sync(self, key: str, limit_spec: LimitSpec) -> bool:
|
|
299
|
+
"""
|
|
300
|
+
Enforces the rate limit synchronously.
|
|
301
|
+
|
|
302
|
+
Args:
|
|
303
|
+
key (str): Unique identifier for the rate limit.
|
|
304
|
+
limit_spec (LimitSpec): Limit specification.
|
|
305
|
+
|
|
306
|
+
Returns:
|
|
307
|
+
bool: True if the request is allowed, False otherwise.
|
|
308
|
+
"""
|
|
309
|
+
try:
|
|
310
|
+
allowed = self.backend.evalsha(
|
|
311
|
+
self._lua_script_hash,
|
|
312
|
+
1,
|
|
313
|
+
self._generate_key(key),
|
|
314
|
+
str(limit_spec.requests),
|
|
315
|
+
str(limit_spec.total_seconds()),
|
|
316
|
+
)
|
|
317
|
+
return bool(allowed)
|
|
318
|
+
except NoScriptError:
|
|
319
|
+
self._ensure_script_loaded_sync()
|
|
320
|
+
try:
|
|
321
|
+
allowed = self.backend.evalsha(
|
|
322
|
+
self._lua_script_hash,
|
|
323
|
+
1,
|
|
324
|
+
self._generate_key(key),
|
|
325
|
+
str(limit_spec.requests),
|
|
326
|
+
str(limit_spec.total_seconds()),
|
|
327
|
+
)
|
|
328
|
+
return bool(allowed)
|
|
329
|
+
except NoScriptError as exc:
|
|
330
|
+
raise RuntimeError(
|
|
331
|
+
"Failed to load Lua script into Redis."
|
|
332
|
+
) from exc
|
|
333
|
+
|
|
334
|
+
async def _enforce_limit_async(
|
|
335
|
+
self, key: str, limit_spec: LimitSpec
|
|
336
|
+
) -> bool:
|
|
337
|
+
"""
|
|
338
|
+
Enforces the rate limit asynchronously.
|
|
339
|
+
|
|
340
|
+
Args:
|
|
341
|
+
key (str): Unique identifier for the rate limit.
|
|
342
|
+
limit_spec (LimitSpec): Limit specification.
|
|
343
|
+
|
|
344
|
+
Returns:
|
|
345
|
+
bool: True if the request is allowed, False otherwise.
|
|
346
|
+
"""
|
|
347
|
+
try:
|
|
348
|
+
allowed = await self.backend.evalsha(
|
|
349
|
+
self._lua_script_hash,
|
|
350
|
+
1,
|
|
351
|
+
self._generate_key(key),
|
|
352
|
+
str(limit_spec.requests),
|
|
353
|
+
str(limit_spec.total_seconds()),
|
|
354
|
+
)
|
|
355
|
+
return bool(allowed)
|
|
356
|
+
except NoScriptError:
|
|
357
|
+
# Attempt to load the script and retry once
|
|
358
|
+
await self._ensure_script_loaded_async()
|
|
359
|
+
try:
|
|
360
|
+
allowed = await self.backend.evalsha(
|
|
361
|
+
self._lua_script_hash,
|
|
362
|
+
1,
|
|
363
|
+
self._generate_key(key),
|
|
364
|
+
str(limit_spec.requests),
|
|
365
|
+
str(limit_spec.total_seconds()),
|
|
366
|
+
)
|
|
367
|
+
return bool(allowed)
|
|
368
|
+
except NoScriptError as exc:
|
|
369
|
+
# If script loading fails again, raise an error
|
|
370
|
+
raise RuntimeError(
|
|
371
|
+
"Failed to load Lua script into Redis."
|
|
372
|
+
) from exc
|
|
373
|
+
|
|
374
|
+
def _generate_key(self, identifier: str) -> str:
|
|
375
|
+
"""
|
|
376
|
+
Generates a unique Redis key for rate limiting.
|
|
377
|
+
|
|
378
|
+
Args:
|
|
379
|
+
identifier (str): Unique identifier for the rate limit.
|
|
380
|
+
|
|
381
|
+
Returns:
|
|
382
|
+
str: Hashed Redis key.
|
|
383
|
+
"""
|
|
384
|
+
return hashlib.sha256(identifier.encode("utf-8")).hexdigest()
|
|
385
|
+
|
|
386
|
+
def _ensure_script_loaded_sync(self) -> None:
|
|
387
|
+
"""Ensures the Lua script is loaded into Redis (synchronously)."""
|
|
388
|
+
if not self.backend.script_exists(self._lua_script_hash)[0]:
|
|
389
|
+
self.backend.script_load(self._lua_script)
|
|
390
|
+
self._script_loaded = True
|
|
391
|
+
|
|
392
|
+
async def _ensure_script_loaded_async(self) -> None:
|
|
393
|
+
"""Ensures the Lua script is loaded into Redis (asynchronously)."""
|
|
394
|
+
if not (await self.backend.script_exists(self._lua_script_hash))[0]:
|
|
395
|
+
await self.backend.script_load(self._lua_script)
|
|
396
|
+
self._script_loaded = True
|