mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -208,6 +208,10 @@ class SystemConfigService:
|
||||
"value": 30,
|
||||
"description": "审计日志保留天数,超过此天数的审计日志将被自动清理",
|
||||
},
|
||||
"enable_db_maintenance": {
|
||||
"value": True,
|
||||
"description": "是否启用数据库表维护任务(定期 VACUUM ANALYZE 防止表和索引膨胀)",
|
||||
},
|
||||
# 系统代理
|
||||
"system_proxy_node_id": {
|
||||
"value": None,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user