forked from xiaohei/taiji-AI-PAD
122 lines
3.7 KiB
Python
122 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
重新初始化数据库表
|
|
用于在数据库表丢失时重新创建所有表结构
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import asyncio
|
|
|
|
# 添加services/mcp-server到路径，以便导入模块
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'services', 'mcp-server'))
|
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy import text
|
|
from models import Base
|
|
from config import settings
|
|
from database import prepare_database_url
|
|
import logging
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def reinit_database():
|
|
"""重新初始化数据库表"""
|
|
try:
|
|
# 准备数据库URL
|
|
database_url = prepare_database_url(settings.database_url)
|
|
|
|
# 确保使用asyncpg驱动
|
|
if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
|
|
database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
|
|
|
|
logger.info(f"正在连接到数据库: {database_url.split('@')[1] if '@' in database_url else '本地数据库'}")
|
|
|
|
# 创建数据库引擎
|
|
engine = create_async_engine(database_url, echo=False)
|
|
|
|
logger.info("开始创建数据库表...")
|
|
|
|
# 创建所有表
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
logger.info("✓ 数据库表创建成功！")
|
|
|
|
# 列出所有创建的表
|
|
async with engine.begin() as conn:
|
|
# 获取所有表名
|
|
if "postgresql" in database_url:
|
|
result = await conn.execute(
|
|
text("""
|
|
SELECT table_name
|
|
FROM information_schema.tables
|
|
WHERE table_schema = 'public'
|
|
ORDER BY table_name
|
|
""")
|
|
)
|
|
else:
|
|
# SQLite
|
|
result = await conn.execute(
|
|
text("""
|
|
SELECT name
|
|
FROM sqlite_master
|
|
WHERE type='table' AND name NOT LIKE 'sqlite_%'
|
|
ORDER BY name
|
|
""")
|
|
)
|
|
|
|
tables = result.fetchall()
|
|
table_names = [row[0] for row in tables]
|
|
|
|
logger.info(f"\n已创建的表列表（共 {len(table_names)} 个）:")
|
|
for i, table_name in enumerate(table_names, 1):
|
|
logger.info(f" {i}. {table_name}")
|
|
|
|
await engine.dispose()
|
|
|
|
logger.info("\n数据库初始化完成！")
|
|
logger.info("注意：此脚本只创建表结构，不会创建初始数据。")
|
|
logger.info("如果需要创建初始数据，请运行服务或使用其他初始化脚本。")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"✗ 数据库初始化失败: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
async def main():
|
|
"""主函数"""
|
|
print("="*80)
|
|
print("重新初始化数据库表")
|
|
print("="*80)
|
|
print()
|
|
|
|
# 确认操作
|
|
print("警告：此操作将创建所有数据库表。")
|
|
print("如果表已存在，SQLAlchemy不会删除或修改现有表。")
|
|
print()
|
|
|
|
success = await reinit_database()
|
|
|
|
if success:
|
|
print("\n" + "="*80)
|
|
print("✓ 数据库表初始化成功！")
|
|
print("="*80)
|
|
sys.exit(0)
|
|
else:
|
|
print("\n" + "="*80)
|
|
print("✗ 数据库表初始化失败！")
|
|
print("="*80)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|
|
|