| import inspect |
| from urllib.parse import urlparse |
| import asyncio |
| import time |
|
|
| import logging |
|
|
| import redis |
|
|
| from open_webui.env import ( |
| REDIS_CLUSTER, |
| REDIS_SOCKET_CONNECT_TIMEOUT, |
| REDIS_SENTINEL_HOSTS, |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| REDIS_SENTINEL_PORT, |
| REDIS_URL, |
| REDIS_RECONNECT_DELAY, |
| ) |
|
|
| log = logging.getLogger(__name__) |
|
|
|
|
| _CONNECTION_CACHE = {} |
|
|
|
|
| class SentinelRedisProxy: |
| def __init__(self, sentinel, service, *, async_mode: bool = True, **kw): |
| self._sentinel = sentinel |
| self._service = service |
| self._kw = kw |
| self._async_mode = async_mode |
|
|
| def _master(self): |
| return self._sentinel.master_for(self._service, **self._kw) |
|
|
| def __getattr__(self, item): |
| master = self._master() |
| orig_attr = getattr(master, item) |
|
|
| if not callable(orig_attr): |
| return orig_attr |
|
|
| FACTORY_METHODS = {"pipeline", "pubsub", "monitor", "client", "transaction"} |
| if item in FACTORY_METHODS: |
| return orig_attr |
|
|
| if self._async_mode: |
| if inspect.isasyncgenfunction(orig_attr): |
|
|
| def _wrapped_iter(*args, **kwargs): |
| async def _iter(): |
| for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT): |
| try: |
| method = getattr(self._master(), item) |
| async for value in method(*args, **kwargs): |
| yield value |
| return |
| except ( |
| redis.exceptions.ConnectionError, |
| redis.exceptions.ReadOnlyError, |
| ) as e: |
| if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1: |
| log.debug( |
| "Redis sentinel fail-over (%s). Retry %s/%s", |
| type(e).__name__, |
| i + 1, |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| ) |
| if REDIS_RECONNECT_DELAY: |
| time.sleep(REDIS_RECONNECT_DELAY / 1000) |
| continue |
| log.error( |
| "Redis operation failed after %s retries: %s", |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| e, |
| ) |
| raise e from e |
|
|
| return _iter() |
|
|
| return _wrapped_iter |
|
|
| async def _wrapped(*args, **kwargs): |
| for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT): |
| try: |
| method = getattr(self._master(), item) |
| result = method(*args, **kwargs) |
| if inspect.iscoroutine(result): |
| return await result |
| return result |
| except ( |
| redis.exceptions.ConnectionError, |
| redis.exceptions.ReadOnlyError, |
| ) as e: |
| if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1: |
| log.debug( |
| "Redis sentinel fail-over (%s). Retry %s/%s", |
| type(e).__name__, |
| i + 1, |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| ) |
| if REDIS_RECONNECT_DELAY: |
| await asyncio.sleep(REDIS_RECONNECT_DELAY / 1000) |
| continue |
| log.error( |
| "Redis operation failed after %s retries: %s", |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| e, |
| ) |
| raise e from e |
|
|
| return _wrapped |
|
|
| else: |
|
|
| def _wrapped(*args, **kwargs): |
| for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT): |
| try: |
| method = getattr(self._master(), item) |
| return method(*args, **kwargs) |
| except ( |
| redis.exceptions.ConnectionError, |
| redis.exceptions.ReadOnlyError, |
| ) as e: |
| if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1: |
| log.debug( |
| "Redis sentinel fail-over (%s). Retry %s/%s", |
| type(e).__name__, |
| i + 1, |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| ) |
| if REDIS_RECONNECT_DELAY: |
| time.sleep(REDIS_RECONNECT_DELAY / 1000) |
| continue |
| log.error( |
| "Redis operation failed after %s retries: %s", |
| REDIS_SENTINEL_MAX_RETRY_COUNT, |
| e, |
| ) |
| raise e from e |
|
|
| return _wrapped |
|
|
|
|
| def parse_redis_service_url(redis_url): |
| parsed_url = urlparse(redis_url) |
| if parsed_url.scheme != "redis" and parsed_url.scheme != "rediss": |
| raise ValueError("Invalid Redis URL scheme. Must be 'redis' or 'rediss'.") |
|
|
| return { |
| "username": parsed_url.username or None, |
| "password": parsed_url.password or None, |
| "service": parsed_url.hostname or "mymaster", |
| "port": parsed_url.port or 6379, |
| "db": int(parsed_url.path.lstrip("/") or 0), |
| } |
|
|
|
|
| def get_redis_client(async_mode=False): |
| try: |
| return get_redis_connection( |
| redis_url=REDIS_URL, |
| redis_sentinels=get_sentinels_from_env( |
| REDIS_SENTINEL_HOSTS, REDIS_SENTINEL_PORT |
| ), |
| redis_cluster=REDIS_CLUSTER, |
| async_mode=async_mode, |
| ) |
| except Exception as e: |
| log.debug(f"Failed to get Redis client: {e}") |
| return None |
|
|
|
|
| def get_redis_connection( |
| redis_url, |
| redis_sentinels, |
| redis_cluster=False, |
| async_mode=False, |
| decode_responses=True, |
| ): |
|
|
| cache_key = ( |
| redis_url, |
| tuple(redis_sentinels) if redis_sentinels else (), |
| async_mode, |
| decode_responses, |
| ) |
|
|
| if cache_key in _CONNECTION_CACHE: |
| return _CONNECTION_CACHE[cache_key] |
|
|
| connection = None |
|
|
| if async_mode: |
| import redis.asyncio as redis |
|
|
| |
| if redis_sentinels: |
| redis_config = parse_redis_service_url(redis_url) |
| sentinel = redis.sentinel.Sentinel( |
| redis_sentinels, |
| port=redis_config["port"], |
| db=redis_config["db"], |
| username=redis_config["username"], |
| password=redis_config["password"], |
| decode_responses=decode_responses, |
| socket_connect_timeout=REDIS_SOCKET_CONNECT_TIMEOUT, |
| ) |
| connection = SentinelRedisProxy( |
| sentinel, |
| redis_config["service"], |
| async_mode=async_mode, |
| ) |
| elif redis_cluster: |
| if not redis_url: |
| raise ValueError("Redis URL must be provided for cluster mode.") |
| return redis.cluster.RedisCluster.from_url( |
| redis_url, decode_responses=decode_responses |
| ) |
| elif redis_url: |
| connection = redis.from_url(redis_url, decode_responses=decode_responses) |
| else: |
| import redis |
|
|
| if redis_sentinels: |
| redis_config = parse_redis_service_url(redis_url) |
| sentinel = redis.sentinel.Sentinel( |
| redis_sentinels, |
| port=redis_config["port"], |
| db=redis_config["db"], |
| username=redis_config["username"], |
| password=redis_config["password"], |
| decode_responses=decode_responses, |
| socket_connect_timeout=REDIS_SOCKET_CONNECT_TIMEOUT, |
| ) |
| connection = SentinelRedisProxy( |
| sentinel, |
| redis_config["service"], |
| async_mode=async_mode, |
| ) |
| elif redis_cluster: |
| if not redis_url: |
| raise ValueError("Redis URL must be provided for cluster mode.") |
| return redis.cluster.RedisCluster.from_url( |
| redis_url, decode_responses=decode_responses |
| ) |
| elif redis_url: |
| connection = redis.Redis.from_url( |
| redis_url, decode_responses=decode_responses |
| ) |
|
|
| _CONNECTION_CACHE[cache_key] = connection |
| return connection |
|
|
|
|
| def get_sentinels_from_env(sentinel_hosts_env, sentinel_port_env): |
| if sentinel_hosts_env: |
| sentinel_hosts = sentinel_hosts_env.split(",") |
| sentinel_port = int(sentinel_port_env) |
| return [(host, sentinel_port) for host in sentinel_hosts] |
| return [] |
|
|
|
|
| def get_sentinel_url_from_env(redis_url, sentinel_hosts_env, sentinel_port_env): |
| redis_config = parse_redis_service_url(redis_url) |
| username = redis_config["username"] or "" |
| password = redis_config["password"] or "" |
| auth_part = "" |
| if username or password: |
| auth_part = f"{username}:{password}@" |
| hosts_part = ",".join( |
| f"{host}:{sentinel_port_env}" for host in sentinel_hosts_env.split(",") |
| ) |
| return f"redis+sentinel://{auth_part}{hosts_part}/{redis_config['db']}/{redis_config['service']}" |
|
|
