refactor: improve Redis wrapper type hints and fix None value handling (#23845)

api/extensions/ext_redis.py

@@ -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(",")
         ]