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()