From 6cb763507c9155fff86ef4d4f8edeb4d814d1634 Mon Sep 17 00:00:00 2001 From: sseide Date: Fri, 11 Aug 2023 20:37:44 +0200 Subject: [PATCH] add basic support for redis cluster server (#9128) This change updates the central utility class to recognize a Redis cluster server after connection and returns an new cluster aware Redis client. The "normal" Redis client would not be able to talk to a cluster node because keys might be stored on other shards of the Redis cluster and therefor not readable or writable. With this patch clients do not need to know what Redis server it is, they just connect though the same API calls for standalone and cluster server. There are no dependencies added due to this MR. Remark - with current redis-py client library (4.6.0) a cluster cannot be used as VectorStore. It can be used for other use-cases. There is a bug / missing feature(?) in the Redis client breaking the VectorStore implementation. I opened an issue at the client library too (redis/redis-py#2888) to fix this. As soon as this is fixed in `redis-py` library it should be usable there too. --------- Co-authored-by: Bagatur --- libs/langchain/langchain/utilities/redis.py | 23 +++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/libs/langchain/langchain/utilities/redis.py b/libs/langchain/langchain/utilities/redis.py index 9255179999..a9535c4425 100644 --- a/libs/langchain/langchain/utilities/redis.py +++ b/libs/langchain/langchain/utilities/redis.py @@ -68,14 +68,17 @@ def get_client(redis_url: str, **kwargs: Any) -> RedisType: # check if normal redis:// or redis+sentinel:// url if redis_url.startswith("redis+sentinel"): redis_client = _redis_sentinel_client(redis_url, **kwargs) - if redis_url.startswith("rediss+sentinel"): # sentinel with TLS support enables + elif redis_url.startswith("rediss+sentinel"): # sentinel with TLS support enables kwargs["ssl"] = True if "ssl_cert_reqs" not in kwargs: kwargs["ssl_cert_reqs"] = "none" redis_client = _redis_sentinel_client(redis_url, **kwargs) else: - # connect to redis server from url + # connect to redis server from url, reconnect with cluster client if needed redis_client = redis.from_url(redis_url, **kwargs) + if _check_for_cluster(redis_client): + redis_client.close() + redis_client = _redis_cluster_client(redis_url, **kwargs) return redis_client @@ -138,3 +141,19 @@ answered NO PASSWORD NEEDED - Please check Sentinel configuration" raise ae return sentinel_client.master_for(service_name) + + +def _check_for_cluster(redis_client: RedisType) -> bool: + import redis + + try: + cluster_info = redis_client.info("cluster") + return cluster_info["cluster_enabled"] == 1 + except redis.exceptions.RedisError: + return False + + +def _redis_cluster_client(redis_url: str, **kwargs: Any) -> RedisType: + from redis.cluster import RedisCluster + + return RedisCluster.from_url(redis_url, **kwargs)