""" 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.")