Allow passing in an explicit connection pool for redis (#110)

* Allow explicit connection pool for redis

doc/source/storage.rst

@@ -152,6 +152,12 @@ or :code:`redis+unix:///path/to/socket?db=n` (for database `n`).
 If the database is password protected the password can be provided in the url, for example
 :code:`redis://:foobared@localhost:6379` or :code:`redis+unix//:foobered/path/to/socket` if using a UDS..
 
+For scenarios where a redis connection pool is already available and can be reused, it can be provided
+in :paramref:`~limits.storage.storage_from_string.options`, for example::
+
+    pool = redis.connections.BlockingConnectionPool.from_url("redis://.....")
+    storage_from_string("redis://", connection_pool=pool)
+    
 Depends on: :pypi:`redis`
 
 Redis over SSL

limits/aio/storage/redis.py

@@ -129,14 +129,18 @@ class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
     """
     DEPENDENCIES = ["coredis"]
 
-    def __init__(self, uri: str, **options) -> None:
+    def __init__(
+        self, uri: str, connection_pool: Optional[Any] = None, **options
+    ) -> None:
         """
         :param uri: uri of the form `async+redis://[:password]@host:port`,
          `async+redis://[:password]@host:port/db`,
          `async+rediss://[:password]@host:port`, `async+unix:///path/to/sock` etc.
          This uri is passed directly to :func:`coredis.StrictRedis.from_url` with
          the initial `a` removed, except for the case of `redis+unix` where it
          is replaced with `unix`.
+        :param connection_pool: if provided, the redis client is initialized with
+         the connection pool and any other params passed as :paramref:`options`
         :param options: all remaining keyword arguments are passed
          directly to the constructor of :class:`coredis.StrictRedis`
         :raise ConfigurationError: when the redis library is not available
@@ -148,7 +152,12 @@ def __init__(self, uri: str, **options) -> None:
         super().__init__()
 
         self.dependency = self.dependencies["coredis"]
-        self.storage = self.dependency.StrictRedis.from_url(uri, **options)
+        if connection_pool:
+            self.storage = self.dependency.StrictRedis(
+                connection_pool=connection_pool, **options
+            )
+        else:
+            self.storage = self.dependency.StrictRedis.from_url(uri, **options)
 
         self.initialize_storage(uri)
 

limits/storage/redis.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Tuple
+from typing import Any, Optional, Tuple
 
 from ..util import get_package_data
 from .base import MovingWindowSupport, Storage
@@ -111,13 +111,20 @@ class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
 
     DEPENDENCIES = ["redis"]
 
-    def __init__(self, uri: str, **options):
+    def __init__(
+        self,
+        uri: str,
+        connection_pool: Optional[Any] = None,
+        **options,
+    ):
         """
         :param uri: uri of the form ``redis://[:password]@host:port``,
          ``redis://[:password]@host:port/db``,
          ``rediss://[:password]@host:port``, ``redis+unix:///path/to/sock`` etc.
          This uri is passed directly to :func:`redis.from_url` except for the
          case of ``redis+unix://`` where it is replaced with ``unix://``.
+        :param connection_pool: if provided, the redis client is initialized with
+         the connection pool and any other params passed as :paramref:`options`
         :param options: all remaining keyword arguments are passed
          directly to the constructor of :class:`redis.Redis`
         :raise ConfigurationError: when the :pypi:`redis` library is not available
@@ -126,7 +133,10 @@ def __init__(self, uri: str, **options):
         redis = self.dependencies["redis"]
         uri = uri.replace("redis+unix", "unix")
 
-        self.storage = redis.from_url(uri, **options)
+        if not connection_pool:
+            self.storage = redis.from_url(uri, **options)
+        else:
+            self.storage = redis.Redis(connection_pool=connection_pool, **options)
         self.initialize_storage(uri)
 
     def initialize_storage(self, _uri: str):

tests/aio/storage/test_redis.py

@@ -59,6 +59,7 @@ async def test_moving_window_clear(self):
         assert await limiter.hit(per_min)
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky
     async def test_moving_window_expiry(self):
         limiter = MovingWindowRateLimiter(self.storage)
         limit = RateLimitItemPerSecond(2)
@@ -97,6 +98,15 @@ async def test_init_options(self, mocker):
             from_url.spy_return.connection_pool.connection_kwargs["stream_timeout"] == 1
         )
 
+    @pytest.mark.asyncio
+    async def test_custom_connection_pool(self):
+        import coredis
+
+        pool = coredis.BlockingConnectionPool.from_url(self.storage_url)
+        storage = storage_from_string("async+redis://", connection_pool=pool)
+
+        assert await storage.check()
+
 
 @pytest.mark.redis
 class TestAsyncRedisAuthStorage(AsyncSharedRedisTests):

tests/storage/test_redis.py

@@ -84,6 +84,12 @@ def test_init_options(self, mocker):
         assert storage_from_string(self.storage_url, socket_timeout=1).check()
         assert from_url.call_args[1]["socket_timeout"] == 1
 
+    def test_custom_connection_pool(self):
+        pool = redis.connection.BlockingConnectionPool.from_url(self.storage_url)
+        storage = storage_from_string("redis://", connection_pool=pool)
+
+        assert storage.check()
+
 
 @pytest.mark.redis
 class TestRedisUnixSocketStorage(SharedRedisTests):

tests/utils.py

@@ -86,7 +86,7 @@ async def async_window(delay_end: float, delay: Optional[float] = None):
             "redis+sentinel://localhost:26379",
             {"service_name": "localhost-redis-sentinel"},
             pytest.lazy_fixture("redis_sentinel"),
-            marks=pytest.mark.redis_sentinel,
+            marks=[pytest.mark.redis_sentinel, pytest.mark.flaky],
         ),
         pytest.param(
             "redis+sentinel://localhost:26379/localhost-redis-sentinel",