重构模型平台前后端并移除Redis依赖
This commit is contained in:
Winnie
@@ -0,0 +1,413 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
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 (
|
||||
AuditLogs,
|
||||
Roles,
|
||||
Users,
|
||||
WorkspaceMembers,
|
||||
Workspaces,
|
||||
)
|
||||
from migrations.data.migrate_system_json import (
|
||||
deterministic_legacy_ulid,
|
||||
set_changed,
|
||||
)
|
||||
|
||||
|
||||
class LegacyWorkspace(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str | None = None
|
||||
userIds: list[str]
|
||||
portOffset: int
|
||||
|
||||
|
||||
WORKSPACE_LIST = TypeAdapter(list[LegacyWorkspace])
|
||||
|
||||
|
||||
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 extract_workspace_definitions(
|
||||
source_path: Path,
|
||||
) -> tuple[list[LegacyWorkspace], str]:
|
||||
raw = source_path.read_bytes()
|
||||
module = ast.parse(
|
||||
raw.decode("utf-8-sig"),
|
||||
filename=str(source_path),
|
||||
)
|
||||
definition: Any | None = None
|
||||
|
||||
for statement in module.body:
|
||||
if not isinstance(statement, ast.Assign):
|
||||
continue
|
||||
if any(
|
||||
isinstance(target, ast.Name)
|
||||
and target.id == "WORKSPACE_DEFINITIONS"
|
||||
for target in statement.targets
|
||||
):
|
||||
definition = ast.literal_eval(statement.value)
|
||||
break
|
||||
|
||||
if definition is None:
|
||||
raise ValueError("WORKSPACE_DEFINITIONS was not found")
|
||||
|
||||
workspaces = WORKSPACE_LIST.validate_python(definition)
|
||||
require_unique([item.id for item in workspaces], "workspace ids")
|
||||
for workspace in workspaces:
|
||||
require_unique(
|
||||
workspace.userIds,
|
||||
f"members of workspace {workspace.id}",
|
||||
)
|
||||
if not workspace.userIds:
|
||||
raise ValueError(
|
||||
f"workspace {workspace.id!r} has no members"
|
||||
)
|
||||
|
||||
return workspaces, hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def new_stats() -> dict[str, int]:
|
||||
return {
|
||||
"workspaces_inserted": 0,
|
||||
"workspaces_updated": 0,
|
||||
"workspace_members_inserted": 0,
|
||||
"workspace_members_updated": 0,
|
||||
"audit_logs_workspace_backfilled": 0,
|
||||
}
|
||||
|
||||
|
||||
async def load_users_and_roles(
|
||||
session: AsyncSession,
|
||||
) -> tuple[dict[str, Users], dict[str, str]]:
|
||||
users = {
|
||||
item.username: item
|
||||
for item in (
|
||||
await session.scalars(select(Users).order_by(Users.username))
|
||||
).all()
|
||||
}
|
||||
roles_by_id = {
|
||||
item.role_id: item.role_code
|
||||
for item in (
|
||||
await session.scalars(select(Roles).order_by(Roles.role_code))
|
||||
).all()
|
||||
}
|
||||
role_codes_by_username = {
|
||||
username: roles_by_id[user.platform_role_id]
|
||||
for username, user in users.items()
|
||||
if user.platform_role_id in roles_by_id
|
||||
}
|
||||
return users, role_codes_by_username
|
||||
|
||||
|
||||
def validate_members(
|
||||
workspaces: list[LegacyWorkspace],
|
||||
users: dict[str, Users],
|
||||
role_codes_by_username: dict[str, str],
|
||||
) -> None:
|
||||
source_members = {
|
||||
username
|
||||
for workspace in workspaces
|
||||
for username in workspace.userIds
|
||||
}
|
||||
missing_users = sorted(source_members - set(users))
|
||||
if missing_users:
|
||||
raise ValueError(
|
||||
f"workspace members were not migrated as users: {missing_users}"
|
||||
)
|
||||
users_without_roles = sorted(
|
||||
source_members - set(role_codes_by_username)
|
||||
)
|
||||
if users_without_roles:
|
||||
raise ValueError(
|
||||
"workspace members have no platform role: "
|
||||
f"{users_without_roles}"
|
||||
)
|
||||
|
||||
|
||||
def workspace_creator(
|
||||
workspace: LegacyWorkspace,
|
||||
users: dict[str, Users],
|
||||
role_codes_by_username: dict[str, str],
|
||||
) -> Users:
|
||||
admin_username = next(
|
||||
(
|
||||
username
|
||||
for username in workspace.userIds
|
||||
if role_codes_by_username[username] == "admin"
|
||||
),
|
||||
workspace.userIds[0],
|
||||
)
|
||||
return users[admin_username]
|
||||
|
||||
|
||||
async def migrate_workspaces(
|
||||
session: AsyncSession,
|
||||
source: list[LegacyWorkspace],
|
||||
users: dict[str, Users],
|
||||
role_codes_by_username: dict[str, str],
|
||||
stats: dict[str, int],
|
||||
) -> dict[str, str]:
|
||||
existing = {
|
||||
item.workspace_code: item
|
||||
for item in (
|
||||
await session.scalars(
|
||||
select(Workspaces).order_by(Workspaces.workspace_code)
|
||||
)
|
||||
).all()
|
||||
}
|
||||
workspace_ids: dict[str, str] = {}
|
||||
|
||||
for legacy in source:
|
||||
workspace = existing.get(legacy.id)
|
||||
if workspace is None:
|
||||
workspace_id = deterministic_legacy_ulid(
|
||||
"workspace", legacy.id
|
||||
)
|
||||
creator = workspace_creator(
|
||||
legacy,
|
||||
users,
|
||||
role_codes_by_username,
|
||||
)
|
||||
workspace = Workspaces(
|
||||
workspace_id=workspace_id,
|
||||
workspace_code=legacy.id,
|
||||
workspace_name=legacy.name,
|
||||
description=legacy.description,
|
||||
active_root_uri=(
|
||||
"file:///workspace/workspaces/"
|
||||
f"{quote(legacy.id, safe='')}"
|
||||
),
|
||||
quota_bytes=0,
|
||||
used_bytes=0,
|
||||
status="active",
|
||||
created_by=creator.user_id,
|
||||
artifact_bucket="model-platform",
|
||||
artifact_prefix=f"workspaces/{workspace_id}",
|
||||
)
|
||||
session.add(workspace)
|
||||
stats["workspaces_inserted"] += 1
|
||||
else:
|
||||
values = {
|
||||
"workspace_name": legacy.name,
|
||||
"description": legacy.description,
|
||||
"active_root_uri": (
|
||||
"file:///workspace/workspaces/"
|
||||
f"{quote(legacy.id, safe='')}"
|
||||
),
|
||||
"artifact_bucket": "model-platform",
|
||||
"artifact_prefix": (
|
||||
f"workspaces/{workspace.workspace_id}"
|
||||
),
|
||||
}
|
||||
if set_changed(workspace, values):
|
||||
stats["workspaces_updated"] += 1
|
||||
workspace_ids[legacy.id] = workspace.workspace_id
|
||||
|
||||
return workspace_ids
|
||||
|
||||
|
||||
async def migrate_members(
|
||||
session: AsyncSession,
|
||||
source: list[LegacyWorkspace],
|
||||
workspace_ids: dict[str, str],
|
||||
users: dict[str, Users],
|
||||
stats: dict[str, int],
|
||||
) -> None:
|
||||
existing = {
|
||||
(item.workspace_id, item.user_id): item
|
||||
for item in (
|
||||
await session.scalars(select(WorkspaceMembers))
|
||||
).all()
|
||||
}
|
||||
|
||||
for legacy in source:
|
||||
workspace_id = workspace_ids[legacy.id]
|
||||
for username in legacy.userIds:
|
||||
user = users[username]
|
||||
pair = (workspace_id, user.user_id)
|
||||
member = existing.get(pair)
|
||||
values = {
|
||||
"role_id": user.platform_role_id,
|
||||
"member_status": "active",
|
||||
}
|
||||
if member is None:
|
||||
member = WorkspaceMembers(
|
||||
workspace_id=workspace_id,
|
||||
user_id=user.user_id,
|
||||
**values,
|
||||
)
|
||||
session.add(member)
|
||||
existing[pair] = member
|
||||
stats["workspace_members_inserted"] += 1
|
||||
elif set_changed(member, values):
|
||||
stats["workspace_members_updated"] += 1
|
||||
|
||||
|
||||
async def backfill_audit_workspaces(
|
||||
session: AsyncSession,
|
||||
source: list[LegacyWorkspace],
|
||||
workspace_ids: dict[str, str],
|
||||
users: dict[str, Users],
|
||||
stats: dict[str, int],
|
||||
) -> None:
|
||||
workspace_codes_by_username: dict[str, set[str]] = defaultdict(set)
|
||||
for workspace in source:
|
||||
for username in workspace.userIds:
|
||||
workspace_codes_by_username[username].add(workspace.id)
|
||||
|
||||
workspace_by_user_id = {
|
||||
users[username].user_id: workspace_ids[next(iter(codes))]
|
||||
for username, codes in workspace_codes_by_username.items()
|
||||
if len(codes) == 1
|
||||
}
|
||||
|
||||
audit_logs = (
|
||||
await session.scalars(
|
||||
select(AuditLogs).where(AuditLogs.workspace_id.is_(None))
|
||||
)
|
||||
).all()
|
||||
for audit_log in audit_logs:
|
||||
if (
|
||||
not isinstance(audit_log.detail_json, dict)
|
||||
or audit_log.detail_json.get("migration_source")
|
||||
!= "platform_data/system.json"
|
||||
):
|
||||
continue
|
||||
workspace_id = workspace_by_user_id.get(
|
||||
audit_log.actor_user_id
|
||||
)
|
||||
if workspace_id is not None:
|
||||
audit_log.workspace_id = workspace_id
|
||||
stats["audit_logs_workspace_backfilled"] += 1
|
||||
|
||||
|
||||
async def run_migration(
|
||||
database_url: str,
|
||||
source: list[LegacyWorkspace],
|
||||
*,
|
||||
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:
|
||||
users, role_codes_by_username = (
|
||||
await load_users_and_roles(session)
|
||||
)
|
||||
validate_members(
|
||||
source,
|
||||
users,
|
||||
role_codes_by_username,
|
||||
)
|
||||
workspace_ids = await migrate_workspaces(
|
||||
session,
|
||||
source,
|
||||
users,
|
||||
role_codes_by_username,
|
||||
stats,
|
||||
)
|
||||
await migrate_members(
|
||||
session,
|
||||
source,
|
||||
workspace_ids,
|
||||
users,
|
||||
stats,
|
||||
)
|
||||
await backfill_audit_workspaces(
|
||||
session,
|
||||
source,
|
||||
workspace_ids,
|
||||
users,
|
||||
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 WORKSPACE_DEFINITIONS into MySQL."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source",
|
||||
required=True,
|
||||
type=Path,
|
||||
help="Path to the legacy server.py 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)
|
||||
workspaces, source_sha256 = extract_workspace_definitions(source_path)
|
||||
stats = await run_migration(
|
||||
os.environ["DATABASE_URL"],
|
||||
workspaces,
|
||||
apply_changes=args.apply,
|
||||
)
|
||||
result = {
|
||||
"mode": "apply" if args.apply else "dry-run",
|
||||
"source": str(source_path),
|
||||
"source_symbol": "WORKSPACE_DEFINITIONS",
|
||||
"source_sha256": source_sha256,
|
||||
"source_counts": {
|
||||
"workspaces": len(workspaces),
|
||||
"workspace_members": sum(
|
||||
len(workspace.userIds) for workspace in workspaces
|
||||
),
|
||||
},
|
||||
"changes": stats,
|
||||
}
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
asyncio.run(async_main())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user