model-platform/migrations/data/migrate_system_json.py

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import os
from pathlib import Path
from typing import Any
from urllib.parse import quote

from pydantic import BaseModel, ConfigDict
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from common.db import create_database_engine, create_session_factory
from common.db.models import (
    Permissions,
    RolePermissions,
    Roles,
    Users,
)

CROCKFORD_BASE32 = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
LEGACY_PASSWORD_HASH = "!legacy-account-without-password!"

PERMISSION_NAMES = {
    "dashboard.view": "查看工作台",
    "script.build": "构建脚本",
    "script.public.manage": "管理公共脚本",
    "schedule.own": "管理本人调度",
    "schedule.all": "管理全部调度",
    "experiment.own": "管理本人实验",
    "experiment.all": "管理全部实验",
    "resource.personal": "管理个人资源",
    "resource.public.upload": "上传公共资源",
    "resource.public.manage": "管理公共资源",
    "system.view": "查看系统管理",
    "system.manage": "管理系统配置",
}


class LegacyUser(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str
    name: str
    role: str
    roleKey: str
    avatar: str | None = None


class LegacyRole(BaseModel):
    model_config = ConfigDict(extra="forbid")

    key: str
    name: str
    description: str | None = None
    permissions: list[str]


class LegacySystem(BaseModel):
    """Legacy platform_data/system.json shape.

    ``extra='ignore'`` so legacy files that include other top-level
    keys (e.g. historical audit-log payloads) can still be parsed —
    only the fields this script actually consumes are listed below.
    """

    model_config = ConfigDict(extra="ignore")

    users: list[LegacyUser]
    roles: list[LegacyRole]


def deterministic_legacy_ulid(entity_type: str, legacy_key: str) -> str:
    """Create a stable ULID-compatible ID with an epoch timestamp prefix."""
    digest = hashlib.sha256(
        f"model-platform-v1:{entity_type}:{legacy_key}".encode("utf-8")
    ).digest()
    value = int.from_bytes(b"\x00" * 6 + digest[:10], byteorder="big")
    encoded = ["0"] * 26
    for index in range(25, -1, -1):
        encoded[index] = CROCKFORD_BASE32[value & 31]
        value >>= 5
    return "".join(encoded)


def require_unique(values: list[str], label: str) -> None:
    duplicates = sorted(
        value for value in set(values) if values.count(value) > 1
    )
    if duplicates:
        raise ValueError(f"duplicate {label}: {duplicates}")


def validate_source(source: LegacySystem) -> None:
    role_codes = [role.key for role in source.roles]
    user_codes = [user.id for user in source.users]
    require_unique(role_codes, "role keys")
    require_unique(user_codes, "user ids")

    role_code_set = set(role_codes)
    unknown_roles = sorted(
        user.roleKey
        for user in source.users
        if user.roleKey not in role_code_set
    )
    if unknown_roles:
        raise ValueError(f"users reference unknown roles: {unknown_roles}")

    for role in source.roles:
        require_unique(role.permissions, f"permissions of role {role.key}")
        for permission_code in role.permissions:
            if "." not in permission_code:
                raise ValueError(
                    f"invalid permission code {permission_code!r}"
                )


def load_source(path: Path) -> tuple[LegacySystem, str]:
    raw = path.read_bytes()
    source = LegacySystem.model_validate_json(raw)
    validate_source(source)
    return source, hashlib.sha256(raw).hexdigest()


def set_changed(instance: Any, values: dict[str, Any]) -> bool:
    changed = False
    for attribute, value in values.items():
        if getattr(instance, attribute) != value:
            setattr(instance, attribute, value)
            changed = True
    return changed


def new_stats() -> dict[str, int]:
    return {
        "roles_inserted": 0,
        "roles_updated": 0,
        "permissions_inserted": 0,
        "permissions_updated": 0,
        "role_permissions_inserted": 0,
        "users_inserted": 0,
        "users_updated": 0,
    }


async def migrate_roles(
    session: AsyncSession,
    source: LegacySystem,
    stats: dict[str, int],
) -> dict[str, str]:
    existing = {
        item.role_code: item
        for item in (
            await session.scalars(select(Roles).order_by(Roles.role_code))
        ).all()
    }
    role_ids: dict[str, str] = {}

    for legacy in source.roles:
        values = {
            "role_name": legacy.name,
            "role_scope": "platform",
            "is_builtin": 1,
            "description": legacy.description,
        }
        role = existing.get(legacy.key)
        if role is None:
            role = Roles(
                role_id=deterministic_legacy_ulid("role", legacy.key),
                role_code=legacy.key,
                **values,
            )
            session.add(role)
            stats["roles_inserted"] += 1
        elif set_changed(role, values):
            stats["roles_updated"] += 1
        role_ids[legacy.key] = role.role_id

    return role_ids


async def migrate_permissions(
    session: AsyncSession,
    source: LegacySystem,
    stats: dict[str, int],
) -> dict[str, str]:
    permission_codes = sorted(
        {
            permission_code
            for role in source.roles
            for permission_code in role.permissions
        }
    )
    existing = {
        item.permission_code: item
        for item in (
            await session.scalars(
                select(Permissions).order_by(Permissions.permission_code)
            )
        ).all()
    }
    permission_ids: dict[str, str] = {}

    for permission_code in permission_codes:
        values = {
            "permission_name": PERMISSION_NAMES.get(
                permission_code, permission_code
            ),
            "module_code": permission_code.split(".", 1)[0],
            "description": f"由旧版 system.json 迁移：{permission_code}",
        }
        permission = existing.get(permission_code)
        if permission is None:
            permission = Permissions(
                permission_id=deterministic_legacy_ulid(
                    "permission", permission_code
                ),
                permission_code=permission_code,
                **values,
            )
            session.add(permission)
            stats["permissions_inserted"] += 1
        elif set_changed(permission, values):
            stats["permissions_updated"] += 1
        permission_ids[permission_code] = permission.permission_id

    return permission_ids


async def migrate_role_permissions(
    session: AsyncSession,
    source: LegacySystem,
    role_ids: dict[str, str],
    permission_ids: dict[str, str],
    stats: dict[str, int],
) -> None:
    existing = set(
        (
            await session.execute(
                select(
                    RolePermissions.role_id,
                    RolePermissions.permission_id,
                )
            )
        ).all()
    )

    for role in source.roles:
        for permission_code in role.permissions:
            pair = (
                role_ids[role.key],
                permission_ids[permission_code],
            )
            if pair not in existing:
                session.add(
                    RolePermissions(
                        role_id=pair[0],
                        permission_id=pair[1],
                    )
                )
                existing.add(pair)
                stats["role_permissions_inserted"] += 1


async def migrate_users(
    session: AsyncSession,
    source: LegacySystem,
    role_ids: dict[str, str],
    stats: dict[str, int],
) -> dict[str, str]:
    existing = {
        item.username: item
        for item in (
            await session.scalars(select(Users).order_by(Users.username))
        ).all()
    }
    user_ids: dict[str, str] = {}

    for legacy in source.users:
        values = {
            "display_name": legacy.name,
            "platform_role_id": role_ids[legacy.roleKey],
            "avatar_uri": (
                f"initial://{quote(legacy.avatar)}"
                if legacy.avatar
                else None
            ),
        }
        user = existing.get(legacy.id)
        if user is None:
            user = Users(
                user_id=deterministic_legacy_ulid("user", legacy.id),
                username=legacy.id,
                password_hash=LEGACY_PASSWORD_HASH,
                status="active",
                **values,
            )
            session.add(user)
            stats["users_inserted"] += 1
        elif set_changed(user, values):
            stats["users_updated"] += 1
        user_ids[legacy.id] = user.user_id

    return user_ids


async def run_migration(
    database_url: str,
    source: LegacySystem,
    *,
    apply_changes: bool,
) -> dict[str, int]:
    engine = create_database_engine(database_url)
    factory = create_session_factory(engine)
    stats = new_stats()

    try:
        async with factory() as session:
            try:
                role_ids = await migrate_roles(session, source, stats)
                permission_ids = await migrate_permissions(
                    session, source, stats
                )
                await migrate_role_permissions(
                    session,
                    source,
                    role_ids,
                    permission_ids,
                    stats,
                )
                await migrate_users(
                    session, source, role_ids, stats
                )

                if apply_changes:
                    await session.commit()
                else:
                    await session.rollback()
            except Exception:
                await session.rollback()
                raise
    finally:
        await engine.dispose()

    return stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Migrate legacy system.json records into MySQL."
    )
    parser.add_argument(
        "--source",
        required=True,
        type=Path,
        help="Path to the legacy platform_data/system.json file.",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Commit changes. Without this flag the transaction is rolled back.",
    )
    return parser.parse_args()


async def async_main() -> None:
    args = parse_args()
    source_path = args.source.resolve(strict=True)
    source, source_sha256 = load_source(source_path)
    database_url = os.environ["DATABASE_URL"]
    stats = await run_migration(
        database_url,
        source,
        apply_changes=args.apply,
    )
    result = {
        "mode": "apply" if args.apply else "dry-run",
        "source": str(source_path),
        "source_sha256": source_sha256,
        "source_counts": {
            "roles": len(source.roles),
            "permissions": len(
                {
                    permission
                    for role in source.roles
                    for permission in role.permissions
                }
            ),
            "role_permissions": sum(
                len(role.permissions) for role in source.roles
            ),
            "users": len(source.users),
        },
        "changes": stats,
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))


def main() -> None:
    asyncio.run(async_main())


if __name__ == "__main__":
    main()