Add REDIS_DB env variable for db selection

Author: notgurev

src/common/config.py

@@ -16,14 +16,16 @@
              "ssl_certfile": os.getenv('REDIS_SSL_CERTFILE', None),
              "ssl_cert_reqs": os.getenv('REDIS_SSL_CERT_REQS', 'required'),
              "ssl_ca_certs": os.getenv('REDIS_SSL_CA_CERTS', None),
-             "cluster_mode": os.getenv('REDIS_CLUSTER_MODE', False) in ('true', '1', 't')}
+             "cluster_mode": os.getenv('REDIS_CLUSTER_MODE', False) in ('true', '1', 't'),
+             "db": int(os.getenv('REDIS_DB', 0))}
 
 
 def generate_redis_uri():
     cfg = REDIS_CFG
     scheme = "rediss" if cfg.get("ssl") else "redis"
     host = cfg.get("host", "127.0.0.1")
     port = cfg.get("port", 6379)
+    db = cfg.get("db", 0)
 
     username = cfg.get("username")
     password = cfg.get("password")
@@ -37,7 +39,7 @@ def generate_redis_uri():
         auth_part = ""
 
     # Base URI
-    base_uri = f"{scheme}://{auth_part}{host}:{port}"
+    base_uri = f"{scheme}://{auth_part}{host}:{port}/{db}"
 
     # Additional SSL query parameters if SSL is enabled
     query_params = {}

src/common/connection.py

@@ -16,23 +16,43 @@ class RedisConnectionManager:
     def get_connection(cls, decode_responses=True) -> Redis:
         if cls._instance is None:
             try:
-                redis_class: Type[Union[Redis, RedisCluster]] = redis.cluster.RedisCluster if REDIS_CFG["cluster_mode"] else redis.Redis
+                if REDIS_CFG["cluster_mode"]:
+                    redis_class: Type[Union[Redis, RedisCluster]] = redis.cluster.RedisCluster
+                    connection_params = {
+                        "host": REDIS_CFG["host"],
+                        "port": REDIS_CFG["port"],
+                        "username": REDIS_CFG["username"],
+                        "password": REDIS_CFG["password"],
+                        "ssl": REDIS_CFG["ssl"],
+                        "ssl_ca_path": REDIS_CFG["ssl_ca_path"],
+                        "ssl_keyfile": REDIS_CFG["ssl_keyfile"],
+                        "ssl_certfile": REDIS_CFG["ssl_certfile"],
+                        "ssl_cert_reqs": REDIS_CFG["ssl_cert_reqs"],
+                        "ssl_ca_certs": REDIS_CFG["ssl_ca_certs"],
+                        "decode_responses": decode_responses,
+                        "lib_name": f"redis-py(mcp-server_v{__version__})",
+                        "max_connections_per_node": 10 
+                    }
+                else:
+                    redis_class: Type[Union[Redis, RedisCluster]] = redis.Redis
+                    connection_params = {
+                        "host": REDIS_CFG["host"],
+                        "port": REDIS_CFG["port"],
+                        "db": REDIS_CFG["db"],
+                        "username": REDIS_CFG["username"],
+                        "password": REDIS_CFG["password"],
+                        "ssl": REDIS_CFG["ssl"],
+                        "ssl_ca_path": REDIS_CFG["ssl_ca_path"],
+                        "ssl_keyfile": REDIS_CFG["ssl_keyfile"],
+                        "ssl_certfile": REDIS_CFG["ssl_certfile"],
+                        "ssl_cert_reqs": REDIS_CFG["ssl_cert_reqs"],
+                        "ssl_ca_certs": REDIS_CFG["ssl_ca_certs"],
+                        "decode_responses": decode_responses,
+                        "lib_name": f"redis-py(mcp-server_v{__version__})",
+                        "max_connections": 10
+                    }
 
-                cls._instance = redis_class(
-                    host=REDIS_CFG["host"],
-                    port=REDIS_CFG["port"],
-                    username=REDIS_CFG["username"],
-                    password=REDIS_CFG["password"],
-                    ssl=REDIS_CFG["ssl"],
-                    ssl_ca_path=REDIS_CFG["ssl_ca_path"],
-                    ssl_keyfile=REDIS_CFG["ssl_keyfile"],
-                    ssl_certfile=REDIS_CFG["ssl_certfile"],
-                    ssl_cert_reqs=REDIS_CFG["ssl_cert_reqs"],
-                    ssl_ca_certs=REDIS_CFG["ssl_ca_certs"],
-                    decode_responses=decode_responses,
-                    lib_name=f"redis-py(mcp-server_v{__version__})",
-                    **({"max_connections_per_node": 10} if REDIS_CFG["cluster_mode"] else {"max_connections": 10})
-                )
+                cls._instance = redis_class(**connection_params)
 
             except redis.exceptions.ConnectionError:
                 print("Failed to connect to Redis server", file=sys.stderr)