didier/database/crud/users.py

42 lines
979 B
Python

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