model-platform/migrations/versions/b71c4f2a9d10_seed_demo_context.py

"""seed the self-hosted demo users and workspaces

Revision ID: b71c4f2a9d10
Revises: 8d86e2f82860
Create Date: 2026-07-31 16:00:00
"""

from collections.abc import Sequence

from alembic import op
import sqlalchemy as sa


revision: str = "b71c4f2a9d10"
down_revision: str | Sequence[str] | None = "8d86e2f82860"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


ADMIN_ROLE_ID = "0000000000000000000000000A"
DEVELOPER_ROLE_ID = "0000000000000000000000000B"

USERS = (
    ("0000000000RF6FG1SDBXG59S13", "admin-zhang", "张三", ADMIN_ROLE_ID),
    ("0000000000H2QYCGPCWQM1JSGS", "admin-li", "李四", ADMIN_ROLE_ID),
    ("0000000000RWG40ESZPGJT629J", "dev-wang", "王五", DEVELOPER_ROLE_ID),
    ("00000000004CQV7WASJA6N6FW4", "dev-zhao", "赵六", DEVELOPER_ROLE_ID),
)

WORKSPACES = (
    ("00000000000BM630VT9ARVFZPC", "model-development", "模型开发 Workspace"),
    ("0000000000AE0NC0V5T424KK86", "risk-validation", "风险验证 Workspace"),
)


def upgrade() -> None:
    roles = sa.table(
        "roles",
        sa.column("role_id", sa.String),
        sa.column("role_code", sa.String),
        sa.column("role_name", sa.String),
        sa.column("role_scope", sa.String),
        sa.column("is_builtin", sa.Integer),
        sa.column("description", sa.String),
    )
    users = sa.table(
        "users",
        sa.column("user_id", sa.String),
        sa.column("username", sa.String),
        sa.column("display_name", sa.String),
        sa.column("password_hash", sa.String),
        sa.column("status", sa.String),
        sa.column("email", sa.String),
        sa.column("platform_role_id", sa.String),
    )
    workspaces = sa.table(
        "workspaces",
        sa.column("workspace_id", sa.String),
        sa.column("workspace_code", sa.String),
        sa.column("workspace_name", sa.String),
        sa.column("active_root_uri", sa.String),
        sa.column("status", sa.String),
        sa.column("created_by", sa.String),
        sa.column("description", sa.String),
    )
    members = sa.table(
        "workspace_members",
        sa.column("workspace_id", sa.String),
        sa.column("user_id", sa.String),
        sa.column("role_id", sa.String),
        sa.column("member_status", sa.String),
    )

    op.bulk_insert(
        roles,
        [
            {
                "role_id": ADMIN_ROLE_ID,
                "role_code": "admin",
                "role_name": "管理员",
                "role_scope": "workspace",
                "is_builtin": 1,
                "description": "Self-hosted workspace administrator",
            },
            {
                "role_id": DEVELOPER_ROLE_ID,
                "role_code": "developer",
                "role_name": "开发人员",
                "role_scope": "workspace",
                "is_builtin": 1,
                "description": "Self-hosted workspace developer",
            },
        ],
    )
    op.bulk_insert(
        users,
        [
            {
                "user_id": user_id,
                "username": username,
                "display_name": display_name,
                "password_hash": "demo-login-disabled",
                "status": "active",
                "email": f"{username}@model-platform.local",
                "platform_role_id": role_id,
            }
            for user_id, username, display_name, role_id in USERS
        ],
    )
    op.bulk_insert(
        workspaces,
        [
            {
                "workspace_id": workspace_id,
                "workspace_code": workspace_code,
                "workspace_name": workspace_name,
                "active_root_uri": f"s3://workspaces/{workspace_id}/",
                "status": "active",
                "created_by": USERS[0][0],
                "description": "Self-hosted demo workspace",
            }
            for workspace_id, workspace_code, workspace_name in WORKSPACES
        ],
    )
    op.bulk_insert(
        members,
        [
            {
                "workspace_id": workspace_id,
                "user_id": user_id,
                "role_id": role_id,
                "member_status": "active",
            }
            for workspace_id, _, _ in WORKSPACES
            for user_id, _, _, role_id in USERS
        ],
    )


def downgrade() -> None:
    connection = op.get_bind()
    workspace_ids = [workspace_id for workspace_id, _, _ in WORKSPACES]
    user_ids = [user_id for user_id, _, _, _ in USERS]
    connection.execute(
        sa.text(
            "DELETE FROM workspace_members "
            "WHERE workspace_id IN :workspace_ids AND user_id IN :user_ids"
        ).bindparams(
            sa.bindparam("workspace_ids", expanding=True),
            sa.bindparam("user_ids", expanding=True),
        ),
        {"workspace_ids": workspace_ids, "user_ids": user_ids},
    )
    connection.execute(
        sa.text("DELETE FROM workspaces WHERE workspace_id IN :workspace_ids")
        .bindparams(sa.bindparam("workspace_ids", expanding=True)),
        {"workspace_ids": workspace_ids},
    )
    connection.execute(
        sa.text("DELETE FROM users WHERE user_id IN :user_ids")
        .bindparams(sa.bindparam("user_ids", expanding=True)),
        {"user_ids": user_ids},
    )
    connection.execute(
        sa.text("DELETE FROM roles WHERE role_id IN :role_ids")
        .bindparams(sa.bindparam("role_ids", expanding=True)),
        {"role_ids": [ADMIN_ROLE_ID, DEVELOPER_ROLE_ID]},
    )