didier/database/crud/users.py

from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from database.models import Bank, NightlyData, User

__all__ = [
    "get_or_add",
]


async def get_or_add(session: AsyncSession, user_id: int) -> User:
    """Get a user's profile
    If it doesn't exist yet, create it (along with all linked datastructures).
    """
    statement = select(User).where(User.user_id == user_id)
    user: Optional[User] = (await session.execute(statement)).scalar_one_or_none()

    # User exists
    if user is not None:
        return user

    # Create new user
    user = User(user_id=user_id)
    session.add(user)
    await session.commit()

    # Add bank & nightly info
    bank = Bank(user_id=user_id)
    nightly_data = NightlyData(user_id=user_id)
    user.bank = bank
    user.nightly_data = nightly_data

    session.add(bank)
    session.add(nightly_data)
    session.add(user)

    await session.commit()

    return user