from typing import Optional from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from database.schemas import Bank, NightlyData, User __all__ = [ "get_or_add_user", ] async def get_or_add_user(session: AsyncSession, user_id: int, *, options: Optional[list] = None) -> User: """Get a user's profile If it doesn't exist yet, create it (along with all linked datastructures) """ if options is None: options = [] statement = select(User).where(User.user_id == user_id).options(*options) 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() await session.refresh(user) return user