mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
import redis
|
|
from redis.connection import Connection, SSLConnection
|
|
from redis.sentinel import Sentinel
|
|
|
|
from configs import dify_config
|
|
|
|
|
|
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 dify_config.REDIS_USE_SSL:
|
|
connection_class = SSLConnection
|
|
|
|
redis_params = {
|
|
"username": dify_config.REDIS_USERNAME,
|
|
"password": dify_config.REDIS_PASSWORD,
|
|
"db": dify_config.REDIS_DB,
|
|
"encoding": "utf-8",
|
|
"encoding_errors": "strict",
|
|
"decode_responses": False,
|
|
}
|
|
|
|
if dify_config.REDIS_USE_SENTINEL:
|
|
sentinel_hosts = [
|
|
(node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")
|
|
]
|
|
sentinel = Sentinel(
|
|
sentinel_hosts,
|
|
sentinel_kwargs={
|
|
"socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
|
|
"username": dify_config.REDIS_SENTINEL_USERNAME,
|
|
"password": dify_config.REDIS_SENTINEL_PASSWORD,
|
|
},
|
|
)
|
|
master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
|
|
redis_client.initialize(master)
|
|
else:
|
|
redis_params.update(
|
|
{
|
|
"host": dify_config.REDIS_HOST,
|
|
"port": dify_config.REDIS_PORT,
|
|
"connection_class": connection_class,
|
|
}
|
|
)
|
|
pool = redis.ConnectionPool(**redis_params)
|
|
redis_client.initialize(redis.Redis(connection_pool=pool))
|
|
|
|
app.extensions["redis"] = redis_client
|