mirror of
https://gitee.com/dify_ai/dify.git
synced 2024-12-05 20:57:46 +08:00
84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
import redis
|
|
from redis.connection import Connection, SSLConnection
|
|
from redis.sentinel import Sentinel
|
|
|
|
|
|
class RedisClientWrapper(redis.Redis):
|
|
"""
|
|
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 (redis.Redis): 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.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._client = None
|
|
|
|
def initialize(self, client):
|
|
if self._client is None:
|
|
self._client = client
|
|
|
|
def __getattr__(self, item):
|
|
if self._client is None:
|
|
raise RuntimeError("Redis client is not initialized. Call init_app first.")
|
|
return getattr(self._client, item)
|
|
|
|
|
|
redis_client = RedisClientWrapper()
|
|
|
|
|
|
def init_app(app):
|
|
global redis_client
|
|
connection_class = Connection
|
|
if app.config.get("REDIS_USE_SSL"):
|
|
connection_class = SSLConnection
|
|
|
|
redis_params = {
|
|
"username": app.config.get("REDIS_USERNAME"),
|
|
"password": app.config.get("REDIS_PASSWORD"),
|
|
"db": app.config.get("REDIS_DB"),
|
|
"encoding": "utf-8",
|
|
"encoding_errors": "strict",
|
|
"decode_responses": False,
|
|
}
|
|
|
|
if app.config.get("REDIS_USE_SENTINEL"):
|
|
sentinel_hosts = [
|
|
(node.split(":")[0], int(node.split(":")[1])) for node in app.config.get("REDIS_SENTINELS").split(",")
|
|
]
|
|
sentinel = Sentinel(
|
|
sentinel_hosts,
|
|
sentinel_kwargs={
|
|
"socket_timeout": app.config.get("REDIS_SENTINEL_SOCKET_TIMEOUT", 0.1),
|
|
"username": app.config.get("REDIS_SENTINEL_USERNAME"),
|
|
"password": app.config.get("REDIS_SENTINEL_PASSWORD"),
|
|
},
|
|
)
|
|
master = sentinel.master_for(app.config.get("REDIS_SENTINEL_SERVICE_NAME"), **redis_params)
|
|
redis_client.initialize(master)
|
|
else:
|
|
redis_params.update(
|
|
{
|
|
"host": app.config.get("REDIS_HOST"),
|
|
"port": app.config.get("REDIS_PORT"),
|
|
"connection_class": connection_class,
|
|
}
|
|
)
|
|
pool = redis.ConnectionPool(**redis_params)
|
|
redis_client.initialize(redis.Redis(connection_pool=pool))
|
|
|
|
app.extensions["redis"] = redis_client
|