Add a configuration for controlling the redis instance / type used for streaming events between celery worker and api (vibe-kanban 08e07904)

Currently, the celery worker executing workflows / chatflows uses redis pubsub to publish events to api.
(See \_topic\_msg\_generator and \_publish\_streaming\_response)

The current implementation uses the default redis client.

For large scale deployment, we need to use a dedicated redis cluster to ensure performance.

To achieve this, you should:

1. introduce a dedicated configuration class to control

  the redis address used for pubsub. (Ideally, there should only be one configuration item such as `pubsub_redis_url`, and its default value should be the original redis confugration.)

2. Add an option to switch between pubsub and sharded pubsub. When shared pubsub is specified, the ShardedRedisBroadcastChannel should be used instead.

COmplete the task above, add some unit tests.

api/tests/unit_tests/configs/test_dify_config.py

@@ -159,6 +159,62 @@ def test_db_extras_options_merging(monkeypatch: pytest.MonkeyPatch):
     assert "timezone=UTC" in options
 
 
+def test_pubsub_redis_url_default(monkeypatch: pytest.MonkeyPatch):
+    os.environ.clear()
+
+    monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
+    monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
+    monkeypatch.setenv("DB_USERNAME", "postgres")
+    monkeypatch.setenv("DB_PASSWORD", "postgres")
+    monkeypatch.setenv("DB_HOST", "localhost")
+    monkeypatch.setenv("DB_PORT", "5432")
+    monkeypatch.setenv("DB_DATABASE", "dify")
+    monkeypatch.setenv("REDIS_HOST", "redis.example.com")
+    monkeypatch.setenv("REDIS_PORT", "6380")
+    monkeypatch.setenv("REDIS_USERNAME", "user")
+    monkeypatch.setenv("REDIS_PASSWORD", "pass@word")
+    monkeypatch.setenv("REDIS_DB", "2")
+    monkeypatch.setenv("REDIS_USE_SSL", "true")
+
+    config = DifyConfig()
+
+    assert config.PUBSUB_REDIS_URL == "rediss://user:pass%40word@redis.example.com:6380/2"
+    assert config.PUBSUB_REDIS_CHANNEL_TYPE == "pubsub"
+
+
+def test_pubsub_redis_url_override(monkeypatch: pytest.MonkeyPatch):
+    os.environ.clear()
+
+    monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
+    monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
+    monkeypatch.setenv("DB_USERNAME", "postgres")
+    monkeypatch.setenv("DB_PASSWORD", "postgres")
+    monkeypatch.setenv("DB_HOST", "localhost")
+    monkeypatch.setenv("DB_PORT", "5432")
+    monkeypatch.setenv("DB_DATABASE", "dify")
+    monkeypatch.setenv("PUBSUB_REDIS_URL", "redis://pubsub-host:6381/5")
+
+    config = DifyConfig()
+
+    assert config.PUBSUB_REDIS_URL == "redis://pubsub-host:6381/5"
+
+
+def test_pubsub_redis_url_required_when_default_unavailable(monkeypatch: pytest.MonkeyPatch):
+    os.environ.clear()
+
+    monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
+    monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
+    monkeypatch.setenv("DB_USERNAME", "postgres")
+    monkeypatch.setenv("DB_PASSWORD", "postgres")
+    monkeypatch.setenv("DB_HOST", "localhost")
+    monkeypatch.setenv("DB_PORT", "5432")
+    monkeypatch.setenv("DB_DATABASE", "dify")
+    monkeypatch.setenv("REDIS_HOST", "")
+
+    with pytest.raises(ValueError, match="PUBSUB_REDIS_URL must be set"):
+        DifyConfig()
+
+
 @pytest.mark.parametrize(
     ("broker_url", "expected_host", "expected_port", "expected_username", "expected_password", "expected_db"),
     [

api/tests/unit_tests/conftest.py

@@ -49,6 +49,8 @@ def _patch_redis_clients_on_loaded_modules():
             continue
         if hasattr(module, "redis_client"):
             module.redis_client = redis_mock
+        if hasattr(module, "pubsub_redis_client"):
+            module.pubsub_redis_client = redis_mock
 
 
 @pytest.fixture
@@ -66,7 +68,10 @@ def _provide_app_context(app: Flask):
 def _patch_redis_clients():
     """Patch redis_client to MagicMock only for unit test executions."""
 
-    with patch.object(ext_redis, "redis_client", redis_mock):
+    with (
+        patch.object(ext_redis, "redis_client", redis_mock),
+        patch.object(ext_redis, "pubsub_redis_client", redis_mock),
+    ):
         _patch_redis_clients_on_loaded_modules()
         yield
 

api/tests/unit_tests/extensions/test_pubsub_channel.py

@@ -0,0 +1,20 @@
+from configs import dify_config
+from extensions import ext_redis
+from libs.broadcast_channel.redis.channel import BroadcastChannel as RedisBroadcastChannel
+from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel
+
+
+def test_get_pubsub_broadcast_channel_defaults_to_pubsub(monkeypatch):
+    monkeypatch.setattr(dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "pubsub")
+
+    channel = ext_redis.get_pubsub_broadcast_channel()
+
+    assert isinstance(channel, RedisBroadcastChannel)
+
+
+def test_get_pubsub_broadcast_channel_sharded(monkeypatch):
+    monkeypatch.setattr(dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "sharded")
+
+    channel = ext_redis.get_pubsub_broadcast_channel()
+
+    assert isinstance(channel, ShardedRedisBroadcastChannel)