Skip to content

Commit 0558d99

Browse files
committed
Merge branch 'obkv-redis' into dev
# Conflicts: # docker/docker-compose.middleware.yaml # docker/docker-compose.yaml
2 parents ecf866d + 3600dac commit 0558d99

File tree

6 files changed

+175
-2
lines changed

6 files changed

+175
-2
lines changed

api/configs/middleware/cache/redis_config.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,3 +93,8 @@ class RedisConfig(BaseSettings):
9393
description="Enable client side cache in redis",
9494
default=False,
9595
)
96+
97+
REDIS_USE_OBKV: bool = Field(
98+
description="Use OBKV instead of Redis",
99+
default=False,
100+
)

api/extensions/ext_redis.py

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
1-
from typing import Any, Union
1+
import logging
2+
from typing import Any, Optional, Union
23

34
import redis
45
from redis.cache import CacheConfig
6+
from redis.client import Pipeline
57
from redis.cluster import ClusterNode, RedisCluster
68
from redis.connection import Connection, SSLConnection
9+
from redis.exceptions import LockNotOwnedError
10+
from redis.lock import Lock
711
from redis.sentinel import Sentinel
812

913
from configs import dify_config
@@ -101,6 +105,22 @@ def init_app(app: DifyApp):
101105
cache_config=clientside_cache_config,
102106
)
103107
)
108+
elif dify_config.REDIS_USE_OBKV:
109+
assert dify_config.REDIS_SERIALIZATION_PROTOCOL < 3, (
110+
"OBKV does not support RESP version 3. "
111+
"Please specify REDIS_SERIALIZATION_PROTOCOL=2 in the configuration file."
112+
)
113+
redis_params.update(
114+
{
115+
"host": dify_config.REDIS_HOST,
116+
"port": dify_config.REDIS_PORT,
117+
"username": dify_config.REDIS_USERNAME,
118+
"password": dify_config.REDIS_PASSWORD,
119+
"connection_class": connection_class,
120+
}
121+
)
122+
pool = redis.ConnectionPool(**redis_params)
123+
redis_client.initialize(OBKVClient(redis.Redis(connection_pool=pool)))
104124
else:
105125
redis_params.update(
106126
{
@@ -115,3 +135,51 @@ def init_app(app: DifyApp):
115135
redis_client.initialize(redis.Redis(connection_pool=pool))
116136

117137
app.extensions["redis"] = redis_client
138+
139+
140+
class OceanBaseRedisLock(Lock):
141+
def do_acquire(self, token: str) -> bool:
142+
try:
143+
result = self.redis.execute_command("SETNX", self.name, token)
144+
except Exception as e:
145+
logging.info(f"Failed to acquire redis lock {self.name}, error: {e}")
146+
result = False
147+
if result == 1:
148+
self.redis.execute_command("EXPIRE", self.name, 60)
149+
return True
150+
else:
151+
return False
152+
153+
def do_release(self, expected_token: str) -> None:
154+
current_value = self.redis.get(self.name)
155+
if current_value == expected_token:
156+
redis_client.delete(self.name)
157+
logging.info(f"Lock released: {self.name}")
158+
else:
159+
raise LockNotOwnedError(
160+
"Cannot release a lock that's no longer owned",
161+
lock_name=self.name,
162+
)
163+
164+
165+
class OceanBasePipeLine(Pipeline):
166+
def execute(self, raise_on_error=True):
167+
"""do nothing, OBKV cannot support redis pipeline"""
168+
logging.info("executing pipeline: do nothing, OBKV cannot support redis pipeline")
169+
170+
171+
class OBKVClient:
172+
def __init__(self, client: redis.Redis):
173+
self._client = client
174+
175+
def lock(self, name: str, timeout: Optional[int] = None, **kwargs):
176+
# Implementation using OceanBaseRedisLock
177+
return self._client.lock(name, timeout=timeout, lock_class=OceanBaseRedisLock, **kwargs)
178+
179+
def pipeline(self, transaction=True, shard_hint=None) -> "Pipeline":
180+
return OceanBasePipeLine(self.connection_pool, self.response_callbacks, transaction, shard_hint)
181+
182+
def __getattr__(self, item):
183+
if self._client is None:
184+
raise RuntimeError("Redis client is not initialized. Call init_app first.")
185+
return getattr(self._client, item)

docker/.env.example

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,10 @@ REDIS_USE_CLUSTERS=false
275275
REDIS_CLUSTERS=
276276
REDIS_CLUSTERS_PASSWORD=
277277

278+
# Whether to use OBKV instead of Redis.
279+
REDIS_USE_OBKV=false
280+
281+
278282
# ------------------------------
279283
# Celery Configuration
280284
# ------------------------------

docker/docker-compose.middleware.yaml

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,43 @@ services:
4343
healthcheck:
4444
test: [ "CMD", "redis-cli", "ping" ]
4545

46+
47+
obkv_redis:
48+
image: oceanbase/oceanbase-ce:4.2.5.3-103000022025033117
49+
container_name: oceanbase
50+
restart: always
51+
environment:
52+
OB_MEMORY_LIMIT: ${OCEANBASE_MEMORY_LIMIT:-6G}
53+
OB_SYS_PASSWORD: ${OCEANBASE_VECTOR_PASSWORD:-difyai123456}
54+
OB_TENANT_PASSWORD: ${OCEANBASE_VECTOR_PASSWORD:-difyai123456}
55+
OB_CLUSTER_NAME: ${OCEANBASE_CLUSTER_NAME:-difyai}
56+
MODE: MINI
57+
volumes:
58+
- ./volumes/oceanbase/obkv_redis/healthcheck.sh:/root/healthcheck.sh
59+
ports:
60+
- "${OBKV_PORT:-2991}:2881"
61+
healthcheck:
62+
test: [ "CMD", "/root/healthcheck.sh" ]
63+
interval: 10s
64+
retries: 30
65+
start_period: 30s
66+
timeout: 10s
67+
68+
obproxy:
69+
image: oceanbase/obproxy-ce:4.3.4.0-1
70+
container_name: obproxy
71+
environment:
72+
APP_NAME: dify_obproxy
73+
PROXYRO_PASSWORD: 123456
74+
RS_LIST: ${OCEANBASE_IP}:${OBKV_PORT:-2991}
75+
OB_CLUSTER: ${OCEANBASE_CLUSTER_NAME:-difyai}
76+
ports:
77+
- "${OBPROXY_SQL_PORT:-2883}:2883"
78+
- "${OBPROXY_RPC_PORT:-2885}:2885"
79+
depends_on:
80+
obkv_redis:
81+
condition: service_healthy
82+
4683
# The DifySandbox
4784
sandbox:
4885
image: langgenius/dify-sandbox:0.2.12
@@ -71,7 +108,7 @@ services:
71108

72109
# plugin daemon
73110
plugin_daemon:
74-
image: langgenius/dify-plugin-daemon:0.1.2-local
111+
image: quay.io/oceanbase-devhub/dify-plugin-daemon:obkv-redis-0.1-local
75112
restart: always
76113
env_file:
77114
- ./middleware.env
@@ -131,6 +168,9 @@ services:
131168
- "${EXPOSE_PLUGIN_DEBUGGING_PORT:-5003}:${PLUGIN_DEBUGGING_PORT:-5003}"
132169
volumes:
133170
- ./volumes/plugin_daemon:/app/storage
171+
depends_on:
172+
obkv_redis:
173+
condition: service_healthy
134174

135175
# ssrf_proxy server
136176
# for more information, please refer to

docker/docker-compose.yaml

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ x-shared-env: &shared-api-worker-env
7575
REDIS_SENTINEL_PASSWORD: ${REDIS_SENTINEL_PASSWORD:-}
7676
REDIS_SENTINEL_SOCKET_TIMEOUT: ${REDIS_SENTINEL_SOCKET_TIMEOUT:-0.1}
7777
REDIS_USE_CLUSTERS: ${REDIS_USE_CLUSTERS:-false}
78+
REDIS_USE_OBKV: ${REDIS_USE_OBKV:-false}
79+
REDIS_SERIALIZATION_PROTOCOL: ${REDIS_SERIALIZATION_PROTOCOL:-2}
7880
REDIS_CLUSTERS: ${REDIS_CLUSTERS:-}
7981
REDIS_CLUSTERS_PASSWORD: ${REDIS_CLUSTERS_PASSWORD:-}
8082
CELERY_BROKER_URL: ${CELERY_BROKER_URL:-redis://:difyai123456@redis:6379/1}
@@ -530,6 +532,8 @@ services:
530532
condition: service_healthy
531533
redis:
532534
condition: service_started
535+
obproxy:
536+
condition: service_started
533537
volumes:
534538
# Mount the storage directory to the container, for storing user files.
535539
- ./volumes/app/storage:/app/api/storage
@@ -557,6 +561,8 @@ services:
557561
condition: service_healthy
558562
redis:
559563
condition: service_started
564+
obproxy:
565+
condition: service_started
560566
volumes:
561567
# Mount the storage directory to the container, for storing user files.
562568
- ./volumes/app/storage:/app/api/storage
@@ -707,6 +713,8 @@ services:
707713
depends_on:
708714
oceanbase:
709715
condition: service_healthy
716+
obproxy:
717+
condition: service_started
710718

711719
# ssrf_proxy server
712720
# for more information, please refer to
@@ -970,6 +978,46 @@ services:
970978
start_period: 30s
971979
timeout: 10s
972980

981+
obkv_redis:
982+
image: oceanbase/oceanbase-ce:4.2.5.3-103000022025033117
983+
container_name: obkv_redis
984+
profiles:
985+
- oceanbase
986+
restart: always
987+
environment:
988+
OB_MEMORY_LIMIT: ${OCEANBASE_MEMORY_LIMIT:-6G}
989+
OB_SYS_PASSWORD: ${OCEANBASE_VECTOR_PASSWORD:-difyai123456}
990+
OB_TENANT_PASSWORD: ${OCEANBASE_VECTOR_PASSWORD:-difyai123456}
991+
OB_CLUSTER_NAME: ${OCEANBASE_CLUSTER_NAME:-difyai}
992+
MODE: MINI
993+
volumes:
994+
- ./volumes/oceanbase/obkv_redis/healthcheck.sh:/root/healthcheck.sh
995+
ports:
996+
- "${OBKV_PORT:-2991}:2881"
997+
healthcheck:
998+
test: [ "CMD", "/root/healthcheck.sh" ]
999+
interval: 10s
1000+
retries: 30
1001+
start_period: 30s
1002+
timeout: 10s
1003+
1004+
obproxy:
1005+
image: oceanbase/obproxy-ce:4.3.4.0-1
1006+
container_name: obproxy
1007+
profiles:
1008+
- oceanbase
1009+
environment:
1010+
APP_NAME: dify_obproxy
1011+
PROXYRO_PASSWORD: 123456
1012+
RS_LIST: ${OCEANBASE_IP}:${OBKV_PORT:-2991}
1013+
OB_CLUSTER: ${OCEANBASE_CLUSTER_NAME:-difyai}
1014+
ports:
1015+
- "${OBPROXY_SQL_PORT:-2883}:2883"
1016+
- "${OBPROXY_RPC_PORT:-2885}:2885"
1017+
depends_on:
1018+
obkv_redis:
1019+
condition: service_healthy
1020+
9731021
# Oracle vector database
9741022
oracle:
9751023
image: container-registry.oracle.com/database/free:latest
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#!/bin/sh
2+
3+
set -x
4+
5+
obclient -h127.0.0.1 -P2881 -uroot@test -p${OB_TENANT_PASSWORD} -e "SELECT 1;"
6+
obclient -h127.0.0.1 -P2881 -uroot@sys -p${OB_SYS_PASSWORD} -e "CREATE USER 'proxyro'@'%' IDENTIFIED BY '123456'; GRANT ALL PRIVILEGES ON *.* TO 'proxyro'@'%'; FLUSH PRIVILEGES;"
7+
# Fixed the bug that obproxy always failed to connect to the test tenant for the first time.
8+
obclient -h127.0.0.1 -P2881 -uroot@sys -p${OB_SYS_PASSWORD} -e "ALTER SYSTEM LOAD MODULE DATA module=redis tenant=sys;"

0 commit comments

Comments
 (0)