Files
model-platform/migrations/data/migrate_system_json.py
T
2026-08-14 19:51:58 +08:00

406 lines
11 KiB
Python

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 common.db import create_database_engine, create_session_factory
from common.db.models import (
Permissions,
RolePermissions,
Roles,
Users,
)
from pydantic import BaseModel, ConfigDict
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
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()
).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()