cloud-reporter/backend/init_db.py

#!/usr/bin/env python
"""
数据库初始化脚本

功能:
- 创建所有数据库表
- 创建默认管理员账户
- 可选: 创建示例数据

用法:
    python init_db.py              # 初始化数据库并创建管理员
    python init_db.py --with-demo  # 初始化并创建示例数据
    python init_db.py --reset      # 重置数据库(删除所有数据后重新初始化)
"""
import os
import sys
import argparse

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app import create_app, db
from app.models import User, AWSCredential, UserCredential, Task, TaskLog, Report


def init_database(reset=False):
    """初始化数据库表"""
    if reset:
        print("⚠️  正在删除所有表...")
        db.drop_all()
        print("✓ 所有表已删除")
    
    print("正在创建数据库表...")
    db.create_all()
    print("✓ 数据库表创建完成")


def create_admin_user():
    """创建默认管理员账户"""
    admin = User.query.filter_by(username='admin').first()
    if admin:
        print("✓ 管理员账户已存在，跳过创建")
        return admin
    
    admin = User(
        username='admin',
        email='admin@example.com',
        role='admin',
        is_active=True
    )
    admin.set_password('admin123')
    
    db.session.add(admin)
    db.session.commit()
    
    print("✓ 管理员账户创建成功")
    print("  用户名: admin")
    print("  密码: admin123")
    print("  ⚠️  请在生产环境中修改默认密码!")
    
    return admin


def create_demo_data(admin_user):
    """创建示例数据"""
    print("\n正在创建示例数据...")
    
    # 创建示例用户
    demo_user = User.query.filter_by(username='demo').first()
    if not demo_user:
        demo_user = User(
            username='demo',
            email='demo@example.com',
            role='user',
            is_active=True
        )
        demo_user.set_password('demo123')
        db.session.add(demo_user)
        print("✓ 示例用户创建成功 (demo/demo123)")
    
    # 创建示例 AWS 凭证 (Assume Role 类型)
    demo_credential = AWSCredential.query.filter_by(name='Demo Account').first()
    if not demo_credential:
        demo_credential = AWSCredential(
            name='Demo Account',
            credential_type='assume_role',
            account_id='123456789012',
            role_arn='arn:aws:iam::123456789012:role/DemoRole',
            external_id='demo-external-id',
            is_active=True
        )
        db.session.add(demo_credential)
        print("✓ 示例 AWS 凭证创建成功")
    
    db.session.commit()
    
    # 关联凭证到管理员
    user_cred = UserCredential.query.filter_by(
        user_id=admin_user.id, 
        credential_id=demo_credential.id
    ).first()
    if not user_cred:
        user_cred = UserCredential(
            user_id=admin_user.id,
            credential_id=demo_credential.id
        )
        db.session.add(user_cred)
        db.session.commit()
        print("✓ 凭证已关联到管理员账户")
    
    print("✓ 示例数据创建完成")


def main():
    parser = argparse.ArgumentParser(description='数据库初始化脚本')
    parser.add_argument('--with-demo', action='store_true', help='创建示例数据')
    parser.add_argument('--reset', action='store_true', help='重置数据库(删除所有数据)')
    parser.add_argument('--env', default='development', help='运行环境 (development/testing/production)')
    args = parser.parse_args()
    
    if args.reset:
        confirm = input("⚠️  确定要删除所有数据吗? 输入 'yes' 确认: ")
        if confirm.lower() != 'yes':
            print("操作已取消")
            return
    
    print(f"\n{'='*50}")
    print(f"数据库初始化 - 环境: {args.env}")
    print(f"{'='*50}\n")
    
    # 创建应用上下文
    app = create_app(args.env)
    
    with app.app_context():
        # 初始化数据库
        init_database(reset=args.reset)
        
        # 创建管理员
        admin = create_admin_user()
        
        # 创建示例数据
        if args.with_demo:
            create_demo_data(admin)
        
        print(f"\n{'='*50}")
        print("✓ 数据库初始化完成!")
        print(f"{'='*50}\n")


if __name__ == '__main__':
    main()