mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
perf: 依赖数据库 CASCADE/SET NULL 替代手动清理关联表,缩短删除事务
- 批量删除移除 cleanup_key_references 手动清理,改为依赖 FK CASCADE/SET NULL - video_tasks.key_id FK 增加 ondelete="SET NULL",附带幂等迁移脚本 - _sync_delete 增加 statement_timeout 和任务级超时保护 - 批量导入在每次 await 前释放闲置 DB 连接,按批次提交写入避免长事务 - 前端轮询改为先查后等,首次查询不再多等一个间隔
This commit is contained in:
@@ -42,6 +42,7 @@ from src.database.database import get_db
|
||||
from src.models.database import Provider, ProviderAPIKey, User
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
from src.services.provider_keys.auth_type import OAUTH_AUTH_TYPES
|
||||
from src.services.scheduling.utils import release_db_connection_before_await
|
||||
from src.utils.auth_utils import require_admin
|
||||
|
||||
router = APIRouter(prefix="/api/admin/provider-oauth", tags=["Provider OAuth"])
|
||||
@@ -129,6 +130,7 @@ _PROVIDER_OAUTH_BATCH_TASK_TTL_SECONDS = 24 * 3600
|
||||
_PROVIDER_OAUTH_BATCH_TASK_MAX_ERROR_SAMPLES = 20
|
||||
_PROVIDER_OAUTH_DEFAULT_TIMEOUT_SECONDS = 30.0
|
||||
_PROVIDER_OAUTH_BATCH_IMPORT_PROXY_TIMEOUT_SECONDS = 60.0
|
||||
_PROVIDER_OAUTH_BATCH_IMPORT_COMMIT_BATCH_SIZE = 25
|
||||
_PROVIDER_OAUTH_BATCH_TASK_ALLOWED_STATUSES = {
|
||||
"submitted",
|
||||
"processing",
|
||||
@@ -1639,6 +1641,27 @@ def _estimate_batch_import_total(provider_type: str, raw_credentials: str) -> in
|
||||
return len(_parse_standard_oauth_import_entries(raw_credentials))
|
||||
|
||||
|
||||
def _release_batch_import_db_connection_before_await(db: Session) -> None:
|
||||
"""Best-effort 释放批量导入任务的只读 DB 连接。
|
||||
|
||||
批量导入会在单个后台任务里执行大量 await(上游 token 校验、邮箱探测、Redis 进度更新)。
|
||||
如果前面做过 Provider 查询而 Session 一直保持事务打开,连接会长时间占着不放,
|
||||
在大批量导入且多数条目最终失败时尤其容易把连接池拖满。
|
||||
|
||||
这里复用调度器已有的 helper:仅在 Session 没有挂起写入时才提前结束事务,
|
||||
避免影响 flush 后尚未提交的数据。
|
||||
"""
|
||||
release_db_connection_before_await(db)
|
||||
|
||||
|
||||
def _commit_batch_import_writes_if_needed(db: Session, pending_writes: int) -> int:
|
||||
"""按固定批次提交导入写入,避免长事务持续占用连接。"""
|
||||
if pending_writes < _PROVIDER_OAUTH_BATCH_IMPORT_COMMIT_BATCH_SIZE:
|
||||
return pending_writes
|
||||
db.commit()
|
||||
return 0
|
||||
|
||||
|
||||
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"):
|
||||
@@ -1901,11 +1924,14 @@ async def _batch_import_standard_oauth_internal(
|
||||
success_count = 0
|
||||
failed_count = 0
|
||||
processed_count = 0
|
||||
pending_success_writes = 0
|
||||
db_lock = asyncio.Lock()
|
||||
sem = asyncio.Semaphore(max(concurrency, 1))
|
||||
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
|
||||
async def _process_entry(idx: int, import_entry: dict[str, Any]) -> None:
|
||||
nonlocal success_count, failed_count, processed_count
|
||||
nonlocal success_count, failed_count, processed_count, pending_success_writes
|
||||
result_item: BatchImportResultItem
|
||||
|
||||
async with sem:
|
||||
@@ -1951,6 +1977,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
json_body = None
|
||||
|
||||
try:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
resp = await post_oauth_token(
|
||||
provider_type=provider_type,
|
||||
token_url=token_url,
|
||||
@@ -1970,6 +1997,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total, processed_count, success_count, failed_count, result_item
|
||||
)
|
||||
@@ -1997,6 +2025,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total, processed_count, success_count, failed_count, result_item
|
||||
)
|
||||
@@ -2016,6 +2045,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total, processed_count, success_count, failed_count, result_item
|
||||
)
|
||||
@@ -2039,6 +2069,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
}
|
||||
|
||||
try:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
auth_config = await enrich_auth_config(
|
||||
provider_type=provider_type,
|
||||
auth_config=auth_config,
|
||||
@@ -2067,6 +2098,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total,
|
||||
processed_count,
|
||||
@@ -2108,6 +2140,11 @@ async def _batch_import_standard_oauth_internal(
|
||||
proxy=key_proxy,
|
||||
)
|
||||
|
||||
pending_success_writes += 1
|
||||
pending_success_writes = _commit_batch_import_writes_if_needed(
|
||||
db, pending_success_writes
|
||||
)
|
||||
|
||||
result_item = BatchImportResultItem(
|
||||
index=idx,
|
||||
status="success",
|
||||
@@ -2129,6 +2166,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(total, processed_count, success_count, failed_count, result_item)
|
||||
|
||||
await asyncio.gather(
|
||||
@@ -2136,7 +2174,7 @@ async def _batch_import_standard_oauth_internal(
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
if success_count > 0:
|
||||
if success_count > 0 and pending_success_writes > 0:
|
||||
db.commit()
|
||||
|
||||
final_results = [r for r in results if r is not None]
|
||||
@@ -2446,11 +2484,14 @@ async def _batch_import_kiro_internal(
|
||||
success_count = 0
|
||||
failed_count = 0
|
||||
processed_count = 0
|
||||
pending_success_writes = 0
|
||||
db_lock = asyncio.Lock()
|
||||
sem = asyncio.Semaphore(max(concurrency, 1))
|
||||
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
|
||||
async def _process_entry(idx: int, cred: dict[str, Any]) -> None:
|
||||
nonlocal success_count, failed_count, processed_count
|
||||
nonlocal success_count, failed_count, processed_count, pending_success_writes
|
||||
result_item: BatchImportResultItem
|
||||
|
||||
async with sem:
|
||||
@@ -2466,6 +2507,7 @@ async def _batch_import_kiro_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total, processed_count, success_count, failed_count, result_item
|
||||
)
|
||||
@@ -2475,6 +2517,7 @@ async def _batch_import_kiro_internal(
|
||||
cfg.provider_type = ProviderType.KIRO.value
|
||||
|
||||
try:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
access_token, new_cfg = await refresh_access_token(
|
||||
cfg,
|
||||
proxy_config=proxy_config,
|
||||
@@ -2490,11 +2533,13 @@ async def _batch_import_kiro_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total, processed_count, success_count, failed_count, result_item
|
||||
)
|
||||
return
|
||||
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
email = await _fetch_kiro_email(new_cfg.to_dict(), proxy_config=proxy_config)
|
||||
if email and not new_cfg.email:
|
||||
new_cfg.email = email
|
||||
@@ -2514,6 +2559,7 @@ async def _batch_import_kiro_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(
|
||||
total,
|
||||
processed_count,
|
||||
@@ -2550,6 +2596,11 @@ async def _batch_import_kiro_internal(
|
||||
proxy=key_proxy,
|
||||
)
|
||||
|
||||
pending_success_writes += 1
|
||||
pending_success_writes = _commit_batch_import_writes_if_needed(
|
||||
db, pending_success_writes
|
||||
)
|
||||
|
||||
result_item = BatchImportResultItem(
|
||||
index=idx,
|
||||
status="success",
|
||||
@@ -2572,6 +2623,7 @@ async def _batch_import_kiro_internal(
|
||||
processed_count += 1
|
||||
results[idx] = result_item
|
||||
if progress_hook is not None:
|
||||
_release_batch_import_db_connection_before_await(db)
|
||||
await progress_hook(total, processed_count, success_count, failed_count, result_item)
|
||||
|
||||
await asyncio.gather(
|
||||
@@ -2580,7 +2632,7 @@ async def _batch_import_kiro_internal(
|
||||
)
|
||||
|
||||
# 提交所有成功的记录
|
||||
if success_count > 0:
|
||||
if success_count > 0 and pending_success_writes > 0:
|
||||
db.commit()
|
||||
|
||||
final_results = [r for r in results if r is not None]
|
||||
|
||||
Reference in New Issue
Block a user