kmtnflower/backend/migrate_roles.py

"""
Migration seed: add missing schema columns/tables, then create the 'admin' role
and assign it to existing is_admin=True users.

Run once after deploying the updated schema. Idempotent — safe to run multiple times.

Usage:
    KMTN_DATABASE_URL=... python migrate_roles.py
    (run from the backend/ directory)
"""

import asyncio
from sqlalchemy import select, text

from app.database import async_session
from app.user_models import User, Role, user_roles


async def ensure_schema(session) -> None:
    """Add missing columns/tables so the User model matches the live database.

    - Adds ``password_hash`` to ``users`` (nullable — existing rows stay NULL).
    - Creates ``roles`` and ``user_roles`` tables if they do not exist.

    All steps use ``IF NOT EXISTS`` / ``NOT EXISTS`` guards so the function is
    fully idempotent.
    """
    conn = await session.connection()

    # 1) Add password_hash column if missing
    result = await conn.execute(
        text(
            "SELECT 1 FROM information_schema.columns "
            "WHERE table_name = 'users' AND column_name = 'password_hash'"
        )
    )
    if not result.scalar():
        await conn.execute(
            text("ALTER TABLE users ADD COLUMN password_hash VARCHAR(255)")
        )
        print("  ✓ Added 'password_hash' column to 'users'")
    else:
        print("  ✓ 'users.password_hash' already exists")

    # 2) Create 'roles' table if missing
    result = await conn.execute(
        text(
            "SELECT 1 FROM information_schema.tables "
            "WHERE table_name = 'roles'"
        )
    )
    if not result.scalar():
        await conn.execute(
            text(
                "CREATE TABLE IF NOT EXISTS roles ("
                "id SERIAL PRIMARY KEY, "
                "name VARCHAR(50) UNIQUE NOT NULL, "
                "description TEXT"
                ")"
            )
        )
        print("  ✓ Created 'roles' table")
    else:
        print("  ✓ 'roles' table already exists")

    # 3) Create 'user_roles' table if missing
    result = await conn.execute(
        text(
            "SELECT 1 FROM information_schema.tables "
            "WHERE table_name = 'user_roles'"
        )
    )
    if not result.scalar():
        await conn.execute(
            text(
                "CREATE TABLE IF NOT EXISTS user_roles ("
                "user_id INTEGER REFERENCES users(id) ON DELETE CASCADE, "
                "role_id INTEGER REFERENCES roles(id) ON DELETE CASCADE, "
                "PRIMARY KEY (user_id, role_id)"
                ")"
            )
        )
        print("  ✓ Created 'user_roles' table")
    else:
        print("  ✓ 'user_roles' table already exists")

    await session.commit()


async def migrate() -> None:
    """Ensure schema is up to date, 'admin' role exists, and legacy admins get the role."""
    async with async_session() as session:
        print("Ensuring schema is up to date...")
        await ensure_schema(session)

        print("Running role migration...")
        # 1) Create 'admin' role if it doesn't exist
        result = await session.execute(select(Role).where(Role.name == "admin"))
        admin_role = result.scalar_one_or_none()
        if admin_role is None:
            admin_role = Role(name="admin", description="Full administrative access")
            session.add(admin_role)
            await session.commit()
            await session.refresh(admin_role)
            print(f"  ✓ Created 'admin' role (id={admin_role.id})")
        else:
            print(f"  ✓ 'admin' role already exists (id={admin_role.id})")

        # 2) Assign 'admin' role to all existing is_admin=True users
        result = await session.execute(
            select(User).where(User.is_admin == True)  # noqa: E712
        )
        admin_users = result.scalars().all()
        migrated = 0
        for user in admin_users:
            if admin_role not in user.roles:
                user.roles.append(admin_role)
                migrated += 1
        if migrated:
            await session.commit()
            print(f"  ✓ Assigned 'admin' role to {migrated} existing admin user(s)")
        else:
            print("  ✓ No new role assignments needed")

        # 3) Summary
        result = await session.execute(
            select(User).where(user_roles.c.role_id == admin_role.id)
        )
        role_admins = result.scalars().all()
        print(f"  → {len(role_admins)} user(s) now have the 'admin' role")


if __name__ == "__main__":
    print("Running role migration...")
    asyncio.run(migrate())
    print("Done.")