Skip to content

Commit

Permalink
support report libname and libver to Redis
Browse files Browse the repository at this point in the history
  • Loading branch information
ztxiuyuan committed Apr 21, 2023
1 parent 7fc4c76 commit 4ebfde6
Show file tree
Hide file tree
Showing 10 changed files with 100 additions and 2 deletions.
2 changes: 2 additions & 0 deletions redis/asyncio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,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,
Expand Down Expand Up @@ -212,6 +213,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
Expand Down
14 changes: 12 additions & 2 deletions redis/asyncio/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, get_lib, str_if_bytes

hiredis = None
if HIREDIS_AVAILABLE:
Expand Down Expand Up @@ -458,6 +458,7 @@ class Connection:
"db",
"username",
"client_name",
"set_lib",
"credential_provider",
"password",
"socket_timeout",
Expand Down Expand Up @@ -504,6 +505,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,
Expand All @@ -522,6 +524,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
Expand Down Expand Up @@ -718,7 +721,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 and self.is_connected is False:
try:
await self.send_command("CLIENT", "SETINFO", "LIB-NAME", "redis-py")
await self.read_response()
await self.send_command("CLIENT", "SETINFO", "LIB-VERSION", get_lib())
await self.read_response()
except Exception:
pass
# if a database is specified, switch to it
if self.db:
await self.send_command("SELECT", self.db)
Expand Down
3 changes: 3 additions & 0 deletions redis/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -745,6 +745,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,
Expand Down Expand Up @@ -938,6 +939,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,
Expand Down Expand Up @@ -988,6 +990,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,
}
Expand Down
1 change: 1 addition & 0 deletions redis/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ class AbstractRedisCluster:
"AUTH",
"CLIENT LIST",
"CLIENT SETNAME",
"CLIENT SETINFO",
"CLIENT GETNAME",
"CONFIG SET",
"CONFIG REWRITE",
Expand Down
8 changes: 8 additions & 0 deletions redis/commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
17 changes: 17 additions & 0 deletions redis/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
CRYPTOGRAPHY_AVAILABLE,
HIREDIS_AVAILABLE,
HIREDIS_PACK_AVAILABLE,
get_lib,
str_if_bytes,
)

Expand Down Expand Up @@ -601,6 +602,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,
Expand All @@ -624,6 +626,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
Expand Down Expand Up @@ -739,6 +742,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):
Expand All @@ -763,12 +770,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_lib())
self.read_response()
except Exception:
pass

# if a database is specified, switch to it
if self.db:
Expand Down
14 changes: 14 additions & 0 deletions redis/utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import sys
from contextlib import contextmanager
from functools import wraps
from typing import Any, Dict, Mapping, Union
Expand All @@ -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):
"""
Expand Down Expand Up @@ -110,3 +116,11 @@ def wrapper(*args, **kwargs):
return wrapper

return decorator


def get_lib():
try:
libver = metadata.version("redis")
except metadata.PackageNotFoundError:
libver = "99.99.99"
return libver
9 changes: 9 additions & 0 deletions tests/test_asyncio/test_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
17 changes: 17 additions & 0 deletions tests/test_asyncio/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
17 changes: 17 additions & 0 deletions tests/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 4ebfde6

Please sign in to comment.