414 lines
12 KiB
Python
414 lines
12 KiB
Python
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()
|