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.

    Steps:
    1. Create default Platform entries (if not exist)
    2. Create Node entries from distinct category values (if not exist)
    3. Update servers: set node_id and platform_id based on category
    """
    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']}")

        await session.commit()

        # ── 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

        await session.commit()

        # ── 3. Backfill server node_id ──
        linux_platform = await session.execute(
            select(Platform).where(Platform.type == "linux")
        )
        linux_plat = linux_platform.scalar_one_or_none()
        linux_platform_id = linux_plat.id if linux_plat else None

        for cat, node_id in category_to_node_id.items():
            await session.execute(
                update(Server)
                .where(Server.category == cat)
                .values(node_id=node_id, platform_id=linux_platform_id)
            )
            logger.info(f"Migrated category '{cat}' → node_id={node_id}, platform_id={linux_platform_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失效)'",
    ]
    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:
                await session.rollback()
                # Column already exists — expected for non-first-run