From aa9b3b3b40432aad988430d422bfb51038c4c254 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=82=A0=E9=A3=8E?= Date: Mon, 10 Apr 2023 16:45:46 +0800 Subject: [PATCH] support report libname and libver to Redis --- redis/asyncio/client.py | 4 ++++ redis/asyncio/connection.py | 14 ++++++++++++-- redis/client.py | 4 ++++ redis/cluster.py | 1 + redis/commands/core.py | 8 ++++++++ redis/connection.py | 17 +++++++++++++++++ redis/utils.py | 14 ++++++++++++++ tests/test_asyncio/test_cluster.py | 9 +++++++++ tests/test_asyncio/test_commands.py | 17 +++++++++++++++++ tests/test_commands.py | 17 +++++++++++++++++ 10 files changed, 103 insertions(+), 2 deletions(-) diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index 3e6626aedf..43f39b9dfb 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -2,6 +2,7 @@ import copy import inspect import re +import sys import warnings from typing import ( TYPE_CHECKING, @@ -59,6 +60,7 @@ from redis.typing import ChannelT, EncodableT, KeyT from redis.utils import safe_str, str_if_bytes + PubSubHandler = Callable[[Dict[str, str]], Awaitable[None]] _KeyT = TypeVar("_KeyT", bound=KeyT) _ArgT = TypeVar("_ArgT", KeyT, EncodableT) @@ -171,6 +173,7 @@ def __init__( single_connection_client: bool = False, health_check_interval: int = 0, client_name: Optional[str] = None, + set_lib: bool = True, username: Optional[str] = None, retry: Optional[Retry] = None, auto_close_connection_pool: bool = True, @@ -212,6 +215,7 @@ def __init__( "max_connections": max_connections, "health_check_interval": health_check_interval, "client_name": client_name, + "set_lib": set_lib, "redis_connect_func": redis_connect_func, } # based on input, setup appropriate connection args diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 58dcd66efb..1eb9fd99a8 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -55,7 +55,7 @@ TimeoutError, ) from redis.typing import EncodableT, EncodedT -from redis.utils import HIREDIS_AVAILABLE, str_if_bytes +from redis.utils import HIREDIS_AVAILABLE, str_if_bytes, get_libver hiredis = None if HIREDIS_AVAILABLE: @@ -457,6 +457,7 @@ class Connection: "db", "username", "client_name", + "set_lib", "credential_provider", "password", "socket_timeout", @@ -503,6 +504,7 @@ def __init__( socket_read_size: int = 65536, health_check_interval: float = 0, client_name: Optional[str] = None, + set_lib: bool = False, username: Optional[str] = None, retry: Optional[Retry] = None, redis_connect_func: Optional[ConnectCallbackT] = None, @@ -521,6 +523,7 @@ def __init__( self.port = int(port) self.db = db self.client_name = client_name + self.set_lib = set_lib, self.credential_provider = credential_provider self.password = password self.username = username @@ -717,7 +720,14 @@ async def on_connect(self) -> None: await self.send_command("CLIENT", "SETNAME", self.client_name) if str_if_bytes(await self.read_response()) != "OK": raise ConnectionError("Error setting client name") - + if self.set_lib is True: + try: + await self.send_command("CLIENT", "SETINFO", "LIB-NAME", "redis-py") + await self.read_response() + await self.send_command("CLIENT", "SETINFO", "LIB-VERSION", get_libver()) + await self.read_response() + except: + pass # if a database is specified, switch to it if self.db: await self.send_command("SELECT", self.db) diff --git a/redis/client.py b/redis/client.py index 1a9b96b83d..1c0edf2b02 100755 --- a/redis/client.py +++ b/redis/client.py @@ -1,6 +1,7 @@ import copy import datetime import re +import sys import threading import time import warnings @@ -745,6 +746,7 @@ class AbstractRedis: "CLIENT LIST": parse_client_list, "CLIENT INFO": parse_client_info, "CLIENT SETNAME": bool_ok, + "CLIENT SETINFO": bool_ok, "CLIENT UNBLOCK": lambda r: r and int(r) == 1 or False, "CLIENT PAUSE": bool_ok, "CLIENT GETREDIR": int, @@ -938,6 +940,7 @@ def __init__( single_connection_client=False, health_check_interval=0, client_name=None, + set_lib=True, username=None, retry=None, redis_connect_func=None, @@ -988,6 +991,7 @@ def __init__( "max_connections": max_connections, "health_check_interval": health_check_interval, "client_name": client_name, + "set_lib": set_lib, "redis_connect_func": redis_connect_func, "credential_provider": credential_provider, } diff --git a/redis/cluster.py b/redis/cluster.py index 5e6e7da546..746f46d35c 100644 --- a/redis/cluster.py +++ b/redis/cluster.py @@ -212,6 +212,7 @@ class AbstractRedisCluster: "AUTH", "CLIENT LIST", "CLIENT SETNAME", + "CLIENT SETINFO", "CLIENT GETNAME", "CONFIG SET", "CONFIG REWRITE", diff --git a/redis/commands/core.py b/redis/commands/core.py index e2cabb85fa..70c3af1931 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -706,6 +706,14 @@ def client_setname(self, name: str, **kwargs) -> ResponseT: """ return self.execute_command("CLIENT SETNAME", name, **kwargs) + def client_setinfo(self, attr: str, value: str, **kwargs) -> ResponseT: + """ + Sets the current connection libname + + For mor information see https://redis.io/commands/client-setinfo + """ + return self.execute_command("CLIENT SETINFO", attr, value, **kwargs) + def client_unblock( self, client_id: int, error: bool = False, **kwargs ) -> ResponseT: diff --git a/redis/connection.py b/redis/connection.py index 162a4c3215..db33ec556e 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -39,6 +39,7 @@ HIREDIS_AVAILABLE, HIREDIS_PACK_AVAILABLE, str_if_bytes, + get_libver, ) try: @@ -600,6 +601,7 @@ def __init__( socket_read_size=65536, health_check_interval=0, client_name=None, + set_lib=False, username=None, retry=None, redis_connect_func=None, @@ -623,6 +625,7 @@ def __init__( self.pid = os.getpid() self.db = db self.client_name = client_name + self.set_lib = set_lib self.credential_provider = credential_provider self.password = password self.username = username @@ -738,6 +741,10 @@ def _error_message(self, exception): def on_connect(self): "Initialize the connection, authenticate and select a database" self._parser.on_connect(self) + if ssl_available: + authened = False + else: + authened = True # if credential provider or username and/or password are set, authenticate if self.credential_provider or (self.username or self.password): @@ -762,12 +769,22 @@ def on_connect(self): if str_if_bytes(auth_response) != "OK": raise AuthenticationError("Invalid Username or Password") + else: + authened = True # if a client_name is given, set it if self.client_name: self.send_command("CLIENT", "SETNAME", self.client_name) if str_if_bytes(self.read_response()) != "OK": raise ConnectionError("Error setting client name") + if self.set_lib is True and authened is True: + try: + self.send_command("CLIENT", "SETINFO", "LIB-NAME", "redis-py") + self.read_response() + self.send_command("CLIENT", "SETINFO", "LIB-VERSION", get_libver()) + self.read_response() + except: + pass # if a database is specified, switch to it if self.db: diff --git a/redis/utils.py b/redis/utils.py index d95e62c042..67d2aa2be6 100644 --- a/redis/utils.py +++ b/redis/utils.py @@ -1,3 +1,4 @@ +import sys from contextlib import contextmanager from functools import wraps from typing import Any, Dict, Mapping, Union @@ -19,6 +20,11 @@ except ImportError: CRYPTOGRAPHY_AVAILABLE = False +if sys.version_info >= (3, 8): + from importlib import metadata +else: + import importlib_metadata as metadata + def from_url(url, **kwargs): """ @@ -110,3 +116,11 @@ def wrapper(*args, **kwargs): return wrapper return decorator + + +def get_libver(): + try: + libver = metadata.version("redis") + except metadata.PackageNotFoundError: + libver = "99.99.99" + return libver diff --git a/tests/test_asyncio/test_cluster.py b/tests/test_asyncio/test_cluster.py index 13e5e26ae3..13278478dc 100644 --- a/tests/test_asyncio/test_cluster.py +++ b/tests/test_asyncio/test_cluster.py @@ -870,6 +870,15 @@ async def test_client_setname(self, r: RedisCluster) -> None: client_name = await r.client_getname(target_nodes=node) assert client_name == "redis_py_test" + @skip_if_server_version_lt("7.2.0") + async def test_client_setinfo(self, r: RedisCluster) -> None: + node = r.get_random_node() + await r.client_setinfo("lib-name", "redis-py", target_nodes=node) + await r.client_setinfo("lib-ver", "4.33", target_nodes=node) + info = await r.client_info(target_nodes=node) + assert "lib-name" in info + assert "lib-ver" in info + async def test_exists(self, r: RedisCluster) -> None: d = {"a": b"1", "b": b"2", "c": b"3", "d": b"4"} await r.mset_nonatomic(d) diff --git a/tests/test_asyncio/test_commands.py b/tests/test_asyncio/test_commands.py index 7c6fd45ab9..096e8bb5df 100644 --- a/tests/test_asyncio/test_commands.py +++ b/tests/test_asyncio/test_commands.py @@ -340,6 +340,15 @@ async def test_client_setname(self, r: redis.Redis): assert await r.client_setname("redis_py_test") assert await r.client_getname() == "redis_py_test" + @skip_if_server_version_lt("7.2.0") + @pytest.mark.onlynoncluster + async def test_client_setinfo(self, r:redis.Redis): + assert await r.client_setinfo("lib-name", "redis-py") + assert await r.client_setinfo("lib-ver", "4.33") + info = await r.client_info() + assert "lib-name" in info + assert "lib-ver" in info + @skip_if_server_version_lt("2.6.9") @pytest.mark.onlynoncluster async def test_client_kill(self, r: redis.Redis, r2): @@ -438,6 +447,14 @@ async def test_client_list_after_client_setname(self, r: redis.Redis): # we don't know which client ours will be assert "redis_py_test" in [c["name"] for c in clients] + @skip_if_server_version_lt("7.2.0") + async def test_client_list_after_client_setinfo(self, r: redis.Redis): + await r.client_setinfo("lib-name", "redis-py") + await r.client_setinfo("lib-ver", "4.33") + clients = await r.client_list() + assert "redis-py" in [c["lib-name"] for c in clients] + assert "4.33" in [c["lib-ver"] for c in clients] + @skip_if_server_version_lt("2.9.50") @pytest.mark.onlynoncluster async def test_client_pause(self, r: redis.Redis): diff --git a/tests/test_commands.py b/tests/test_commands.py index 94249e9419..172daf58b5 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -530,6 +530,15 @@ def test_client_setname(self, r): assert r.client_setname("redis_py_test") assert r.client_getname() == "redis_py_test" + @pytest.mark.onlynoncluster + @skip_if_server_version_lt("7.2.0") + def test_client_setinfo(self, r): + assert r.client_setinfo("lib-name", "redis-py") + assert r.client_setinfo("lib-ver", "4.33") + info = r.client_info() + assert "lib-name" in info + assert "lib-ver" in info + @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.6.9") def test_client_kill(self, r, r2): @@ -628,6 +637,14 @@ def test_client_list_after_client_setname(self, r): # we don't know which client ours will be assert "redis_py_test" in [c["name"] for c in clients] + @skip_if_server_version_lt("7.2.0") + def test_client_list_after_client_setinfo(self, r): + r.client_setinfo("lib-name", "redis-py") + r.client_setinfo("lib-ver", "4.33") + clients = r.client_list() + assert "redis-py" in [c["lib-name"] for c in clients] + assert "4.33" in [c["lib-ver"] for c in clients] + @skip_if_server_version_lt("6.2.0") def test_client_kill_filter_by_laddr(self, r, r2): r.client_setname("redis-py-c1")