Add Redis.from_pool() class method, for explicitly owning and closing a ConnectionPool (#2913)

kristjanvalur · web-flow · commit 012f7cf7d295 · 2023-09-19T02:09:41.000+03:00
* Add the `from_pool` argument to asyncio.Redis

* Add tests for the `from_pool` argument

* Add a "from_pool" argument for the sync client too

* Add automatic connection pool close for redis.sentinel

* use from_pool() class method instead

* re-add the auto_close_connection_pool arg to Connection.from_url()

* Deprecate the "auto_close_connection_pool" argument to Redis() and Redis.from_url()
diff --git a/docs/examples/asyncio_examples.ipynb b/docs/examples/asyncio_examples.ipynb
@@ -44,6 +44,26 @@
     "await connection.close()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If you create custom `ConnectionPool` for the `Redis` instance to use alone, use the `from_pool` class method to create it.  This will cause the pool to be disconnected along with the Redis instance. Disconnecting the connection pool simply disconnects all connections hosted in the pool."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import redis.asyncio as redis\n",
+    "\n",
+    "pool = redis.ConnectionPool.from_url(\"redis://localhost\")\n",
+    "connection = redis.Redis.from_pool(pool)\n",
+    "await connection.close()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -53,7 +73,8 @@
     }
    },
    "source": [
-    "If you supply a custom `ConnectionPool` that is supplied to several `Redis` instances, you may want to disconnect the connection pool explicitly. Disconnecting the connection pool simply disconnects all connections hosted in the pool."
+    "\n",
+    "However, If you supply a `ConnectionPool` that is shared several `Redis` instances, you may want to disconnect the connection pool explicitly. use the `connection_pool` argument in that case."
    ]
   },
   {
@@ -69,10 +90,12 @@
    "source": [
     "import redis.asyncio as redis\n",
     "\n",
-    "connection = redis.Redis(auto_close_connection_pool=False)\n",
-    "await connection.close()\n",
-    "# Or: await connection.close(close_connection_pool=False)\n",
-    "await connection.connection_pool.disconnect()"
+    "pool = redis.ConnectionPool.from_url(\"redis://localhost\")\n",
+    "connection1 = redis.Redis(connection_pool=pool)\n",
+    "connection2 = redis.Redis(connection_pool=pool)\n",
+    "await connection1.close()\n",
+    "await connection2.close()\n",
+    "await pool.disconnect()"
    ]
   },
   {
diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py
@@ -114,7 +114,7 @@ def from_url(
         cls,
         url: str,
         single_connection_client: bool = False,
-        auto_close_connection_pool: bool = True,
+        auto_close_connection_pool: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -160,12 +160,39 @@ class initializer. In the case of conflicting arguments, querystring
 
         """
         connection_pool = ConnectionPool.from_url(url, **kwargs)
-        redis = cls(
+        client = cls(
             connection_pool=connection_pool,
             single_connection_client=single_connection_client,
         )
-        redis.auto_close_connection_pool = auto_close_connection_pool
-        return redis
+        if auto_close_connection_pool is not None:
+            warnings.warn(
+                DeprecationWarning(
+                    '"auto_close_connection_pool" is deprecated '
+                    "since version 5.0.0. "
+                    "Please create a ConnectionPool explicitly and "
+                    "provide to the Redis() constructor instead."
+                )
+            )
+        else:
+            auto_close_connection_pool = True
+        client.auto_close_connection_pool = auto_close_connection_pool
+        return client
+
+    @classmethod
+    def from_pool(
+        cls: Type["Redis"],
+        connection_pool: ConnectionPool,
+    ) -> "Redis":
+        """
+        Return a Redis client from the given connection pool.
+        The Redis client will take ownership of the connection pool and
+        close it when the Redis client is closed.
+        """
+        client = cls(
+            connection_pool=connection_pool,
+        )
+        client.auto_close_connection_pool = True
+        return client
 
     def __init__(
         self,
@@ -200,7 +227,8 @@ def __init__(
         lib_version: Optional[str] = get_lib_version(),
         username: Optional[str] = None,
         retry: Optional[Retry] = None,
-        auto_close_connection_pool: bool = True,
+        # deprecated. create a pool and use connection_pool instead
+        auto_close_connection_pool: Optional[bool] = None,
         redis_connect_func=None,
         credential_provider: Optional[CredentialProvider] = None,
         protocol: Optional[int] = 2,
@@ -213,14 +241,21 @@ def __init__(
         To retry on TimeoutError, `retry_on_timeout` can also be set to `True`.
         """
         kwargs: Dict[str, Any]
-        # auto_close_connection_pool only has an effect if connection_pool is
-        # None. This is a similar feature to the missing __del__ to resolve #1103,
-        # but it accounts for whether a user wants to manually close the connection
-        # pool, as a similar feature to ConnectionPool's __del__.
-        self.auto_close_connection_pool = (
-            auto_close_connection_pool if connection_pool is None else False
-        )
+
+        if auto_close_connection_pool is not None:
+            warnings.warn(
+                DeprecationWarning(
+                    '"auto_close_connection_pool" is deprecated '
+                    "since version 5.0.0. "
+                    "Please create a ConnectionPool explicitly and "
+                    "provide to the Redis() constructor instead."
+                )
+            )
+        else:
+            auto_close_connection_pool = True
+
         if not connection_pool:
+            # Create internal connection pool, expected to be closed by Redis instance
             if not retry_on_error:
                 retry_on_error = []
             if retry_on_timeout is True:
@@ -277,7 +312,13 @@ def __init__(
                             "ssl_check_hostname": ssl_check_hostname,
                         }
                     )
+            # This arg only used if no pool is passed in
+            self.auto_close_connection_pool = auto_close_connection_pool
             connection_pool = ConnectionPool(**kwargs)
+        else:
+            # If a pool is passed in, do not close it
+            self.auto_close_connection_pool = False
+
         self.connection_pool = connection_pool
         self.single_connection_client = single_connection_client
         self.connection: Optional[Connection] = None
diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py
@@ -1106,8 +1106,8 @@ class BlockingConnectionPool(ConnectionPool):
     """
     A blocking connection pool::
 
-        >>> from redis.asyncio.client import Redis
-        >>> client = Redis(connection_pool=BlockingConnectionPool())
+        >>> from redis.asyncio import Redis, BlockingConnectionPool
+        >>> client = Redis.from_pool(BlockingConnectionPool())
 
     It performs the same function as the default
     :py:class:`~redis.asyncio.ConnectionPool` implementation, in that,
diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py
@@ -340,12 +340,7 @@ def master_for(
 
         connection_pool = connection_pool_class(service_name, self, **connection_kwargs)
         # The Redis object "owns" the pool
-        auto_close_connection_pool = True
-        client = redis_class(
-            connection_pool=connection_pool,
-        )
-        client.auto_close_connection_pool = auto_close_connection_pool
-        return client
+        return redis_class.from_pool(connection_pool)
 
     def slave_for(
         self,
@@ -377,9 +372,4 @@ def slave_for(
 
         connection_pool = connection_pool_class(service_name, self, **connection_kwargs)
         # The Redis object "owns" the pool
-        auto_close_connection_pool = True
-        client = redis_class(
-            connection_pool=connection_pool,
-        )
-        client.auto_close_connection_pool = auto_close_connection_pool
-        return client
+        return redis_class.from_pool(connection_pool)
diff --git a/redis/client.py b/redis/client.py
@@ -4,7 +4,7 @@
 import time
 import warnings
 from itertools import chain
-from typing import Optional
+from typing import Optional, Type
 
 from redis._parsers.helpers import (
     _RedisCallbacks,
@@ -136,10 +136,28 @@ class initializer. In the case of conflicting arguments, querystring
         """
         single_connection_client = kwargs.pop("single_connection_client", False)
         connection_pool = ConnectionPool.from_url(url, **kwargs)
-        return cls(
+        client = cls(
             connection_pool=connection_pool,
             single_connection_client=single_connection_client,
         )
+        client.auto_close_connection_pool = True
+        return client
+
+    @classmethod
+    def from_pool(
+        cls: Type["Redis"],
+        connection_pool: ConnectionPool,
+    ) -> "Redis":
+        """
+        Return a Redis client from the given connection pool.
+        The Redis client will take ownership of the connection pool and
+        close it when the Redis client is closed.
+        """
+        client = cls(
+            connection_pool=connection_pool,
+        )
+        client.auto_close_connection_pool = True
+        return client
 
     def __init__(
         self,
@@ -275,6 +293,10 @@ def __init__(
                         }
                     )
             connection_pool = ConnectionPool(**kwargs)
+            self.auto_close_connection_pool = True
+        else:
+            self.auto_close_connection_pool = False
+
         self.connection_pool = connection_pool
         self.connection = None
         if single_connection_client:
@@ -477,6 +499,9 @@ def close(self):
             self.connection = None
             self.connection_pool.release(conn)
 
+        if self.auto_close_connection_pool:
+            self.connection_pool.disconnect()
+
     def _send_command_parse_response(self, conn, command_name, *args, **options):
         """
         Send a command and parse the response
diff --git a/redis/sentinel.py b/redis/sentinel.py
@@ -353,10 +353,8 @@ def master_for(
         kwargs["is_master"] = True
         connection_kwargs = dict(self.connection_kwargs)
         connection_kwargs.update(kwargs)
-        return redis_class(
-            connection_pool=connection_pool_class(
-                service_name, self, **connection_kwargs
-            )
+        return redis_class.from_pool(
+            connection_pool_class(service_name, self, **connection_kwargs)
         )
 
     def slave_for(
@@ -386,8 +384,6 @@ def slave_for(
         kwargs["is_master"] = False
         connection_kwargs = dict(self.connection_kwargs)
         connection_kwargs.update(kwargs)
-        return redis_class(
-            connection_pool=connection_pool_class(
-                service_name, self, **connection_kwargs
-            )
+        return redis_class.from_pool(
+            connection_pool_class(service_name, self, **connection_kwargs)
         )
diff --git a/tests/test_asyncio/test_connection.py b/tests/test_asyncio/test_connection.py
diff --git a/tests/test_connection.py b/tests/test_connection.py
diff --git a/tests/test_sentinel.py b/tests/test_sentinel.py
