Underground-Digital/Workflow-Engine

1import redis2from redis.connection import Connection, SSLConnection3from redis.sentinel import Sentinel4 5from configs import dify_config6 7 8class RedisClientWrapper(redis.Redis):9    """10    A wrapper class for the Redis client that addresses the issue where the global11    `redis_client` variable cannot be updated when a new Redis instance is returned12    by Sentinel.13 14    This class allows for deferred initialization of the Redis client, enabling the15    client to be re-initialized with a new instance when necessary. This is particularly16    useful in scenarios where the Redis instance may change dynamically, such as during17    a failover in a Sentinel-managed Redis setup.18 19    Attributes:20        _client (redis.Redis): The actual Redis client instance. It remains None until21                               initialized with the `initialize` method.22 23    Methods:24        initialize(client): Initializes the Redis client if it hasn't been initialized already.25        __getattr__(item): Delegates attribute access to the Redis client, raising an error26                           if the client is not initialized.27    """28 29    def __init__(self):30        self._client = None31 32    def initialize(self, client):33        if self._client is None:34            self._client = client35 36    def __getattr__(self, item):37        if self._client is None:38            raise RuntimeError("Redis client is not initialized. Call init_app first.")39        return getattr(self._client, item)40 41 42redis_client = RedisClientWrapper()43 44 45def init_app(app):46    global redis_client47    connection_class = Connection48    if dify_config.REDIS_USE_SSL:49        connection_class = SSLConnection50 51    redis_params = {52        "username": dify_config.REDIS_USERNAME,53        "password": dify_config.REDIS_PASSWORD,54        "db": dify_config.REDIS_DB,55        "encoding": "utf-8",56        "encoding_errors": "strict",57        "decode_responses": False,58    }59 60    if dify_config.REDIS_USE_SENTINEL:61        sentinel_hosts = [62            (node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")63        ]64        sentinel = Sentinel(65            sentinel_hosts,66            sentinel_kwargs={67                "socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,68                "username": dify_config.REDIS_SENTINEL_USERNAME,69                "password": dify_config.REDIS_SENTINEL_PASSWORD,70            },71        )72        master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)73        redis_client.initialize(master)74    else:75        redis_params.update(76            {77                "host": dify_config.REDIS_HOST,78                "port": dify_config.REDIS_PORT,79                "connection_class": connection_class,80            }81        )82        pool = redis.ConnectionPool(**redis_params)83        redis_client.initialize(redis.Redis(connection_pool=pool))84 85    app.extensions["redis"] = redis_client86