Nexus/server/infrastructure/database/migrations.py

"""Nexus — Data Migration: category → Node, backfill platform_id + schema migrations

D8: Migrate existing server category field to the new Node tree structure,
and create default Platform entries for existing server types.

Also handles safe schema migrations for existing databases (adding new columns).
"""

import logging
from sqlalchemy import select, update, text
from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server, Node, Platform, Base
from server.infrastructure.database.session import AsyncSessionLocal

logger = logging.getLogger("nexus.migration")


# Default platforms to create
DEFAULT_PLATFORMS = [
    {"name": "Linux 服务器", "category": "host", "type": "linux",
     "default_protocols": [{"name": "ssh", "port": 22, "primary": True}]},
    {"name": "Windows 服务器", "category": "host", "type": "windows",
     "default_protocols": [{"name": "rdp", "port": 3389, "primary": True}, {"name": "ssh", "port": 22}]},
    {"name": "MySQL", "category": "database", "type": "mysql",
     "default_protocols": [{"name": "mysql", "port": 3306, "primary": True}]},
    {"name": "网络设备", "category": "device", "type": "switch",
     "default_protocols": [{"name": "ssh", "port": 22, "primary": True}]},
]


async def migrate_category_to_node():
    """Migrate server.category values to Node tree + backfill platform_id.

    All steps in a single transaction — single commit at the end for atomicity.
    """
    async with AsyncSessionLocal() as session:
        # ── 1. Create default platforms ──
        for plat_data in DEFAULT_PLATFORMS:
            existing = await session.execute(
                select(Platform).where(Platform.name == plat_data["name"])
            )
            if not existing.scalar_one_or_none():
                platform = Platform(**plat_data)
                session.add(platform)
                logger.info(f"Created platform: {plat_data['name']}")

        # ── 2. Create nodes from existing categories ──
        result = await session.execute(
            select(Server.category).distinct().where(Server.category.isnot(None))
        )
        categories = [row[0] for row in result.all() if row[0]]

        category_to_node_id = {}
        for cat in categories:
            existing = await session.execute(
                select(Node).where(Node.name == cat)
            )
            node = existing.scalar_one_or_none()
            if not node:
                node = Node(name=cat, sort_order=0)
                session.add(node)
                await session.flush()
                logger.info(f"Created node: {cat}")
            category_to_node_id[cat] = node.id

        # ── 3. Backfill server node_id + platform_id ──
        # Build lookup: platform type → platform id
        all_platforms = await session.execute(select(Platform))
        type_to_platform_id = {p.type: p.id for p in all_platforms.scalars().all()}

        # Map category patterns to platform types
        def _resolve_platform_type(category: str) -> str | None:
            cat_lower = category.lower().strip()
            if cat_lower in type_to_platform_id:
                return cat_lower
            if cat_lower in ("linux", "centos", "ubuntu", "debian", "redhat"):
                return "linux"
            if cat_lower in ("windows", "win"):
                return "windows"
            if cat_lower in ("mysql", "mariadb", "database", "postgresql"):
                return "mysql"
            if cat_lower in ("switch", "router", "network", "cisco"):
                return "switch"
            return None

        for cat, node_id in category_to_node_id.items():
            plat_type = _resolve_platform_type(cat)
            plat_id = type_to_platform_id.get(plat_type) if plat_type else None
            await session.execute(
                update(Server)
                .where(Server.category == cat)
                .values(node_id=node_id, platform_id=plat_id)
            )
            logger.info(f"Migrated category '{cat}' → node_id={node_id}, platform_id={plat_id}")

        await session.commit()
        logger.info(f"Migration complete: {len(category_to_node_id)} categories migrated")


async def run_migrations():
    """Run all data migrations (called during startup after init_db)"""
    try:
        await migrate_category_to_node()
    except Exception as e:
        logger.error(f"Data migration failed: {e}")
        # Non-fatal — tables exist, migration can be retried

    try:
        await run_schema_migrations()
    except Exception as e:
        logger.error(f"Schema migration failed: {e}")


