Merge pull request #148 from JWCook/per-host-buckets

Fix per-host rate-limiting for Redis and Postgres backends

Author: JWCook

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.8
+    rev: v0.15.9
     hooks:
       - id: ruff
       - id: ruff-format

HISTORY.md

@@ -1,8 +1,12 @@
 # History
 
+## 0.10.0 (Unreleased)
+* Fix per-host rate-limiting for Redis and Postgres backends
+* If both `per_host=True` and a `bucket_name` is specified, use `bucket_name` as a bucket prefix
+* Add warning if a custom Limiter object is passed with `per_host=True` and no HostBucketFactory
+
 ## 0.9.3 (2026-04-02)
-* Fix compatibility with `RedisBucket`
-* Fix compatibility with `PostgresBucket`
+* Fix bucket initialization for `RedisBucket` and `PostgresBucket`
 * Use built-in support for pickling `Limiter` from pyrate-limiter 4.1.0
 
 ## 0.9.2 (2026-02-27)

README.md

@@ -123,16 +123,20 @@ The following parameters are available for the most common rate limit intervals:
 <!-- TODO: Section explaining burst rate limit -->
 
 ### Advanced Settings
-If you need to define more complex rate limits, you can create a `Limiter` object instead:
+If you need to define more complex rate limits, you can use `Limiter` directly.
+Note that it must be used with `HostBucketFactory` if you want per-host rate-limiting.
+
 ```python
-from pyrate_limiter import Duration, RequestRate, Limiter
-from requests_ratelimiter import LimiterSession
+from pyrate_limiter import Duration, Rate, Limiter
+from requests_ratelimiter import LimiterSession, HostBucketFactory
 
-nanocentury_rate = RequestRate(10, Duration.SECOND * 3.156)
-fortnight_rate = RequestRate(1000, Duration.DAY * 14)
-trimonthly_rate = RequestRate(10000, Duration.MONTH * 3)
-limiter = Limiter(nanocentury_rate, fortnight_rate, trimonthly_rate)
+nanocentury_rate = Rate(10, Duration.SECOND * 3.156)
+fortnight_rate = Rate(1000, Duration.DAY * 14)
+trimonthly_rate = Rate(10000, Duration.DAY * 90)
 
+# This factory object is required for per-host rate-limiting
+factory = HostBucketFactory(rates=[nanocentury_rate, fortnight_rate, trimonthly_rate])
+limiter = Limiter(factory)
 session = LimiterSession(limiter=limiter)
 ```
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'requests-ratelimiter'
-version = '0.9.3'
+version = '0.10.0'
 description = 'Rate-limiting for the requests library'
 authors = [{name = 'Jordan Cook'}]
 license = 'MIT'

requests_ratelimiter/buckets.py

@@ -1,3 +1,4 @@
+import re
 from typing import Dict, Optional, Type
 
 from pyrate_limiter import InMemoryBucket, PostgresBucket, Rate, RedisBucket, SQLiteBucket
@@ -33,53 +34,54 @@ def get(self, item: RateItem) -> AbstractBucket:
         """Get or create a bucket for the given item name"""
         if item.name not in self.buckets:
             # Create new bucket for this name
-            bucket = self._create_bucket()
+            bucket = self._create_bucket(item.name)
             self.schedule_leak(bucket)
             self.buckets[item.name] = bucket
 
         return self.buckets[item.name]
 
-    def _create_bucket(self) -> AbstractBucket:
-        """Create a new bucket instance with the configured bucket class"""
+    def _create_bucket(self, name: str) -> AbstractBucket:
+        """Create a new bucket instance, and handle per-host naming for each supported backend"""
         if self.bucket_class == InMemoryBucket:
             return InMemoryBucket(self.rates)
         elif self.bucket_class == SQLiteBucket:
-            kwargs = prepare_sqlite_kwargs(self.bucket_init_kwargs, self.bucket_name)
+            kwargs = prepare_sqlite_kwargs(self.bucket_init_kwargs, name)
             return SQLiteBucket.init_from_file(rates=self.rates, **kwargs)
         elif self.bucket_class == RedisBucket:
             kwargs = self.bucket_init_kwargs.copy()
-            bucket_key = kwargs.pop('bucket_key', self.bucket_name or 'default')
             redis = kwargs.pop('redis')
