feat(codex): 增强 OAuth 导入解析与账号信息提取,优化维护调度器线程模型

Codex OAuth:
- 导入解析支持附加账号字段(account_id/plan_type/user_id/email)
- enrich_codex 扩展从 access_token 和直接字段提取账号信息
- parse_codex_id_token 支持 JWT/JSON 字符串/dict 三种输入格式
- request patching 新增 openai:compact 格式支持
- codex_usage_parser 新增 credits_unlimited 字段解析

维护调度器:
- 同步 DB 操作迁移到线程池执行,避免阻塞事件循环
- 新增 request_candidates 定期清理任务
- 新增每周 VACUUM ANALYZE 数据库表维护任务
- 新增 enable_db_maintenance 配置项

前端:
- ElapsedTimeText 从 setInterval 改为 requestAnimationFrame
- 时间线过滤 available/unused 占位记录
- Usage 页面默认关闭全局自动刷新
- 移除号池管理中的配额更新时间显示
This commit is contained in:
fawney19
2026-03-04 10:00:38 +08:00
parent e181329a81
commit 82a9fb3c39
17 changed files with 822 additions and 317 deletions

View File

@@ -1301,62 +1301,115 @@ async def complete_provider_oauth(
# ==============================================================================
def _parse_tokens_input(raw_input: str) -> list[str]:
def _coerce_import_str(value: Any) -> str | None:
if not isinstance(value, str):
return None
normalized = value.strip()
return normalized or None
def _extract_standard_oauth_import_entry(item: Any) -> dict[str, str] | None:
if isinstance(item, str):
token = _coerce_import_str(item)
if token:
return {"refresh_token": token}
return None
if not isinstance(item, dict):
return None
refresh_token = _coerce_import_str(item.get("refresh_token")) or _coerce_import_str(
item.get("refreshToken")
)
if not refresh_token:
return None
entry: dict[str, str] = {"refresh_token": refresh_token}
account_id = (
_coerce_import_str(item.get("account_id"))
or _coerce_import_str(item.get("accountId"))
or _coerce_import_str(item.get("chatgpt_account_id"))
or _coerce_import_str(item.get("chatgptAccountId"))
)
if account_id:
entry["account_id"] = account_id
plan_type = (
_coerce_import_str(item.get("plan_type"))
or _coerce_import_str(item.get("planType"))
or _coerce_import_str(item.get("chatgpt_plan_type"))
or _coerce_import_str(item.get("chatgptPlanType"))
)
if plan_type:
entry["plan_type"] = plan_type.lower()
user_id = (
_coerce_import_str(item.get("user_id"))
or _coerce_import_str(item.get("userId"))
or _coerce_import_str(item.get("chatgpt_user_id"))
or _coerce_import_str(item.get("chatgptUserId"))
)
if user_id:
entry["user_id"] = user_id
email = _coerce_import_str(item.get("email"))
if email:
entry["email"] = email
return entry
def _parse_standard_oauth_import_entries(raw_input: str) -> list[dict[str, str]]:
"""
解析通用 Token 导入输入,支持多种格式
解析标准 OAuth 导入输入,保留 refresh_token 及可用账号提示字段
支持的格式:
1. 单个 Token 字符串
2. JSON 字符串数组: ["token1", "token2", ...]
3. JSON 对象数组: [{"refresh_token": "token1", ...}, ...]
3. JSON 对象数组: [{"refresh_token": "token1", "account_id": "...", ...}, ...]
4. 单个 JSON 对象: {"refresh_token": "token1", ...}
5. 纯 Token 导入(一行一个): "token1\\ntoken2\\ntoken3"
返回: Token 字符串列表
"""
raw = raw_input.strip()
if not raw:
return []
result: list[str] = []
result: list[dict[str, str]] = []
# 尝试解析为 JSON 数组
if raw.startswith("["):
try:
parsed = json.loads(raw)
if isinstance(parsed, list):
for item in parsed:
if isinstance(item, str) and item.strip():
result.append(item.strip())
elif isinstance(item, dict):
token = item.get("refresh_token", "")
if isinstance(token, str) and token.strip():
result.append(token.strip())
return result
except json.JSONDecodeError:
pass # 不是有效 JSON继续尝试其他格式
parsed = None
if isinstance(parsed, list):
for item in parsed:
entry = _extract_standard_oauth_import_entry(item)
if entry:
result.append(entry)
return result
# 单个 JSON 对象
if raw.startswith("{"):
try:
parsed = json.loads(raw)
if isinstance(parsed, dict):
token = parsed.get("refresh_token", "")
if isinstance(token, str) and token.strip():
return [token.strip()]
except json.JSONDecodeError:
pass
parsed = None
if isinstance(parsed, dict):
entry = _extract_standard_oauth_import_entry(parsed)
return [entry] if entry else []
# 纯 Token 导入(一行一个)
lines = raw.splitlines()
for line in lines:
for line in raw.splitlines():
token = line.strip()
if token and not token.startswith("#"): # 忽略空行和注释行
result.append(token)
if token and not token.startswith("#"):
result.append({"refresh_token": token})
return result
def _parse_tokens_input(raw_input: str) -> list[str]:
"""兼容旧逻辑:仅返回 refresh_token 列表。"""
return [entry["refresh_token"] for entry in _parse_standard_oauth_import_entries(raw_input)]
def _parse_kiro_import_input(raw_input: str) -> list[dict[str, Any]]:
"""
解析 Kiro 凭据导入输入。
@@ -1581,7 +1634,15 @@ def _task_state_to_response(state: dict[str, Any]) -> BatchImportTaskStatusRespo
def _estimate_batch_import_total(provider_type: str, raw_credentials: str) -> int:
if provider_type == ProviderType.KIRO.value:
return len(_parse_kiro_import_input(raw_credentials))
return len(_parse_tokens_input(raw_credentials))
return len(_parse_standard_oauth_import_entries(raw_credentials))
def _apply_codex_import_hints(auth_config: dict[str, Any], import_entry: dict[str, str]) -> None:
"""将导入文件中可用的 Codex 账号信息作为兜底补全(不覆盖已有值)。"""
for field in ("account_id", "plan_type", "user_id", "email"):
value = import_entry.get(field)
if value and not auth_config.get(field):
auth_config[field] = value
@router.post(
@@ -1823,8 +1884,8 @@ async def _batch_import_standard_oauth_internal(
template = _require_oauth_template(provider_type)
timeout_seconds = _resolve_batch_import_timeout_seconds(proxy_config)
tokens = _parse_tokens_input(raw_credentials)
if not tokens:
import_entries = _parse_standard_oauth_import_entries(raw_credentials)
if not import_entries:
raise InvalidRequestException("未找到有效的 Token 数据")
api_formats = _get_provider_api_formats(provider)
@@ -1836,7 +1897,8 @@ async def _batch_import_standard_oauth_internal(
success_count = 0
failed_count = 0
for idx, refresh_token in enumerate(tokens):
for idx, import_entry in enumerate(import_entries):
refresh_token = import_entry.get("refresh_token", "")
result_item: BatchImportResultItem
try:
if not refresh_token or len(refresh_token) < 10:
@@ -1895,7 +1957,7 @@ async def _batch_import_standard_oauth_internal(
results.append(result_item)
if progress_hook is not None:
await progress_hook(
len(tokens),
len(import_entries),
idx + 1,
success_count,
failed_count,
@@ -1923,7 +1985,7 @@ async def _batch_import_standard_oauth_internal(
results.append(result_item)
if progress_hook is not None:
await progress_hook(
len(tokens),
len(import_entries),
idx + 1,
success_count,
failed_count,
@@ -1945,7 +2007,7 @@ async def _batch_import_standard_oauth_internal(
results.append(result_item)
if progress_hook is not None:
await progress_hook(
len(tokens),
len(import_entries),
idx + 1,
success_count,
failed_count,
@@ -1981,6 +2043,9 @@ async def _batch_import_standard_oauth_internal(
except Exception as exc:
logger.warning("批量导入: enrich_auth_config 失败 (index={}): {}", idx, exc)
if provider_type == ProviderType.CODEX.value:
_apply_codex_import_hints(auth_config, import_entry)
try:
existing_key = _check_duplicate_oauth_account(db, provider_id, auth_config)
except InvalidRequestException as exc:
@@ -1993,7 +2058,7 @@ async def _batch_import_standard_oauth_internal(
results.append(result_item)
if progress_hook is not None:
await progress_hook(
len(tokens),
len(import_entries),
idx + 1,
success_count,
failed_count,
@@ -2053,7 +2118,9 @@ async def _batch_import_standard_oauth_internal(
results.append(result_item)
if progress_hook is not None:
await progress_hook(len(tokens), idx + 1, success_count, failed_count, result_item)
await progress_hook(
len(import_entries), idx + 1, success_count, failed_count, result_item
)
if success_count > 0:
db.commit()
@@ -2063,12 +2130,12 @@ async def _batch_import_standard_oauth_internal(
provider_id,
provider_type,
success_count,
len(tokens),
len(import_entries),
failed_count,
)
return BatchImportResponse(
total=len(tokens),
total=len(import_entries),
success=success_count,
failed=failed_count,
results=results,

View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import json
from typing import Any, Awaitable, Callable
from urllib.parse import urlsplit, urlunsplit
@@ -208,51 +209,130 @@ async def post_oauth_token(
)
def parse_codex_id_token(id_token: str | None) -> dict[str, Any]:
"""Parse Codex id_token WITHOUT signature verification.
def _as_non_empty_str(value: Any) -> str | None:
if not isinstance(value, str):
return None
text = value.strip()
return text or None
Extract from claim `https://api.openai.com/auth`:
- email: claim `email`
- account_id: `chatgpt_account_id`
- plan_type: `chatgpt_plan_type` (e.g. "plus", "free", "team", "enterprise")
- user_id: `chatgpt_user_id`
Return dict with extracted fields. On any failure returns empty dict.
"""
def _first_non_empty_str(values: list[Any]) -> str | None:
for value in values:
text = _as_non_empty_str(value)
if text:
return text
return None
if not id_token:
return {}
def _decode_unverified_jwt_payload(token: str) -> dict[str, Any] | None:
try:
claims = jwt.decode(
id_token,
token,
options={
"verify_signature": False,
"verify_aud": False,
},
)
result: dict[str, Any] = {}
email = claims.get("email")
if isinstance(email, str) and email:
result["email"] = email
auth_info = claims.get("https://api.openai.com/auth") or {}
if isinstance(auth_info, dict):
account_id = auth_info.get("chatgpt_account_id")
if isinstance(account_id, str) and account_id:
result["account_id"] = account_id
plan_type = auth_info.get("chatgpt_plan_type")
if isinstance(plan_type, str) and plan_type:
result["plan_type"] = plan_type
user_id = auth_info.get("chatgpt_user_id")
if isinstance(user_id, str) and user_id:
result["user_id"] = user_id
return result
except Exception:
return None
return claims if isinstance(claims, dict) else None
def _extract_codex_fields_from_claims(claims: dict[str, Any]) -> dict[str, Any]:
auth_info = claims.get("https://api.openai.com/auth")
auth = auth_info if isinstance(auth_info, dict) else {}
result: dict[str, Any] = {}
email = _first_non_empty_str(
[
claims.get("email"),
auth.get("email"),
]
)
if email:
result["email"] = email
account_id = _first_non_empty_str(
[
auth.get("chatgpt_account_id"),
auth.get("chatgptAccountId"),
auth.get("account_id"),
auth.get("accountId"),
claims.get("chatgpt_account_id"),
claims.get("chatgptAccountId"),
claims.get("account_id"),
claims.get("accountId"),
]
)
if account_id:
result["account_id"] = account_id
plan_type = _first_non_empty_str(
[
auth.get("chatgpt_plan_type"),
auth.get("chatgptPlanType"),
auth.get("plan_type"),
auth.get("planType"),
claims.get("chatgpt_plan_type"),
claims.get("chatgptPlanType"),
claims.get("plan_type"),
claims.get("planType"),
]
)
if plan_type:
result["plan_type"] = plan_type
user_id = _first_non_empty_str(
[
auth.get("chatgpt_user_id"),
auth.get("chatgptUserId"),
auth.get("user_id"),
auth.get("userId"),
claims.get("chatgpt_user_id"),
claims.get("chatgptUserId"),
claims.get("user_id"),
claims.get("userId"),
claims.get("sub"),
]
)
if user_id:
result["user_id"] = user_id
return result
def parse_codex_id_token(id_token: Any) -> dict[str, Any]:
"""Parse Codex token payload without signature verification.
Supports:
- JWT string (typical id_token / sometimes access_token)
- JSON string containing claims
- Already-decoded dict payload
"""
claims: dict[str, Any] | None = None
if isinstance(id_token, dict):
claims = id_token
else:
token_text = _as_non_empty_str(id_token)
if not token_text:
return {}
if token_text.startswith("{"):
try:
payload = json.loads(token_text)
except Exception:
payload = None
if isinstance(payload, dict):
claims = payload
if claims is None:
claims = _decode_unverified_jwt_payload(token_text)
if not isinstance(claims, dict):
return {}
return _extract_codex_fields_from_claims(claims)
async def fetch_google_email(

View File

@@ -82,17 +82,68 @@ async def enrich_codex(
"""Codex auth_config enrichment: parse id_token -> email/account_id/plan_type/user_id."""
from src.core.provider_oauth_utils import parse_codex_id_token
id_token = token_response.get("id_token")
def _read_non_empty_str(*values: Any) -> str | None:
for value in values:
if isinstance(value, str):
normalized = value.strip()
if normalized:
return normalized
return None
# Prefer explicit fields if token endpoint returns them directly.
direct_account_id = _read_non_empty_str(
token_response.get("account_id"),
token_response.get("accountId"),
token_response.get("chatgpt_account_id"),
token_response.get("chatgptAccountId"),
)
if direct_account_id and not auth_config.get("account_id"):
auth_config["account_id"] = direct_account_id
direct_plan_type = _read_non_empty_str(
token_response.get("plan_type"),
token_response.get("planType"),
token_response.get("chatgpt_plan_type"),
token_response.get("chatgptPlanType"),
)
if direct_plan_type and not auth_config.get("plan_type"):
auth_config["plan_type"] = direct_plan_type
direct_user_id = _read_non_empty_str(
token_response.get("user_id"),
token_response.get("userId"),
token_response.get("chatgpt_user_id"),
token_response.get("chatgptUserId"),
)
if direct_user_id and not auth_config.get("user_id"):
auth_config["user_id"] = direct_user_id
direct_email = _read_non_empty_str(token_response.get("email"))
if direct_email and not auth_config.get("email"):
auth_config["email"] = direct_email
logger.debug(
"Codex enrich_auth_config: id_token_present={} token_keys={}",
bool(id_token),
"Codex enrich_auth_config: id_token_present={} access_token_present={} token_keys={}",
bool(token_response.get("id_token") or token_response.get("idToken")),
bool(token_response.get("access_token") or token_response.get("accessToken")),
list(token_response.keys()),
)
# parse_codex_id_token 仅返回非空有效字段,直接 update 即可
codex_info = parse_codex_id_token(id_token)
if codex_info:
logger.debug("Codex parsed id_token fields: {}", list(codex_info.keys()))
auth_config.update(codex_info)
token_candidates = [
token_response.get("id_token"),
token_response.get("idToken"),
token_response.get("access_token"),
token_response.get("accessToken"),
]
for token_payload in token_candidates:
codex_info = parse_codex_id_token(token_payload)
if not codex_info:
continue
logger.debug("Codex parsed token fields: {}", list(codex_info.keys()))
for key, value in codex_info.items():
if not auth_config.get(key):
auth_config[key] = value
return auth_config

View File

@@ -40,7 +40,7 @@ def maybe_patch_request_for_codex(
"""
if (provider_type or "").lower() != ProviderType.CODEX:
return request_body
if (provider_api_format or "").lower() not in {"openai:cli"}:
if (provider_api_format or "").lower() not in {"openai:cli", "openai:compact"}:
return request_body
if not isinstance(request_body, dict):
return request_body

View File

@@ -264,9 +264,14 @@ def parse_codex_wham_usage_response(data: dict[str, Any]) -> dict[str, Any] | No
has_credits = credits.get("has_credits")
if has_credits is not None:
result["has_credits"] = _coerce_bool(has_credits, "credits.has_credits")
balance = credits.get("balance")
if balance is not None:
result["credits_balance"] = _coerce_float(balance, "credits.balance")
credits_balance = _coerce_optional_float(credits.get("balance"), "credits.balance")
if credits_balance is not None:
result["credits_balance"] = credits_balance
credits_unlimited = _coerce_optional_bool(credits.get("unlimited"), "credits.unlimited")
if credits_unlimited is not None:
result["credits_unlimited"] = credits_unlimited
# 添加更新时间戳
if result:

View File

@@ -208,6 +208,10 @@ class SystemConfigService:
"value": 30,
"description": "审计日志保留天数,超过此天数的审计日志将被自动清理",
},
"enable_db_maintenance": {
"value": True,
"description": "是否启用数据库表维护任务(定期 VACUUM ANALYZE 防止表和索引膨胀)",
},
# 系统代理
"system_proxy_node_id": {
"value": None,

View File

@@ -9,6 +9,8 @@
- 连接池监控:定期检查数据库连接池状态
- Pending 状态清理:清理异常的 Pending 状态记录
- Gemini 文件映射清理:清理过期的 Gemini 文件→Key 映射
- 请求候选记录清理:定期清理过期的 request_candidates 记录
- 数据库表维护:定期 VACUUM ANALYZE 防止表和索引膨胀
使用 APScheduler 进行任务调度,支持时区配置。
"""
@@ -19,12 +21,11 @@ import asyncio
from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import delete
from sqlalchemy.orm import Session
from sqlalchemy import delete, text
from src.core.logger import logger
from src.database import create_session
from src.models.database import ApiKey, AuditLog, Provider, Usage
from src.models.database import ApiKey, AuditLog, Provider, RequestCandidate, Usage
from src.services.provider_ops.service import ProviderOpsService
from src.services.system.config import SystemConfigService
from src.services.system.scheduler import get_scheduler
@@ -280,6 +281,25 @@ class MaintenanceScheduler:
name="Gemini文件映射清理",
)
# 请求候选记录清理 - 凌晨 3:30 执行
scheduler.add_cron_job(
self._scheduled_candidate_cleanup,
hour=3,
minute=30,
job_id="candidate_cleanup",
name="请求候选记录清理",
)
# 数据库表维护 - 每周日凌晨 5 点执行 VACUUM ANALYZE
scheduler.add_cron_job(
self._scheduled_db_maintenance,
day_of_week="sun",
hour=5,
minute=0,
job_id="db_maintenance",
name="数据库表维护",
)
# Antigravity User-Agent 版本刷新 - 每 6 小时
scheduler.add_interval_job(
self._scheduled_antigravity_ua_refresh,
@@ -389,6 +409,14 @@ class MaintenanceScheduler:
"""审计日志清理任务(定时调用)"""
await self._perform_audit_cleanup()
async def _scheduled_candidate_cleanup(self) -> None:
"""请求候选记录清理任务(定时调用)"""
await self._perform_candidate_cleanup()
async def _scheduled_db_maintenance(self) -> None:
"""数据库表维护任务(定时调用)"""
await self._perform_db_maintenance()
async def _scheduled_gemini_file_mapping_cleanup(self) -> None:
"""Gemini 文件映射清理任务(定时调用)"""
await self._perform_gemini_file_mapping_cleanup()
@@ -614,92 +642,99 @@ class MaintenanceScheduler:
async def _perform_pending_cleanup(self) -> None:
"""执行 pending 状态清理"""
db = create_session()
try:
from src.services.usage.service import UsageService
# 获取配置的超时时间(默认 10 分钟)
timeout_minutes = SystemConfigService.get_config(
db, "pending_request_timeout_minutes", 10
)
def _do_pending_cleanup() -> int:
db = create_session()
try:
from src.services.usage.service import UsageService
# 执行清理
cleaned_count = UsageService.cleanup_stale_pending_requests(
db, timeout_minutes=timeout_minutes
)
timeout_minutes = SystemConfigService.get_config(
db, "pending_request_timeout_minutes", 10
)
return UsageService.cleanup_stale_pending_requests(
db, timeout_minutes=timeout_minutes
)
except Exception as e:
logger.exception(f"清理 pending 请求失败: {e}")
db.rollback()
return 0
finally:
db.close()
if cleaned_count > 0:
logger.info(f"清理了 {cleaned_count} 条超时的 pending/streaming 请求")
except Exception as e:
logger.exception(f"清理 pending 请求失败: {e}")
db.rollback()
finally:
db.close()
loop = asyncio.get_running_loop()
cleaned_count = await loop.run_in_executor(None, _do_pending_cleanup)
if cleaned_count > 0:
logger.info(f"清理了 {cleaned_count} 条超时的 pending/streaming 请求")
async def _perform_audit_cleanup(self) -> None:
"""执行审计日志清理任务"""
db = create_session()
try:
# 检查是否启用自动清理
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
logger.info("自动清理已禁用,跳过审计日志清理")
return
# 获取审计日志保留天数(默认 30 天,最少 7 天)
audit_retention_days = max(
SystemConfigService.get_config(db, "audit_log_retention_days", 30),
7, # 最少保留 7 天,防止误配置删除所有审计日志
)
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
cutoff_time = datetime.now(timezone.utc) - timedelta(days=audit_retention_days)
logger.info(f"开始清理 {audit_retention_days} 天前的审计日志...")
total_deleted = 0
while True:
# 先查询要删除的记录 ID分批
records_to_delete = (
db.query(AuditLog.id)
.filter(AuditLog.created_at < cutoff_time)
.limit(batch_size)
.all()
)
if not records_to_delete:
break
record_ids = [r.id for r in records_to_delete]
# 执行删除
result = db.execute(
delete(AuditLog)
.where(AuditLog.id.in_(record_ids))
.execution_options(synchronize_session=False)
)
rows_deleted = result.rowcount
db.commit()
total_deleted += rows_deleted
logger.debug(f"已删除 {rows_deleted} 条审计日志,累计 {total_deleted}")
await asyncio.sleep(0.1)
if total_deleted > 0:
logger.info(f"审计日志清理完成,共删除 {total_deleted} 条记录")
else:
logger.info("无需清理的审计日志")
except Exception as e:
logger.exception(f"审计日志清理失败: {e}")
def _do_audit_cleanup() -> int:
db = create_session()
try:
db.rollback()
except Exception:
pass
finally:
db.close()
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
logger.info("自动清理已禁用,跳过审计日志清理")
return 0
audit_retention_days = max(
SystemConfigService.get_config(db, "audit_log_retention_days", 30),
7,
)
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
cutoff_time = datetime.now(timezone.utc) - timedelta(days=audit_retention_days)
logger.info(f"开始清理 {audit_retention_days} 天前的审计日志...")
total_deleted = 0
while True:
batch_db = create_session()
try:
records_to_delete = (
batch_db.query(AuditLog.id)
.filter(AuditLog.created_at < cutoff_time)
.limit(batch_size)
.all()
)
if not records_to_delete:
break
record_ids = [r.id for r in records_to_delete]
result = batch_db.execute(
delete(AuditLog)
.where(AuditLog.id.in_(record_ids))
.execution_options(synchronize_session=False)
)
rows_deleted = result.rowcount
batch_db.commit()
total_deleted += rows_deleted
logger.debug(f"已删除 {rows_deleted} 条审计日志,累计 {total_deleted}")
except Exception as e:
logger.exception(f"删除审计日志批次失败: {e}")
try:
batch_db.rollback()
except Exception:
pass
break
finally:
batch_db.close()
return total_deleted
except Exception as e:
logger.exception(f"审计日志清理失败: {e}")
return 0
finally:
db.close()
loop = asyncio.get_running_loop()
total_deleted = await loop.run_in_executor(None, _do_audit_cleanup)
if total_deleted > 0:
logger.info(f"审计日志清理完成,共删除 {total_deleted} 条记录")
else:
logger.info("无需清理的审计日志")
async def _perform_gemini_file_mapping_cleanup(self) -> None:
"""清理过期的 Gemini 文件映射记录"""
@@ -1038,82 +1073,208 @@ class MaintenanceScheduler:
finally:
db.close()
async def _perform_cleanup(self) -> None:
"""执行清理任务"""
async def _perform_candidate_cleanup(self) -> None:
"""清理过期的 request_candidates 记录"""
def _do_candidate_cleanup() -> int:
db = create_session()
try:
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
logger.info("自动清理已禁用,跳过候选记录清理")
return 0
retention_days = max(
SystemConfigService.get_config(db, "detail_log_retention_days", 7),
3,
)
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
cutoff_time = datetime.now(timezone.utc) - timedelta(days=retention_days)
logger.info(f"开始清理 {retention_days} 天前的请求候选记录...")
except Exception as e:
logger.exception(f"候选记录清理配置读取失败: {e}")
return 0
finally:
db.close()
total_deleted = 0
while True:
batch_db = create_session()
try:
records_to_delete = (
batch_db.query(RequestCandidate.id)
.filter(RequestCandidate.created_at < cutoff_time)
.limit(batch_size)
.all()
)
if not records_to_delete:
break
record_ids = [r.id for r in records_to_delete]
result = batch_db.execute(
delete(RequestCandidate)
.where(RequestCandidate.id.in_(record_ids))
.execution_options(synchronize_session=False)
)
rows_deleted = result.rowcount
batch_db.commit()
total_deleted += rows_deleted
logger.debug(f"已删除 {rows_deleted} 条候选记录,累计 {total_deleted}")
except Exception as e:
logger.exception(f"删除候选记录批次失败: {e}")
try:
batch_db.rollback()
except Exception:
pass
break
finally:
batch_db.close()
return total_deleted
loop = asyncio.get_running_loop()
total_deleted = await loop.run_in_executor(None, _do_candidate_cleanup)
if total_deleted > 0:
logger.info(f"请求候选记录清理完成,共删除 {total_deleted} 条记录")
else:
logger.info("无需清理的候选记录")
async def _perform_db_maintenance(self) -> None:
"""执行数据库表维护VACUUM ANALYZE
对大表执行 VACUUM ANALYZE防止表和索引膨胀。
VACUUM 不能在事务内执行,需要使用 autocommit 连接。
使用线程池执行,避免阻塞事件循环。
"""
db = create_session()
try:
# 检查是否启用自动清理
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
logger.info("自动清理已禁用,跳过清理任务")
if not SystemConfigService.get_config(db, "enable_db_maintenance", True):
logger.info("数据库维护已禁用,跳过")
return
logger.info("开始执行使用记录分级清理...")
# 获取配置参数
detail_retention = SystemConfigService.get_config(db, "detail_log_retention_days", 7)
compressed_retention = SystemConfigService.get_config(
db, "compressed_log_retention_days", 30
)
header_retention = SystemConfigService.get_config(db, "header_retention_days", 90)
log_retention = SystemConfigService.get_config(db, "log_retention_days", 365)
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
now = datetime.now(timezone.utc)
# 1. 压缩详细日志 (body 字段 -> 压缩字段)
detail_cutoff = now - timedelta(days=detail_retention)
body_compressed = await self._cleanup_body_fields(db, detail_cutoff, batch_size)
# 2. 清理压缩字段90天后
compressed_cutoff = now - timedelta(days=compressed_retention)
compressed_cleaned = await self._cleanup_compressed_fields(
db, compressed_cutoff, batch_size
)
# 3. 清理请求头
header_cutoff = now - timedelta(days=header_retention)
header_cleaned = await self._cleanup_header_fields(db, header_cutoff, batch_size)
# 4. 删除过期记录
log_cutoff = now - timedelta(days=log_retention)
records_deleted = await self._delete_old_records(db, log_cutoff, batch_size)
# 5. 清理过期的API Keys
auto_delete = SystemConfigService.get_config(db, "auto_delete_expired_keys", False)
keys_cleaned = ApiKeyService.cleanup_expired_keys(db, auto_delete=auto_delete)
logger.info(
f"清理完成: 压缩 {body_compressed} 条, "
f"清理压缩字段 {compressed_cleaned} 条, "
f"清理header {header_cleaned} 条, "
f"删除记录 {records_deleted} 条, "
f"清理过期Keys {keys_cleaned}"
)
except Exception as e:
logger.exception(f"清理任务执行失败: {e}")
db.rollback()
logger.exception(f"读取数据库维护配置失败: {e}")
return
finally:
db.close()
async def _cleanup_body_fields(
self, db: Session, cutoff_time: datetime, batch_size: int
) -> int:
tables = ["usage", "request_candidates", "audit_logs"]
logger.info(f"开始数据库表维护VACUUM ANALYZE目标表: {', '.join(tables)}")
from src.database.database import _ensure_engine
engine = _ensure_engine()
def _vacuum_table(table_name: str) -> tuple[str, bool, str]:
"""在线程池中执行 VACUUM ANALYZE同步阻塞操作"""
try:
with engine.connect() as raw_conn:
conn = raw_conn.execution_options(isolation_level="AUTOCOMMIT")
conn.execute(text(f"VACUUM ANALYZE {table_name}"))
return table_name, True, ""
except Exception as e:
return table_name, False, str(e)
loop = asyncio.get_running_loop()
for table in tables:
table_name, success, error = await loop.run_in_executor(None, _vacuum_table, table)
if success:
logger.info(f"VACUUM ANALYZE {table_name} 完成")
else:
logger.warning(f"VACUUM ANALYZE {table_name} 失败: {error}")
logger.info("数据库表维护完成")
async def _perform_cleanup(self) -> None:
"""执行清理任务(在线程池中运行,避免阻塞事件循环)"""
def _do_cleanup() -> None:
db = create_session()
try:
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
logger.info("自动清理已禁用,跳过清理任务")
return
logger.info("开始执行使用记录分级清理...")
detail_retention = SystemConfigService.get_config(
db, "detail_log_retention_days", 7
)
compressed_retention = SystemConfigService.get_config(
db, "compressed_log_retention_days", 30
)
header_retention = SystemConfigService.get_config(db, "header_retention_days", 90)
log_retention = SystemConfigService.get_config(db, "log_retention_days", 365)
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
auto_delete = SystemConfigService.get_config(db, "auto_delete_expired_keys", False)
except Exception as e:
logger.exception(f"清理任务配置读取失败: {e}")
return
finally:
db.close()
try:
now = datetime.now(timezone.utc)
# 1. 压缩详细日志 (body 字段 -> 压缩字段)
detail_cutoff = now - timedelta(days=detail_retention)
body_compressed = self._cleanup_body_fields(detail_cutoff, batch_size)
# 2. 清理压缩字段
compressed_cutoff = now - timedelta(days=compressed_retention)
compressed_cleaned = self._cleanup_compressed_fields(compressed_cutoff, batch_size)
# 3. 清理请求头
header_cutoff = now - timedelta(days=header_retention)
header_cleaned = self._cleanup_header_fields(header_cutoff, batch_size)
# 4. 删除过期记录
log_cutoff = now - timedelta(days=log_retention)
records_deleted = self._delete_old_records(log_cutoff, batch_size)
# 5. 清理过期的API Keys
keys_db = create_session()
try:
keys_cleaned = ApiKeyService.cleanup_expired_keys(
keys_db, auto_delete=auto_delete
)
except Exception as e:
logger.exception(f"清理过期 Keys 失败: {e}")
keys_cleaned = 0
finally:
keys_db.close()
logger.info(
f"清理完成: 压缩 {body_compressed} 条, "
f"清理压缩字段 {compressed_cleaned} 条, "
f"清理header {header_cleaned} 条, "
f"删除记录 {records_deleted} 条, "
f"清理过期Keys {keys_cleaned}"
)
except Exception as e:
logger.exception(f"清理任务执行失败: {e}")
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, _do_cleanup)
def _cleanup_body_fields(self, cutoff_time: datetime, batch_size: int) -> int:
"""压缩 request_body 和 response_body 字段到压缩字段
逐条处理,确保每条记录都正确更新
逐条处理,确保每条记录都正确更新(同步方法,在线程池中调用)
"""
from sqlalchemy import null, update
total_compressed = 0
no_progress_count = 0 # 连续无进展计数
processed_ids: set = set() # 记录已处理的 ID防止重复处理
no_progress_count = 0
processed_ids: set = set()
while True:
batch_db = create_session()
try:
# 1. 查询需要压缩的记录
# 注意:排除已经是 NULL 或 JSON null 的记录
records = (
batch_db.query(
Usage.id,
@@ -1136,7 +1297,6 @@ class MaintenanceScheduler:
if not records:
break
# 过滤掉实际值为 None 的记录JSON null 被解析为 Python None
valid_records = [
r
for r in records
@@ -1147,7 +1307,6 @@ class MaintenanceScheduler:
]
if not valid_records:
# 所有记录都是 JSON null需要清理它们
logger.warning(
f"检测到 {len(records)} 条记录的 body 字段为 JSON null进行清理"
)
@@ -1165,7 +1324,6 @@ class MaintenanceScheduler:
batch_db.commit()
continue
# 检测是否有重复的 ID说明更新未生效
current_ids = {r.id for r in valid_records}
repeated_ids = current_ids & processed_ids
if repeated_ids:
@@ -1177,10 +1335,8 @@ class MaintenanceScheduler:
batch_success = 0
# 2. 逐条更新(确保每条都正确处理)
for r in valid_records:
try:
# 使用 null() 确保设置的是 SQL NULL 而不是 JSON null
result = batch_db.execute(
update(Usage)
.where(Usage.id == r.id)
@@ -1216,7 +1372,6 @@ class MaintenanceScheduler:
batch_db.commit()
# 3. 检查是否有实际进展
if batch_success == 0:
no_progress_count += 1
if no_progress_count >= 3:
@@ -1226,15 +1381,13 @@ class MaintenanceScheduler:
)
break
else:
no_progress_count = 0 # 重置计数
no_progress_count = 0
total_compressed += batch_success
logger.debug(
f"已压缩 {batch_success} 条记录的 body 字段,累计 {total_compressed}"
)
await asyncio.sleep(0.1)
except Exception as e:
logger.exception(f"压缩 body 字段失败: {e}")
try:
@@ -1247,12 +1400,10 @@ class MaintenanceScheduler:
return total_compressed
async def _cleanup_compressed_fields(
self, db: Session, cutoff_time: datetime, batch_size: int
) -> int:
"""清理压缩字段90天后删除压缩的body
def _cleanup_compressed_fields(self, cutoff_time: datetime, batch_size: int) -> int:
"""清理压缩字段删除压缩的body
每批使用短生命周期 session,避免 ORM 缓存问题
每批使用短生命周期 session(同步方法,在线程池中调用)
"""
from sqlalchemy import null, update
@@ -1261,7 +1412,6 @@ class MaintenanceScheduler:
while True:
batch_db = create_session()
try:
# 查询需要清理压缩字段的记录
records_to_clean = (
batch_db.query(Usage.id)
.filter(Usage.created_at < cutoff_time)
@@ -1280,7 +1430,6 @@ class MaintenanceScheduler:
record_ids = [r.id for r in records_to_clean]
# 批量更新,使用 null() 确保设置 SQL NULL
result = batch_db.execute(
update(Usage)
.where(Usage.id.in_(record_ids))
@@ -1302,8 +1451,6 @@ class MaintenanceScheduler:
total_cleaned += rows_updated
logger.debug(f"已清理 {rows_updated} 条记录的压缩字段,累计 {total_cleaned}")
await asyncio.sleep(0.1)
except Exception as e:
logger.exception(f"清理压缩字段失败: {e}")
try:
@@ -1316,12 +1463,10 @@ class MaintenanceScheduler:
return total_cleaned
async def _cleanup_header_fields(
self, db: Session, cutoff_time: datetime, batch_size: int
) -> int:
def _cleanup_header_fields(self, cutoff_time: datetime, batch_size: int) -> int:
"""清理 request_headers, response_headers 和 provider_request_headers 字段
每批使用短生命周期 session,避免 ORM 缓存问题
每批使用短生命周期 session(同步方法,在线程池中调用)
"""
from sqlalchemy import null, update
@@ -1330,7 +1475,6 @@ class MaintenanceScheduler:
while True:
batch_db = create_session()
try:
# 先查询需要清理的记录ID分批
records_to_clean = (
batch_db.query(Usage.id)
.filter(Usage.created_at < cutoff_time)
@@ -1348,7 +1492,6 @@ class MaintenanceScheduler:
record_ids = [r.id for r in records_to_clean]
# 批量更新,使用 null() 确保设置 SQL NULL
result = batch_db.execute(
update(Usage)
.where(Usage.id.in_(record_ids))
@@ -1369,8 +1512,6 @@ class MaintenanceScheduler:
total_cleaned += rows_updated
logger.debug(f"已清理 {rows_updated} 条记录的 header 字段,累计 {total_cleaned}")
await asyncio.sleep(0.1)
except Exception as e:
logger.exception(f"清理 header 字段失败: {e}")
try:
@@ -1383,15 +1524,15 @@ class MaintenanceScheduler:
return total_cleaned
async def _delete_old_records(self, db: Session, cutoff_time: datetime, batch_size: int) -> int:
"""删除过期的完整记录"""
def _delete_old_records(self, cutoff_time: datetime, batch_size: int) -> int:
"""删除过期的完整记录(同步方法,在线程池中调用)"""
total_deleted = 0
while True:
batch_db = create_session()
try:
# 查询要删除的记录ID分批
records_to_delete = (
db.query(Usage.id)
batch_db.query(Usage.id)
.filter(Usage.created_at < cutoff_time)
.limit(batch_size)
.all()
@@ -1402,28 +1543,27 @@ class MaintenanceScheduler:
record_ids = [r.id for r in records_to_delete]
# 执行删除
result = db.execute(
result = batch_db.execute(
delete(Usage)
.where(Usage.id.in_(record_ids))
.execution_options(synchronize_session=False)
)
rows_deleted = result.rowcount
db.commit()
batch_db.commit()
total_deleted += rows_deleted
logger.debug(f"已删除 {rows_deleted} 条过期记录,累计 {total_deleted}")
await asyncio.sleep(0.1)
except Exception as e:
logger.exception(f"删除过期记录失败: {e}")
try:
db.rollback()
batch_db.rollback()
except Exception:
pass
break
finally:
batch_db.close()
return total_deleted

View File

@@ -44,6 +44,7 @@ class TaskScheduler:
func: Callable[..., Any],
hour: int | str,
minute: int = 0,
day_of_week: str | int | None = None,
job_id: str | None = None,
name: str | None = None,
timezone: str | None = None,
@@ -56,12 +57,15 @@ class TaskScheduler:
func: 要执行的函数
hour: 执行时间(小时),使用业务时区
minute: 执行时间(分钟)
day_of_week: 星期几执行(如 "sun", "mon" 或 0-6
job_id: 任务ID
name: 任务名称(用于日志)
**kwargs: 传递给任务函数的参数
"""
trigger_timezone = timezone or APP_TIMEZONE
trigger = CronTrigger(hour=hour, minute=minute, timezone=trigger_timezone)
trigger = CronTrigger(
day_of_week=day_of_week, hour=hour, minute=minute, timezone=trigger_timezone
)
job_id = job_id or func.__name__
display_name = name or job_id