diff --git a/api/.env.example b/api/.env.example index 2cc6410cdd..208821cdc7 100644 --- a/api/.env.example +++ b/api/.env.example @@ -33,6 +33,7 @@ REDIS_USERNAME= REDIS_PASSWORD=difyai123456 REDIS_USE_SSL=false REDIS_DB=0 +REDIS_KEY_PREFIX= # redis Sentinel configuration. REDIS_USE_SENTINEL=false diff --git a/api/configs/middleware/cache/redis_config.py b/api/configs/middleware/cache/redis_config.py index 2e98c31ec3..55f436f8bf 100644 --- a/api/configs/middleware/cache/redis_config.py +++ b/api/configs/middleware/cache/redis_config.py @@ -34,6 +34,11 @@ class RedisConfig(BaseSettings): default=0, ) + REDIS_KEY_PREFIX: Optional[str] = Field( + description="Redis global string key prefix (if required)", + default=None, + ) + REDIS_USE_SSL: bool = Field( description="Enable SSL/TLS for the Redis connection", default=False, diff --git a/api/core/app/features/rate_limiting/rate_limit.py b/api/core/app/features/rate_limiting/rate_limit.py index 632f35d106..10ffa67d57 100644 --- a/api/core/app/features/rate_limiting/rate_limit.py +++ b/api/core/app/features/rate_limiting/rate_limit.py @@ -5,6 +5,7 @@ from collections.abc import Generator, Mapping from datetime import timedelta from typing import Any, Optional, Union +from configs import dify_config from core.errors.error import AppInvokeQuotaExceededError from extensions.ext_redis import redis_client @@ -12,8 +13,9 @@ logger = logging.getLogger(__name__) class RateLimit: - _MAX_ACTIVE_REQUESTS_KEY = "dify:rate_limit:{}:max_active_requests" - _ACTIVE_REQUESTS_KEY = "dify:rate_limit:{}:active_requests" + _KEY_PREFIX = dify_config.REDIS_KEY_PREFIX if dify_config.REDIS_KEY_PREFIX is not None else "dify" + _MAX_ACTIVE_REQUESTS_KEY = _KEY_PREFIX + ":rate_limit:{}:max_active_requests" + _ACTIVE_REQUESTS_KEY = _KEY_PREFIX + ":rate_limit:{}:active_requests" _UNLIMITED_REQUEST_ID = "unlimited_request_id" _REQUEST_MAX_ALIVE_TIME = 10 * 60 # 10 minutes _ACTIVE_REQUESTS_COUNT_FLUSH_INTERVAL = 5 * 60 # recalculate request_count from request_detail every 5 minutes diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index f8679f7e4b..b97a59ff62 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -1,4 +1,5 @@ -from typing import Any, Union +import functools +from typing import Any, Callable, Generic, TypeVar, Union import redis from redis.cluster import ClusterNode, RedisCluster @@ -8,6 +9,44 @@ from redis.sentinel import Sentinel from configs import dify_config from dify_app import DifyApp +T = TypeVar("T") + + +class KeyPrefixMethodProxy(Generic[T]): + """ + KeyPrefixMethodProxy is a generic class used to proxy method calls. + It can preprocess parameters before calling a method, such as prefixing specific keys. + Key features include: + If a method's argument contains'key', prefix it. + If the method's positional parameter contains a key of type string, prefix it. + """ + + def __init__(self, client: T, prefix: str): + self._client = client + self._prefix = prefix + + def __getattr__(self, item): + attr = getattr(self._client, item) + if callable(attr): + return self._wrap_method(attr) + return attr + + def _wrap_method(self, method: Callable) -> Callable: + @functools.wraps(method) + def wrapper(*args, **kwargs): + if "key" in kwargs: + if isinstance(kwargs["key"], str): + kwargs["key"] = self._add_prefix(kwargs["key"]) + elif args: + if args and isinstance(args[0], str): + args = (self._add_prefix(args[0]),) + args[1:] + return method(*args, **kwargs) + + return wrapper + + def _add_prefix(self, key: str) -> str: + return f"{self._prefix}:{key}" if self._prefix else key + class RedisClientWrapper: """ @@ -32,14 +71,19 @@ class RedisClientWrapper: def __init__(self): self._client = None + self._prefix = None - def initialize(self, client): + def initialize(self, client, prefix=None): if self._client is None: self._client = client + if prefix is not None: + self._prefix = prefix def __getattr__(self, item): if self._client is None: raise RuntimeError("Redis client is not initialized. Call init_app first.") + if self._prefix is not None: + return getattr(KeyPrefixMethodProxy(self._client, self._prefix), item) return getattr(self._client, item) @@ -75,7 +119,7 @@ def init_app(app: DifyApp): }, ) master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params) - redis_client.initialize(master) + redis_client.initialize(master, prefix=dify_config.REDIS_KEY_PREFIX) elif dify_config.REDIS_USE_CLUSTERS: assert dify_config.REDIS_CLUSTERS is not None, "REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True" nodes = [ @@ -83,7 +127,10 @@ def init_app(app: DifyApp): for node in dify_config.REDIS_CLUSTERS.split(",") ] # FIXME: mypy error here, try to figure out how to fix it - redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD)) # type: ignore + redis_client.initialize( + RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), # type: ignore + prefix=dify_config.REDIS_KEY_PREFIX, + ) else: redis_params.update( { @@ -93,6 +140,6 @@ def init_app(app: DifyApp): } ) pool = redis.ConnectionPool(**redis_params) - redis_client.initialize(redis.Redis(connection_pool=pool)) + redis_client.initialize(redis.Redis(connection_pool=pool), prefix=dify_config.REDIS_KEY_PREFIX) app.extensions["redis"] = redis_client diff --git a/docker/.env.example b/docker/.env.example index f1ea72d8cc..49fae31e41 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -247,6 +247,7 @@ REDIS_USERNAME= REDIS_PASSWORD=difyai123456 REDIS_USE_SSL=false REDIS_DB=0 +REDIS_KEY_PREFIX= # Whether to use Redis Sentinel mode. # If set to true, the application will automatically discover and connect to the master node through Sentinel. diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index aee36b3986..0ba3a6e3cc 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -65,6 +65,7 @@ x-shared-env: &shared-api-worker-env REDIS_PASSWORD: ${REDIS_PASSWORD:-difyai123456} REDIS_USE_SSL: ${REDIS_USE_SSL:-false} REDIS_DB: ${REDIS_DB:-0} + REDIS_KEY_PREFIX: ${REDIS_KEY_PREFIX:-} REDIS_USE_SENTINEL: ${REDIS_USE_SENTINEL:-false} REDIS_SENTINELS: ${REDIS_SENTINELS:-} REDIS_SENTINEL_SERVICE_NAME: ${REDIS_SENTINEL_SERVICE_NAME:-}