model-platform/migrations/data/migrate_legacy_workspaces.py

from __future__ import annotations

import argparse
import ast
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, 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 (
    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,
    }


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

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