async def run_schema_migrations():
    """Safe schema migrations — add new columns to existing tables.
    Uses ALTER TABLE with try/except so it's a no-op if the column already exists.
    """
    migrations = [
        "ALTER TABLE push_schedules ADD COLUMN sync_mode VARCHAR(20) DEFAULT 'incremental' COMMENT '同步模式'",
        "ALTER TABLE admins ADD COLUMN token_version INT DEFAULT 0 COMMENT '令牌版本(重用时递增使旧token失效)'",
        "ALTER TABLE script_executions ADD COLUMN credential_id INT NULL COMMENT 'DB凭据替换$DB_*变量'",
        # Alert log table
        """CREATE TABLE IF NOT EXISTS `alert_logs` (
          `id` INT AUTO_INCREMENT PRIMARY KEY,
          `server_id` INT NOT NULL,
          `server_name` VARCHAR(100) NULL,
          `alert_type` VARCHAR(20) NOT NULL,
          `value` VARCHAR(50) NULL,
          `is_recovery` TINYINT(1) DEFAULT 0,
          `created_at` DATETIME DEFAULT CURRENT_TIMESTAMP,
          INDEX `idx_alert_logs_server_id` (`server_id`),
          INDEX `idx_alert_logs_created_at` (`created_at`),
          INDEX `idx_alert_logs_type` (`alert_type`)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",
        # Script-schedule columns (push_schedules extended to support script execution)
        "ALTER TABLE push_schedules ADD COLUMN schedule_type VARCHAR(20) DEFAULT 'push' COMMENT '调度类型: push|script'",
        "ALTER TABLE push_schedules ADD COLUMN script_id INT NULL COMMENT '脚本库引用'",
        "ALTER TABLE push_schedules ADD COLUMN script_content TEXT NULL COMMENT '临时命令内容'",
        "ALTER TABLE push_schedules ADD COLUMN exec_timeout INT DEFAULT 60 COMMENT '执行超时(秒)'",
        "ALTER TABLE push_schedules ADD COLUMN long_task TINYINT(1) DEFAULT 0 COMMENT '长任务nohup模式'",
        "ALTER TABLE push_schedules MODIFY COLUMN source_path VARCHAR(500) NULL COMMENT '推送源目录(push类型)'",
        # run_mode + fire_at: support one-time schedules (default) vs recurring cron
        "ALTER TABLE push_schedules ADD COLUMN run_mode VARCHAR(10) DEFAULT 'once' COMMENT '运行模式: once|cron'",
        "ALTER TABLE push_schedules ADD COLUMN fire_at DATETIME NULL COMMENT '一次性执行时间(once模式)'",
        "ALTER TABLE push_schedules MODIFY COLUMN cron_expr VARCHAR(100) NULL COMMENT 'cron表达式(cron模式)'",
        # ── Server table: columns added after initial schema ──
        "ALTER TABLE servers ADD COLUMN platform_id INT NULL COMMENT '平台类型ID'",
        "ALTER TABLE servers ADD COLUMN node_id INT NULL COMMENT '所属节点ID'",
        "ALTER TABLE servers ADD COLUMN protocols JSON NULL COMMENT '实例级协议覆盖'",
        "ALTER TABLE servers ADD COLUMN extra_attrs JSON NULL COMMENT '扩展属性'",
        "ALTER TABLE servers ADD COLUMN connectivity VARCHAR(20) DEFAULT 'unknown' COMMENT '连接状态'",
        "ALTER TABLE servers ADD COLUMN last_checked_at DATETIME NULL COMMENT '上次连接测试时间'",
        "ALTER TABLE servers ADD COLUMN ssh_key_configured TINYINT(1) DEFAULT 0 COMMENT '是否已配置SSH Key'",
        # Fix: ssh_key columns were VARCHAR(500) in legacy schema — change to TEXT for large encrypted keys
        "ALTER TABLE servers MODIFY COLUMN ssh_key_private TEXT NULL COMMENT '加密后的私钥内容(Fernet)'",
        "ALTER TABLE servers MODIFY COLUMN ssh_key_public TEXT NULL COMMENT '公钥内容'",
    ]
    async with AsyncSessionLocal() as session:
        for sql in migrations:
            try:
                await session.execute(text(sql))
                await session.commit()
                logger.info(f"Schema migration applied: {sql[:60]}...")
            except Exception as e:
                await session.rollback()
                err = str(e).lower()
                if "duplicate column" in err or "1060" in err:
                    logger.debug("Schema migration skipped (already applied): %s", sql[:60])
                else:
                    logger.error("Schema migration failed: %s — %s", sql[:60], e)
                    raise