arbret/backend/seed.py

"""Seed the database with roles, permissions, and dev users."""
import asyncio
import os
from typing import List
from sqlalchemy import select

from database import engine, async_session, Base
from models import User, Role, Permission, role_permissions, ROLE_DEFINITIONS
from auth import get_password_hash

DEV_USER_EMAIL = os.environ["DEV_USER_EMAIL"]
DEV_USER_PASSWORD = os.environ["DEV_USER_PASSWORD"]
DEV_ADMIN_EMAIL = os.environ["DEV_ADMIN_EMAIL"]
DEV_ADMIN_PASSWORD = os.environ["DEV_ADMIN_PASSWORD"]


async def upsert_role(db, name: str, description: str, permissions: List[Permission]) -> Role:
    """Create or update a role with the given permissions."""
    result = await db.execute(select(Role).where(Role.name == name))
    role = result.scalar_one_or_none()
    
    if role:
        role.description = description
        print(f"Updated role: {name}")
    else:
        role = Role(name=name, description=description)
        db.add(role)
        await db.flush()  # Get the role ID
        print(f"Created role: {name}")
    
    # Set permissions for the role
    await role.set_permissions(db, permissions)
    print(f"  Permissions: {', '.join(p.value for p in permissions)}")
    
    return role


async def upsert_user(db, email: str, password: str, role_names: List[str]) -> User:
    """Create or update a user with the given credentials and roles."""
    result = await db.execute(select(User).where(User.email == email))
    user = result.scalar_one_or_none()
    
    # Get roles
    roles = []
    for role_name in role_names:
        result = await db.execute(select(Role).where(Role.name == role_name))
        role = result.scalar_one_or_none()
        if role:
            roles.append(role)
    
    if user:
        user.hashed_password = get_password_hash(password)
        user.roles = roles
        print(f"Updated user: {email} with roles: {role_names}")
    else:
        user = User(
            email=email,
            hashed_password=get_password_hash(password),
            roles=roles,
        )
        db.add(user)
        print(f"Created user: {email} with roles: {role_names}")
    
    return user


async def seed():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    async with async_session() as db:
        print("\n=== Seeding Roles ===")
        for role_name, role_config in ROLE_DEFINITIONS.items():
            await upsert_role(
                db,
                role_name,
                role_config["description"],
                role_config["permissions"],
            )
        
        print("\n=== Seeding Users ===")
        # Create regular dev user
        await upsert_user(db, DEV_USER_EMAIL, DEV_USER_PASSWORD, ["regular"])
        
        # Create admin dev user
        await upsert_user(db, DEV_ADMIN_EMAIL, DEV_ADMIN_PASSWORD, ["admin"])
        
        await db.commit()
        print("\n=== Seeding Complete ===\n")


if __name__ == "__main__":
    asyncio.run(seed())