api/extensions/ext_redis.py
Rama main · réplica de lectura del repositorio original.
import functools
import logging
import socket
import ssl
import sys
from collections.abc import Callable
from datetime import timedelta
from typing import Any, Protocol, cast, runtime_checkable
import redis
from redis import RedisError
from redis.backoff import ExponentialWithJitterBackoff # type: ignore
from redis.cache import CacheConfig
from redis.client import PubSub
from redis.cluster import ClusterNode, RedisCluster
from redis.connection import Connection, SSLConnection
from redis.exceptions import ConnectionError, TimeoutError
from redis.retry import Retry
from redis.sentinel import Sentinel
from typing_extensions import TypedDict
from configs import dify_config
from dify_app import DifyApp
from extensions.redis_names import (
normalize_redis_key_prefix,
serialize_redis_name,
serialize_redis_name_arg,
serialize_redis_name_args,
)
from libs.broadcast_channel.channel import BroadcastChannel as BroadcastChannelProtocol
from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel
from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel
from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel
logger = logging.getLogger(__name__)
_normalize_redis_key_prefix = normalize_redis_key_prefix
_serialize_redis_name = serialize_redis_name
_serialize_redis_name_arg = serialize_redis_name_arg
_serialize_redis_name_args = serialize_redis_name_args
@runtime_checkable
class _ScriptRegistrar(Protocol):
def register_script(self, script: str) -> Callable[..., Any]: ...
_INCR_IF_EXISTS_LUA = """
-- Exactly 1 key and no arguments.
-- KEYS[1]: Existing integer counter key, including the deployment prefix.
-- Returns the incremented integer, or nil (Lua false) if the key is missing.
-- INCR preserves TTL; checking existence atomically avoids recreating an
-- expired/deleted counter without an expiry.
if redis.call('EXISTS', KEYS[1]) == 0 then return false end
return redis.call('INCR', KEYS[1])
"""
class RedisClientWrapper:
"""
A wrapper class for the Redis client that addresses the issue where the global
`redis_client` variable cannot be updated when a new Redis instance is returned
by Sentinel.
This class allows for deferred initialization of the Redis client, enabling the
client to be re-initialized with a new instance when necessary. This is particularly
useful in scenarios where the Redis instance may change dynamically, such as during
a failover in a Sentinel-managed Redis setup.
Attributes:
_client: The actual Redis client instance. It remains None until
initialized with the `initialize` method.
Methods:
initialize(client): Initializes the Redis client if it hasn't been initialized already.
__getattr__(item): Delegates attribute access to the Redis client, raising an error
if the client is not initialized.
"""
_client: redis.Redis | RedisCluster | None
def __init__(self) -> None:
self._client = None
def initialize(self, client: redis.Redis | RedisCluster) -> None:
if self._client is None:
self._client = client
def _require_client(self) -> redis.Redis | RedisCluster:
if self._client is None:
raise RuntimeError("Redis client is not initialized. Call init_app first.")
return self._client
def _get_prefix(self) -> str:
return dify_config.REDIS_KEY_PREFIX
def get(self, name: str | bytes) -> Any:
return self._require_client().get(_serialize_redis_name_arg(name, self._get_prefix()))
def set(
self,
name: str | bytes,
value: Any,
ex: int | None = None,
px: int | None = None,
nx: bool = False,
xx: bool = False,
keepttl: bool = False,
get: bool = False,
exat: int | None = None,
pxat: int | None = None,
) -> Any:
return self._require_client().set(
_serialize_redis_name_arg(name, self._get_prefix()),
value,
ex=ex,
px=px,
nx=nx,
xx=xx,
keepttl=keepttl,
get=get,
exat=exat,
pxat=pxat,
)
def setex(self, name: str | bytes, time: int | timedelta, value: Any) -> Any:
return self._require_client().setex(_serialize_redis_name_arg(name, self._get_prefix()), time, value)
def setnx(self, name: str | bytes, value: Any) -> Any:
return self._require_client().setnx(_serialize_redis_name_arg(name, self._get_prefix()), value)
def delete(self, *names: str | bytes) -> Any:
return self._require_client().delete(*_serialize_redis_name_args(names, self._get_prefix()))
def incr(self, name: str | bytes, amount: int = 1) -> Any:
return self._require_client().incr(_serialize_redis_name_arg(name, self._get_prefix()), amount)
def expire(
self,
name: str | bytes,
time: int | timedelta,
nx: bool = False,
xx: bool = False,
gt: bool = False,
lt: bool = False,
) -> Any:
return self._require_client().expire(
_serialize_redis_name_arg(name, self._get_prefix()),
time,
nx=nx,
xx=xx,
gt=gt,
lt=lt,
)
def exists(self, *names: str | bytes) -> Any:
return self._require_client().exists(*_serialize_redis_name_args(names, self._get_prefix()))
def ttl(self, name: str | bytes) -> Any:
return self._require_client().ttl(_serialize_redis_name_arg(name, self._get_prefix()))
def pttl(self, name: str | bytes) -> Any:
return self._require_client().pttl(_serialize_redis_name_arg(name, self._get_prefix()))
def getdel(self, name: str | bytes) -> Any:
return self._require_client().getdel(_serialize_redis_name_arg(name, self._get_prefix()))
def lock(
self,
name: str,
timeout: float | None = None,
sleep: float = 0.1,
blocking: bool = True,
blocking_timeout: float | None = None,
thread_local: bool = True,
) -> Any:
return self._require_client().lock(
_serialize_redis_name(name, self._get_prefix()),
timeout=timeout,
sleep=sleep,
blocking=blocking,
blocking_timeout=blocking_timeout,
thread_local=thread_local,
)
def hset(self, name: str | bytes, *args: Any, **kwargs: Any) -> Any:
return self._require_client().hset(_serialize_redis_name_arg(name, self._get_prefix()), *args, **kwargs)
def incr_if_exists(self, name: str) -> int | None:
"""Increment an existing integer counter, preserving TTL; return None if absent."""
count = self.eval(_INCR_IF_EXISTS_LUA, 1, _serialize_redis_name(name, self._get_prefix()))
if count is not None and not isinstance(count, int):
raise ValueError("Invalid Redis counter response")
return count
def hgetall(self, name: str | bytes) -> Any:
return self._require_client().hgetall(_serialize_redis_name_arg(name, self._get_prefix()))
def hkeys(self, name: str | bytes) -> Any:
return self._require_client().hkeys(_serialize_redis_name_arg(name, self._get_prefix()))
def hexists(self, name: str | bytes, key: str | bytes) -> Any:
return self._require_client().hexists(_serialize_redis_name_arg(name, self._get_prefix()), key)
def hdel(self, name: str | bytes, *keys: str | bytes) -> Any:
return self._require_client().hdel(_serialize_redis_name_arg(name, self._get_prefix()), *keys)
def hlen(self, name: str | bytes) -> Any:
return self._require_client().hlen(_serialize_redis_name_arg(name, self._get_prefix()))
def zadd(
self,
name: str | bytes,
mapping: dict[str | bytes | int | float, float | int | str | bytes],
nx: bool = False,
xx: bool = False,
ch: bool = False,
incr: bool = False,
gt: bool = False,
lt: bool = False,
) -> Any:
return self._require_client().zadd(
_serialize_redis_name_arg(name, self._get_prefix()),
cast(Any, mapping),
nx=nx,
xx=xx,
ch=ch,
incr=incr,
gt=gt,
lt=lt,
)
def zrem(self, name: str | bytes, *values: str | bytes | int | float) -> Any:
return self._require_client().zrem(_serialize_redis_name_arg(name, self._get_prefix()), *values)
def zremrangebyscore(self, name: str | bytes, min: float | str, max: float | str) -> Any:
return self._require_client().zremrangebyscore(_serialize_redis_name_arg(name, self._get_prefix()), min, max)
def zcard(self, name: str | bytes) -> Any:
return self._require_client().zcard(_serialize_redis_name_arg(name, self._get_prefix()))
def pubsub(self) -> PubSub:
return self._require_client().pubsub()
def eval(self, script: str, numkeys: int, *args: str | int) -> Any:
# Lua callers supply physical names, matching the historically delegated API.
# Redis script responses are dynamically shaped at this SDK boundary.
return self._require_client().execute_command("EVAL", script, numkeys, *args)
def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> Any:
return self._require_client().pipeline(transaction=transaction, shard_hint=shard_hint)
def register_script(self, script: str) -> Callable[..., Any]:
"""Register a Lua script whose key arguments use logical Redis names."""
client = self._require_client()
if not isinstance(client, _ScriptRegistrar):
raise RuntimeError("Redis client does not support Lua script registration")
registered_script = client.register_script(script)
def execute(
keys: list[str | bytes] | tuple[str | bytes, ...] | None = None,
args: list[Any] | tuple[Any, ...] | None = None,
client: RedisClientWrapper | redis.Redis | RedisCluster | None = None,
) -> Any:
redis_keys = _serialize_redis_name_args(tuple(keys or ()), self._get_prefix())
execution_client = client._require_client() if isinstance(client, RedisClientWrapper) else client
return registered_script(
keys=redis_keys,
args=args,
client=execution_client or self._require_client(),
)
return execute
def __getattr__(self, item: str) -> Any:
return getattr(self._require_client(), item)
redis_client: RedisClientWrapper = RedisClientWrapper()
_pubsub_redis_client: redis.Redis | RedisCluster | None = None
class RedisSSLParamsDict(TypedDict):
ssl_cert_reqs: int
ssl_ca_certs: str | None
ssl_certfile: str | None
ssl_keyfile: str | None
class RedisHealthParamsDict(TypedDict):
retry: Retry
socket_timeout: float | None
socket_connect_timeout: float | None
health_check_interval: int | None
socket_keepalive: bool
socket_keepalive_options: dict[int, int]
class RedisClusterHealthParamsDict(TypedDict):
retry: Retry
socket_timeout: float | None
socket_connect_timeout: float | None
socket_keepalive: bool
socket_keepalive_options: dict[int, int]
class RedisBaseParamsDict(TypedDict):
username: str | None
password: str | None
db: int
encoding: str
encoding_errors: str
decode_responses: bool
protocol: int
cache_config: CacheConfig | None
retry: Retry
socket_timeout: float | None
socket_connect_timeout: float | None
health_check_interval: int | None
socket_keepalive: bool
socket_keepalive_options: dict[int, int]
def _get_ssl_configuration() -> tuple[type[Connection | SSLConnection], dict[str, Any]]:
"""Get SSL configuration for Redis connection."""
if not dify_config.REDIS_USE_SSL:
return Connection, {}
cert_reqs_map = {
"CERT_NONE": ssl.CERT_NONE,
"CERT_OPTIONAL": ssl.CERT_OPTIONAL,
"CERT_REQUIRED": ssl.CERT_REQUIRED,
}
ssl_cert_reqs = cert_reqs_map.get(dify_config.REDIS_SSL_CERT_REQS, ssl.CERT_NONE)
ssl_kwargs = {
"ssl_cert_reqs": ssl_cert_reqs,
"ssl_ca_certs": dify_config.REDIS_SSL_CA_CERTS,
"ssl_certfile": dify_config.REDIS_SSL_CERTFILE,
"ssl_keyfile": dify_config.REDIS_SSL_KEYFILE,
}
return SSLConnection, ssl_kwargs
def _get_cache_configuration() -> CacheConfig | None:
"""Get client-side cache configuration if enabled."""
if not dify_config.REDIS_ENABLE_CLIENT_SIDE_CACHE:
return None
resp_protocol = dify_config.REDIS_SERIALIZATION_PROTOCOL
if resp_protocol < 3:
raise ValueError("Client side cache is only supported in RESP3")
return CacheConfig()
def _get_retry_policy() -> Retry:
"""Build the shared retry policy for Redis connections."""
return Retry(
backoff=ExponentialWithJitterBackoff(
base=dify_config.REDIS_RETRY_BACKOFF_BASE,
cap=dify_config.REDIS_RETRY_BACKOFF_CAP,
),
retries=dify_config.REDIS_RETRY_RETRIES,
supported_errors=(
ConnectionError,
TimeoutError,
BrokenPipeError,
OSError,
),
)
def _get_connection_health_params() -> RedisHealthParamsDict:
"""Get connection health and retry parameters for standalone and Sentinel Redis clients."""
socket_keepalive_options: dict[int, int] = {}
if sys.platform == "linux":
socket_keepalive_options[socket.TCP_KEEPIDLE] = dify_config.REDIS_KEEPALIVE_IDLE
socket_keepalive_options[socket.TCP_KEEPINTVL] = dify_config.REDIS_KEEPALIVE_INTERVAL
socket_keepalive_options[socket.TCP_KEEPCNT] = dify_config.REDIS_KEEPALIVE_COUNT
elif sys.platform == "darwin":
socket_keepalive_options[socket.TCP_KEEPALIVE] = dify_config.REDIS_KEEPALIVE_IDLE
return RedisHealthParamsDict(
retry=_get_retry_policy(),
socket_timeout=dify_config.REDIS_SOCKET_TIMEOUT,
socket_connect_timeout=dify_config.REDIS_SOCKET_CONNECT_TIMEOUT,
health_check_interval=dify_config.REDIS_HEALTH_CHECK_INTERVAL,
socket_keepalive=dify_config.REDIS_KEEPALIVE,
socket_keepalive_options=socket_keepalive_options,
)
def _get_cluster_connection_health_params() -> RedisClusterHealthParamsDict:
"""Get retry and timeout parameters for Redis Cluster clients.
RedisCluster does not support ``health_check_interval`` as a constructor
keyword (it is silently stripped by ``cleanup_kwargs``), so it is excluded
here. Only ``retry``, ``socket_timeout``, and ``socket_connect_timeout``
are passed through.
"""
health_params = _get_connection_health_params()
result: RedisClusterHealthParamsDict = {
"retry": health_params["retry"],
"socket_timeout": health_params["socket_timeout"],
"socket_connect_timeout": health_params["socket_connect_timeout"],
"socket_keepalive": health_params["socket_keepalive"],
"socket_keepalive_options": health_params["socket_keepalive_options"],
}
return result
def _get_base_redis_params() -> RedisBaseParamsDict:
"""Get base Redis connection parameters including retry and health policy."""
return RedisBaseParamsDict(
username=dify_config.REDIS_USERNAME,
password=dify_config.REDIS_PASSWORD or None,
db=dify_config.REDIS_DB,
encoding="utf-8",
encoding_errors="strict",
decode_responses=False,
protocol=dify_config.REDIS_SERIALIZATION_PROTOCOL,
cache_config=_get_cache_configuration(),
**_get_connection_health_params(),
)
def _parse_redis_nodes(value: str) -> list[tuple[str, int]]:
nodes = []
for raw_node in value.split(","):
host, separator, port = raw_node.strip().rpartition(":")
if not separator or not host or not port:
raise ValueError(f"Invalid Redis node: {raw_node}")
if host.startswith("[") and host.endswith("]"):
host = host[1:-1]
nodes.append((host, int(port)))
return nodes
def _create_sentinel_client(redis_params: RedisBaseParamsDict) -> redis.Redis | RedisCluster:
"""Create Redis client using Sentinel configuration."""
if not dify_config.REDIS_SENTINELS:
raise ValueError("REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True")
if not dify_config.REDIS_SENTINEL_SERVICE_NAME:
raise ValueError("REDIS_SENTINEL_SERVICE_NAME must be set when REDIS_USE_SENTINEL is True")
sentinel_hosts = _parse_redis_nodes(dify_config.REDIS_SENTINELS)
health_params = _get_connection_health_params()
sentinel_kwargs = {
"socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
"username": dify_config.REDIS_SENTINEL_USERNAME,
"password": dify_config.REDIS_SENTINEL_PASSWORD,
"socket_keepalive": health_params["socket_keepalive"],
"socket_keepalive_options": health_params["socket_keepalive_options"],
}
if dify_config.REDIS_MAX_CONNECTIONS:
sentinel_kwargs["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS
sentinel = Sentinel(
sentinel_hosts,
sentinel_kwargs=sentinel_kwargs,
)
params: dict[str, Any] = {**redis_params}
master: redis.Redis = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **params)
return master
def _create_cluster_client() -> redis.Redis | RedisCluster:
"""Create Redis cluster client."""
if not dify_config.REDIS_CLUSTERS:
raise ValueError("REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True")
nodes = [ClusterNode(host=host, port=port) for host, port in _parse_redis_nodes(dify_config.REDIS_CLUSTERS)]
cluster_kwargs: dict[str, Any] = {
"startup_nodes": nodes,
"password": dify_config.REDIS_CLUSTERS_PASSWORD,
"protocol": dify_config.REDIS_SERIALIZATION_PROTOCOL,
"cache_config": _get_cache_configuration(),
**_get_cluster_connection_health_params(),
}
if dify_config.REDIS_MAX_CONNECTIONS:
cluster_kwargs["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS
cluster: RedisCluster = RedisCluster(**cluster_kwargs)
return cluster
def _create_standalone_client(redis_params: RedisBaseParamsDict) -> redis.Redis | RedisCluster:
"""Create standalone Redis client."""
connection_class, ssl_kwargs = _get_ssl_configuration()
params: dict[str, Any] = {
**redis_params,
"host": dify_config.REDIS_HOST,
"port": dify_config.REDIS_PORT,
"connection_class": connection_class,
}
if dify_config.REDIS_MAX_CONNECTIONS:
params["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS
if ssl_kwargs:
params.update(ssl_kwargs)
pool = redis.ConnectionPool(**params)
client: redis.Redis = redis.Redis(connection_pool=pool)
return client
def _create_pubsub_client(pubsub_url: str, use_clusters: bool) -> redis.Redis | RedisCluster:
max_conns = dify_config.REDIS_MAX_CONNECTIONS
if use_clusters:
health_params = _get_cluster_connection_health_params()
kwargs: dict[str, Any] = {**health_params}
if max_conns:
kwargs["max_connections"] = max_conns
return RedisCluster.from_url(pubsub_url, **kwargs)
standalone_health_params: dict[str, Any] = dict(_get_connection_health_params())
kwargs = {**standalone_health_params}
if max_conns:
kwargs["max_connections"] = max_conns
return redis.Redis.from_url(pubsub_url, **kwargs)
def init_app(app: DifyApp) -> None:
"""Attach the cache client and select the global event-transport client.
Requires valid settings for the selected cache topology. Explicit PubSub
URLs take precedence; otherwise Sentinel is reused without requiring
standalone Redis host settings.
"""
global redis_client
# Determine Redis mode and create appropriate client
if dify_config.REDIS_USE_SENTINEL:
redis_params = _get_base_redis_params()
client = _create_sentinel_client(redis_params)
elif dify_config.REDIS_USE_CLUSTERS:
client = _create_cluster_client()
else:
redis_params = _get_base_redis_params()
client = _create_standalone_client(redis_params)
# Initialize the wrapper and attach to app
redis_client.initialize(client)
app.extensions["redis"] = redis_client
global _pubsub_redis_client
_pubsub_redis_client = client
pubsub_redis_url = (dify_config.PUBSUB_REDIS_URL or "").strip()
if pubsub_redis_url or not dify_config.REDIS_USE_SENTINEL:
_pubsub_redis_client = _create_pubsub_client(
dify_config.normalized_pubsub_redis_url, dify_config.PUBSUB_REDIS_USE_CLUSTERS
)
def get_pubsub_broadcast_channel() -> BroadcastChannelProtocol:
assert _pubsub_redis_client is not None, "PubSub redis Client should be initialized here."
if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "sharded":
return ShardedRedisBroadcastChannel(_pubsub_redis_client)
if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "streams":
return StreamsBroadcastChannel(
_pubsub_redis_client,
retention_seconds=dify_config.PUBSUB_STREAMS_RETENTION_SECONDS,
)
return RedisBroadcastChannel(_pubsub_redis_client)
def redis_fallback[T](default_return: T | None = None): # type: ignore
"""
decorator to handle Redis operation exceptions and return a default value when Redis is unavailable.
Args:
default_return: The value to return when a Redis operation fails. Defaults to None.
"""
def decorator[**P, R](func: Callable[P, R]) -> Callable[P, R | T | None]:
@functools.wraps(func)
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | T | None:
try:
return func(*args, **kwargs)
except RedisError as e:
func_name = getattr(func, "__name__", "Unknown")
logger.warning("Redis operation failed in %s: %s", func_name, str(e), exc_info=True)
return default_return
return wrapper
return decorator