klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' of https://github.com/dataease/SQLBot

Browse Source
This commit is contained in:
dataeaseShu
2025-07-23 17:59:09 +08:00
parent dc41cab372 46523fa3ef
commit f0bf7584d5
7 changed files with 159 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/api/datasource.py
+2 -2
View File
@@ -99,8 +99,8 @@ async def edit_field(session: SessionDep, field: CoreField):
@router.post("/previewData/{id}")
async def edit_local(session: SessionDep, id: int, data: TableObj):
return preview(session, id, data)
async def edit_local(session: SessionDep, current_user: CurrentUser, id: int, data: TableObj):
return preview(session, current_user, id, data)
@router.post("/fieldEnum/{id}")
backend/apps/datasource/crud/datasource.py
+18 -2
View File
@@ -237,11 +237,27 @@ def updateField(session: SessionDep, field: CoreField):
update_field(session, field)
def preview(session: SessionDep, id: int, data: TableObj):
def preview(session: SessionDep, current_user: CurrentUser, id: int, data: TableObj):
if data.fields is None or len(data.fields) == 0:
return {"fields": [], "data": [], "sql": ''}
fields = [f.field_name for f in data.fields if f.checked]
# column is checked, and, column permission for data.fields
f_list = [f for f in data.fields if f.checked]
column_permissions = session.query(DsPermission).filter(
and_(DsPermission.table_id == data.table.id, DsPermission.type == 'column')).all()
if column_permissions is not None:
for permission in column_permissions:
# check permission and user in same rules
obj = session.query(DsRules).filter(
and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)),
or_(DsRules.user_list.op('@>')(cast([f'{current_user.id}'], JSONB)),
DsRules.user_list.op('@>')(cast([current_user.id], JSONB))))
).first()
if obj is not None:
permission_list = json.loads(permission.permissions)
f_list = filter_list(f_list, permission_list)
fields = [f.field_name for f in f_list]
if fields is None or len(fields) == 0:
return {"fields": [], "data": [], "sql": ''}
backend/apps/system/api/user.py
+2 -5
View File
@@ -1,7 +1,7 @@
from typing import Optional
from fastapi import APIRouter, HTTPException, Query
from sqlmodel import func, or_, select, delete as sqlmodel_delete
from apps.system.crud.user import clean_user_cache, get_db_user, single_delete, user_ws_options
from apps.system.crud.user import get_db_user, single_delete, user_ws_options
from apps.system.models.system_model import UserWsModel
from apps.system.models.user import UserModel
from apps.system.schemas.auth import CacheName, CacheNamespace
@@ -95,14 +95,13 @@ async def ws_options(session: SessionDep, current_user: CurrentUser, trans: Tran
return await user_ws_options(session, current_user.id, trans)
@router.put("/ws/{oid}")
# @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
async def ws_change(session: SessionDep, current_user: CurrentUser, oid: int):
ws_list: list[UserWs] = await user_ws_options(session, current_user.id)
if not any(x.id == oid for x in ws_list):
raise HTTPException(f"oid [{oid}] is invalid!")
user_model: UserModel = get_db_user(session = session, user_id = current_user.id)
user_model.oid = oid
await clean_user_cache(user_model.id)
session.add(user_model)
session.commit()
@@ -181,7 +180,6 @@ async def langChange(session: SessionDep, current_user: CurrentUser, language: U
return {"message": "Language not supported"}
db_user: UserModel = get_db_user(session=session, user_id=current_user.id)
db_user.language = lang
await clean_user_cache(db_user.id)
session.add(db_user)
session.commit()
@@ -202,6 +200,5 @@ async def pwdUpdate(session: SessionDep, current_user: CurrentUser, editor: PwdE
if not verify_md5pwd(editor.pwd, db_user.password):
raise HTTPException("pwd error")
db_user.password = md5pwd(editor.new_pwd)
await clean_user_cache(db_user.id)
session.add(db_user)
session.commit()
backend/apps/system/middleware/auth.py
+2 -2
View File
@@ -71,7 +71,7 @@ class TokenMiddleware(BaseHTTPMiddleware):
) """
return True, session_user
except Exception as e:
SQLBotLogUtil.exception(f"Token validation error: {str(e)}", exc_info=True)
SQLBotLogUtil.exception(f"Token validation error: {str(e)}")
return False, e
@@ -97,6 +97,6 @@ class TokenMiddleware(BaseHTTPMiddleware):
assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
return True, session_user, assistant_info
except Exception as e:
SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}", exc_info=True)
SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}")
# Return False and the exception message
return False, e
backend/common/core/response_middleware.py
+2 -2
View File
@@ -68,7 +68,7 @@ class ResponseMiddleware(BaseHTTPMiddleware):
class exception_handler():
@staticmethod
async def http_exception_handler(request: Request, exc: HTTPException):
SQLBotLogUtil.exception(f"HTTP Exception: {exc.detail}", exc_info=True)
SQLBotLogUtil.error(f"HTTP Exception: {exc.detail}", exc_info=True)
return JSONResponse(
status_code=exc.status_code,
content={
@@ -82,7 +82,7 @@ class exception_handler():
@staticmethod
async def global_exception_handler(request: Request, exc: Exception):
SQLBotLogUtil.exception(f"Unhandled Exception: {str(exc)}", exc_info=True)
SQLBotLogUtil.error(f"Unhandled Exception: {str(exc)}", exc_info=True)
return JSONResponse(
status_code=500,
content={
backend/common/core/sqlbot_cache.py
+102 -108
View File
@@ -1,35 +1,50 @@
import contextvars
from fastapi_cache import FastAPICache
from fastapi_cache.decorator import cache as original_cache
from functools import partial, wraps
from typing import Optional, Set, Any, Dict, Tuple
from inspect import Parameter, signature
from typing import Optional, Any, Dict, Tuple
from inspect import signature
from contextlib import asynccontextmanager
import asyncio
import random
from collections import defaultdict
from common.utils.utils import SQLBotLogUtil
# 锁管理
_cache_locks = {}
# 使用contextvar来跟踪当前线程已持有的锁
_held_locks = contextvars.ContextVar('held_locks', default=set())
# 高效锁管理器
class LockManager:
_locks = defaultdict(asyncio.Lock)
@classmethod
def get_lock(cls, key: str) -> asyncio.Lock:
return cls._locks[key]
@asynccontextmanager
async def _get_cache_lock(key: str):
lock = _cache_locks.setdefault(key, asyncio.Lock())
async with lock:
try:
yield
finally:
if key in _cache_locks and not lock.locked():
del _cache_locks[key]
def should_skip_param(param: Parameter) -> bool:
"""判断参数是否应该被忽略(依赖注入参数)"""
return (
param.kind == Parameter.VAR_KEYWORD or # **kwargs
param.kind == Parameter.VAR_POSITIONAL or # *args
hasattr(param.annotation, "__module__") and
param.annotation.__module__.startswith(('fastapi', 'starlette', "sqlmodel.orm.session"))
)
# 获取当前已持有的锁集合
current_locks = _held_locks.get()
# 如果已经持有这个锁，直接yield（锁传递）
if key in current_locks:
yield
return
# 否则获取锁并添加到当前上下文中
lock = LockManager.get_lock(key)
try:
await lock.acquire()
# 更新当前持有的锁集合
new_locks = current_locks | {key}
token = _held_locks.set(new_locks)
yield
finally:
# 恢复之前的锁集合
_held_locks.reset(token)
if lock.locked():
lock.release()
def custom_key_builder(
func: Any,
@@ -37,137 +52,116 @@ def custom_key_builder(
*,
args: Tuple[Any, ...] = (),
kwargs: Dict[str, Any],
additional_skip_args: Optional[Set[str]] = None,
cacheName: Optional[str] = None,
cacheName: str,
keyExpression: Optional[str] = None,
) -> str:
"""完全兼容FastAPICache的键生成器"""
if cacheName:
try:
base_key = f"{namespace}:{cacheName}:"
if keyExpression:
try:
sig = signature(func)
bound_args = sig.bind_partial(*args, **kwargs)
bound_args.apply_defaults()
if keyExpression.startswith("args["):
import re
match = re.match(r"args\[(\d+)\]", keyExpression)
if match:
index = int(match.group(1))
value = bound_args.args[index]
base_key += f"{value}:"
else:
parts = keyExpression.split('.')
value = bound_args.arguments[parts[0]]
for part in parts[1:]:
value = getattr(value, part)
base_key += f"{value}:"
except (IndexError, KeyError, AttributeError) as e:
SQLBotLogUtil.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
sig = signature(func)
bound_args = sig.bind_partial(*args, **kwargs)
bound_args.apply_defaults()
# 支持args[0]格式
if keyExpression.startswith("args["):
import re
if match := re.match(r"args\[(\d+)\]", keyExpression):
index = int(match.group(1))
value = bound_args.args[index]
return f"{base_key}{value}"
# 支持属性路径格式
parts = keyExpression.split('.')
value = bound_args.arguments[parts[0]]
for part in parts[1:]:
value = getattr(value, part)
return f"{base_key}{value}"
return base_key
sig = signature(func)
auto_skip_args = {
name for name, param in sig.parameters.items()
if should_skip_param(param)
}
skip_args = auto_skip_args.union(additional_skip_args or set())
filtered_kwargs = {
k: v for k, v in kwargs.items() if k not in skip_args
}
bound_args = sig.bind_partial(*args, **kwargs)
bound_args.apply_defaults()
filtered_args = []
for i, (name, value) in enumerate(bound_args.arguments.items()):
if i < len(args) and name not in skip_args:
filtered_args.append(value)
filtered_args = tuple(filtered_args)
default_key_builder = FastAPICache.get_key_builder()
return default_key_builder(
func=func,
namespace=namespace,
args=filtered_args,
kwargs=filtered_kwargs,
)
# 默认使用第一个参数作为key
return f"{base_key}{args[0] if args else 'default'}"
except Exception as e:
SQLBotLogUtil.error(f"Key builder error: {str(e)}")
raise ValueError(f"Invalid cache key generation: {e}") from e
def cache(
expire: Optional[int] = 60 * 60 * 24,
namespace: Optional[str] = None,
key_builder: Optional[Any] = None,
expire: int = 60 * 60 * 24,
namespace: str = "",
*,
additional_skip_args: Optional[Set[str]] = None,
cacheName: Optional[str] = None,
cacheName: str, # 必须提供cacheName
keyExpression: Optional[str] = None,
jitter: int = 60, # 默认抖动60秒
):
"""完全兼容的缓存装饰器"""
def decorator(func):
if key_builder is None:
used_key_builder = partial(
custom_key_builder,
additional_skip_args=additional_skip_args,
cacheName=cacheName,
keyExpression=keyExpression
)
else:
used_key_builder = key_builder
# 预先生成key builder
used_key_builder = partial(
custom_key_builder,
cacheName=cacheName,
keyExpression=keyExpression
)
@wraps(func)
async def wrapper(*args, **kwargs):
# 生成缓存键
cache_key = used_key_builder(
func=func,
namespace=namespace or "",
namespace=namespace,
args=args,
kwargs=kwargs
)
# 防击穿锁
async with _get_cache_lock(cache_key):
SQLBotLogUtil.info(f"Using cache key: {cache_key}")
backend = FastAPICache.get_backend()
cached_value = await backend.get(cache_key)
if cached_value is not None:
SQLBotLogUtil.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
return cached_value
# 双重检查
if (cached := await backend.get(cache_key)) is not None:
SQLBotLogUtil.debug(f"Cache hit: {cache_key}")
return cached
# 执行函数并缓存结果
result = await func(*args, **kwargs)
await backend.set(cache_key, result, expire)
SQLBotLogUtil.info(f"Cache miss for key: {cache_key}, result cached.")
actual_expire = expire + random.randint(-jitter, jitter)
await backend.set(cache_key, result, actual_expire)
SQLBotLogUtil.debug(f"Cache set: {cache_key} (expire: {actual_expire}s)")
return result
return wrapper
return decorator
def clear_cache(
namespace: Optional[str] = None,
cacheName: Optional[str] = None,
namespace: str = "",
*,
cacheName: str,
keyExpression: Optional[str] = None,
):
"""清除缓存的装饰器"""
"""精确清除单个缓存项的装饰器"""
def decorator(func):
@wraps(func)
async def wrapper(*args, **kwargs):
cache_key = custom_key_builder(
func=func,
namespace=namespace or "",
namespace=namespace,
args=args,
kwargs=kwargs,
cacheName=cacheName,
keyExpression=keyExpression,
)
# 加锁防止竞争
async with _get_cache_lock(cache_key):
if await FastAPICache.get_backend().get(cache_key):
await FastAPICache.clear(key=cache_key)
backend = FastAPICache.get_backend()
result = None
if await backend.get(cache_key):
await backend.clear(cache_key)
result = await func(*args, **kwargs)
SQLBotLogUtil.info(f"Clearing cache for key: {cache_key}")
return result
return await func(*args, **kwargs)
if await backend.get(cache_key):
await backend.clear(cache_key)
SQLBotLogUtil.info(f"Cache cleared: {cache_key}")
return result
return wrapper
return decorator
frontend/src/utils/request.ts
+31 -1
View File
@@ -308,6 +308,36 @@ class HttpService {
}
public loadRemoteScript(url: string, id?: string, cb?: any): Promise<HTMLElement> {
if (!url) {
return Promise.reject(new Error('URL is required to load remote script'))
}
if (id && document.getElementById(id)) {
return Promise.resolve(document.getElementById(id) as HTMLElement)
}
if (url.startsWith('/')) {
const real_url = import.meta.env.VITE_API_BASE_URL.replace('/api/v1', '')
url = real_url + url
}
return new Promise<HTMLElement>((resolve, reject) => {
// 改用传统的script标签加载方式
const script = document.createElement('script')
script.src = url
script.id = id || `remote-script-${Date.now()}`
script.onload = () => {
if (cb) cb()
resolve(script)
}
script.onerror = (error) => {
console.error(`Failed to load script from ${url}:`, error)
reject(new Error(`Failed to load script from ${url}`))
}
document.head.appendChild(script)
})
}
/* public loadRemoteScript(url: string, id?: string, cb?: any): Promise<HTMLElement> {
if (!url) {
return Promise.reject(new Error('URL is required to load remote script'))
}
@@ -337,7 +367,7 @@ class HttpService {
reject(new Error(`Failed to load script from ${url}: ${error.message}`))
})
})
}
} */
}
// Create singleton instance