"""Seed the database with roles, permissions, and dev users.""" import asyncio import os from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from database import engine, async_session, Base from models import User, Role, Permission, role_permissions, ROLE_DEFINITIONS, ROLE_REGULAR, ROLE_ADMIN 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: AsyncSession, 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: AsyncSession, 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 not role: raise ValueError(f"Role '{role_name}' not found") roles.append(role) if user: user.hashed_password = get_password_hash(password) user.roles = roles # type: ignore[assignment] 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() -> None: 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, [ROLE_REGULAR]) # Create admin dev user await upsert_user(db, DEV_ADMIN_EMAIL, DEV_ADMIN_PASSWORD, [ROLE_ADMIN]) await db.commit() print("\n=== Seeding Complete ===\n") if __name__ == "__main__": asyncio.run(seed())