refactor: improve Redis wrapper type hints and fix None value handling (#23845)
This commit is contained in:
@@ -1,18 +1,23 @@
|
||||
import functools
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Union
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
import redis
|
||||
from redis import RedisError
|
||||
from redis.cache import CacheConfig
|
||||
from redis.cluster import ClusterNode, RedisCluster
|
||||
from redis.connection import Connection, SSLConnection
|
||||
from redis.lock import Lock
|
||||
from redis.sentinel import Sentinel
|
||||
|
||||
from configs import dify_config
|
||||
from dify_app import DifyApp
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis.lock import Lock
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -28,8 +33,8 @@ class RedisClientWrapper:
|
||||
a failover in a Sentinel-managed Redis setup.
|
||||
|
||||
Attributes:
|
||||
_client (redis.Redis): The actual Redis client instance. It remains None until
|
||||
initialized with the `initialize` method.
|
||||
_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.
|
||||
@@ -37,20 +42,78 @@ class RedisClientWrapper:
|
||||
if the client is not initialized.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
_client: Union[redis.Redis, RedisCluster, None]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = None
|
||||
|
||||
def initialize(self, client):
|
||||
def initialize(self, client: Union[redis.Redis, RedisCluster]) -> None:
|
||||
if self._client is None:
|
||||
self._client = client
|
||||
|
||||
def __getattr__(self, item):
|
||||
if TYPE_CHECKING:
|
||||
# Type hints for IDE support and static analysis
|
||||
# These are not executed at runtime but provide type information
|
||||
def get(self, name: str | bytes) -> Any: ...
|
||||
|
||||
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: ...
|
||||
|
||||
def setex(self, name: str | bytes, time: int | timedelta, value: Any) -> Any: ...
|
||||
def setnx(self, name: str | bytes, value: Any) -> Any: ...
|
||||
def delete(self, *names: str | bytes) -> Any: ...
|
||||
def incr(self, name: str | bytes, amount: int = 1) -> Any: ...
|
||||
def expire(
|
||||
self,
|
||||
name: str | bytes,
|
||||
time: int | timedelta,
|
||||
nx: bool = False,
|
||||
xx: bool = False,
|
||||
gt: bool = False,
|
||||
lt: bool = False,
|
||||
) -> Any: ...
|
||||
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,
|
||||
) -> Lock: ...
|
||||
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: ...
|
||||
def zremrangebyscore(self, name: str | bytes, min: float | str, max: float | str) -> Any: ...
|
||||
def zcard(self, name: str | bytes) -> Any: ...
|
||||
def getdel(self, name: str | bytes) -> Any: ...
|
||||
|
||||
def __getattr__(self, item: str) -> Any:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Redis client is not initialized. Call init_app first.")
|
||||
return getattr(self._client, item)
|
||||
|
||||
|
||||
redis_client = RedisClientWrapper()
|
||||
redis_client: RedisClientWrapper = RedisClientWrapper()
|
||||
|
||||
|
||||
def init_app(app: DifyApp):
|
||||
@@ -80,6 +143,9 @@ def init_app(app: DifyApp):
|
||||
|
||||
if dify_config.REDIS_USE_SENTINEL:
|
||||
assert dify_config.REDIS_SENTINELS is not None, "REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True"
|
||||
assert dify_config.REDIS_SENTINEL_SERVICE_NAME is not None, (
|
||||
"REDIS_SENTINEL_SERVICE_NAME must be set when REDIS_USE_SENTINEL is True"
|
||||
)
|
||||
sentinel_hosts = [
|
||||
(node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")
|
||||
]
|
||||
|
@@ -162,6 +162,7 @@ dev = [
|
||||
"pandas-stubs~=2.2.3",
|
||||
"scipy-stubs>=1.15.3.0",
|
||||
"types-python-http-client>=3.3.7.20240910",
|
||||
"types-redis>=4.6.0.20241004",
|
||||
]
|
||||
|
||||
############################################################
|
||||
|
@@ -24,9 +24,20 @@ def queue_monitor_task():
|
||||
queue_name = "dataset"
|
||||
threshold = dify_config.QUEUE_MONITOR_THRESHOLD
|
||||
|
||||
if threshold is None:
|
||||
logging.warning(click.style("QUEUE_MONITOR_THRESHOLD is not configured, skipping monitoring", fg="yellow"))
|
||||
return
|
||||
|
||||
try:
|
||||
queue_length = celery_redis.llen(f"{queue_name}")
|
||||
logging.info(click.style(f"Start monitor {queue_name}", fg="green"))
|
||||
|
||||
if queue_length is None:
|
||||
logging.error(
|
||||
click.style(f"Failed to get queue length for {queue_name} - Redis may be unavailable", fg="red")
|
||||
)
|
||||
return
|
||||
|
||||
logging.info(click.style(f"Queue length: {queue_length}", fg="green"))
|
||||
|
||||
if queue_length >= threshold:
|
||||
|
17
api/uv.lock
generated
17
api/uv.lock
generated
@@ -1,5 +1,5 @@
|
||||
version = 1
|
||||
revision = 3
|
||||
revision = 2
|
||||
requires-python = ">=3.11, <3.13"
|
||||
resolution-markers = [
|
||||
"python_full_version >= '3.12.4' and platform_python_implementation != 'PyPy' and sys_platform == 'linux'",
|
||||
@@ -1371,6 +1371,7 @@ dev = [
|
||||
{ name = "types-python-http-client" },
|
||||
{ name = "types-pywin32" },
|
||||
{ name = "types-pyyaml" },
|
||||
{ name = "types-redis" },
|
||||
{ name = "types-regex" },
|
||||
{ name = "types-requests" },
|
||||
{ name = "types-requests-oauthlib" },
|
||||
@@ -1557,6 +1558,7 @@ dev = [
|
||||
{ name = "types-python-http-client", specifier = ">=3.3.7.20240910" },
|
||||
{ name = "types-pywin32", specifier = "~=310.0.0" },
|
||||
{ name = "types-pyyaml", specifier = "~=6.0.12" },
|
||||
{ name = "types-redis", specifier = ">=4.6.0.20241004" },
|
||||
{ name = "types-regex", specifier = "~=2024.11.6" },
|
||||
{ name = "types-requests", specifier = "~=2.32.0" },
|
||||
{ name = "types-requests-oauthlib", specifier = "~=2.0.0" },
|
||||
@@ -6064,6 +6066,19 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/99/5f/e0af6f7f6a260d9af67e1db4f54d732abad514252a7a378a6c4d17dd1036/types_pyyaml-6.0.12.20250516-py3-none-any.whl", hash = "sha256:8478208feaeb53a34cb5d970c56a7cd76b72659442e733e268a94dc72b2d0530", size = 20312, upload-time = "2025-05-16T03:08:04.019Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "types-redis"
|
||||
version = "4.6.0.20241004"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cryptography" },
|
||||
{ name = "types-pyopenssl" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/3a/95/c054d3ac940e8bac4ca216470c80c26688a0e79e09f520a942bb27da3386/types-redis-4.6.0.20241004.tar.gz", hash = "sha256:5f17d2b3f9091ab75384153bfa276619ffa1cf6a38da60e10d5e6749cc5b902e", size = 49679, upload-time = "2024-10-04T02:43:59.224Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/55/82/7d25dce10aad92d2226b269bce2f85cfd843b4477cd50245d7d40ecf8f89/types_redis-4.6.0.20241004-py3-none-any.whl", hash = "sha256:ef5da68cb827e5f606c8f9c0b49eeee4c2669d6d97122f301d3a55dc6a63f6ed", size = 58737, upload-time = "2024-10-04T02:43:57.968Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "types-regex"
|
||||
version = "2024.11.6.20250403"
|
||||
|
Reference in New Issue
Block a user