forked from xiaohei/taiji-AI-PAD
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
||
"""修复已存在的渠道配额记录
|
||
|
||
对于已经使用 customAgentResources 分配过资源的渠道，
|
||
但还没有 ChannelCustomAgentQuota 记录的，创建对应的记录。
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/../services/mcp-server')
|
||
|
||
from sqlalchemy import select
|
||
from database import engine, AsyncSessionLocal
|
||
from models import Channel, ChannelCustomAgentQuota
|
||
import asyncio
|
||
|
||
async def fix_channel_quotas():
|
||
"""修复渠道配额记录"""
|
||
async with AsyncSessionLocal() as db:
|
||
# 查询所有渠道
|
||
result = await db.execute(select(Channel))
|
||
channels = result.scalars().all()
|
||
|
||
print(f"找到 {len(channels)} 个渠道")
|
||
print("=" * 60)
|
||
|
||
fixed_count = 0
|
||
for channel in channels:
|
||
# 检查是否有旧格式的配额（custom_agent_cpu 或 custom_agent_memory > 0）
|
||
has_old_quota = (
|
||
(channel.custom_agent_cpu and channel.custom_agent_cpu > 0) or
|
||
(channel.custom_agent_memory and channel.custom_agent_memory > 0)
|
||
)
|
||
|
||
if not has_old_quota:
|
||
continue
|
||
|
||
# 检查是否已经有 ChannelCustomAgentQuota 记录
|
||
quota_result = await db.execute(
|
||
select(ChannelCustomAgentQuota).where(
|
||
ChannelCustomAgentQuota.channel_id == channel.id
|
||
)
|
||
)
|
||
existing_quota = quota_result.scalar_one_or_none()
|
||
|
||
if existing_quota:
|
||
print(f"✅ 渠道 {channel.name} ({channel.id}) 已有配额记录，跳过")
|
||
continue
|
||
|
||
# 创建新的配额记录
|
||
cpu_quota = float(channel.custom_agent_cpu or 0)
|
||
memory_quota = float(channel.custom_agent_memory or 0)
|
||
|
||
print(f"\n🔧 修复渠道 {channel.name} ({channel.id})")
|
||
print(f" 旧格式配额: CPU={cpu_quota}核, Memory={memory_quota}GB")
|
||
|
||
# 同时更新新格式字段
|
||
channel.custom_agent_cpu_quota = cpu_quota
|
||
channel.custom_agent_memory_quota = memory_quota
|
||
|
||
# 创建配额记录
|
||
new_quota = ChannelCustomAgentQuota(
|
||
channel_id=channel.id,
|
||
cpu_quota=cpu_quota,
|
||
memory_quota=memory_quota,
|
||
cpu_allocated=0,
|
||
memory_allocated=0,
|
||
)
|
||
db.add(new_quota)
|
||
|
||
print(f" ✅ 创建 ChannelCustomAgentQuota 记录")
|
||
fixed_count += 1
|
||
|
||
if fixed_count > 0:
|
||
await db.commit()
|
||
print("\n" + "=" * 60)
|
||
print(f"✅ 修复完成，共修复 {fixed_count} 个渠道")
|
||
else:
|
||
print("\n" + "=" * 60)
|
||
print("✅ 没有需要修复的渠道")
|
||
|
||
print("=" * 60)
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(fix_channel_quotas())
|
||
except Exception as e:
|
||
print(f"\n❌ 修复失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
sys.exit(1)
|