-            return RedisBucket.init(rates=self.rates, redis=redis, bucket_key=bucket_key)
+            bucket_key = kwargs.pop('bucket_key', _sanitize_name(name))
+            return RedisBucket.init(rates=self.rates, redis=redis, bucket_key=bucket_key, **kwargs)
         elif self.bucket_class == PostgresBucket:
             kwargs = self.bucket_init_kwargs.copy()
             pool = kwargs.pop('pool')
-            table = kwargs.pop('table', self.bucket_name or 'default')
+            table = kwargs.pop('table', _sanitize_name(name))
             return PostgresBucket(pool=pool, table=table, rates=self.rates)
         else:
-            # Generic bucket creation - pass rates as first arg
             return self.bucket_class(self.rates, **self.bucket_init_kwargs)
 
     def __getitem__(self, name: str) -> AbstractBucket:
         """Dict-like access for backward compatibility with _fill_bucket() method"""
         if name not in self.buckets:
-            # Create bucket on access
             temp_item = RateItem(name, 0, 1)
             return self.get(temp_item)
         return self.buckets[name]
 
 
 def prepare_sqlite_kwargs(bucket_kwargs: Dict, bucket_name: Optional[str] = None) -> Dict:
-    """Prepare SQLiteBucket kwargs for v4 compatibility"""
     kwargs = bucket_kwargs.copy()
     if 'path' in kwargs:
         kwargs['db_path'] = str(kwargs.pop('path'))
 
     # If bucket_name is specified, use it as the table name to ensure separation
     # This allows multiple sessions with different bucket_names to share a db file
     if bucket_name and 'table' not in kwargs:
-        kwargs['table'] = f'bucket_{bucket_name}'
+        kwargs['table'] = f'bucket_{_sanitize_name(bucket_name)}'
 
     # Filter to only supported parameters for SQLiteBucket.init_from_file
     supported_params = {'table', 'db_path', 'create_new_table', 'use_file_lock'}
     return {k: v for k, v in kwargs.items() if k in supported_params}
+
+
+def _sanitize_name(name: str) -> str:
+    return re.sub(r'[^a-zA-Z0-9_]', '_', name)

requests_ratelimiter/requests_ratelimiter.py

