update
This commit is contained in:
108
team_service.py
108
team_service.py
@@ -104,8 +104,17 @@ def fetch_account_id(team: dict, silent: bool = False) -> str:
|
|||||||
return ""
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
def _is_shutdown_requested() -> bool:
|
||||||
|
"""检查是否收到停止请求"""
|
||||||
|
try:
|
||||||
|
import run
|
||||||
|
return run._shutdown_requested
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def preload_all_account_ids() -> tuple[int, int]:
|
def preload_all_account_ids() -> tuple[int, int]:
|
||||||
"""预加载所有 Team 的 account_id
|
"""预加载所有 Team 的 account_id (并行处理)
|
||||||
|
|
||||||
在程序启动时调用,避免后续重复获取
|
在程序启动时调用,避免后续重复获取
|
||||||
只处理有 token 的 Team,没有 token 的跳过
|
只处理有 token 的 Team,没有 token 的跳过
|
||||||
@@ -114,6 +123,7 @@ def preload_all_account_ids() -> tuple[int, int]:
|
|||||||
tuple: (success_count, fail_count)
|
tuple: (success_count, fail_count)
|
||||||
"""
|
"""
|
||||||
import sys
|
import sys
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
|
||||||
success_count = 0
|
success_count = 0
|
||||||
fail_count = 0
|
fail_count = 0
|
||||||
@@ -122,83 +132,71 @@ def preload_all_account_ids() -> tuple[int, int]:
|
|||||||
teams_with_token = [t for t in TEAMS if t.get("auth_token")]
|
teams_with_token = [t for t in TEAMS if t.get("auth_token")]
|
||||||
teams_need_fetch = [t for t in teams_with_token if not t.get("account_id")]
|
teams_need_fetch = [t for t in teams_with_token if not t.get("account_id")]
|
||||||
|
|
||||||
|
# 已缓存的数量
|
||||||
|
cached_count = len(teams_with_token) - len(teams_need_fetch)
|
||||||
|
|
||||||
if not teams_need_fetch:
|
if not teams_need_fetch:
|
||||||
if teams_with_token:
|
if teams_with_token:
|
||||||
log.success(f"所有 Team account_id 已缓存 ({len(teams_with_token)} 个)")
|
log.success(f"所有 Team account_id 已缓存 ({len(teams_with_token)} 个)")
|
||||||
return len(teams_with_token), 0
|
return len(teams_with_token), 0
|
||||||
|
|
||||||
total = len(teams_with_token)
|
total = len(teams_need_fetch)
|
||||||
log.info(f"预加载 {len(teams_need_fetch)} 个 Team 的 account_id...", icon="sync")
|
log.info(f"并行预加载 {total} 个 Team 的 account_id...", icon="sync")
|
||||||
|
|
||||||
need_save = False
|
need_save = False
|
||||||
failed_teams = []
|
failed_teams = []
|
||||||
|
stopped = False
|
||||||
|
|
||||||
# 检测是否为 TTY 环境
|
def fetch_one(team):
|
||||||
is_tty = sys.stdout.isatty()
|
"""获取单个 Team 的 account_id"""
|
||||||
|
if _is_shutdown_requested():
|
||||||
if is_tty:
|
return team, None, "stopped"
|
||||||
# TTY 环境: 使用 rich 进度条
|
|
||||||
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
|
|
||||||
|
|
||||||
with Progress(
|
|
||||||
SpinnerColumn(),
|
|
||||||
TextColumn("[progress.description]{task.description}"),
|
|
||||||
BarColumn(),
|
|
||||||
TaskProgressColumn(),
|
|
||||||
TextColumn("{task.fields[status]}"),
|
|
||||||
) as progress:
|
|
||||||
task = progress.add_task("加载中", total=total, status="")
|
|
||||||
|
|
||||||
for team in teams_with_token:
|
|
||||||
progress.update(task, description=f"[cyan]{team['name']}", status="")
|
|
||||||
|
|
||||||
if team.get("account_id"):
|
|
||||||
success_count += 1
|
|
||||||
progress.update(task, advance=1, status="[green]✓ 已缓存")
|
|
||||||
continue
|
|
||||||
|
|
||||||
account_id = fetch_account_id(team, silent=True)
|
account_id = fetch_account_id(team, silent=True)
|
||||||
if account_id:
|
return team, account_id, "ok" if account_id else "failed"
|
||||||
|
|
||||||
|
# 使用线程池并行获取 (最多 10 个并发)
|
||||||
|
max_workers = min(10, total)
|
||||||
|
completed = 0
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||||
|
futures = {executor.submit(fetch_one, team): team for team in teams_need_fetch}
|
||||||
|
|
||||||
|
for future in as_completed(futures):
|
||||||
|
# 检查停止标志
|
||||||
|
if _is_shutdown_requested():
|
||||||
|
stopped = True
|
||||||
|
# 取消剩余任务
|
||||||
|
for f in futures:
|
||||||
|
f.cancel()
|
||||||
|
break
|
||||||
|
|
||||||
|
team, account_id, status = future.result()
|
||||||
|
completed += 1
|
||||||
|
|
||||||
|
if status == "stopped":
|
||||||
|
stopped = True
|
||||||
|
break
|
||||||
|
elif status == "ok":
|
||||||
success_count += 1
|
success_count += 1
|
||||||
progress.update(task, advance=1, status="[green]✓")
|
|
||||||
if team.get("format") == "new":
|
if team.get("format") == "new":
|
||||||
need_save = True
|
need_save = True
|
||||||
|
log.info(f"预加载 [{completed}/{total}] {team['name']}: ✓ {account_id}")
|
||||||
else:
|
else:
|
||||||
fail_count += 1
|
fail_count += 1
|
||||||
failed_teams.append(team['name'])
|
failed_teams.append(team['name'])
|
||||||
progress.update(task, advance=1, status="[red]✗")
|
log.warning(f"预加载 [{completed}/{total}] {team['name']}: ✗ 失败")
|
||||||
else:
|
|
||||||
# 非 TTY 环境 (systemd/journalctl): 使用普通日志输出
|
|
||||||
for idx, team in enumerate(teams_with_token, 1):
|
|
||||||
team_name = team['name']
|
|
||||||
|
|
||||||
if team.get("account_id"):
|
# 加上已缓存的数量
|
||||||
success_count += 1
|
success_count += cached_count
|
||||||
log.info(f"预加载 [{idx}/{total}] {team_name}: ✓ 已缓存")
|
|
||||||
continue
|
|
||||||
|
|
||||||
account_id = fetch_account_id(team, silent=True)
|
|
||||||
if account_id:
|
|
||||||
success_count += 1
|
|
||||||
log.info(f"预加载 [{idx}/{total}] {team_name}: ✓ {account_id}")
|
|
||||||
if team.get("format") == "new":
|
|
||||||
need_save = True
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
failed_teams.append(team_name)
|
|
||||||
log.warning(f"预加载 [{idx}/{total}] {team_name}: ✗ 失败")
|
|
||||||
|
|
||||||
# 输出失败的 team (仅 TTY 环境,非 TTY 已在循环中输出)
|
|
||||||
if is_tty:
|
|
||||||
for name in failed_teams:
|
|
||||||
log.warning(f"Team {name}: 获取 account_id 失败")
|
|
||||||
|
|
||||||
# 持久化到 team.json
|
# 持久化到 team.json
|
||||||
if need_save:
|
if need_save:
|
||||||
if save_team_json():
|
if save_team_json():
|
||||||
log.success(f"account_id 已保存到 team.json")
|
log.success(f"account_id 已保存到 team.json")
|
||||||
|
|
||||||
if fail_count == 0 and success_count > 0:
|
if stopped:
|
||||||
|
log.warning(f"预加载已停止 (已完成 {completed} 个)")
|
||||||
|
elif fail_count == 0 and success_count > 0:
|
||||||
log.success(f"所有 Team account_id 加载完成 ({success_count} 个)")
|
log.success(f"所有 Team account_id 加载完成 ({success_count} 个)")
|
||||||
elif fail_count > 0:
|
elif fail_count > 0:
|
||||||
log.warning(f"account_id 加载: 成功 {success_count}, 失败 {fail_count}")
|
log.warning(f"account_id 加载: 成功 {success_count}, 失败 {fail_count}")
|
||||||
|
|||||||
Reference in New Issue
Block a user