arbret/backend/seed.py

94 lines
3.1 KiB
Python

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