@@ -66,6 +66,11 @@ def __init__(
         if limiter:
             self.limiter = limiter
             self._custom_limiter = True
+            if per_host and not isinstance(limiter.bucket_factory, HostBucketFactory):
+                logger.warning(
+                    'Custom limiter does not use HostBucketFactory; per-host rate limiting will '
+                    'not work. Use HostBucketFactory to enable per-host rate limiting.'
+                )
         else:
             factory = HostBucketFactory(
                 rates=rates,
@@ -101,10 +106,11 @@ def send(self, request: PreparedRequest, **kwargs) -> Response:
 
     def _bucket_name(self, request):
         """Get a bucket name for the given request"""
-        if self.bucket_name:
+        if self.per_host:
+            host = urlparse(request.url).netloc
+            return f'{self.bucket_name}:{host}' if self.bucket_name else host
+        elif self.bucket_name:
             return self.bucket_name
-        elif self.per_host:
-            return urlparse(request.url).netloc
         else:
             return self._default_bucket
 
@@ -188,8 +194,9 @@ class LimiterSession(LimiterMixin, Session):
             :py:class:`~pyrate_limiter.buckets.redis_bucket.RedisBucket`
         bucket_kwargs: Bucket backend keyword arguments
         limiter: An existing Limiter object to use instead of the above params
-        per_host: Track request rate limits separately for each host
         limit_statuses: Alternative HTTP status codes that indicate a rate limit was exceeded
+        per_host: Track request rate limits separately for each host
+        bucket_name: Override default bucket name. In per-host mode, this sets the bucket prefix.
     """
 
     __attrs__ = Session.__attrs__ + [

test/test_buckets.py

@@ -52,14 +52,14 @@ def test_separate_buckets_per_name():
 
 
 @pytest.mark.parametrize(
-    'extra_init_kwargs, expected_key',
+    'extra_init_kwargs, identity, expected_key',
     [
-        ({'bucket_key': 'test_bucket'}, 'test_bucket'),
-        ({}, 'default'),
+        ({}, 'api.example.com', 'api_example_com'),
+        ({'bucket_key': 'override'}, 'api.example.com', 'override'),
+        ({}, 'myapp:api.example.com', 'myapp_api_example_com'),
     ],
-    ids=['explicit_key', 'default_key'],
 )
-def test_redis_bucket(extra_init_kwargs, expected_key):
+def test_redis_bucket(extra_init_kwargs, identity, expected_key):
     mock_redis = MagicMock()
     mock_redis.script_load.return_value = 'fake_sha1'
     factory = HostBucketFactory(
@@ -69,7 +69,7 @@ def test_redis_bucket(extra_init_kwargs, expected_key):
     )
 
     with patch.object(RedisBucket, 'init', wraps=RedisBucket.init) as mock_init:
-        factory._create_bucket()
+        factory._create_bucket(identity)
 
     mock_init.assert_called_once_with(
         rates=factory.rates, redis=mock_redis, bucket_key=expected_key
@@ -82,20 +82,19 @@ def _postgres_init_stub(self, pool, table, rates):
 
 
 @pytest.mark.parametrize(
-    'extra_init_kwargs, factory_kwargs, expected_table',
+    'extra_init_kwargs, identity, expected_table',
     [
-        ({'table': 'test_table'}, {}, 'test_table'),
-        ({}, {'bucket_name': 'my_table'}, 'my_table'),
+        ({}, 'api.example.com', 'api_example_com'),
+        ({'table': 'override'}, 'api.example.com', 'override'),
+        ({}, 'myapp:api.example.com', 'myapp_api_example_com'),
     ],
-    ids=['explicit_table', 'default_from_bucket_name'],
 )
-def test_postgres_bucket(extra_init_kwargs, factory_kwargs, expected_table):
+def test_postgres_bucket(extra_init_kwargs, identity, expected_table):
     mock_pool = MagicMock()
     factory = HostBucketFactory(
         rates=[Rate(5, 1000)],
         bucket_class=PostgresBucket,
         bucket_init_kwargs={'pool': mock_pool, **extra_init_kwargs},
-        **factory_kwargs,
     )
 
     with patch.object(
@@ -105,7 +104,7 @@ def test_postgres_bucket(extra_init_kwargs, factory_kwargs, expected_table):
         return_value=None,
         side_effect=_postgres_init_stub,
     ) as mock_init:
-        factory._create_bucket()
+        factory._create_bucket(identity)
 
     _, kwargs = mock_init.call_args
     assert kwargs == {'pool': mock_pool, 'table': expected_table, 'rates': factory.rates}
@@ -119,7 +118,7 @@ class CustomBucket(InMemoryBucket):
         rates=[Rate(5, 1000)],
         bucket_class=CustomBucket,
     )
-    bucket = factory._create_bucket()
+    bucket = factory._create_bucket('test_host')
     assert isinstance(bucket, CustomBucket)
 
 

test/test_requests_ratelimiter.py

@@ -18,7 +18,7 @@
 from requests import PreparedRequest, Session
 from requests_cache import CacheMixin
 
-from requests_ratelimiter import LimiterMixin, LimiterSession
+from requests_ratelimiter import HostBucketFactory, LimiterMixin, LimiterSession
 from requests_ratelimiter.requests_ratelimiter import _convert_rate, _get_valid_kwargs
 from test.conftest import (
     MOCKED_URL,
@@ -48,7 +48,6 @@ def __init__(self, *args, flag: bool = False, **kwargs):
         lambda: LimiterSession(per_second=5),
         lambda: CustomSession(per_second=5, flag=True),
     ],
-    ids=['LimiterSession', 'CustomSession'],
 )
 def test_rate_limit_enforcement(mock_sleep, session_factory):
     session = mount_mock_adapter(session_factory())
@@ -92,6 +91,40 @@ def test_custom_limiter(mock_sleep):
     assert mock_sleep.called is True
 
 
+@patch_sleep
+def test_custom_limiter__per_host(mock_sleep):
+    factory = HostBucketFactory(rates=[Rate(5, Duration.SECOND)])
+    limiter = Limiter(factory)
+    session = get_mock_session(limiter=limiter, per_host=True)
+
+    for _ in range(5):
+        session.get(MOCKED_URL)
+    assert mock_sleep.called is False
+
+    # A different host should not be affected
+    session.get(MOCKED_URL_ALT_HOST)
+    assert mock_sleep.called is False
+
+    session.get(MOCKED_URL)
+    assert mock_sleep.called is True
+
+
+@pytest.mark.parametrize(
+    'session_kwargs, expect_warning',
+    [
+        ({'per_host': True}, True),
+        ({'per_host': False}, False),
+    ],
+)
+def test_custom_limiter__per_host_warning(caplog, session_kwargs, expect_warning):
+    """Warn when a custom limiter without HostBucketFactory is used with per_host=True"""
+    bucket = InMemoryBucket([Rate(5, Duration.SECOND)])
+    limiter = Limiter(bucket)
+    with caplog.at_level('WARNING', logger='requests_ratelimiter'):
+        LimiterSession(limiter=limiter, **session_kwargs)
+    assert ('HostBucketFactory' in caplog.text) == expect_warning
+
+
 @patch_sleep
 @pytest.mark.parametrize(
     'url, session_kwargs, expect_sleep',
@@ -299,7 +332,7 @@ def test_close_before_any_request_and_idempotent():
 def test_fill_bucket_with_custom_limiter(mock_sleep):
     bucket = InMemoryBucket([Rate(5, Duration.SECOND)])
     limiter = Limiter(bucket)
-    session = get_mock_session(limiter=limiter)
+    session = get_mock_session(limiter=limiter, per_host=False)
     session.get(MOCKED_URL_429)
     session.get(MOCKED_URL_429)
     assert mock_sleep.called is True
@@ -335,11 +368,11 @@ class MyBucket(InMemoryBucket):
     session.close()
 
 
-def test_bucket_name_overrides_per_host():
-    session = LimiterSession(per_second=5, bucket_name='fixed', per_host=True)
+def test_bucket_name_prefixes_per_host():
+    session = LimiterSession(per_second=5, bucket_name='myapp', per_host=True)
     req = PreparedRequest()
     req.url = MOCKED_URL
-    assert session._bucket_name(req) == 'fixed'
+    assert session._bucket_name(req) == 'myapp:requests-ratelimiter.com'
 
 
 def test_max_delay_logs_warning(caplog):
@@ -365,7 +398,6 @@ def test_burst_allows_consecutive_requests(mock_sleep):
         (InMemoryBucket, {}),  # InMemoryBucket
         (SQLiteBucket, None),  # SQLiteBucket (will use fixture to provide kwargs)
     ],
-    ids=['in_memory', 'sqlite'],
 )
 def test_pickling(mock_sleep, bucket_class, bucket_kwargs, tmp_path):
     if bucket_class == SQLiteBucket:
@@ -434,18 +466,19 @@ def test_no_rate_limits_no_limiter():
 
 
 @pytest.mark.parametrize(
-    'url, expected_name',
+    'url, bucket_name, expected_name',
     [
-        ('http+mock://example.com/path', 'example.com'),
-        ('http+mock://example.com:8080/path', 'example.com:8080'),
-        ('http+mock://192.168.1.1/path', '192.168.1.1'),
-        ('http+mock://[::1]/path', '[::1]'),
-        ('http+mock://[::1]:8080/path', '[::1]:8080'),
+        ('http+mock://example.com/path', None, 'example.com'),
+        ('http+mock://example.com:8080/path', None, 'example.com:8080'),
+        ('http+mock://192.168.1.1/path', None, '192.168.1.1'),
+        ('http+mock://[::1]/path', None, '[::1]'),
+        ('http+mock://[::1]:8080/path', None, '[::1]:8080'),
+        ('http+mock://example.com/path', 'myapp', 'myapp:example.com'),
     ],
 )
-def test_bucket_name_from_url(url, expected_name):
+def test_bucket_name_from_url(url, bucket_name, expected_name):
     """per_host bucket names are derived from URL netloc, including ports and IPs"""
-    session = LimiterSession(per_second=5, per_host=True)
+    session = LimiterSession(per_second=5, per_host=True, bucket_name=bucket_name)
     req = PreparedRequest()
     req.url = url
     assert session._bucket_name(req) == expected_name
@@ -454,7 +487,7 @@ def test_bucket_name_from_url(url, expected_name):
 def test_custom_limiter_close_does_not_stop_factory():
     bucket = InMemoryBucket([Rate(5, Duration.SECOND)])
     limiter = Limiter(bucket)
-    session = get_mock_session(limiter=limiter)
+    session = get_mock_session(limiter=limiter, per_host=False)
     session.get(MOCKED_URL)
     session.close()