Source code for baseplate.lib.ratelimit.ratelimit
from baseplate import Span
from baseplate.clients import ContextFactory
from baseplate.lib.ratelimit.backends import RateLimitBackend
[docs]class RateLimitExceededException(Exception):
"""This exception gets raised whenever a rate limit is exceeded."""
[docs]class RateLimiterContextFactory(ContextFactory):
"""RateLimiter context factory.
:param backend_factory: An instance of
:py:class:`baseplate.clients.ContextFactory`. The context factory must
return an instance of :py:class:`baseplate.lib.ratelimit.backends.RateLimitBackend`
:param allowance: The maximum allowance allowed per key.
:param interval: The interval (in seconds) to reset allowances.
"""
def __init__(self, backend_factory: ContextFactory, allowance: int, interval: int):
if allowance < 1:
raise ValueError("minimum allowance is 1")
if interval < 1:
raise ValueError("minimum interval is 1")
if not isinstance(backend_factory, ContextFactory):
raise TypeError("backend_factory must be an instance of ContextFactory")
self.backend_factory = backend_factory
self.allowance = allowance
self.interval = interval
[docs] def make_object_for_context(self, name: str, span: Span) -> "RateLimiter":
backend = self.backend_factory.make_object_for_context(name, span)
return RateLimiter(backend, self.allowance, self.interval)
[docs]class RateLimiter:
"""A class for rate limiting actions.
:param backend: The backend to use for storing rate limit counters.
:param allowance: The maximum allowance allowed per key.
:param interval: The interval (in seconds) to reset allowances.
"""
def __init__(self, backend: RateLimitBackend, allowance: int, interval: int):
if allowance < 1:
raise ValueError("minimum allowance is 1")
if interval < 1:
raise ValueError("minimum interval is 1")
if not isinstance(backend, RateLimitBackend):
raise TypeError("backend must be an instance of RateLimitBackend")
self.backend = backend
self.allowance = allowance
self.interval = interval
[docs] def consume(self, key: str, amount: int = 1) -> None:
"""Consume the given `amount` from the allowance for the given `key`.
This will raise
:py:class:`baseplate.lib.ratelimit.RateLimitExceededException` if the
allowance for `key` is exhausted.
:param key: The name of the rate limit bucket to consume from.
:param amount: The amount to consume from the rate limit bucket.
"""
if not self.backend.consume(key, amount, self.allowance, self.interval):
raise RateLimitExceededException("Rate limit exceeded.")