#!/usr/bin/env python
"""
应用数据库迁移脚本

用法:
    python apply_migration.py
"""
import os
import sys

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

from app import create_app, db
from sqlalchemy import text

def apply_migration():
    """应用account_id字段可空的迁移"""
    app = create_app()
    
    with app.app_context():
        try:
            print("正在应用数据库迁移...")
            
            # 检查当前表结构
            result = db.engine.execute(text("PRAGMA table_info(aws_credentials)"))
            columns = result.fetchall()
            
            account_id_nullable = False
            for column in columns:
                if column[1] == 'account_id':  # column[1] is column name
                    account_id_nullable = column[3] == 0  # column[3] is notnull (0=nullable, 1=not null)
                    break
            
            if account_id_nullable:
                print("✓ account_id字段已经是可空的")
                return
            
            print("正在修改account_id字段为可空...")
            
            # SQLite doesn't support ALTER COLUMN directly, so we need to recreate the table
            # First, create a backup of the data
            db.engine.execute(text("""
                CREATE TABLE aws_credentials_backup AS 
                SELECT * FROM aws_credentials
            """))
            
            # Drop the original table
            db.engine.execute(text("DROP TABLE aws_credentials"))
            
            # Recreate the table with nullable account_id
            db.engine.execute(text("""
                CREATE TABLE aws_credentials (
                    id INTEGER PRIMARY KEY,
                    name VARCHAR(100) NOT NULL,
                    credential_type VARCHAR(20) NOT NULL,
                    account_id VARCHAR(12),  -- Now nullable
                    role_arn VARCHAR(255),
                    external_id VARCHAR(255),
                    access_key_id VARCHAR(255),
                    secret_access_key_encrypted TEXT,
                    created_at DATETIME,
                    is_active BOOLEAN,
                    CHECK (credential_type IN ('assume_role', 'access_key'))
                )
            """))
            
            # Create index on account_id
            db.engine.execute(text("CREATE INDEX ix_aws_credentials_account_id ON aws_credentials (account_id)"))
            
            # Restore the data
            db.engine.execute(text("""
                INSERT INTO aws_credentials 
                SELECT * FROM aws_credentials_backup
            """))
            
            # Drop the backup table
            db.engine.execute(text("DROP TABLE aws_credentials_backup"))
            
            print("✓ 数据库迁移完成")
            
        except Exception as e:
            print(f"❌ 迁移失败: {e}")
            # Try to restore from backup if it exists
            try:
                db.engine.execute(text("DROP TABLE IF EXISTS aws_credentials"))
                db.engine.execute(text("ALTER TABLE aws_credentials_backup RENAME TO aws_credentials"))
                print("已恢复原始数据")
            except:
                pass
            raise

if __name__ == '__main__':
    apply_migration()