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:
fawney19
2026-04-03 16:26:16 +08:00
parent 8f26e1a31f
commit 1d9c77522a
868 changed files with 1735 additions and 2433 deletions

View File

@@ -0,0 +1,3 @@
from .timeout import AsyncTimeoutError, run_with_timeout, with_timeout
__all__ = ["with_timeout", "run_with_timeout", "AsyncTimeoutError"]

View 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

View 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

View 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] 是 selfargs[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

View 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

View 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"
)

View 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)

View 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"

View 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]

View 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

View 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

View 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}")

View 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)}")

View 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")
)