Skip to main content

A fast token bucket rate limiter

Project description

FastLimiter

PyPI version License: MIT

A fast and flexible token bucket rate limiter for FastAPI applications.

Table of Contents

Features

  • Flexible Rate Limiting: Configure rate limits based on tokens, capacity, burst size, and time intervals to suit different scenarios.
  • Token Bucket Algorithm: Implements the efficient token bucket algorithm, which is ideal for handling bursts and smoothing out request rates.
  • FastAPI Integration: Seamlessly integrates with FastAPI applications through middleware or route decorators.
  • Per-Client Limiting: Rate limits are applied per client, identified by IP address or custom keys.
  • Statistics Collection: Optionally collect detailed statistics, including allowed and denied requests, which can be used for monitoring and analytics.
  • Customizable Callbacks: Add custom functions to execute after each request, allowing for logging, alerting, or other side effects.
  • Dynamic Configuration: Update rate limiting parameters at runtime without restarting the application.

Installation

Prerequisites

  • Python 3.8 or higher
  • FastAPI 0.65.0 or higher

Install via Pip

Install FastLimiter from PyPI:

pip install fastlimiter

Install from Source

Alternatively, you can clone the repository and install it:

git clone https://github.com/anto18671/fastlimiter.git
cd fastlimiter
pip install .

Quick Start

Here's how you can quickly get started with FastLimiter in your FastAPI application:

from fastapi import FastAPI, Request
from fastlimiter import RateLimiter, setup_rate_limiter

app = FastAPI()

# Initialize the rate limiter
rate_limiter = RateLimiter(rate=10, capacity=100, seconds=60)

# Setup the rate limiter middleware
setup_rate_limiter(app, rate_limiter)

@app.get("/items")
async def read_items(request: Request):
    return {"message": "Hello, World!"}

In this example, the rate limiter allows up to 100 requests per client per 60 seconds, refilling at a rate of 10 tokens per 60 seconds.

Usage

Basic Usage

As Middleware

To apply rate limiting to all incoming requests, use the middleware:

from fastapi import FastAPI
from fastlimiter import RateLimiter, setup_rate_limiter

app = FastAPI()

rate_limiter = RateLimiter(rate=5, capacity=10, seconds=60)
setup_rate_limiter(app, rate_limiter)

As Decorator

To apply rate limiting to specific routes, use the limit decorator:

from fastapi import FastAPI, Request
from fastlimiter import RateLimiter

app = FastAPI()

rate_limiter = RateLimiter(rate=5, capacity=10, seconds=60)

@app.get("/limited")
@rate_limiter.limit()
async def limited_endpoint(request: Request):
    return {"message": "This endpoint is rate-limited."}

Advanced Usage

Custom Key Function

You can customize how the rate limiter identifies clients by overriding get_key_from_request:

def custom_key_function(self, request: Request) -> str:
    return request.headers.get("X-API-Key", "default")

rate_limiter.get_key_from_request = custom_key_function.__get__(rate_limiter, RateLimiter)

Collecting Statistics

Enable statistics collection to monitor rate limiting:

rate_limiter.enable_stats_collection()

stats = rate_limiter.get_stats("client_key")
print(stats)

Updating Parameters at Runtime

You can dynamically update the rate limiter's parameters at runtime:

rate_limiter.update_rate(new_rate=20)
rate_limiter.update_capacity(new_capacity=50)
rate_limiter.update_burst(new_burst=10)
rate_limiter.update_time(seconds=30)
rate_limiter.update_stats_window(new_stats_window=120)

API Reference

RateLimiter Class

Initialization

RateLimiter(
    rate: int,
    capacity: int = 1024,
    burst: int = None,
    stats_window: int = 60,
    enable_stats: bool = True,
    seconds: int = 0,
    minutes: int = 0,
    hours: int = 0
)
  • rate (int): The rate at which tokens are added to the bucket.
  • capacity (int): Maximum number of tokens in the bucket.
  • burst (int, optional): Extra tokens allowed during bursts. Defaults to 0 if not provided.
  • stats_window (int, optional): Time window for collecting statistics in seconds.
  • enable_stats (bool, optional): Enable or disable statistics collection.
  • seconds, minutes, hours: Define the time interval over which tokens are refilled.

Methods

allow_request(key: str) -> bool

Determines if a request identified by key is allowed based on the rate limit.

  • Parameters:
    • key (str): The unique identifier for the client/request.
  • Returns:
    • bool: True if the request is allowed, False otherwise.
  • Raises:
    • HTTPException: Raises an HTTP 429 exception if the request is denied.
get_wait_time(key: str) -> float

Returns the time in seconds a client identified by key needs to wait before making a new request.

  • Parameters:
    • key (str): The unique identifier for the client/request.
  • Returns:
    • float: Time in seconds to wait before the next allowed request.
update_stats(key: str, allowed: bool, timestamp: float)

Updates the statistics for a given key.

  • Parameters:
    • key (str): The unique identifier for the client/request.
    • allowed (bool): Whether the request was allowed.
    • timestamp (float): The time when the request was processed.
get_stats(key: str) -> Dict

Retrieves the statistics for a given key.

  • Parameters:
    • key (str): The unique identifier for the client/request.
  • Returns:
    • Dict: A dictionary containing statistical data such as total allowed, total denied, etc.
reset(key: str = None)

Resets the rate limiter state for a specific key or all keys if key is None.

  • Parameters:
    • key (str, optional): The unique identifier for the client/request. If None, resets all keys.
update_capacity(new_capacity: int)

