Skip to content

Commit e739284

Browse files
perf: Cache config
1 parent 19c0f94 commit e739284

File tree

5 files changed

+36
-16
lines changed

5 files changed

+36
-16
lines changed

.env.example

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ LOG_DIR="/opt/sqlbot/logs"
2121
LOG_FORMAT="%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"
2222
SQL_DEBUG=False
2323

24+
CACHE_TYPE="memory"
25+
CACHE_REDIS_URL="redis://127.0.0.1:6379"
26+
2427
# Postgres
2528
POSTGRES_SERVER=localhost
2629
POSTGRES_PORT=5432

backend/common/core/config.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,9 @@ def all_cors_origins(self) -> list[str]:
6161
DEFAULT_PWD: str = "SQLBot@123456"
6262
ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN"
6363

64+
CACHE_TYPE: Literal["redis", "memory", "None"] = "memory"
65+
CACHE_REDIS_URL: str | None = None # Redis URL, e.g., "redis://[[username]:[password]]@localhost:6379/0"
66+
6467
LOG_LEVEL: str = "INFO" # DEBUG, INFO, WARNING, ERROR
6568
LOG_DIR: str = "logs"
6669
LOG_FORMAT: str = "%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"

backend/common/core/sqlbot_cache.py

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,9 @@
77
import asyncio
88
import random
99
from collections import defaultdict
10-
11-
from apps.system.schemas.auth import CacheName
12-
from apps.system.schemas.system_schema import UserInfoDTO
10+
from common.core.config import settings
1311
from common.utils.utils import SQLBotLogUtil
14-
12+
from fastapi_cache.backends.inmemory import InMemoryBackend
1513
# 使用contextvar来跟踪当前线程已持有的锁
1614
_held_locks = contextvars.ContextVar('held_locks', default=set())
1715
# 高效锁管理器
@@ -105,6 +103,8 @@ def decorator(func):
105103

106104
@wraps(func)
107105
async def wrapper(*args, **kwargs):
106+
if not settings.CACHE_TYPE or settings.CACHE_TYPE.lower() == "none":
107+
return await func(*args, **kwargs)
108108
# 生成缓存键
109109
cache_key = used_key_builder(
110110
func=func,
@@ -116,13 +116,9 @@ async def wrapper(*args, **kwargs):
116116
# 防击穿锁
117117
async with _get_cache_lock(cache_key):
118118
backend = FastAPICache.get_backend()
119-
120119
# 双重检查
121120
if (cached := await backend.get(cache_key)) is not None:
122121
SQLBotLogUtil.debug(f"Cache hit: {cache_key}")
123-
if CacheName.USER_INFO.value in cache_key:
124-
user = UserInfoDTO.model_validate(cached)
125-
SQLBotLogUtil.info(f"User cache hit: [uid: {user.id}, account: {user.account}, oid: {user.oid}]")
126122
return cached
127123

128124
# 执行函数并缓存结果
@@ -134,9 +130,6 @@ async def wrapper(*args, **kwargs):
134130
await backend.set(cache_key, result, actual_expire)
135131

136132
SQLBotLogUtil.debug(f"Cache set: {cache_key} (expire: {actual_expire}s)")
137-
if CacheName.USER_INFO.value in cache_key:
138-
user = UserInfoDTO.model_validate(result)
139-
SQLBotLogUtil.info(f"User cache set: [uid: {user.id}, account: {user.account}, oid: {user.oid}]")
140133
return result
141134

142135
return wrapper
@@ -152,6 +145,8 @@ def clear_cache(
152145
def decorator(func):
153146
@wraps(func)
154147
async def wrapper(*args, **kwargs):
148+
if not settings.CACHE_TYPE or settings.CACHE_TYPE.lower() == "none":
149+
return await func(*args, **kwargs)
155150
cache_key = custom_key_builder(
156151
func=func,
157152
namespace=namespace,
@@ -170,8 +165,27 @@ async def wrapper(*args, **kwargs):
170165
result = await func(*args, **kwargs)
171166
if await backend.get(cache_key):
172167
await backend.clear(cache_key)
173-
SQLBotLogUtil.info(f"Cache cleared: {cache_key}")
168+
SQLBotLogUtil.debug(f"Cache cleared: {cache_key}")
174169
return result
175170

176171
return wrapper
177-
return decorator
172+
return decorator
173+
174+
175+
def init_sqlbot_cache():
176+
cache_type: str = settings.CACHE_TYPE
177+
if cache_type == "memory":
178+
FastAPICache.init(InMemoryBackend())
179+
SQLBotLogUtil.info("SQLBot 使用内存缓存, 仅支持单进程模式")
180+
elif cache_type == "redis":
181+
from fastapi_cache.backends.redis import RedisBackend
182+
import redis.asyncio as redis
183+
from redis.asyncio.connection import ConnectionPool
184+
redis_url = settings.CACHE_REDIS_URL or "redis://localhost:6379/0"
185+
pool = ConnectionPool.from_url(url=redis_url)
186+
redis_client = redis.Redis(connection_pool=pool)
187+
FastAPICache.init(RedisBackend(redis_client), prefix="sqlbot-cache")
188+
SQLBotLogUtil.info(f"SQLBot 使用Redis缓存, 可使用多进程模式")
189+
else:
190+
SQLBotLogUtil.warning("SQLBot 未启用缓存, 可使用多进程模式")
191+

backend/main.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,9 @@
1414
from fastapi_mcp import FastApiMCP
1515
from fastapi.staticfiles import StaticFiles
1616
import sqlbot_xpack
17-
from fastapi_cache import FastAPICache
18-
from fastapi_cache.backends.inmemory import InMemoryBackend
1917

2018
from common.utils.utils import SQLBotLogUtil
19+
from common.core.sqlbot_cache import init_sqlbot_cache
2120

2221
def run_migrations():
2322
alembic_cfg = Config("alembic.ini")
@@ -27,7 +26,7 @@ def run_migrations():
2726
@asynccontextmanager
2827
async def lifespan(app: FastAPI):
2928
run_migrations()
30-
FastAPICache.init(InMemoryBackend())
29+
init_sqlbot_cache()
3130
init_dynamic_cors(app)
3231
SQLBotLogUtil.info("✅ SQLBot 初始化完成")
3332
yield

backend/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ dependencies = [
3939
"sqlbot-xpack==0.0.3.7",
4040
"fastapi-cache2>=0.2.2",
4141
"sqlparse>=0.5.3",
42+
"redis>=6.2.0",
4243
]
4344
[[tool.uv.index]]
4445
name = "default"

0 commit comments

Comments
 (0)