didier/database/crud/users.py

52 lines
1.2 KiB
Python

from typing import Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from database.schemas import Bank, BankSavings, 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
savings = BankSavings(user_id=user_id)
user.savings = savings
session.add(bank)
session.add(nightly_data)
session.add(savings)
session.add(user)
await session.commit()
await session.refresh(user)
return user