Updates the capacity of the token bucket.

  • Parameters:
    • new_capacity (int): The new maximum number of tokens.
update_burst(new_burst: int)

Updates the burst size.

  • Parameters:
    • new_burst (int): The new burst size.
update_stats_window(new_stats_window: int)

Updates the time window for statistics collection.

  • Parameters:
    • new_stats_window (int): The new statistics window in seconds.
update_time(seconds: int = 0, minutes: int = 0, hours: int = 0)

Updates the time interval over which tokens are refilled.

  • Parameters:
    • seconds, minutes, hours: Time components for the new interval.
update_rate(new_rate: int)

Updates the rate at which tokens are added to the bucket.

  • Parameters:
    • new_rate (int): The new token addition rate.
get_key_from_request(request: Request) -> str

Retrieves a unique key from the incoming request. By default, uses the client's IP address.

  • Parameters:
    • request (Request): The incoming FastAPI request.
  • Returns:
    • str: The unique key for rate limiting.
enable_stats_collection()

Enables statistics collection.

disable_stats_collection()

Disables statistics collection.

add_callback(callback: Callable)

Adds a callback function to be executed after each request is processed.

  • Parameters:
    • callback (Callable): The function to call after processing a request.
limit()

Creates a rate-limiting decorator for FastAPI routes.

  • Returns:
    • Callable: A decorator that can be applied to FastAPI route handlers.
fastapi_middleware(request: Request, call_next: Callable)

Middleware function to apply rate limiting to incoming FastAPI requests.

  • Parameters:
    • request (Request): The incoming FastAPI request.
    • call_next (Callable): The next middleware or route handler.

Configuration

Customize the rate limiter by adjusting the parameters:

  • Rate: The number of tokens added to the bucket over the specified time interval.
  • Capacity: The maximum number of tokens the bucket can hold.
  • Burst: Additional tokens allowed for handling bursts of requests.
  • Time Interval: The period over which tokens are refilled (specified in seconds, minutes, and/or hours).

Example:

rate_limiter = RateLimiter(
    rate=100,
    capacity=200,
    burst=50,
    seconds=60,
    enable_stats=True,
    stats_window=300
)

Examples

Limiting Based on API Key

def api_key_key_function(self, request: Request) -> str:
    return request.headers.get("X-API-Key", "anonymous")

rate_limiter.get_key_from_request = api_key_key_function.__get__(rate_limiter, RateLimiter)

Custom Callback Function

def log_request(allowed: bool, key: str):
    status = "allowed" if allowed else "denied"
    emoji = "😊" if allowed else "😞"
    print(f"Request from {key} was {status}. {emoji}")

rate_limiter.add_callback(log_request)

Resetting Rate Limiter for a Specific Client

rate_limiter.reset(key="client_ip")

Applying Different Limits to Different Endpoints

from fastapi import FastAPI, Request
from fastlimiter import RateLimiter

app = FastAPI()

# General rate limiter for most endpoints
general_limiter = RateLimiter(rate=100, capacity=200, seconds=60)

# Specific rate limiter for a heavy endpoint
heavy_limiter = RateLimiter(rate=10, capacity=20, seconds=60)

@app.get("/general")
@general_limiter.limit()
async def general_endpoint(request: Request):
    return {"message": "This is a general endpoint."}

@app.get("/heavy")
@heavy_limiter.limit()
async def heavy_endpoint(request: Request):
    return {"message": "This endpoint has stricter rate limiting."}

Limiting by User ID

If your application uses authentication, you might want to rate limit based on user ID:

def user_id_key_function(self, request: Request) -> str:
    user = request.state.user  # Assuming user is stored in request.state
    return str(user.id) if user else "anonymous"

rate_limiter.get_key_from_request = user_id_key_function.__get__(rate_limiter, RateLimiter)

Customizing Rate Limit Exceeded Response

Customize the response when the rate limit is exceeded:

from fastapi.responses import PlainTextResponse

async def rate_limit_middleware(request: Request, call_next: Callable):
    key = rate_limiter.get_key_from_request(request)
    try:
        await rate_limiter.allow_request(key)
        response = await call_next(request)
        return response
    except HTTPException as exc:
        return PlainTextResponse(
            "You have exceeded your request rate limit. Please try again later.",
            status_code=exc.status_code
        )

app.middleware("http")(rate_limit_middleware)

Contributing

We welcome contributions to improve FastLimiter! Here's how you can contribute:

Reporting Bugs

If you find a bug, please open an issue on GitHub with detailed steps to reproduce the problem.

Feature Requests

Have an idea for a new feature? Open an issue with the tag "enhancement" to discuss it.

Pull Requests

  1. Fork the repository on GitHub.
  2. Clone your forked repository to your local machine.
  3. Create a new branch for your changes: git checkout -b feature/your-feature-name.
  4. Make your changes and commit them with clear messages.
  5. Push your changes to your fork: git push origin feature/your-feature-name.
  6. Open a pull request on the main repository.

Code Style

Please follow the PEP 8 style guide for Python code.

Testing

Ensure that all existing tests pass and write new tests for your code.

  • Run tests with pytest.
  • For asynchronous code, use pytest_asyncio.

License

This project is licensed under the MIT License. See the LICENSE file for details.

Contact

For questions or suggestions, please open an issue on GitHub:

  • Author: Anthony Therrien

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

fastlimiter-0.1.0.tar.gz (20.1 kB view hashes)

Uploaded Source

Built Distribution

fastlimiter-0.1.0-py3-none-any.whl (11.1 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page