bigscreen/init_database.py

#!/usr/bin/env python3
"""
PostgreSQL 数据库初始化脚本
为 FastAPI 应用创建所有表结构
"""

import asyncio
import sys
import os
from pathlib import Path

# 添加项目路径
sys.path.insert(0, str(Path(__file__).parent / "backend"))

# 设置环境变量（从 .env 读取）
from dotenv import load_dotenv
load_dotenv()

async def init_database():
    """初始化数据库"""
    try:
        from app.core.config import settings
        from app.db.base import Base
        from app.core.database import engine

        # 导入所有模型以确保它们被注册
        from app.models import (
            health, alert, device, user, message,
            org, config, audit, meeting
        )

        print("=" * 70)
        print("PostgreSQL 数据库初始化开始")
        print("=" * 70)
        print(f"📦 应用: {settings.app_name} v{settings.app_version}")
        print(f"🗄️  数据库: {settings.database_name}")
        print(f"🔗 主机: {settings.database_host}:{settings.database_port}")
        print(f"👤 用户: {settings.database_user}")
        print("=" * 70)

        # 创建所有表
        print("\n正在创建数据库表...")
        async with engine.begin() as conn:
            # 删除所有表（如果存在）
            print("  ⚠️  删除旧表（如果存在）...")
            await conn.run_sync(Base.metadata.drop_all)

            # 创建所有表
            print("  ✨ 创建新表...")
            await conn.run_sync(Base.metadata.create_all)

        print("\n✅ 数据库表创建完成！\n")

        # 验证表创建
        from sqlalchemy import inspect, text
        async with engine.connect() as conn:
            # 获取所有表名
            def get_tables(connection):
                inspector = inspect(connection)
                return inspector.get_table_names()

            tables = await conn.run_sync(get_tables)

            print(f"📊 共创建 {len(tables)} 个表：")
            print("-" * 70)

            # 按类别分组显示表
            categories = {
                "健康数据": ["t_user_health_data", "t_user_health_data_daily", "t_user_health_data_weekly", "t_health_summary_daily"],
                "告警管理": ["t_alert_info", "t_alert_log", "t_alert_rules", "t_user_alerts", "t_health_anomaly", "t_event_alarm_queue"],
                "设备管理": ["t_device_info", "t_device_user", "t_device_config", "t_device_status_log"],
                "用户管理": ["t_user", "t_user_profile"],
                "组织架构": ["t_org_info", "t_user_org", "t_position", "t_user_position", "t_org_closure", "t_org_manager_cache"],
                "消息通知": ["t_device_message", "t_notification_config", "t_notification_log"],
                "系统配置": ["t_customer", "t_customer_config", "t_system_config", "t_api_key"],
                "审计日志": ["t_audit_log", "t_audit_operation_config"],
                "会议室": ["t_meeting_room"]
            }

            categorized = set()
            for category, table_list in categories.items():
                found = [t for t in tables if t in table_list]
                if found:
                    print(f"\n{category}:")
                    for table in sorted(found):
                        print(f"  ✓ {table}")
                        categorized.add(table)

            # 显示未分类的表
            uncategorized = set(tables) - categorized
            if uncategorized:
                print(f"\n其他:")
                for table in sorted(uncategorized):
                    print(f"  ✓ {table}")

            # 测试数据库连接
            print("\n" + "-" * 70)
            print("正在测试数据库连接...")
            result = await conn.execute(text("SELECT version();"))
            row = result.fetchone()
            print(f"  PostgreSQL 版本: {row[0][:60]}...")

        print("\n" + "=" * 70)
        print("✅ 数据库初始化成功完成！")
        print("=" * 70)
        print("\n💡 提示:")
        print("  - 可以使用 start_app.sh 启动应用")
        print("  - API 文档: http://localhost:8008/docs")
        print("  - 大屏地址: http://localhost:8118/bigscreen")
        print()

    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(init_database())