mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本,归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块,新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块,引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程,移除 base image 二级构建,统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
3
_deprecated_py_src/utils/__init__.py
Normal file
3
_deprecated_py_src/utils/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
from .timeout import AsyncTimeoutError, run_with_timeout, with_timeout
|
||||
|
||||
__all__ = ["with_timeout", "run_with_timeout", "AsyncTimeoutError"]
|
||||
66
_deprecated_py_src/utils/async_utils.py
Normal file
66
_deprecated_py_src/utils/async_utils.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
异步工具函数
|
||||
|
||||
提供在异步上下文中安全执行同步函数的工具,避免阻塞事件循环。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable, Coroutine
|
||||
from functools import partial, wraps
|
||||
from typing import Any, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# 全局 task 引用集合,防止 fire-and-forget task 被 GC 回收
|
||||
_background_tasks: set[asyncio.Task[Any]] = set()
|
||||
|
||||
|
||||
def safe_create_task(coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any] | None:
|
||||
"""创建后台 task 并持有引用,防止被 GC 回收。
|
||||
|
||||
用于 fire-and-forget 场景(如缓存失效、异步指标上报等),
|
||||
替代裸 ``asyncio.create_task()`` 调用。
|
||||
|
||||
Returns:
|
||||
创建的 Task 对象;若没有运行中的事件循环则返回 None。
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
task = loop.create_task(coro)
|
||||
_background_tasks.add(task)
|
||||
task.add_done_callback(_background_tasks.discard)
|
||||
return task
|
||||
|
||||
|
||||
async def run_in_executor(func: Callable[..., T], *args: Any, **kwargs: Any) -> T:
|
||||
"""
|
||||
在线程池中运行同步函数,避免阻塞事件循环。
|
||||
|
||||
用法:
|
||||
result = await run_in_executor(some_sync_function, arg1, arg2)
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
bound = partial(func, *args, **kwargs)
|
||||
return await loop.run_in_executor(None, bound)
|
||||
|
||||
|
||||
def async_wrap_sync(func: Callable[..., T]) -> Callable[..., Coroutine[Any, Any, T]]:
|
||||
"""
|
||||
装饰器:将同步函数包装成异步函数(在线程池中执行)。
|
||||
|
||||
用法:
|
||||
@async_wrap_sync
|
||||
def do_sync(...): ...
|
||||
|
||||
result = await do_sync(...)
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> T:
|
||||
return await run_in_executor(func, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
209
_deprecated_py_src/utils/auth_utils.py
Normal file
209
_deprecated_py_src/utils/auth_utils.py
Normal file
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
认证工具函数
|
||||
提供统一的用户认证和授权功能
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Any
|
||||
|
||||
from fastapi import Depends, Header, HTTPException, Request, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ManagementToken
|
||||
from src.services.auth.service import AuthService
|
||||
from src.utils.request_utils import get_client_ip
|
||||
|
||||
from ..core.exceptions import ForbiddenException
|
||||
from ..database import get_db
|
||||
from ..models.database import User, UserRole
|
||||
|
||||
security = HTTPBearer()
|
||||
|
||||
|
||||
async def authenticate_user_from_bearer_token(
|
||||
token: str,
|
||||
db: Session,
|
||||
request: Request | None = None,
|
||||
) -> User:
|
||||
if token.startswith(ManagementToken.TOKEN_PREFIX):
|
||||
client_ip = get_client_ip(request) if request is not None else "unknown"
|
||||
result = await AuthService.authenticate_management_token(db, token, client_ip)
|
||||
if not result:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的Token")
|
||||
|
||||
user, management_token = result
|
||||
if request is not None:
|
||||
request.state.user_id = user.id
|
||||
request.state.management_token_id = management_token.id
|
||||
return user
|
||||
|
||||
# 验证Token格式和签名
|
||||
try:
|
||||
payload = await AuthService.verify_token(token, token_type="access")
|
||||
except HTTPException as token_error:
|
||||
token_fp = hashlib.sha256(token.encode()).hexdigest()[:12]
|
||||
logger.error(
|
||||
"Token验证失败: {}: {}, token_fp={}",
|
||||
token_error.status_code,
|
||||
token_error.detail,
|
||||
token_fp,
|
||||
)
|
||||
raise
|
||||
except Exception as token_error:
|
||||
token_fp = hashlib.sha256(token.encode()).hexdigest()[:12]
|
||||
logger.error("Token验证失败: {}, token_fp={}", token_error, token_fp)
|
||||
raise ForbiddenException("无效的Token")
|
||||
|
||||
user_id = payload.get("user_id")
|
||||
|
||||
if not user_id:
|
||||
logger.error("Token缺少user_id字段: payload={}", payload)
|
||||
raise ForbiddenException("无效的认证凭据")
|
||||
|
||||
token_fp = hashlib.sha256(token.encode()).hexdigest()[:12]
|
||||
|
||||
if not isinstance(user_id, str):
|
||||
logger.error("Token中user_id格式错误: {} - {}", type(user_id), user_id)
|
||||
raise ForbiddenException("认证信息格式错误,请重新登录")
|
||||
|
||||
try:
|
||||
from src.services.user.service import UserService
|
||||
|
||||
user = UserService.get_user(db, user_id)
|
||||
except Exception as db_error:
|
||||
logger.error("数据库查询失败: user_id={}, error={}", user_id, db_error)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="数据库查询失败,请稍后重试",
|
||||
)
|
||||
|
||||
if not user:
|
||||
logger.error("用户不存在: user_id={}", user_id)
|
||||
raise ForbiddenException("用户不存在或已禁用")
|
||||
|
||||
if not user.is_active:
|
||||
logger.error("用户已禁用: user_id={}", user_id)
|
||||
raise ForbiddenException("用户不存在或已禁用")
|
||||
|
||||
if user.is_deleted:
|
||||
logger.error("用户已删除: user_id={}", user_id)
|
||||
raise ForbiddenException("用户不存在或已禁用")
|
||||
|
||||
if not AuthService.token_identity_matches_user(payload, user):
|
||||
logger.error("Token身份校验失败: user_id={}, token_fp={}", user_id, token_fp)
|
||||
raise ForbiddenException("身份验证失败")
|
||||
|
||||
if request is not None:
|
||||
request.state.user_id = user.id
|
||||
if hasattr(request.state, "management_token_id"):
|
||||
request.state.management_token_id = None
|
||||
|
||||
return user
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
request: Request,
|
||||
credentials: HTTPAuthorizationCredentials = Depends(security),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User:
|
||||
"""
|
||||
获取当前登录用户
|
||||
统一的认证依赖函数
|
||||
|
||||
Args:
|
||||
credentials: Bearer token 凭据
|
||||
db: 数据库会话
|
||||
|
||||
Returns:
|
||||
User: 当前用户对象
|
||||
|
||||
Raises:
|
||||
HTTPException: 认证失败时抛出
|
||||
"""
|
||||
try:
|
||||
return await authenticate_user_from_bearer_token(credentials.credentials, db, request)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("认证失败,未预期的错误: {}", e)
|
||||
# 返回500而不是401,避免触发前端的退出逻辑
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="认证服务暂时不可用"
|
||||
)
|
||||
|
||||
|
||||
async def get_current_user_from_header(
|
||||
request: Request,
|
||||
authorization: str | None = Header(None),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User:
|
||||
"""
|
||||
从Header中获取当前用户(兼容性函数)
|
||||
|
||||
Args:
|
||||
authorization: Authorization header
|
||||
db: 数据库会话
|
||||
|
||||
Returns:
|
||||
User: 当前用户对象
|
||||
|
||||
Raises:
|
||||
HTTPException: 认证失败时抛出
|
||||
"""
|
||||
if not authorization or not authorization.startswith("Bearer "):
|
||||
raise ForbiddenException("未提供认证令牌")
|
||||
|
||||
try:
|
||||
return await authenticate_user_from_bearer_token(
|
||||
authorization.replace("Bearer ", ""),
|
||||
db,
|
||||
request,
|
||||
)
|
||||
except HTTPException:
|
||||
# 保持原始的HTTPException (包括401)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"认证失败: {e}")
|
||||
raise ForbiddenException("认证失败")
|
||||
|
||||
|
||||
def require_admin(current_user: User = Depends(get_current_user)) -> User:
|
||||
"""
|
||||
要求管理员权限
|
||||
|
||||
Args:
|
||||
current_user: 当前用户
|
||||
|
||||
Returns:
|
||||
User: 管理员用户对象
|
||||
|
||||
Raises:
|
||||
HTTPException: 非管理员时抛出403错误
|
||||
"""
|
||||
if current_user.role != UserRole.ADMIN:
|
||||
raise ForbiddenException("需要管理员权限")
|
||||
return current_user
|
||||
|
||||
|
||||
def require_role(required_role: UserRole) -> Any:
|
||||
"""
|
||||
要求特定角色权限的装饰器工厂
|
||||
|
||||
Args:
|
||||
required_role: 需要的用户角色
|
||||
|
||||
Returns:
|
||||
依赖函数
|
||||
"""
|
||||
|
||||
def check_role(current_user: User = Depends(get_current_user)) -> User:
|
||||
if current_user.role != required_role:
|
||||
raise ForbiddenException(f"需要{required_role.value}权限")
|
||||
return current_user
|
||||
|
||||
return check_role
|
||||
155
_deprecated_py_src/utils/cache_decorator.py
Normal file
155
_deprecated_py_src/utils/cache_decorator.py
Normal file
@@ -0,0 +1,155 @@
|
||||
"""缓存装饰器工具"""
|
||||
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from src.clients.redis_client import get_redis_client_sync
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
def _is_adapter_instance(obj: Any) -> bool:
|
||||
"""检查对象是否是 ApiAdapter 的实例(延迟导入避免循环依赖)"""
|
||||
try:
|
||||
from src.api.base.adapter import ApiAdapter
|
||||
|
||||
return isinstance(obj, ApiAdapter)
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def _is_api_context(obj: Any) -> bool:
|
||||
"""检查对象是否是 ApiRequestContext(通过 duck typing)"""
|
||||
return hasattr(obj, "user") and hasattr(obj, "db")
|
||||
|
||||
|
||||
def _resolve_attr(obj: Any, dotted_name: str) -> tuple[bool, Any]:
|
||||
"""Resolve a possibly dotted attribute path (e.g. 'time_range.start_date').
|
||||
|
||||
Returns (found, value). When any segment along the chain is missing or
|
||||
the intermediate value is None the lookup stops and returns (False, None).
|
||||
"""
|
||||
current = obj
|
||||
for part in dotted_name.split("."):
|
||||
if current is None:
|
||||
return False, None
|
||||
if not hasattr(current, part):
|
||||
return False, None
|
||||
current = getattr(current, part)
|
||||
return True, current
|
||||
|
||||
|
||||
def _hash_vary(vary: dict[str, Any]) -> str:
|
||||
"""Build a short stable hash for cache key variations."""
|
||||
try:
|
||||
raw = json.dumps(vary, sort_keys=True, ensure_ascii=False, default=str)
|
||||
except Exception:
|
||||
raw = str(vary)
|
||||
return hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def cache_result(
|
||||
key_prefix: str,
|
||||
ttl: int = 60,
|
||||
user_specific: bool = True,
|
||||
*,
|
||||
vary_by: list[str] | None = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
缓存函数结果的装饰器
|
||||
|
||||
Args:
|
||||
key_prefix: 缓存键前缀
|
||||
ttl: 缓存过期时间(秒)
|
||||
user_specific: 是否针对用户缓存(从 context.user.id 获取)
|
||||
"""
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
redis_client = get_redis_client_sync()
|
||||
|
||||
# 如果 Redis 不可用,直接执行原函数
|
||||
if redis_client is None:
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
# 构建缓存键
|
||||
try:
|
||||
# 从 args 中获取 context
|
||||
# 对于实例方法,args[0] 是 self,args[1] 才是 context
|
||||
# 对于普通函数,args[0] 是 context
|
||||
context = None
|
||||
adapter_self = None
|
||||
|
||||
if len(args) >= 2 and _is_adapter_instance(args[0]) and _is_api_context(args[1]):
|
||||
# 实例方法: handle(self, context)
|
||||
adapter_self = args[0]
|
||||
context = args[1]
|
||||
elif len(args) >= 1 and _is_api_context(args[0]):
|
||||
# 普通函数或 context 在第一个位置
|
||||
context = args[0]
|
||||
elif len(args) >= 1 and _is_adapter_instance(args[0]):
|
||||
# 实例方法但 context 可能在 kwargs 中
|
||||
adapter_self = args[0]
|
||||
context = kwargs.get("context")
|
||||
|
||||
if user_specific and context and hasattr(context, "user") and context.user:
|
||||
cache_key = f"{key_prefix}:user:{context.user.id}"
|
||||
else:
|
||||
cache_key = f"{key_prefix}:global"
|
||||
|
||||
# If there are extra parameters, include them into the key.
|
||||
# - When vary_by is provided: hash the selected attributes to keep key short.
|
||||
# - Otherwise keep backward-compatible "days/limit" suffix behavior.
|
||||
if adapter_self and hasattr(adapter_self, "__dict__"):
|
||||
if vary_by:
|
||||
vary: dict[str, Any] = {}
|
||||
for attr_name in vary_by:
|
||||
found, value = _resolve_attr(adapter_self, attr_name)
|
||||
if found:
|
||||
vary[attr_name] = value
|
||||
if vary:
|
||||
cache_key += f":v:{_hash_vary(vary)}"
|
||||
else:
|
||||
for attr_name in ["days", "limit"]:
|
||||
if hasattr(adapter_self, attr_name):
|
||||
attr_value = getattr(adapter_self, attr_name)
|
||||
cache_key += f":{attr_name}:{attr_value}"
|
||||
|
||||
# 尝试从缓存获取
|
||||
cached = await redis_client.get(cache_key)
|
||||
if cached:
|
||||
try:
|
||||
result = json.loads(cached)
|
||||
logger.trace(f"缓存命中: {cache_key}")
|
||||
return result
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"缓存解析失败,删除损坏缓存: {cache_key}, 错误: {e}")
|
||||
try:
|
||||
await redis_client.delete(cache_key)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 执行原函数
|
||||
result = await func(*args, **kwargs)
|
||||
|
||||
# 保存到缓存
|
||||
try:
|
||||
await redis_client.setex(
|
||||
cache_key, ttl, json.dumps(result, ensure_ascii=False, default=str)
|
||||
)
|
||||
logger.trace(f"缓存已保存: {cache_key}, TTL: {ttl}s")
|
||||
except Exception as e:
|
||||
logger.warning(f"保存缓存失败: {e}")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"缓存处理出错: {e}, 直接执行原函数")
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
75
_deprecated_py_src/utils/compression.py
Normal file
75
_deprecated_py_src/utils/compression.py
Normal file
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
数据压缩/解压工具
|
||||
|
||||
提供JSON数据的gzip压缩和解压功能
|
||||
"""
|
||||
|
||||
import gzip
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
|
||||
def compress_json(data: Any) -> bytes | None:
|
||||
"""
|
||||
将JSON数据压缩为gzip格式的字节
|
||||
|
||||
Args:
|
||||
data: 任意可JSON序列化的数据
|
||||
|
||||
Returns:
|
||||
gzip压缩后的字节,如果输入为None则返回None
|
||||
"""
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# 转换为JSON字符串
|
||||
json_str = json.dumps(data, ensure_ascii=False)
|
||||
# gzip压缩
|
||||
compressed = gzip.compress(json_str.encode("utf-8"), compresslevel=6)
|
||||
return compressed
|
||||
except Exception:
|
||||
# 如果压缩失败,返回None
|
||||
return None
|
||||
|
||||
|
||||
def decompress_json(compressed_data: bytes | None) -> Any | None:
|
||||
"""
|
||||
解压gzip格式的字节为JSON数据
|
||||
|
||||
Args:
|
||||
compressed_data: gzip压缩的字节数据
|
||||
|
||||
Returns:
|
||||
解压后的JSON数据,如果输入为None或解压失败则返回None
|
||||
"""
|
||||
if compressed_data is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# gzip解压
|
||||
json_str = gzip.decompress(compressed_data).decode("utf-8")
|
||||
# 解析JSON
|
||||
data = json.loads(json_str)
|
||||
return data
|
||||
except Exception:
|
||||
# 如果解压失败,返回None
|
||||
return None
|
||||
|
||||
|
||||
def get_body_size(data: Any) -> int:
|
||||
"""
|
||||
获取JSON数据序列化后的字节大小
|
||||
|
||||
Args:
|
||||
data: 任意可JSON序列化的数据
|
||||
|
||||
Returns:
|
||||
字节大小
|
||||
"""
|
||||
if data is None:
|
||||
return 0
|
||||
try:
|
||||
return len(json.dumps(data, ensure_ascii=False).encode("utf-8"))
|
||||
except Exception:
|
||||
return 0
|
||||
118
_deprecated_py_src/utils/database_helpers.py
Normal file
118
_deprecated_py_src/utils/database_helpers.py
Normal file
@@ -0,0 +1,118 @@
|
||||
"""
|
||||
数据库方言兼容性辅助函数
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
|
||||
|
||||
def escape_like_pattern(pattern: str) -> str:
|
||||
"""
|
||||
转义 SQL LIKE 语句中的特殊字符(%、_、\\)
|
||||
|
||||
Args:
|
||||
pattern: 原始搜索模式
|
||||
|
||||
Returns:
|
||||
转义后的模式,可安全用于 LIKE 查询(需配合 escape="\\\\")
|
||||
|
||||
Examples:
|
||||
>>> escape_like_pattern("hello_world%test")
|
||||
'hello\\\\_world\\\\%test'
|
||||
"""
|
||||
return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
|
||||
|
||||
def safe_truncate_escaped(escaped: str, max_len: int) -> str:
|
||||
"""
|
||||
安全截断已转义的字符串,避免截断在转义序列中间
|
||||
|
||||
转义后的字符串中,反斜杠总是成对出现(\\\\)或作为转义符(\\%, \\_)。
|
||||
如果在某个位置截断导致末尾有奇数个反斜杠,说明截断发生在转义序列中间,
|
||||
需要去掉最后一个反斜杠以保持转义完整性。
|
||||
|
||||
Args:
|
||||
escaped: 已经过 escape_like_pattern 处理的字符串
|
||||
max_len: 最大长度
|
||||
|
||||
Returns:
|
||||
截断后的字符串,保证不会破坏转义序列
|
||||
"""
|
||||
if len(escaped) <= max_len:
|
||||
return escaped
|
||||
|
||||
truncated = escaped[:max_len]
|
||||
|
||||
# 统计末尾连续的反斜杠数量
|
||||
trailing_backslashes = 0
|
||||
for i in range(len(truncated) - 1, -1, -1):
|
||||
if truncated[i] == "\\":
|
||||
trailing_backslashes += 1
|
||||
else:
|
||||
break
|
||||
|
||||
# 如果末尾反斜杠数量为奇数,说明截断在转义序列中间
|
||||
# 需要去掉最后一个反斜杠
|
||||
if trailing_backslashes % 2 == 1:
|
||||
truncated = truncated[:-1]
|
||||
|
||||
return truncated
|
||||
|
||||
|
||||
def date_trunc_portable(dialect_name: str, interval: str, column: Any) -> Any:
|
||||
"""
|
||||
跨数据库的日期截断函数
|
||||
|
||||
Args:
|
||||
dialect_name: 数据库方言名称 ('postgresql', 'sqlite', 'mysql')
|
||||
interval: 时间间隔 ('day', 'week', 'month', 'year')
|
||||
column: 日期列
|
||||
|
||||
Returns:
|
||||
SQLAlchemy ClauseElement
|
||||
|
||||
Raises:
|
||||
NotImplementedError: 不支持的数据库方言
|
||||
|
||||
Examples:
|
||||
>>> # PostgreSQL
|
||||
>>> period_func = date_trunc_portable('postgresql', 'week', Usage.created_at)
|
||||
>>> # 等价于: func.date_trunc('week', Usage.created_at)
|
||||
|
||||
>>> # SQLite
|
||||
>>> period_func = date_trunc_portable('sqlite', 'month', Usage.created_at)
|
||||
>>> # 等价于: func.strftime("%Y-%m", Usage.created_at)
|
||||
"""
|
||||
if dialect_name == "postgresql":
|
||||
# PostgreSQL 使用 date_trunc 函数
|
||||
return func.date_trunc(interval, column)
|
||||
|
||||
elif dialect_name == "sqlite":
|
||||
# SQLite 使用 strftime 函数
|
||||
format_map = {
|
||||
"year": "%Y",
|
||||
"month": "%Y-%m",
|
||||
"week": "%Y-%W",
|
||||
"day": "%Y-%m-%d",
|
||||
}
|
||||
if interval not in format_map:
|
||||
raise ValueError(f"Unsupported interval for SQLite: {interval}")
|
||||
return func.strftime(format_map[interval], column)
|
||||
|
||||
elif dialect_name == "mysql":
|
||||
# MySQL 使用 date_format 函数
|
||||
format_map = {
|
||||
"year": "%Y",
|
||||
"month": "%Y-%m",
|
||||
"day": "%Y-%m-%d",
|
||||
}
|
||||
if interval not in format_map:
|
||||
raise ValueError(f"Unsupported interval for MySQL: {interval}")
|
||||
return func.date_format(column, format_map[interval])
|
||||
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported database dialect: {dialect_name}. "
|
||||
f"Supported dialects: postgresql, sqlite, mysql"
|
||||
)
|
||||
135
_deprecated_py_src/utils/perf.py
Normal file
135
_deprecated_py_src/utils/perf.py
Normal file
@@ -0,0 +1,135 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class PerfRecorder:
|
||||
"""轻量性能记录器(可选启用)"""
|
||||
|
||||
@staticmethod
|
||||
def enabled() -> bool:
|
||||
return bool(config.perf_metrics_enabled or config.perf_log_slow_ms > 0)
|
||||
|
||||
@staticmethod
|
||||
def start(force: bool = False) -> float | None:
|
||||
if not force and not PerfRecorder.enabled():
|
||||
return None
|
||||
return time.perf_counter()
|
||||
|
||||
@staticmethod
|
||||
def stop(
|
||||
start: float | None,
|
||||
name: str,
|
||||
labels: dict[str, str] | None = None,
|
||||
*,
|
||||
sample_rate: float | None = None,
|
||||
log_hint: str | None = None,
|
||||
) -> float | None:
|
||||
if start is None:
|
||||
return None
|
||||
duration = time.perf_counter() - start
|
||||
PerfRecorder.record_timing(
|
||||
name,
|
||||
duration,
|
||||
labels=labels,
|
||||
sample_rate=sample_rate,
|
||||
log_hint=log_hint,
|
||||
)
|
||||
return duration
|
||||
|
||||
@staticmethod
|
||||
def record_timing(
|
||||
name: str,
|
||||
duration: float,
|
||||
labels: dict[str, str] | None = None,
|
||||
*,
|
||||
sample_rate: float | None = None,
|
||||
log_hint: str | None = None,
|
||||
) -> None:
|
||||
if not PerfRecorder.enabled():
|
||||
return
|
||||
if not PerfRecorder._should_sample(sample_rate):
|
||||
return
|
||||
|
||||
duration_ms = duration * 1000.0
|
||||
if config.perf_log_slow_ms > 0 and duration_ms >= float(config.perf_log_slow_ms):
|
||||
hint = f" | {log_hint}" if log_hint else ""
|
||||
logger.info("[PERF] {} took {:.2f}ms{}", name, duration_ms, hint)
|
||||
|
||||
if not config.perf_metrics_enabled:
|
||||
return
|
||||
|
||||
plugin = PerfRecorder._get_monitor_plugin()
|
||||
if not plugin:
|
||||
return
|
||||
PerfRecorder._create_task(plugin.timing(PerfRecorder._metric_name(name), duration, labels))
|
||||
|
||||
@staticmethod
|
||||
def record_counter(
|
||||
name: str,
|
||||
value: float = 1,
|
||||
labels: dict[str, str] | None = None,
|
||||
*,
|
||||
sample_rate: float | None = None,
|
||||
) -> None:
|
||||
if not config.perf_metrics_enabled:
|
||||
return
|
||||
if not PerfRecorder._should_sample(sample_rate):
|
||||
return
|
||||
|
||||
plugin = PerfRecorder._get_monitor_plugin()
|
||||
if not plugin:
|
||||
return
|
||||
PerfRecorder._create_task(
|
||||
plugin.increment(PerfRecorder._metric_name(name), value=value, labels=labels)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def should_store() -> bool:
|
||||
return bool(getattr(config, "perf_store_enabled", False))
|
||||
|
||||
@staticmethod
|
||||
def should_store_sample() -> bool:
|
||||
if not PerfRecorder.should_store():
|
||||
return False
|
||||
rate = float(getattr(config, "perf_store_sample_rate", 1.0))
|
||||
return PerfRecorder._should_sample(rate)
|
||||
|
||||
@staticmethod
|
||||
def _should_sample(sample_rate: float | None) -> bool:
|
||||
rate = float(sample_rate if sample_rate is not None else config.perf_sample_rate)
|
||||
if rate >= 1:
|
||||
return True
|
||||
if rate <= 0:
|
||||
return False
|
||||
return random.random() < rate
|
||||
|
||||
@staticmethod
|
||||
def _get_monitor_plugin() -> Any | None:
|
||||
# Lazy import to avoid circular deps during app startup.
|
||||
try:
|
||||
from src.plugins.manager import get_plugin_manager
|
||||
except Exception:
|
||||
return None
|
||||
try:
|
||||
plugin = get_plugin_manager().get_plugin("monitor")
|
||||
except Exception:
|
||||
return None
|
||||
if plugin and getattr(plugin, "enabled", True):
|
||||
return plugin
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _metric_name(name: str) -> str:
|
||||
return name if name.startswith("perf_") else f"perf_{name}"
|
||||
|
||||
@staticmethod
|
||||
def _create_task(coro: Any) -> None:
|
||||
from src.utils.async_utils import safe_create_task
|
||||
|
||||
safe_create_task(coro)
|
||||
184
_deprecated_py_src/utils/request_utils.py
Normal file
184
_deprecated_py_src/utils/request_utils.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
请求处理工具函数
|
||||
提供统一的HTTP请求信息提取功能
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
TRACE_ID_HEADER = "x-trace-id"
|
||||
_MISSING = object()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RequestIdentityMetadata:
|
||||
"""请求身份元数据(最小集合)。"""
|
||||
|
||||
request_id: str | None
|
||||
client_ip: str
|
||||
user_agent: str
|
||||
|
||||
|
||||
def get_client_ip(request: Request) -> str:
|
||||
"""
|
||||
获取客户端真实IP地址
|
||||
|
||||
按优先级检查:
|
||||
1. X-Real-IP 头(最可靠,由最外层可信 Nginx 直接设置)
|
||||
2. X-Forwarded-For 头的第一个 IP(原始客户端)
|
||||
3. 直接客户端IP
|
||||
|
||||
安全说明:
|
||||
- X-Real-IP 优先级最高,因为它通常由最外层 Nginx 设置为 $remote_addr,
|
||||
Nginx 会直接覆盖这个头,不会传递客户端伪造的值
|
||||
- 只要最外层 Nginx 配置了 proxy_set_header X-Real-IP $remote_addr; 即可正确获取真实 IP
|
||||
|
||||
Args:
|
||||
request: FastAPI Request 对象
|
||||
|
||||
Returns:
|
||||
str: 客户端IP地址,如果无法获取则返回 "unknown"
|
||||
"""
|
||||
# 优先检查 X-Real-IP 头(由最外层 Nginx 设置,最可靠)
|
||||
real_ip = request.headers.get("X-Real-IP")
|
||||
if real_ip:
|
||||
return real_ip.strip()
|
||||
|
||||
# 检查 X-Forwarded-For 头,取第一个 IP(原始客户端)
|
||||
forwarded_for = request.headers.get("X-Forwarded-For")
|
||||
if forwarded_for:
|
||||
# X-Forwarded-For 格式: "client, proxy1, proxy2"
|
||||
ips = [ip.strip() for ip in forwarded_for.split(",") if ip.strip()]
|
||||
if ips:
|
||||
return ips[0]
|
||||
|
||||
# 回退到直接客户端IP
|
||||
if request.client and request.client.host:
|
||||
return request.client.host
|
||||
|
||||
return "unknown"
|
||||
|
||||
|
||||
def get_user_agent(request: Request) -> str:
|
||||
"""
|
||||
获取用户代理字符串
|
||||
|
||||
Args:
|
||||
request: FastAPI Request 对象
|
||||
|
||||
Returns:
|
||||
str: User-Agent 字符串,如果无法获取则返回 "unknown"
|
||||
"""
|
||||
return request.headers.get("User-Agent", "unknown")
|
||||
|
||||
|
||||
def get_request_id(request: Request) -> str | None:
|
||||
"""
|
||||
获取请求ID(如果存在)
|
||||
|
||||
Args:
|
||||
request: FastAPI Request 对象
|
||||
|
||||
Returns:
|
||||
Optional[str]: 请求ID,如果不存在则返回 None
|
||||
"""
|
||||
request_id = getattr(request.state, "request_id", None)
|
||||
if request_id:
|
||||
return request_id
|
||||
|
||||
trace_id = request.headers.get(TRACE_ID_HEADER)
|
||||
if trace_id:
|
||||
return trace_id.strip() or None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def update_request_state(
|
||||
request: Request,
|
||||
*,
|
||||
request_id: object = _MISSING,
|
||||
user_id: object = _MISSING,
|
||||
api_key_id: object = _MISSING,
|
||||
management_token_id: object = _MISSING,
|
||||
user_session_id: object = _MISSING,
|
||||
prefetched_balance_remaining: object = _MISSING,
|
||||
gateway_execution_path: object = _MISSING,
|
||||
rate_limit_scope: object = _MISSING,
|
||||
) -> None:
|
||||
"""集中维护请求级 runtime state,减少散点赋值。"""
|
||||
|
||||
if request_id is not _MISSING:
|
||||
request.state.request_id = request_id
|
||||
if user_id is not _MISSING:
|
||||
request.state.user_id = user_id
|
||||
if api_key_id is not _MISSING:
|
||||
request.state.api_key_id = api_key_id
|
||||
if management_token_id is not _MISSING:
|
||||
request.state.management_token_id = management_token_id
|
||||
if user_session_id is not _MISSING:
|
||||
request.state.user_session_id = user_session_id
|
||||
if prefetched_balance_remaining is not _MISSING:
|
||||
request.state.prefetched_balance_remaining = prefetched_balance_remaining
|
||||
if gateway_execution_path is not _MISSING:
|
||||
request.state.gateway_execution_path = gateway_execution_path
|
||||
if rate_limit_scope is not _MISSING:
|
||||
request.state.rate_limit_scope = rate_limit_scope
|
||||
|
||||
|
||||
def get_request_metadata(request: Request) -> dict:
|
||||
"""
|
||||
获取请求的完整元数据
|
||||
|
||||
Args:
|
||||
request: FastAPI Request 对象
|
||||
|
||||
Returns:
|
||||
dict: 包含请求元数据的字典
|
||||
"""
|
||||
identity = get_request_identity_metadata(request)
|
||||
return {
|
||||
**asdict(identity),
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"query_params": str(request.query_params) if request.query_params else None,
|
||||
"content_type": request.headers.get("Content-Type"),
|
||||
"content_length": request.headers.get("Content-Length"),
|
||||
}
|
||||
|
||||
|
||||
def get_request_identity_metadata(request: Request) -> RequestIdentityMetadata:
|
||||
"""集中读取 request_id/client_ip/user_agent,避免散点访问。"""
|
||||
|
||||
return RequestIdentityMetadata(
|
||||
request_id=get_request_id(request),
|
||||
client_ip=get_client_ip(request),
|
||||
user_agent=get_user_agent(request),
|
||||
)
|
||||
|
||||
|
||||
def extract_ip_from_headers(headers: dict) -> str:
|
||||
"""
|
||||
从HTTP头字典中提取IP地址(用于中间件等场景)
|
||||
|
||||
Args:
|
||||
headers: HTTP头字典
|
||||
|
||||
Returns:
|
||||
str: 客户端IP地址
|
||||
"""
|
||||
# 优先检查 X-Real-IP(由最外层 Nginx 设置,最可靠)
|
||||
real_ip = headers.get("x-real-ip", "")
|
||||
if real_ip:
|
||||
return real_ip.strip()
|
||||
|
||||
# 检查 X-Forwarded-For,取第一个 IP
|
||||
forwarded_for = headers.get("x-forwarded-for", "")
|
||||
if forwarded_for:
|
||||
ips = [ip.strip() for ip in forwarded_for.split(",") if ip.strip()]
|
||||
if ips:
|
||||
return ips[0]
|
||||
|
||||
return "unknown"
|
||||
109
_deprecated_py_src/utils/sse_parser.py
Normal file
109
_deprecated_py_src/utils/sse_parser.py
Normal file
@@ -0,0 +1,109 @@
|
||||
class SSEEventParser:
|
||||
"""轻量SSE解析器,按行接收输入并输出完整事件。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._reset_buffer()
|
||||
|
||||
def _reset_buffer(self) -> None:
|
||||
self._buffer: dict[str, str | None | list[str]] = {
|
||||
"event": None,
|
||||
"data": [],
|
||||
"id": None,
|
||||
"retry": None,
|
||||
}
|
||||
|
||||
def _finalize_event(self) -> dict[str, str | None] | None:
|
||||
data_lines = self._buffer.get("data", [])
|
||||
if not isinstance(data_lines, list) or not data_lines:
|
||||
self._reset_buffer()
|
||||
return None
|
||||
|
||||
data_str = "\n".join(data_lines)
|
||||
event_val = self._buffer.get("event")
|
||||
id_val = self._buffer.get("id")
|
||||
retry_val = self._buffer.get("retry")
|
||||
event: dict[str, str | None] = {
|
||||
"event": event_val if isinstance(event_val, str) else None,
|
||||
"data": data_str,
|
||||
"id": id_val if isinstance(id_val, str) else None,
|
||||
"retry": retry_val if isinstance(retry_val, str) else None,
|
||||
}
|
||||
|
||||
self._reset_buffer()
|
||||
return event
|
||||
|
||||
def feed_line(self, line: str | None) -> list[dict[str, str | None]]:
|
||||
"""处理单行SSE文本,返回所有完成的事件。"""
|
||||
|
||||
normalized_line = (line or "").rstrip("\r")
|
||||
events: list[dict[str, str | None]] = []
|
||||
|
||||
# 空行表示事件结束
|
||||
if normalized_line == "":
|
||||
event = self._finalize_event()
|
||||
if event:
|
||||
events.append(event)
|
||||
return events
|
||||
|
||||
# 注释行直接忽略
|
||||
if normalized_line.startswith(":") and not normalized_line.startswith("::"):
|
||||
return events
|
||||
|
||||
if normalized_line.startswith("event:"):
|
||||
_, rest = normalized_line.split(":", 1)
|
||||
value = rest.lstrip()
|
||||
|
||||
if " data:" in value:
|
||||
event_part, data_part = value.split("data:", 1)
|
||||
event_name = event_part.strip() or None
|
||||
data_value = data_part.lstrip()
|
||||
self._buffer["event"] = event_name
|
||||
if data_value:
|
||||
self._append_data_line(data_value)
|
||||
event = self._finalize_event()
|
||||
if event:
|
||||
events.append(event)
|
||||
else:
|
||||
event_name = value.strip() or None
|
||||
self._buffer["event"] = event_name
|
||||
return events
|
||||
|
||||
if normalized_line.startswith("data:"):
|
||||
# 如果已经有缓存的 data,先完成上一个事件
|
||||
# 这样可以处理没有空行分隔的连续 data 行
|
||||
existing_data = self._buffer.get("data", [])
|
||||
if existing_data and len(existing_data) > 0:
|
||||
event = self._finalize_event()
|
||||
if event:
|
||||
events.append(event)
|
||||
|
||||
_, rest = normalized_line.split(":", 1)
|
||||
self._append_data_line(rest[1:] if rest.startswith(" ") else rest)
|
||||
return events
|
||||
|
||||
if normalized_line.startswith("id:"):
|
||||
_, rest = normalized_line.split(":", 1)
|
||||
self._buffer["id"] = rest.strip() or None
|
||||
return events
|
||||
|
||||
if normalized_line.startswith("retry:"):
|
||||
_, rest = normalized_line.split(":", 1)
|
||||
self._buffer["retry"] = rest.strip() or None
|
||||
return events
|
||||
|
||||
# 未知行:视作数据追加(部分实现会缺少 data: 前缀)
|
||||
self._append_data_line(normalized_line)
|
||||
return events
|
||||
|
||||
def flush(self) -> list[dict[str, str | None]]:
|
||||
"""在流结束时调用,输出尚未完成的事件。"""
|
||||
|
||||
event = self._finalize_event()
|
||||
return [event] if event else []
|
||||
|
||||
def _append_data_line(self, value: str) -> None:
|
||||
data_lines = self._buffer.get("data")
|
||||
if isinstance(data_lines, list):
|
||||
data_lines.append(value)
|
||||
else:
|
||||
self._buffer["data"] = [value]
|
||||
117
_deprecated_py_src/utils/ssl_utils.py
Normal file
117
_deprecated_py_src/utils/ssl_utils.py
Normal file
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
SSL 工具函数
|
||||
提供统一的 SSL 上下文创建功能
|
||||
"""
|
||||
|
||||
import ssl
|
||||
|
||||
from loguru import logger
|
||||
|
||||
|
||||
def _create_default_ssl_context() -> ssl.SSLContext:
|
||||
try:
|
||||
import certifi
|
||||
|
||||
return ssl.create_default_context(cafile=certifi.where())
|
||||
except ImportError:
|
||||
return ssl.create_default_context()
|
||||
|
||||
|
||||
try:
|
||||
_SSL_CONTEXT = _create_default_ssl_context()
|
||||
except Exception:
|
||||
_SSL_CONTEXT = ssl.create_default_context()
|
||||
|
||||
_PROXY_SSL_CONTEXT: ssl.SSLContext | None = None
|
||||
_PROFILE_SSL_CONTEXTS: dict[str, ssl.SSLContext] = {}
|
||||
|
||||
|
||||
def get_ssl_context() -> ssl.SSLContext:
|
||||
"""
|
||||
获取 SSL 上下文
|
||||
|
||||
优先使用 certifi 证书包,如果未安装则使用系统默认证书。
|
||||
返回模块级缓存的 SSL 上下文实例。
|
||||
|
||||
Returns:
|
||||
ssl.SSLContext: SSL 上下文
|
||||
"""
|
||||
return _SSL_CONTEXT
|
||||
|
||||
|
||||
def get_proxy_ssl_context(expected_fingerprint: str | None = None) -> ssl.SSLContext:
|
||||
"""
|
||||
获取用于代理连接的 SSL 上下文(连接 aether-proxy TLS 端口)
|
||||
|
||||
当前使用 CERT_NONE(不验证证书),因为 aether-proxy 使用自签名证书。
|
||||
expected_fingerprint 参数预留供未来实现指纹校验。
|
||||
|
||||
Args:
|
||||
expected_fingerprint: 预期的证书 SHA-256 指纹(hex,预留参数)
|
||||
|
||||
Returns:
|
||||
ssl.SSLContext: 代理专用 SSL 上下文
|
||||
"""
|
||||
global _PROXY_SSL_CONTEXT
|
||||
|
||||
if expected_fingerprint:
|
||||
logger.warning("TLS 证书指纹校验尚未实现, fingerprint={} 被忽略", expected_fingerprint)
|
||||
|
||||
if _PROXY_SSL_CONTEXT is None:
|
||||
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
ctx.check_hostname = False
|
||||
ctx.verify_mode = ssl.CERT_NONE
|
||||
_PROXY_SSL_CONTEXT = ctx
|
||||
# TODO: 实现基于 expected_fingerprint 的证书指纹校验
|
||||
return _PROXY_SSL_CONTEXT
|
||||
|
||||
|
||||
def _build_claude_code_ssl_context() -> ssl.SSLContext:
|
||||
"""构建 Claude Code best-effort TLS 配置。
|
||||
|
||||
说明:Python/OpenSSL 无法完整模拟 Node.js ClientHello。
|
||||
这里仅做可控项的尽力对齐(ALPN/TLS 版本/常见 cipher 偏好)。
|
||||
"""
|
||||
ctx = _create_default_ssl_context()
|
||||
ctx.minimum_version = ssl.TLSVersion.TLSv1_2
|
||||
try:
|
||||
ctx.maximum_version = ssl.TLSVersion.TLSv1_3
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
ctx.set_alpn_protocols(["h2", "http/1.1"])
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
ctx.set_ciphers(
|
||||
"ECDHE-ECDSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-RSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-ECDSA-AES256-GCM-SHA384:"
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:"
|
||||
"ECDHE-ECDSA-CHACHA20-POLY1305:"
|
||||
"ECDHE-RSA-CHACHA20-POLY1305"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return ctx
|
||||
|
||||
|
||||
def get_ssl_context_for_profile(tls_profile: str | None = None) -> ssl.SSLContext:
|
||||
"""按 profile 返回 SSL 上下文。"""
|
||||
profile = str(tls_profile or "").strip().lower()
|
||||
if not profile:
|
||||
return get_ssl_context()
|
||||
|
||||
if profile in _PROFILE_SSL_CONTEXTS:
|
||||
logger.debug("复用 TLS profile SSL context: {}", profile)
|
||||
return _PROFILE_SSL_CONTEXTS[profile]
|
||||
|
||||
if profile == "claude_code_nodejs":
|
||||
logger.info("启用 TLS profile: {}(best-effort)", profile)
|
||||
ctx = _build_claude_code_ssl_context()
|
||||
else:
|
||||
logger.warning("未知 TLS profile: {},回退默认 SSL context", profile)
|
||||
ctx = get_ssl_context()
|
||||
|
||||
_PROFILE_SSL_CONTEXTS[profile] = ctx
|
||||
return ctx
|
||||
241
_deprecated_py_src/utils/task_coordinator.py
Normal file
241
_deprecated_py_src/utils/task_coordinator.py
Normal file
@@ -0,0 +1,241 @@
|
||||
"""分布式任务协调器,确保仅有一个 worker 执行特定任务
|
||||
|
||||
锁清理策略:
|
||||
- 单实例模式(默认):启动时使用原子操作清理旧锁并获取新锁
|
||||
- 多实例模式:使用 NX 选项竞争锁,依赖 TTL 处理异常退出
|
||||
|
||||
使用方式:
|
||||
- 默认行为:启动时清理旧锁(适用于单机部署)
|
||||
- 多实例部署:设置 SINGLE_INSTANCE_MODE=false 禁用启动清理
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import pathlib
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
try:
|
||||
import fcntl # type: ignore
|
||||
except ImportError: # pragma: no cover - Windows 环境
|
||||
fcntl = None
|
||||
|
||||
|
||||
class StartupTaskCoordinator:
|
||||
"""利用 Redis 或文件锁,保证任务只在单个进程/实例中运行"""
|
||||
|
||||
# 类级别标记:在当前进程中是否已尝试过启动清理
|
||||
# 注意:这在 fork 模式下每个 worker 都是独立的
|
||||
_startup_cleanup_attempted = False
|
||||
|
||||
def __init__(self, redis_client: Any | None = None, lock_dir: str | None = None) -> None:
|
||||
self.redis = redis_client
|
||||
self._tokens: dict[str, str] = {}
|
||||
self._file_handles: dict[str, object] = {}
|
||||
self._refresh_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._lock_lost_callbacks: dict[str, Callable[[str], Awaitable[None] | None]] = {}
|
||||
self._lock_dir = pathlib.Path(lock_dir or os.getenv("TASK_LOCK_DIR", "./.locks"))
|
||||
if not self._lock_dir.exists():
|
||||
self._lock_dir.mkdir(parents=True, exist_ok=True)
|
||||
# 单实例模式:启动时清理旧锁(适用于单机部署,避免残留锁问题)
|
||||
self._single_instance_mode = os.getenv("SINGLE_INSTANCE_MODE", "true").lower() == "true"
|
||||
|
||||
def _redis_key(self, name: str) -> str:
|
||||
return f"task_lock:{name}"
|
||||
|
||||
async def acquire(self, name: str, ttl: int | None = None) -> bool:
|
||||
ttl = ttl or int(os.getenv("TASK_COORDINATOR_LOCK_TTL", "86400"))
|
||||
|
||||
if self.redis:
|
||||
token = str(uuid.uuid4())
|
||||
try:
|
||||
if self._single_instance_mode:
|
||||
# 单实例模式:使用 Lua 脚本原子性地"清理旧锁 + 竞争获取"
|
||||
# 只有当锁不存在或成功获取时才返回 1
|
||||
# 这样第一个执行的 worker 会清理旧锁并获取,后续 worker 会正常竞争
|
||||
script = """
|
||||
local key = KEYS[1]
|
||||
local token = ARGV[1]
|
||||
local ttl = tonumber(ARGV[2])
|
||||
local startup_key = KEYS[1] .. ':startup'
|
||||
|
||||
-- 检查是否已有 worker 执行过启动清理
|
||||
local cleaned = redis.call('GET', startup_key)
|
||||
if not cleaned then
|
||||
-- 第一个 worker:删除旧锁,标记已清理
|
||||
redis.call('DEL', key)
|
||||
redis.call('SET', startup_key, '1', 'EX', 60)
|
||||
end
|
||||
|
||||
-- 尝试获取锁(NX 模式)
|
||||
local result = redis.call('SET', key, token, 'NX', 'EX', ttl)
|
||||
if result then
|
||||
return 1
|
||||
end
|
||||
return 0
|
||||
"""
|
||||
result = await self.redis.eval(
|
||||
script, 2, self._redis_key(name), self._redis_key(name), token, ttl
|
||||
)
|
||||
if result == 1:
|
||||
self._tokens[name] = token
|
||||
self._start_refresh_task(name, ttl)
|
||||
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
||||
return True
|
||||
return False
|
||||
else:
|
||||
# 多实例模式:直接使用 NX 选项竞争锁
|
||||
acquired = await self.redis.set(self._redis_key(name), token, nx=True, ex=ttl)
|
||||
if acquired:
|
||||
self._tokens[name] = token
|
||||
self._start_refresh_task(name, ttl)
|
||||
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
||||
return True
|
||||
return False
|
||||
except Exception as exc: # pragma: no cover - Redis 异常回退
|
||||
logger.warning(f"Redis 锁获取失败,回退到文件锁: {exc}")
|
||||
|
||||
return await self._acquire_file_lock(name)
|
||||
|
||||
async def release(self, name: str) -> Any:
|
||||
refresh_task = self._refresh_tasks.pop(name, None)
|
||||
if refresh_task is not None:
|
||||
refresh_task.cancel()
|
||||
self._lock_lost_callbacks.pop(name, None)
|
||||
|
||||
if self.redis and name in self._tokens:
|
||||
token = self._tokens.pop(name)
|
||||
script = """
|
||||
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||||
return redis.call('DEL', KEYS[1])
|
||||
end
|
||||
return 0
|
||||
"""
|
||||
try:
|
||||
await self.redis.eval(script, 1, self._redis_key(name), token)
|
||||
except Exception as exc: # pragma: no cover
|
||||
logger.warning(f"释放 Redis 锁失败: {exc}")
|
||||
|
||||
handle = self._file_handles.pop(name, None)
|
||||
if handle and fcntl:
|
||||
try:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
finally:
|
||||
handle.close()
|
||||
|
||||
async def _acquire_file_lock(self, name: str) -> bool:
|
||||
if fcntl is None:
|
||||
# 在不支持 fcntl 的平台上退化为单进程锁
|
||||
if name in self._file_handles:
|
||||
return False
|
||||
self._file_handles[name] = object()
|
||||
logger.warning("操作系统不支持文件锁,任务锁仅在当前进程生效")
|
||||
return True
|
||||
|
||||
lock_path = self._lock_dir / f"{name}.lock"
|
||||
handle = open(lock_path, "a+")
|
||||
try:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
self._file_handles[name] = handle
|
||||
logger.info(f"任务 {name} 使用文件锁独占执行")
|
||||
return True
|
||||
except BlockingIOError:
|
||||
handle.close()
|
||||
return False
|
||||
|
||||
def register_lock_lost_callback(
|
||||
self,
|
||||
name: str,
|
||||
callback: Callable[[str], Awaitable[None] | None],
|
||||
) -> None:
|
||||
self._lock_lost_callbacks[name] = callback
|
||||
|
||||
def _start_refresh_task(self, name: str, ttl: int) -> None:
|
||||
if not self.redis:
|
||||
return
|
||||
|
||||
existing_task = self._refresh_tasks.pop(name, None)
|
||||
if existing_task is not None:
|
||||
existing_task.cancel()
|
||||
|
||||
self._refresh_tasks[name] = asyncio.create_task(self._refresh_lock_loop(name, ttl))
|
||||
|
||||
async def _refresh_lock(self, name: str, ttl: int) -> bool:
|
||||
token = self._tokens.get(name)
|
||||
if not self.redis or token is None:
|
||||
return False
|
||||
|
||||
script = """
|
||||
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||||
return redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]))
|
||||
end
|
||||
return 0
|
||||
"""
|
||||
result = await self.redis.eval(script, 1, self._redis_key(name), token, ttl)
|
||||
return result == 1
|
||||
|
||||
async def _refresh_lock_loop(self, name: str, ttl: int) -> None:
|
||||
interval = max(1, ttl // 3)
|
||||
max_consecutive_failures = 5
|
||||
consecutive_failures = 0
|
||||
try:
|
||||
while name in self._tokens:
|
||||
await asyncio.sleep(interval)
|
||||
if name not in self._tokens:
|
||||
break
|
||||
try:
|
||||
refreshed = await self._refresh_lock(name, ttl)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
consecutive_failures += 1
|
||||
if consecutive_failures >= max_consecutive_failures:
|
||||
logger.error(
|
||||
f"续租任务锁 {name} 连续失败 {consecutive_failures} 次,视为锁丢失: {exc}"
|
||||
)
|
||||
self._tokens.pop(name, None)
|
||||
await self._notify_lock_lost(name)
|
||||
break
|
||||
backoff = min(interval, 2**consecutive_failures)
|
||||
logger.warning(
|
||||
f"续租任务锁 {name} 失败 ({consecutive_failures}/{max_consecutive_failures}),"
|
||||
f"{backoff}s 后重试: {exc}"
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
|
||||
consecutive_failures = 0
|
||||
if not refreshed:
|
||||
logger.warning(f"任务 {name} 的 Redis 锁已失效,停止续租")
|
||||
self._tokens.pop(name, None)
|
||||
await self._notify_lock_lost(name)
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
|
||||
async def _notify_lock_lost(self, name: str) -> None:
|
||||
callback = self._lock_lost_callbacks.pop(name, None)
|
||||
if callback is None:
|
||||
return
|
||||
|
||||
try:
|
||||
result = callback(name)
|
||||
if result is not None:
|
||||
await result
|
||||
except Exception as exc: # pragma: no cover - 回调失败仅记录日志
|
||||
logger.exception("任务 {} 的失锁回调执行失败: {}", name, exc)
|
||||
|
||||
|
||||
async def ensure_singleton_task(
|
||||
name: str, redis_client: Any | None = None, ttl: int | None = None
|
||||
) -> Any:
|
||||
"""便捷协程,返回 (coordinator, acquired)"""
|
||||
|
||||
coordinator = StartupTaskCoordinator(redis_client)
|
||||
acquired = await coordinator.acquire(name, ttl=ttl)
|
||||
return coordinator, acquired
|
||||
223
_deprecated_py_src/utils/timeout.py
Normal file
223
_deprecated_py_src/utils/timeout.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
超时保护工具
|
||||
|
||||
为异步函数和操作提供超时保护
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import Any, TypeVar
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class AsyncTimeoutError(TimeoutError):
|
||||
"""异步操作超时错误"""
|
||||
|
||||
def __init__(self, message: str, operation: str, timeout: float):
|
||||
super().__init__(message)
|
||||
self.operation = operation
|
||||
self.timeout = timeout
|
||||
|
||||
|
||||
def with_timeout(seconds: float, operation_name: str | None = None) -> Any:
|
||||
"""
|
||||
装饰器:为异步函数添加超时保护
|
||||
|
||||
Args:
|
||||
seconds: 超时时间(秒)
|
||||
operation_name: 操作名称(用于日志,默认使用函数名)
|
||||
|
||||
Usage:
|
||||
@with_timeout(30.0)
|
||||
async def my_async_function():
|
||||
...
|
||||
|
||||
@with_timeout(60.0, operation_name="API请求")
|
||||
async def api_call():
|
||||
...
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
@wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
op_name = operation_name or func.__name__
|
||||
try:
|
||||
return await asyncio.wait_for(func(*args, **kwargs), timeout=seconds)
|
||||
except TimeoutError:
|
||||
logger.warning(f"操作超时: {op_name} (timeout={seconds}s)")
|
||||
raise AsyncTimeoutError(
|
||||
f"{op_name} 操作超时({seconds}秒)",
|
||||
operation=op_name,
|
||||
timeout=seconds,
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
async def run_with_timeout(
|
||||
coro: Any,
|
||||
timeout: float,
|
||||
operation_name: str = "operation",
|
||||
default: T | None = None,
|
||||
raise_on_timeout: bool = True,
|
||||
) -> T:
|
||||
"""
|
||||
为协程添加超时保护(函数式调用)
|
||||
|
||||
Args:
|
||||
coro: 协程对象
|
||||
timeout: 超时时间(秒)
|
||||
operation_name: 操作名称(用于日志)
|
||||
default: 超时时返回的默认值(仅在 raise_on_timeout=False 时有效)
|
||||
raise_on_timeout: 超时时是否抛出异常
|
||||
|
||||
Returns:
|
||||
协程的返回值,或超时时的默认值
|
||||
|
||||
Usage:
|
||||
result = await run_with_timeout(
|
||||
my_async_function(),
|
||||
timeout=30.0,
|
||||
operation_name="API请求"
|
||||
)
|
||||
"""
|
||||
try:
|
||||
return await asyncio.wait_for(coro, timeout=timeout)
|
||||
except TimeoutError:
|
||||
logger.warning(f"操作超时: {operation_name} (timeout={timeout}s)")
|
||||
if raise_on_timeout:
|
||||
raise AsyncTimeoutError(
|
||||
f"{operation_name} 操作超时({timeout}秒)",
|
||||
operation=operation_name,
|
||||
timeout=timeout,
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
class TimeoutContext:
|
||||
"""
|
||||
超时上下文管理器
|
||||
|
||||
Usage:
|
||||
async with TimeoutContext(30.0, "数据库查询") as ctx:
|
||||
result = await db.query(...)
|
||||
# 如果超过30秒会抛出 AsyncTimeoutError
|
||||
"""
|
||||
|
||||
def __init__(self, timeout: float, operation_name: str = "operation"):
|
||||
self.timeout = timeout
|
||||
self.operation_name = operation_name
|
||||
self._task: asyncio.Task | None = None
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
# asyncio.timeout 在 Python 3.11+ 可用
|
||||
# 这里使用更通用的方式
|
||||
pass
|
||||
|
||||
|
||||
async def with_timeout_context(timeout: float, operation_name: str = "operation") -> Any:
|
||||
"""
|
||||
超时上下文管理器(Python 3.11+ asyncio.timeout 的替代)
|
||||
|
||||
Usage:
|
||||
async with with_timeout_context(30.0, "API请求"):
|
||||
result = await api_call()
|
||||
"""
|
||||
try:
|
||||
# Python 3.11+ 使用内置的 asyncio.timeout
|
||||
return asyncio.timeout(timeout)
|
||||
except AttributeError:
|
||||
# Python 3.10 及以下版本的兼容实现
|
||||
# 注意:这个简单实现不支持嵌套取消
|
||||
pass
|
||||
|
||||
|
||||
async def read_first_chunk_with_ttfb_timeout(
|
||||
byte_iterator: Any,
|
||||
timeout: float,
|
||||
request_id: str,
|
||||
provider_name: str,
|
||||
) -> tuple[bytes, Any]:
|
||||
"""
|
||||
读取流的首字节并应用 TTFB 超时检测
|
||||
|
||||
首字节超时(Time To First Byte)用于检测慢响应的 Provider,
|
||||
超时时触发故障转移到其他可用的 Provider。
|
||||
|
||||
Args:
|
||||
byte_iterator: 异步字节流迭代器
|
||||
timeout: TTFB 超时时间(秒)
|
||||
request_id: 请求 ID(用于日志)
|
||||
provider_name: Provider 名称(用于日志和异常)
|
||||
|
||||
Returns:
|
||||
(first_chunk, aiter): 首个字节块和异步迭代器
|
||||
|
||||
Raises:
|
||||
ProviderTimeoutException: 如果首字节超时
|
||||
"""
|
||||
from src.core.exceptions import ProviderTimeoutException
|
||||
|
||||
aiter = byte_iterator.__aiter__()
|
||||
|
||||
try:
|
||||
first_chunk = await asyncio.wait_for(aiter.__anext__(), timeout=timeout)
|
||||
return first_chunk, aiter
|
||||
except TimeoutError:
|
||||
# 完整的资源清理:先关闭迭代器,再关闭底层响应
|
||||
await _cleanup_iterator_resources(aiter, request_id)
|
||||
logger.warning(
|
||||
f" [{request_id}] 流首字节超时 (TTFB): "
|
||||
f"Provider={provider_name}, timeout={timeout}s"
|
||||
)
|
||||
raise ProviderTimeoutException(
|
||||
provider_name=provider_name,
|
||||
timeout=int(timeout),
|
||||
)
|
||||
|
||||
|
||||
async def _cleanup_iterator_resources(aiter: Any, request_id: str) -> None:
|
||||
"""
|
||||
清理异步迭代器及其底层资源
|
||||
|
||||
确保在 TTFB 超时后正确释放 HTTP 连接,避免连接泄漏。
|
||||
|
||||
Args:
|
||||
aiter: 异步迭代器
|
||||
request_id: 请求 ID(用于日志)
|
||||
"""
|
||||
# 1. 关闭迭代器本身
|
||||
if hasattr(aiter, "aclose"):
|
||||
try:
|
||||
await aiter.aclose()
|
||||
except Exception as e:
|
||||
logger.debug(f" [{request_id}] 关闭迭代器失败: {e}")
|
||||
|
||||
# 2. 关闭底层响应对象(httpx.Response)
|
||||
# 迭代器可能持有 _response 属性指向底层响应
|
||||
response = getattr(aiter, "_response", None)
|
||||
if response is not None and hasattr(response, "aclose"):
|
||||
try:
|
||||
await response.aclose()
|
||||
except Exception as e:
|
||||
logger.debug(f" [{request_id}] 关闭底层响应失败: {e}")
|
||||
|
||||
# 3. 尝试关闭 httpx 流(如果迭代器是 httpx 的 aiter_bytes)
|
||||
# httpx 的 Response.aiter_bytes() 返回的生成器可能有 _stream 属性
|
||||
stream = getattr(aiter, "_stream", None)
|
||||
if stream is not None and hasattr(stream, "aclose"):
|
||||
try:
|
||||
await stream.aclose()
|
||||
except Exception as e:
|
||||
logger.debug(f" [{request_id}] 关闭流对象失败: {e}")
|
||||
308
_deprecated_py_src/utils/transaction_manager.py
Normal file
308
_deprecated_py_src/utils/transaction_manager.py
Normal file
@@ -0,0 +1,308 @@
|
||||
"""
|
||||
数据库事务管理工具
|
||||
提供事务装饰器和事务上下文管理器
|
||||
支持同步和异步函数
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.exc import DatabaseError, IntegrityError
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class TransactionError(Exception):
|
||||
"""事务处理异常"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def _find_db_session(args: Any, kwargs: Any) -> Session | None:
|
||||
"""从参数中查找数据库会话"""
|
||||
# 从位置参数中查找Session
|
||||
for arg in args:
|
||||
if isinstance(arg, Session):
|
||||
return arg
|
||||
|
||||
# 从关键字参数中查找Session
|
||||
for value in kwargs.values():
|
||||
if isinstance(value, Session):
|
||||
return value
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def transactional(commit: bool = True, rollback_on_error: bool = True) -> Any:
|
||||
"""
|
||||
事务装饰器,支持同步和异步函数
|
||||
|
||||
Args:
|
||||
commit: 是否在成功时自动提交,默认True
|
||||
rollback_on_error: 是否在错误时自动回滚,默认True
|
||||
|
||||
Usage:
|
||||
@transactional()
|
||||
def create_user_with_api_key(db: Session, ...):
|
||||
# 同步方法会在事务中执行
|
||||
pass
|
||||
|
||||
@transactional()
|
||||
async def create_user_async(db: Session, ...):
|
||||
# 异步方法也会在事务中执行
|
||||
pass
|
||||
"""
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
# 检查是否是异步函数
|
||||
if inspect.iscoroutinefunction(func):
|
||||
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
db_session = _find_db_session(args, kwargs)
|
||||
|
||||
if not db_session:
|
||||
raise TransactionError(
|
||||
f"No SQLAlchemy Session found in arguments for {func.__name__}"
|
||||
)
|
||||
|
||||
# 检查是否已经在事务中
|
||||
if db_session.in_transaction():
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
transaction_id = f"{func.__module__}.{func.__name__}"
|
||||
logger.debug(f"开始异步事务: {transaction_id}")
|
||||
|
||||
try:
|
||||
result = await func(*args, **kwargs)
|
||||
|
||||
if commit:
|
||||
db_session.commit()
|
||||
logger.debug(f"异步事务提交成功: {transaction_id}")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
if rollback_on_error:
|
||||
try:
|
||||
db_session.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(
|
||||
f"异步事务回滚: {transaction_id} - {type(e).__name__}: {str(e)}"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"异步事务异常(未回滚): {transaction_id} - {type(e).__name__}: {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
return async_wrapper
|
||||
else:
|
||||
|
||||
@functools.wraps(func)
|
||||
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
db_session = _find_db_session(args, kwargs)
|
||||
|
||||
if not db_session:
|
||||
raise TransactionError(
|
||||
f"No SQLAlchemy Session found in arguments for {func.__name__}"
|
||||
)
|
||||
|
||||
# 检查是否已经在事务中
|
||||
if db_session.in_transaction():
|
||||
return func(*args, **kwargs)
|
||||
|
||||
transaction_id = f"{func.__module__}.{func.__name__}"
|
||||
logger.debug(f"开始事务: {transaction_id}")
|
||||
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
|
||||
if commit:
|
||||
db_session.commit()
|
||||
logger.debug(f"事务提交成功: {transaction_id}")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
if rollback_on_error:
|
||||
try:
|
||||
db_session.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"事务回滚: {transaction_id} - {type(e).__name__}: {str(e)}")
|
||||
else:
|
||||
logger.error(
|
||||
f"事务异常(未回滚): {transaction_id} - {type(e).__name__}: {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
return sync_wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@contextmanager
|
||||
def transaction_scope(
|
||||
db: Session,
|
||||
commit_on_success: bool = True,
|
||||
rollback_on_error: bool = True,
|
||||
operation_name: str | None = None,
|
||||
) -> Generator[Session]:
|
||||
"""
|
||||
事务上下文管理器
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
commit_on_success: 成功时是否自动提交
|
||||
rollback_on_error: 失败时是否自动回滚
|
||||
operation_name: 操作名称,用于日志
|
||||
|
||||
Usage:
|
||||
with transaction_scope(db, operation_name="create_user") as tx:
|
||||
user = User(...)
|
||||
tx.add(user)
|
||||
# 自动提交或回滚
|
||||
"""
|
||||
operation_name = operation_name or "database_operation"
|
||||
|
||||
# 检查是否已经在事务中
|
||||
if db.in_transaction():
|
||||
# 已经在事务中,直接返回session
|
||||
logger.debug(f"使用现有事务: {operation_name}")
|
||||
yield db
|
||||
return
|
||||
|
||||
logger.debug(f"开始事务范围: {operation_name}")
|
||||
|
||||
try:
|
||||
yield db
|
||||
|
||||
if commit_on_success:
|
||||
db.commit()
|
||||
logger.debug(f"事务范围提交成功: {operation_name}")
|
||||
|
||||
except Exception as e:
|
||||
if rollback_on_error:
|
||||
db.rollback()
|
||||
logger.error(f"事务范围回滚: {operation_name} - {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def retry_on_database_error(max_retries: int = 3, delay: float = 0.1) -> Any:
|
||||
"""
|
||||
数据库错误重试装饰器
|
||||
|
||||
Args:
|
||||
max_retries: 最大重试次数
|
||||
delay: 重试延迟(秒)
|
||||
"""
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
import random
|
||||
import time
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
except (DatabaseError, IntegrityError) as e:
|
||||
if attempt < max_retries - 1:
|
||||
# 随机化延迟,避免多个请求同时重试
|
||||
actual_delay = delay * (2**attempt) + random.uniform(0, 0.1)
|
||||
logger.warning(
|
||||
f"数据库操作失败,{actual_delay:.2f}秒后重试 (尝试 {attempt + 1}/{max_retries}): {str(e)}"
|
||||
)
|
||||
time.sleep(actual_delay)
|
||||
continue
|
||||
else:
|
||||
logger.error(
|
||||
f"数据库操作最终失败,已达最大重试次数({max_retries}): {func.__name__} - {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class BatchOperation:
|
||||
"""
|
||||
批量操作管理器
|
||||
用于处理大量数据插入/更新操作
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session, batch_size: int = 100):
|
||||
self.db = db
|
||||
self.batch_size = batch_size
|
||||
self.operations = []
|
||||
self.operation_count = 0
|
||||
|
||||
def add(self, obj: Any) -> Any:
|
||||
"""添加对象到批处理"""
|
||||
self.operations.append(("add", obj))
|
||||
self.operation_count += 1
|
||||
|
||||
if self.operation_count >= self.batch_size:
|
||||
self.flush()
|
||||
|
||||
def update(self, obj: Any) -> Any:
|
||||
"""添加更新操作到批处理"""
|
||||
self.operations.append(("merge", obj))
|
||||
self.operation_count += 1
|
||||
|
||||
if self.operation_count >= self.batch_size:
|
||||
self.flush()
|
||||
|
||||
def flush(self) -> Any:
|
||||
"""执行当前批次的所有操作"""
|
||||
if not self.operations:
|
||||
return
|
||||
|
||||
logger.debug(f"执行批量操作: {len(self.operations)} 项")
|
||||
|
||||
try:
|
||||
for operation, obj in self.operations:
|
||||
if operation == "add":
|
||||
self.db.add(obj)
|
||||
elif operation == "merge":
|
||||
self.db.merge(obj)
|
||||
|
||||
self.db.flush() # 只flush,不提交
|
||||
logger.debug(f"批量操作flush完成: {len(self.operations)} 项")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"批量操作失败({len(self.operations)}项): {type(e).__name__}: {str(e)}")
|
||||
raise
|
||||
|
||||
finally:
|
||||
# 清空操作列表
|
||||
self.operations.clear()
|
||||
self.operation_count = 0
|
||||
|
||||
def commit(self) -> Any:
|
||||
"""提交所有操作"""
|
||||
self.flush() # 确保所有操作都已flush
|
||||
self.db.commit()
|
||||
logger.debug("批量操作提交完成")
|
||||
|
||||
def __enter__(self) -> None:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
if exc_type is None:
|
||||
# 正常退出,提交事务
|
||||
self.commit()
|
||||
else:
|
||||
# 异常退出,回滚事务
|
||||
self.db.rollback()
|
||||
logger.error(f"批量操作异常退出,已回滚: {str(exc_val)}")
|
||||
40
_deprecated_py_src/utils/url_utils.py
Normal file
40
_deprecated_py_src/utils/url_utils.py
Normal file
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
URL 处理工具函数
|
||||
|
||||
提供 URL 模式检测和处理功能。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
def is_official_openai_api_url(base_url: str | None) -> bool:
|
||||
"""判断是否为 OpenAI 官方 API 端点。"""
|
||||
value = str(base_url or "").strip()
|
||||
if not value:
|
||||
return False
|
||||
|
||||
parsed = urlparse(value if "://" in value else f"https://{value}")
|
||||
host = str(parsed.hostname or "").strip().lower()
|
||||
return host == "api.openai.com"
|
||||
|
||||
|
||||
def is_codex_url(base_url: str) -> bool:
|
||||
"""判断是否是 Codex OAuth 端点。
|
||||
|
||||
Codex OAuth 端点(如 chatgpt.com/backend-api/codex)不走标准 /v1 前缀,
|
||||
直接使用 /responses 而非 /v1/responses。
|
||||
|
||||
Args:
|
||||
base_url: 端点基础 URL
|
||||
|
||||
Returns:
|
||||
bool: 是否是 Codex 端点
|
||||
"""
|
||||
url = base_url.rstrip("/")
|
||||
return (
|
||||
"/backend-api/codex" in url
|
||||
or "/backendapi/codex" in url
|
||||
or url.endswith("/codex")
|
||||
)
|
||||
Reference in New Issue
Block a user