146 lines
4.4 KiB
Python
146 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
完整的数据库初始化脚本
|
|
包括创建表和插入测试数据
|
|
"""
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
from sqlalchemy import text
|
|
from app.database import engine, Base, AsyncSessionLocal
|
|
from loguru import logger
|
|
|
|
# 导入初始化数据函数
|
|
from init_test_data import init_test_data
|
|
|
|
logger.remove()
|
|
logger.add(sys.stderr, level="INFO", format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | {message}")
|
|
|
|
|
|
async def test_database_connection():
|
|
"""测试数据库连接"""
|
|
logger.info("正在测试数据库连接...")
|
|
|
|
try:
|
|
async with engine.connect() as conn:
|
|
result = await conn.execute(text("SELECT version()"))
|
|
version = result.fetchone()
|
|
logger.info(f"✅ 数据库连接成功")
|
|
logger.info(f"数据库版本: {version[0] if version else 'N/A'}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"❌ 数据库连接失败: {e}")
|
|
logger.error("\n请检查:")
|
|
logger.error("1. PostgreSQL 服务是否运行")
|
|
logger.error("2. 数据库 'deeprisk_db' 是否存在")
|
|
logger.error("3. 用户名密码是否正确")
|
|
logger.error(f"\n当前数据库URL配置:")
|
|
logger.error(f"postgresql+asyncpg://postgres:***@localhost:15432/deeprisk_db")
|
|
return False
|
|
|
|
|
|
async def init_tables():
|
|
"""初始化数据库表"""
|
|
logger.info("=" * 60)
|
|
logger.info("正在初始化数据库表结构...")
|
|
logger.info("=" * 60)
|
|
|
|
try:
|
|
async with engine.begin() as conn:
|
|
# 创建所有表
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
logger.info("✅ 数据库表创建成功")
|
|
|
|
# 列出已创建的表
|
|
result = await conn.execute(
|
|
text("""
|
|
SELECT table_name
|
|
FROM information_schema.tables
|
|
WHERE table_schema = 'public'
|
|
ORDER BY table_name
|
|
""")
|
|
)
|
|
tables = result.fetchall()
|
|
if tables:
|
|
logger.info("\n已创建的表:")
|
|
for table in tables:
|
|
logger.info(f" ✓ {table[0]}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ 创建数据库表失败: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
async def show_data_summary():
|
|
"""显示数据统计"""
|
|
logger.info("\n" + "=" * 60)
|
|
logger.info("数据库数据统计")
|
|
logger.info("=" * 60)
|
|
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
# 检查各表的数据量
|
|
tables = [
|
|
"mcn_agency",
|
|
"streamer_info",
|
|
"platform_recharge",
|
|
"contract",
|
|
"order",
|
|
"settlement",
|
|
"expense",
|
|
"tax_declaration",
|
|
"bank_transaction",
|
|
"invoice",
|
|
"detection_rule",
|
|
]
|
|
|
|
for table in tables:
|
|
try:
|
|
result = await session.execute(text(f"SELECT COUNT(*) FROM {table}"))
|
|
count = result.scalar()
|
|
logger.info(f" {table:30} : {count:6} 条记录")
|
|
except Exception as e:
|
|
logger.warning(f" {table:30} : 查询失败 ({str(e)[:30]})")
|
|
|
|
except Exception as e:
|
|
logger.error(f"统计查询失败: {e}")
|
|
|
|
|
|
async def main():
|
|
"""主函数"""
|
|
logger.info("\n" + "=" * 60)
|
|
logger.info("🎯 DeepRisk 数据库完整初始化")
|
|
logger.info("=" * 60 + "\n")
|
|
|
|
# 1. 测试数据库连接
|
|
if not await test_database_connection():
|
|
sys.exit(1)
|
|
|
|
# 2. 初始化表结构
|
|
await init_tables()
|
|
|
|
# 3. 初始化测试数据
|
|
try:
|
|
await init_test_data()
|
|
except Exception as e:
|
|
logger.error(f"初始化测试数据失败: {e}")
|
|
logger.warning("表创建成功，但数据插入失败")
|
|
logger.warning("可以稍后手动运行 init_test_data.py 插入数据")
|
|
|
|
# 4. 显示数据统计
|
|
await show_data_summary()
|
|
|
|
logger.info("\n" + "=" * 60)
|
|
logger.info("🎉 数据库初始化完成！")
|
|
logger.info("=" * 60)
|
|
logger.info("\n可以使用以下方式启动应用:")
|
|
logger.info(" uvicorn app.main:app --reload --host 0.0.0.0 --port 8000\n")
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|